From e6dee738a84239ebe038a709a2c4245b2896d04e Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 12:55:41 -0400 Subject: [PATCH 01/18] refactor: re-apply Phase 1 perf/best-practices cleanup (cosmetic + fail-fast) Re-applies the previously-reverted Phase 1 batch from the perf/best-practices audit: fix the FP8 Q/DQ gate comment to match the actual threshold, convert a dict() call to a dict literal (clears the one pre-existing ruff C408 on a touched file), require an explicit opt_batch_size in EngineBuilder.build() and its 5 compile_* wrappers (real callers already always pass it), drop a dead try/except around a single logger.info call, and tighten two type hints in StreamDiffusion.pipeline (Optional/Union/Tuple). No behavior change except the opt_batch_size fail-fast, which only affects callers that previously relied on an unused default of 1. --- .../acceleration/tensorrt/__init__.py | 10 ++++---- .../acceleration/tensorrt/builder.py | 22 +++++++++-------- src/streamdiffusion/pipeline.py | 24 +++++++------------ src/streamdiffusion/wrapper.py | 16 ++++++------- 4 files changed, 34 insertions(+), 38 deletions(-) diff --git a/src/streamdiffusion/acceleration/tensorrt/__init__.py b/src/streamdiffusion/acceleration/tensorrt/__init__.py index 7a74b65a2..b4c7a7ca2 100644 --- a/src/streamdiffusion/acceleration/tensorrt/__init__.py +++ b/src/streamdiffusion/acceleration/tensorrt/__init__.py @@ -66,7 +66,7 @@ def compile_vae_encoder( onnx_path: str, onnx_opt_path: str, engine_path: str, - opt_batch_size: int = 1, + opt_batch_size: Optional[int] = None, engine_build_options: Optional[dict] = None, ): vae = vae.to(torch.device("cuda")) @@ -86,7 +86,7 @@ def compile_vae_decoder( onnx_path: str, onnx_opt_path: str, engine_path: str, - opt_batch_size: int = 1, + opt_batch_size: Optional[int] = None, engine_build_options: Optional[dict] = None, ): vae = vae.to(torch.device("cuda")) @@ -106,7 +106,7 @@ def compile_safety_checker( onnx_path: str, onnx_opt_path: str, engine_path: str, - opt_batch_size: int = 1, + opt_batch_size: Optional[int] = None, engine_build_options: Optional[dict] = None, ): safety_checker = safety_checker.to(torch.device("cuda")) @@ -126,7 +126,7 @@ def compile_unet( onnx_path: str, onnx_opt_path: str, engine_path: str, - opt_batch_size: int = 1, + opt_batch_size: Optional[int] = None, engine_build_options: Optional[dict] = None, ): # Extract FP8-specific options before passing the rest to EngineBuilder.build(). @@ -177,7 +177,7 @@ def compile_controlnet( onnx_path: str, onnx_opt_path: str, engine_path: str, - opt_batch_size: int = 1, + opt_batch_size: Optional[int] = None, engine_build_options: Optional[dict] = None, ): build_options = dict(engine_build_options or {}) diff --git a/src/streamdiffusion/acceleration/tensorrt/builder.py b/src/streamdiffusion/acceleration/tensorrt/builder.py index 2b6672bde..045dd7390 100644 --- a/src/streamdiffusion/acceleration/tensorrt/builder.py +++ b/src/streamdiffusion/acceleration/tensorrt/builder.py @@ -93,7 +93,7 @@ def build( engine_path: str, opt_image_height: int = 512, opt_image_width: int = 512, - opt_batch_size: int = 1, + opt_batch_size: Optional[int] = None, min_image_resolution: int = 256, max_image_resolution: int = 1024, build_enable_refit: bool = False, @@ -118,6 +118,8 @@ def build( is_controlnet: bool = False, artifact_prefix: str = "unet", ): + if opt_batch_size is None: + raise ValueError("build() requires an explicit opt_batch_size") build_total_start = time.perf_counter() engine_name = Path(engine_path).parent.name engine_filename = Path(engine_path).name @@ -146,14 +148,14 @@ def build( else: print(f"Exporting model: {onnx_path}") t0 = time.perf_counter() - _export_kwargs = dict( - onnx_path=onnx_path, - model_data=self.model, - opt_image_height=opt_image_height, - opt_image_width=opt_image_width, - opt_batch_size=opt_batch_size, - onnx_opset=onnx_opset, - ) + _export_kwargs = { + "onnx_path": onnx_path, + "model_data": self.model, + "opt_image_height": opt_image_height, + "opt_image_width": opt_image_width, + "opt_batch_size": opt_batch_size, + "onnx_opset": onnx_opset, + } export_onnx(self.network, **_export_kwargs) elapsed = time.perf_counter() - t0 stats["stages"]["onnx_export"] = {"status": "built", "elapsed_s": round(elapsed, 2)} @@ -304,7 +306,7 @@ def _quant_fn(): stats["stages"]["trt_build"] = {"status": "built", "elapsed_s": round(elapsed, 2)} _build_logger.info(f"[BUILD] TRT engine build ({engine_filename}): {elapsed:.1f}s") - # --- FP8 Q/DQ layer count (sanity gate: < 100 means quantization is inactive) --- + # --- FP8 Q/DQ layer count (sanity gate: < 500 means quantization is inactive) --- if fp8 and os.path.exists(engine_path): try: import json as _json diff --git a/src/streamdiffusion/pipeline.py b/src/streamdiffusion/pipeline.py index 7fa0c6f48..f2771a3ad 100644 --- a/src/streamdiffusion/pipeline.py +++ b/src/streamdiffusion/pipeline.py @@ -631,11 +631,7 @@ def prepare( _use_seq_loop = not (self.use_denoising_batch and isinstance(self.scheduler, LCMScheduler)) if _use_seq_loop and self.sub_timesteps_tensor.dim() >= 1: self._sub_timesteps_expanded = ( - self.sub_timesteps_tensor - .view(-1) - .unsqueeze(1) - .expand(-1, self.frame_bff_size) - .contiguous() + self.sub_timesteps_tensor.view(-1).unsqueeze(1).expand(-1, self.frame_bff_size).contiguous() ) # shape [loop_steps, frame_bff_size] else: self._sub_timesteps_expanded = None @@ -726,14 +722,12 @@ def _refresh_derived_tensors(self) -> None: dtype=self.dtype, device=self.device, ) - self._cfg_t_buf = torch.empty( - cfg_batch, dtype=self.sub_timesteps_tensor.dtype, device=self.device - ) + self._cfg_t_buf = torch.empty(cfg_batch, dtype=self.sub_timesteps_tensor.dtype, device=self.device) else: self._cfg_latent_buf = None self._cfg_t_buf = None - def _get_scheduler_scalings(self, timestep): + def _get_scheduler_scalings(self, timestep: Union[int, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]: """Get LCM/TCD-specific scaling factors for boundary conditions.""" if isinstance(self.scheduler, LCMScheduler): c_skip, c_out = self.scheduler.get_scalings_for_boundary_condition_discrete(timestep) @@ -1146,8 +1140,10 @@ def predict_x0_batch(self, x_t_latent: torch.Tensor) -> torch.Tensor: # Resolve scalar timestep tensor on device. # Avoid per-step t.view(1).repeat() allocations by using the # precomputed _sub_timesteps_expanded table from prepare(). - t = timestep if isinstance(timestep, torch.Tensor) else torch.tensor( - timestep, device=self.device, dtype=torch.long + t = ( + timestep + if isinstance(timestep, torch.Tensor) + else torch.tensor(timestep, device=self.device, dtype=torch.long) ) if self._sub_timesteps_expanded is not None: t_expanded = self._sub_timesteps_expanded[idx] # [frame_bff_size], no alloc @@ -1188,7 +1184,7 @@ def predict_x0_batch(self, x_t_latent: torch.Tensor) -> torch.Tensor: return x_0_pred_out @torch.inference_mode() - def __call__(self, x: Union[torch.Tensor, PIL.Image.Image, np.ndarray] = None) -> torch.Tensor: + def __call__(self, x: Optional[Union[torch.Tensor, PIL.Image.Image, np.ndarray]] = None) -> torch.Tensor: start = self._timing_start end = self._timing_end if self.similar_image_filter: @@ -1241,9 +1237,7 @@ def __call__(self, x: Union[torch.Tensor, PIL.Image.Image, np.ndarray] = None) - torch.cuda.empty_cache() raise if "expanded size" in _msg and "must match" in _msg: - logger.error( - "StreamDiffusion.__call__: tensor size mismatch — attempting scheduler rebuild: %s", _e - ) + logger.error("StreamDiffusion.__call__: tensor size mismatch — attempting scheduler rebuild: %s", _e) try: self._param_updater._update_timestep_calculations() self._refresh_derived_tensors() diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py index f17a57422..f0579fce1 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -737,8 +737,7 @@ def update_stream_params( _normalized = [(str(p), float(w)) for p, w in prompt_list] _current = self.stream._param_updater.get_current_prompts() _neg_unchanged = ( - negative_prompt is None - or negative_prompt == self.stream._param_updater._current_negative_prompt + negative_prompt is None or negative_prompt == self.stream._param_updater._current_negative_prompt ) if _normalized == _current and _neg_unchanged: logger.debug( @@ -2100,10 +2099,7 @@ def _load_model( num_ip_layers = getattr(temp_wrapped_unet.ipadapter_wrapper, "num_ip_layers", None) if not isinstance(num_ip_layers, int) or num_ip_layers <= 0: raise RuntimeError("Failed to determine num_ip_layers for IP-Adapter") - try: - logger.info(f"compile_and_load_engine: discovered num_ip_layers={num_ip_layers}") - except Exception: - pass + logger.info(f"compile_and_load_engine: discovered num_ip_layers={num_ip_layers}") unet_model = UNet( stream.unet, @@ -2358,7 +2354,9 @@ def _install_cached_proc(attn_module): try: del stream.unet except Exception as del_error: - logger.debug(f"Failed to delete stream.unet during OOM fallback: {del_error}", exc_info=True) + logger.debug( + f"Failed to delete stream.unet during OOM fallback: {del_error}", exc_info=True + ) self.cleanup_gpu_memory() @@ -2416,7 +2414,9 @@ def _install_cached_proc(attn_module): try: del stream.vae except Exception as del_error: - logger.debug(f"Failed to delete stream.vae during OOM fallback: {del_error}", exc_info=True) + logger.debug( + f"Failed to delete stream.vae during OOM fallback: {del_error}", exc_info=True + ) self.cleanup_gpu_memory() From 1070583b4e748de7112af43f725c01e3da539172 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 13:08:29 -0400 Subject: [PATCH 02/18] refactor: Phase 2 exception hygiene and observability fixes - Add shared _is_oom_error() helper (typed torch.cuda.OutOfMemoryError + string-heuristic fallback) and use it at both TensorRT UNet/VAE engine build fallback sites instead of duplicated inline substring checks. - Log swallowed LoRA candidate-weight-name failures in _load_lora_with_offline_fallback instead of silently continuing. - Log the inner fallback failure in advanced-model-detection instead of a bare except/pass. - Hoist the optional diffusers_ipadapter helper imports in IPAdapterModule.build_unet_hook out of the per-frame hook into a single build-time resolution, and log all previously-silent per-step fallback branches for observability. - Make _load_model fail fast (raise RuntimeError) when an SDXL pipeline retry after a type mismatch also fails, instead of silently continuing with the wrong pipeline type; also null out the stale mismatched `pipe` reference so a later loading-method failure can't let it slip through the final success check. - Add regression tests for _is_oom_error and the SDXL fail-fast path. --- .../modules/ipadapter_module.py | 62 ++++--- src/streamdiffusion/pipeline.py | 1 + src/streamdiffusion/wrapper.py | 46 +++--- tests/unit/test_wrapper_exception_hygiene.py | 152 ++++++++++++++++++ 4 files changed, 220 insertions(+), 41 deletions(-) create mode 100644 tests/unit/test_wrapper_exception_hygiene.py diff --git a/src/streamdiffusion/modules/ipadapter_module.py b/src/streamdiffusion/modules/ipadapter_module.py index 5bc10579a..caca6c3c1 100644 --- a/src/streamdiffusion/modules/ipadapter_module.py +++ b/src/streamdiffusion/modules/ipadapter_module.py @@ -431,6 +431,26 @@ def build_unet_hook(self, stream) -> UnetHook: """ _last_enabled_state = None # Track previous enabled state to avoid redundant updates + # Resolve the optional diffusers_ipadapter helpers once, at hook-build time, + # instead of importing them on every frame inside the per-step hook below. + build_time_weight_factor = None + build_layer_weights = None + try: + from diffusers_ipadapter.ip_adapter.attention_processor import ( + build_layer_weights as _build_layer_weights, + ) + from diffusers_ipadapter.ip_adapter.attention_processor import ( + build_time_weight_factor as _build_time_weight_factor, + ) + + build_time_weight_factor = _build_time_weight_factor + build_layer_weights = _build_layer_weights + except Exception as import_error: + logger.debug( + f"IPAdapterModule.build_unet_hook: diffusers_ipadapter helpers unavailable: {import_error}", + exc_info=True, + ) + def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta: # If no IP-Adapter installed, do nothing if not hasattr(stream, "ipadapter") or stream.ipadapter is None: @@ -442,7 +462,8 @@ def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta: # Read base weight and weight type from IPAdapter instance try: base_weight = float(getattr(stream.ipadapter, "scale", 1.0)) if enabled else 0.0 - except Exception: + except Exception as scale_error: + logger.debug(f"IPAdapterModule: failed to read ipadapter.scale: {scale_error}", exc_info=True) base_weight = 0.0 if not enabled else 1.0 weight_type = getattr(stream.ipadapter, "weight_type", None) @@ -453,41 +474,38 @@ def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta: total_steps = int(stream.denoising_steps_num) elif hasattr(stream, "t_list") and stream.t_list is not None: total_steps = len(stream.t_list) - except Exception: + except Exception as steps_error: + logger.debug(f"IPAdapterModule: failed to determine total_steps: {steps_error}", exc_info=True) total_steps = None time_factor = 1.0 - if total_steps is not None and ctx.step_index is not None: + if total_steps is not None and ctx.step_index is not None and build_time_weight_factor is not None: try: - from diffusers_ipadapter.ip_adapter.attention_processor import build_time_weight_factor - time_factor = float(build_time_weight_factor(weight_type, int(ctx.step_index), int(total_steps))) - except Exception: + except Exception as scale_factor_error: # Do not add fallback mechanisms - pass + logger.debug( + f"IPAdapterModule: build_time_weight_factor failed: {scale_factor_error}", exc_info=True + ) # TensorRT engine path: supply ipadapter_scale vector via extra kwargs try: is_trt_unet = ( hasattr(stream, "unet") and hasattr(stream.unet, "engine") and hasattr(stream.unet, "stream") ) - except Exception: + except Exception as trt_check_error: + logger.debug(f"IPAdapterModule: TensorRT UNet detection failed: {trt_check_error}", exc_info=True) is_trt_unet = False if is_trt_unet and getattr(stream.unet, "use_ipadapter", False): - try: - from diffusers_ipadapter.ip_adapter.attention_processor import build_layer_weights - except Exception: - # If helper unavailable, do not construct weights here - build_layer_weights = None # type: ignore - num_ip_layers = getattr(stream.unet, "num_ip_layers", None) if isinstance(num_ip_layers, int) and num_ip_layers > 0: weights_tensor = None try: if build_layer_weights is not None: weights_tensor = build_layer_weights(num_ip_layers, float(base_weight), weight_type) - except Exception: + except Exception as weights_error: + logger.debug(f"IPAdapterModule: build_layer_weights failed: {weights_error}", exc_info=True) weights_tensor = None if weights_tensor is None: weights_tensor = torch.full( @@ -496,8 +514,11 @@ def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta: # Apply per-step time factor try: weights_tensor = weights_tensor * float(time_factor) - except Exception: - pass + except Exception as time_scale_error: + logger.debug( + f"IPAdapterModule: applying time_factor to weights_tensor failed: {time_scale_error}", + exc_info=True, + ) return UnetKwargsDelta(extra_unet_kwargs={"ipadapter_scale": weights_tensor}) # PyTorch UNet path: modulate installed processor scales by time factor and enabled state @@ -513,8 +534,11 @@ def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta: # Apply both enabled state and time factor final_scale = float(base_val) * float(time_factor) if enabled else 0.0 proc.scale = final_scale - except Exception: - pass + except Exception as processor_scale_error: + logger.debug( + f"IPAdapterModule: PyTorch UNet processor scale update failed: {processor_scale_error}", + exc_info=True, + ) return UnetKwargsDelta() diff --git a/src/streamdiffusion/pipeline.py b/src/streamdiffusion/pipeline.py index f2771a3ad..e8964e763 100644 --- a/src/streamdiffusion/pipeline.py +++ b/src/streamdiffusion/pipeline.py @@ -365,6 +365,7 @@ def _load_lora_with_offline_fallback( return except Exception as e: last_err = e + logger.debug(f"load_lora_weights: candidate {weight_name!r} failed: {e}", exc_info=True) continue if last_err is not None: diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py index f0579fce1..1e9c3463c 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -17,6 +17,19 @@ logger = logging.getLogger(__name__) + +def _is_oom_error(exc: BaseException) -> bool: + """Detect CUDA out-of-memory errors, including ones surfaced as generic + RuntimeErrors by third-party code (e.g. TensorRT) rather than the typed + torch.cuda.OutOfMemoryError.""" + if isinstance(exc, torch.cuda.OutOfMemoryError): + return True + error_msg = str(exc).lower() + return ( + "out of memory" in error_msg or "outofmemory" in error_msg or "oom" in error_msg or "cuda error" in error_msg + ) + + # Text-encoder CPU offload frees ~1.6 GB VRAM but each prompt update pays a # CPU<->GPU round-trip plus torch.cuda.empty_cache() — a measurable stall # mid-stream on high-VRAM GPUs. Default off (encoders stay resident on GPU); @@ -1509,8 +1522,13 @@ def _load_model( pipe = StableDiffusionXLPipeline.from_single_file(model_id_or_path).to(dtype=self.dtype) logger.info("_load_model: Successfully loaded using SDXL pipeline on retry") except Exception as retry_error: - logger.warning(f"_load_model: SDXL pipeline retry failed: {retry_error}") - # Continue with the originally loaded pipeline + # Discard the mismatched-type pipe so a subsequent loading-method + # failure can't leave a wrong-type pipe looking like a success. + pipe = None + raise RuntimeError( + f"_load_model: SDXL model detected but pipeline retry with " + f"StableDiffusionXLPipeline also failed: {retry_error}" + ) from retry_error break except Exception as e: @@ -1851,8 +1869,8 @@ def _load_model( unet_arch = extract_unet_architecture(stream.unet) unet_arch = validate_architecture(unet_arch, model_type) use_controlnet_trt = True - except Exception: - pass + except Exception as fallback_error: + logger.error(f"Basic fallback detection also failed: {fallback_error}", exc_info=True) if not use_controlnet_trt and not self.use_controlnet: logger.info("ControlNet not enabled, building engines without ControlNet support") @@ -2336,15 +2354,7 @@ def _install_cached_proc(attn_module): ) except Exception as e: - error_msg = str(e).lower() - is_oom_error = ( - "out of memory" in error_msg - or "outofmemory" in error_msg - or "oom" in error_msg - or "cuda error" in error_msg - ) - - if is_oom_error: + if _is_oom_error(e): logger.error(f"TensorRT UNet engine OOM: {e}") logger.info("Falling back to PyTorch UNet (no TensorRT acceleration)") logger.info("This will be slower but should work with less memory") @@ -2396,15 +2406,7 @@ def _install_cached_proc(attn_module): logger.info("TensorRT VAE engines loaded successfully") except Exception as e: - error_msg = str(e).lower() - is_oom_error = ( - "out of memory" in error_msg - or "outofmemory" in error_msg - or "oom" in error_msg - or "cuda error" in error_msg - ) - - if is_oom_error: + if _is_oom_error(e): logger.error(f"TensorRT VAE engine OOM: {e}") logger.info("Falling back to PyTorch VAE (no TensorRT acceleration)") logger.info("This will be slower but should work with less memory") diff --git a/tests/unit/test_wrapper_exception_hygiene.py b/tests/unit/test_wrapper_exception_hygiene.py new file mode 100644 index 000000000..af10d1e68 --- /dev/null +++ b/tests/unit/test_wrapper_exception_hygiene.py @@ -0,0 +1,152 @@ +""" +Regression tests for wrapper.py exception-hygiene fixes (Phase 2 of the +perf/best-practices remediation, docs/perf_bestpractices_audit_2026-07-10.md). + +Covers two independent fixes: + +1. `_is_oom_error` — a shared OOM-detection helper that recognizes both the + typed `torch.cuda.OutOfMemoryError` and the string-matching heuristic used + by third-party code (e.g. TensorRT) that surfaces OOM as a generic + RuntimeError. Pure function, CPU-only. + +2. `_load_model`'s SDXL pipeline-type-mismatch handling — previously, when an + SDXL model was detected but loaded with the wrong pipeline type and the + explicit `StableDiffusionXLPipeline` retry also failed, the code silently + logged a warning and continued with the mismatched pipeline. The fix + raises RuntimeError on retry failure instead, refusing to proceed with a + known-wrong pipeline type. That raise is caught by the enclosing + `for method in loading_methods` loop's own except/continue, so it falls + through to the next loading method rather than crashing outright — this + test drives all three loading methods to fail and asserts the final error + carries evidence that the fail-fast fired (rather than the old silent + continuation). + +Both tests are deliberately CPU-only and model-free (no real diffusers model +is loaded), following the object.__new__ shell pattern used elsewhere in +tests/unit/ (see test_safety_checker.py, test_normal_bae_fallback.py). +""" + +from unittest.mock import patch + +import pytest +import torch + +from streamdiffusion import wrapper as wrapper_module +from streamdiffusion.wrapper import StreamDiffusionWrapper, _is_oom_error + + +# --------------------------------------------------------------------------- +# _is_oom_error +# --------------------------------------------------------------------------- + + +class TestIsOomError: + def test_typed_cuda_oom_is_detected(self): + exc = torch.cuda.OutOfMemoryError("CUDA out of memory. Tried to allocate 2.00 GiB") + assert _is_oom_error(exc) is True + + @pytest.mark.parametrize( + "message", + [ + "CUDA out of memory. Tried to allocate 20.00 MiB", + "torch.OutOfMemoryError: OutOfMemory", + "generic oom while building engine", + "CUDA error: an illegal memory access was encountered", + ], + ) + def test_string_heuristic_matches_known_oom_substrings(self, message): + assert _is_oom_error(RuntimeError(message)) is True + + def test_unrelated_error_is_not_oom(self): + assert _is_oom_error(RuntimeError("shape mismatch: expected [1, 4, 64, 64]")) is False + assert _is_oom_error(ValueError("invalid literal for int()")) is False + + +# --------------------------------------------------------------------------- +# SDXL pipeline-mismatch fail-fast +# --------------------------------------------------------------------------- + + +class _FakeWrongTypePipe: + """Stand-in for a successfully-loaded but non-SDXL pipeline object.""" + + def to(self, *args, **kwargs): + return self + + +def _make_wrapper_shell() -> StreamDiffusionWrapper: + """Construct a minimal StreamDiffusionWrapper without model loading.""" + w = object.__new__(StreamDiffusionWrapper) + w.device = torch.device("cpu") + w.dtype = torch.float32 + w.cleanup_gpu_memory = lambda: None + return w + + +class TestSdxlPipelineMismatchFailFast: + def test_retry_failure_raises_instead_of_continuing_silently(self, caplog): + """ + Path contains 'sdxl' + .safetensors -> loading_methods tries + StableDiffusionXLPipeline.from_single_file first. Its first call + succeeds but returns a non-SDXL-typed pipe, triggering the mismatch + branch; the retry (second call to the same mocked method) fails. + The other two loading methods are also made to fail so _load_model + exhausts all methods and raises its final RuntimeError. + + Two things must hold: + - the mismatch branch's failure is logged as a "pipeline retry ... + also failed" warning (proving the fail-fast fired instead of the + old silent "continue with the originally loaded pipeline"), and + - _load_model does not crash trying to use the discarded + wrong-typed pipe as if it were a valid load (regression guard for + the stale `pipe` reference the fail-fast must clear). + """ + xl_call_count = [0] + + def fake_xl_from_single_file(path, *args, **kwargs): + xl_call_count[0] += 1 + if xl_call_count[0] == 1: + return _FakeWrongTypePipe() + raise RuntimeError("SDXL pipeline retry boom") + + def fake_auto_from_pretrained(path, *args, **kwargs): + raise RuntimeError("AutoPipeline boom") + + def fake_sd_from_single_file(path, *args, **kwargs): + raise RuntimeError("SD pipeline boom") + + w = _make_wrapper_shell() + + with ( + patch.object( + wrapper_module.StableDiffusionXLPipeline, + "from_single_file", + staticmethod(fake_xl_from_single_file), + ), + patch.object( + wrapper_module.AutoPipelineForText2Image, + "from_pretrained", + staticmethod(fake_auto_from_pretrained), + ), + patch.object( + wrapper_module.StableDiffusionPipeline, + "from_single_file", + staticmethod(fake_sd_from_single_file), + ), + caplog.at_level("WARNING", logger="streamdiffusion.wrapper"), + ): + with pytest.raises(RuntimeError) as exc_info: + w._load_model("fake_sdxl_model.safetensors", t_index_list=[0]) + + # The XL method must have been called twice: once for the initial + # attempt (wrong-typed success) and once for the mismatch retry. + assert xl_call_count[0] == 2, "expected exactly one retry attempt after the type mismatch" + + # The final error is "all methods exhausted" (the last method's error) - + # that's expected once every loading method has failed. What matters is + # that the mismatch retry's own failure was surfaced, not swallowed. + assert "all loading methods failed" in str(exc_info.value).lower() + assert any( + "pipeline retry" in record.message.lower() and "also failed" in record.message.lower() + for record in caplog.records + ), "expected a logged warning proving the SDXL retry fail-fast fired, not a silent continue" From 1559d6987bbde4d5a02ade429c5a25bf54ee9c0d Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 13:16:53 -0400 Subject: [PATCH 03/18] chore: migrate deprecated top-level ruff settings to [tool.ruff.lint] Ruff warns that top-level ignore/select/isort/per-file-ignores are deprecated in favour of the lint section. Move them under [tool.ruff.lint] (and .isort / .per-file-ignores); line-length stays at the top level. Verified against cuda-link's pyproject.toml, which already uses this layout. No behavior change: ruff check/format report identical results before and after (same 479 pre-existing repo-wide findings, 0 on the files touched by the Phase 1/2 commits). Investigated adding a matching [tool.pyrefly] section for parity with cuda-link, but did not adopt it: pyrefly falls back to a bundled 'basic' preset when no pyrefly config exists, and that preset is substantially more lenient than explicit config. Adding even an empty [tool.pyrefly] section opts out of it and jumps reported errors from 72 to ~900 (mostly dynamically-set attributes on StreamDiffusion / StreamDiffusionWrapper that were never checked before). Doing this properly needs the same per-module triage cuda-link performed (docs/adr/0005-static-typing-hardening.md) before turning on strict project-includes checking, which is out of scope here. --- pyproject.toml | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5fd693ae1..3b4b63bd0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,14 +1,16 @@ [tool.ruff] +line-length = 119 + +[tool.ruff.lint] # Never enforce `E501` (line length violations). ignore = ["C901", "E501", "E741", "F402", "F403", "F405", "F823"] select = ["C", "E", "F", "I", "W"] -line-length = 119 # Ignore import violations in all `__init__.py` files. -[tool.ruff.per-file-ignores] +[tool.ruff.lint.per-file-ignores] "__init__.py" = ["E402", "F401", "F811"] -[tool.ruff.isort] +[tool.ruff.lint.isort] lines-after-imports = 2 known-first-party = ["streamdiffusion"] @@ -23,4 +25,4 @@ indent-style = "space" skip-magic-trailing-comma = false # Like Black, automatically detect the appropriate line ending. -line-ending = "auto" \ No newline at end of file +line-ending = "auto" From cf3df187a4baad5ad7cbc1989eb46aa8996ba568 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 14:04:33 -0400 Subject: [PATCH 04/18] fix: sync _make_fake_engine test double with TensorRTEngine.__init__ _make_fake_engine (tests/unit/test_trt_engine_guards.py) built engines via TensorRTEngine.__new__() to bypass __init__, but never set _dedicated_stream, _pre_exec_event, _post_exec_event, or _buf_cache -- attributes __init__ has initialized since infer() gained cross-stream sync and LRU shape-cache support. This caused 3 AttributeErrors long treated as "pre-existing, out of scope" failures in every prior phase gate. Not a production bug: __init__ correctly sets all four to None/empty, and infer() guards each with `is not None`. Fix mirrors __init__ in the fake. Full unit suite now fully green (91 passed, 0 failed) with no waived tests. --- tests/unit/test_trt_engine_guards.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/unit/test_trt_engine_guards.py b/tests/unit/test_trt_engine_guards.py index 435ce5f2a..759bc252a 100644 --- a/tests/unit/test_trt_engine_guards.py +++ b/tests/unit/test_trt_engine_guards.py @@ -59,6 +59,13 @@ def _make_fake_engine( eng = TensorRTEngine.__new__(TensorRTEngine) eng.engine_path = "/fake/test.engine" eng._cuda_stream = None + # Mirror the remaining TensorRTEngine.__init__ defaults that infer() relies + # on (pre-activate() state: no dedicated stream/events yet, empty LRU cache). + # Keep this in sync with __init__ -- infer() will AttributeError otherwise. + eng._dedicated_stream = None + eng._pre_exec_event = None + eng._post_exec_event = None + eng._buf_cache = OrderedDict() eng.engine = MagicMock() eng.context = MagicMock() From 37443dbfac35f8216ce8081ed8edf1dcca65741b Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 14:05:46 -0400 Subject: [PATCH 05/18] fix: Phase 3 correctness and resource-safety fixes Four targeted fixes from the perf/best-practices audit, quick-wins #4 and #5: - StreamDiffusion.prepare(): generator defaulted to a pre-constructed torch.Generator(), a mutable default evaluated once at def-time so every instance sharing the default shared one Generator. Default is now None, constructed fresh in the method body. - StreamDiffusionWrapper.prepare(): both runtime stream.prepare() calls (single-prompt and prompt-blending paths) omitted seed, so every runtime prompt change silently reset the RNG to stream.prepare()'s own default. Now forwards seed=getattr(self.stream, "current_seed", 2) to preserve the active seed across prompt changes. - postprocess_image()'s "latent" branch returned an internal decode buffer by reference; callers retaining it across frames would see it mutate. Now clones at this public API boundary. The "np" pinned-buffer alias and the internal __call__/txt2img fast-path aliases are documented as intentional, not cloned (would defeat the pinned-buffer/reuse optimizations). - cleanup_engines_and_rebuild's hardcoded engines_dir = "engines" now honors self._engine_dir when set, matching the existing idiom elsewhere in the file. - cleanup_gpu_memory dropped two explicit Engine.__del__() calls; the method already does del + triple gc.collect() + empty_cache() + ipc_collect() at its tail, and Engine.__del__ is self-guarding, so the explicit dunder calls were redundant. Adds tests/unit/test_phase3_correctness.py (5 regression tests, model-free CPU-only). Full unit suite: 91 passed, 0 failed (no pre-existing failures remain -- see cf3df18). --- src/streamdiffusion/pipeline.py | 12 ++- src/streamdiffusion/wrapper.py | 35 +++--- tests/unit/test_phase3_correctness.py | 146 ++++++++++++++++++++++++++ 3 files changed, 179 insertions(+), 14 deletions(-) create mode 100644 tests/unit/test_phase3_correctness.py diff --git a/src/streamdiffusion/pipeline.py b/src/streamdiffusion/pipeline.py index e8964e763..e9b32faa7 100644 --- a/src/streamdiffusion/pipeline.py +++ b/src/streamdiffusion/pipeline.py @@ -401,9 +401,11 @@ def prepare( num_inference_steps: int = 50, guidance_scale: float = 1.2, delta: float = 1.0, - generator: Optional[torch.Generator] = torch.Generator(), + generator: Optional[torch.Generator] = None, seed: int = 2, ) -> None: + if generator is None: + generator = torch.Generator() self.generator = generator self.generator.manual_seed(seed) self.current_seed = seed @@ -1292,6 +1294,10 @@ def __call__(self, x: Optional[Union[torch.Tensor, PIL.Image.Image, np.ndarray]] inference_time = start.elapsed_time(end) / 1000 self.inference_time_ema = 0.9 * self.inference_time_ema + 0.1 * inference_time + # NOTE: x_output aliases self._image_decode_buf, a persistent buffer reused every + # frame to avoid a per-frame .clone() on this hot path (see buffer init above). + # Intentional: callers needing an independent copy must clone it themselves. + # StreamDiffusionWrapper.postprocess_image() clones at the public API boundary. return x_output # ========================================================================= @@ -1384,6 +1390,10 @@ def txt2img(self, batch_size: int = 1) -> torch.Tensor: # IMAGE POSTPROCESSING HOOKS: After VAE decoding, before final output x_output = self._apply_image_postprocessing_hooks(x_output) + # NOTE: x_output aliases self._image_decode_buf, a persistent buffer reused every + # call to avoid a per-frame .clone() (see buffer init above). Intentional: callers + # needing an independent copy must clone it themselves. + # StreamDiffusionWrapper.postprocess_image() clones at the public API boundary. return x_output @torch.inference_mode() diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py index 1e9c3463c..885b47bf9 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -511,6 +511,9 @@ def prepare( num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, delta=delta, + # Preserve the active seed across prompt changes -- stream.prepare()'s + # own default (seed=2) would otherwise silently reset the RNG here. + seed=getattr(self.stream, "current_seed", 2), ) finally: self._offload_text_encoders() @@ -537,6 +540,9 @@ def prepare( num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, delta=delta, + # Preserve the active seed across prompt changes -- stream.prepare()'s + # own default (seed=2) would otherwise silently reset the RNG here. + seed=getattr(self.stream, "current_seed", 2), ) finally: self._offload_text_encoders() @@ -1023,7 +1029,10 @@ def postprocess_image( # Fast paths for non-PIL outputs (avoid unnecessary conversions) if output_type == "latent": - return image_tensor + # Clone: image_tensor may alias an internal decode buffer reused across + # frames (see StreamDiffusion.__call__/txt2img). Callers of this public + # API must get an independent tensor, not a view that mutates next frame. + return image_tensor.clone() elif output_type == "pt": # Denormalize on GPU, return tensor return self._denormalize_on_gpu(image_tensor) @@ -1045,6 +1054,9 @@ def postprocess_image( with profiler.region("d2h_sync"): self._d2h_event.record() self._d2h_event.synchronize() + # NOTE: this numpy array is a view of `_output_pin_buf`, a pinned host + # buffer reused every frame (deliberate DMA optimization). Callers that + # retain the returned array across frames must copy it themselves. out = self._output_pin_buf.numpy() return out if self.frame_buffer_size > 1 else out[0] @@ -2905,15 +2917,10 @@ def cleanup_gpu_memory(self) -> None: unet_engine = self.stream.unet logger.info(" Cleaning up TensorRT UNet engine...") - # Check if it's a TensorRT engine and cleanup properly - if hasattr(unet_engine, "engine") and hasattr(unet_engine.engine, "__del__"): - try: - # Call the engine's destructor explicitly - unet_engine.engine.__del__() - except Exception: - pass - - # Clear all engine-related attributes + # Clear all engine-related attributes. Engine.__del__ is self-guarding + # (safe to invoke twice), so dropping these references is sufficient -- + # GC + the empty_cache()/ipc_collect() below reclaim the memory. No need + # to call __del__ explicitly first. if hasattr(unet_engine, "context"): try: del unet_engine.context @@ -2938,9 +2945,11 @@ def cleanup_gpu_memory(self) -> None: for engine_name in ["vae_encoder", "vae_decoder"]: if hasattr(vae_engine, engine_name): engine = getattr(vae_engine, engine_name) - if hasattr(engine, "engine") and hasattr(engine.engine, "__del__"): + # Drop the inner Engine reference so its self-guarding __del__ + # runs via GC (no need to invoke the dunder explicitly). + if hasattr(engine, "engine"): try: - engine.engine.__del__() + del engine.engine except Exception: pass try: @@ -3062,7 +3071,7 @@ def cleanup_engines_and_rebuild(self, reduce_batch_size: bool = True, reduce_res self.cleanup_gpu_memory() # Remove engines directory - engines_dir = "engines" + engines_dir = str(getattr(self, "_engine_dir", "engines")) if os.path.exists(engines_dir): try: shutil.rmtree(engines_dir) diff --git a/tests/unit/test_phase3_correctness.py b/tests/unit/test_phase3_correctness.py new file mode 100644 index 000000000..49459b63e --- /dev/null +++ b/tests/unit/test_phase3_correctness.py @@ -0,0 +1,146 @@ +""" +Regression tests for Phase 3 correctness & resource-safety fixes +(docs/perf_bestpractices_audit_2026-07-10.md quick-wins #4 and #5). + +Covers two independent fixes: + +1. `StreamDiffusion.prepare()`'s `generator` parameter previously defaulted to + `torch.Generator()` -- a mutable default evaluated once at function-definition + time, so every `StreamDiffusion` instance that didn't pass an explicit + `generator` shared the SAME `torch.Generator` object. The fix changes the + default to `None` and constructs a fresh `torch.Generator()` inside the + method body when `generator is None`. + +2. `StreamDiffusionWrapper.prepare()`'s two runtime `self.stream.prepare(...)` + calls (single-prompt and prompt-blending paths) previously omitted `seed`, + so `StreamDiffusion.prepare()`'s own default (`seed=2`) silently reset the + RNG on every runtime prompt change. The fix forwards + `seed=getattr(self.stream, "current_seed", 2)` so the active seed persists + across prompt changes. + +3. `StreamDiffusionWrapper.postprocess_image()`'s `"latent"` output branch + previously returned the internal decode buffer by reference (aliased + across frames). The fix clones at this public API boundary. + +All tests are deliberately CPU-only and model-free (no real diffusers model +is loaded), following the object.__new__ shell pattern used elsewhere in +tests/unit/ (see test_wrapper_exception_hygiene.py, test_safety_checker.py). +""" + +import inspect + +import torch + +from streamdiffusion.pipeline import StreamDiffusion +from streamdiffusion.wrapper import StreamDiffusionWrapper + + +# --------------------------------------------------------------------------- +# StreamDiffusion.prepare() -- mutable default removed +# --------------------------------------------------------------------------- + + +class TestPrepareGeneratorDefault: + def test_generator_default_is_none_not_a_shared_instance(self): + """ + The `generator` parameter's default must be `None`, not a + pre-constructed `torch.Generator()` -- the latter is a classic mutable + default footgun: the SAME Generator instance would be shared and + mutated by every StreamDiffusion instance that omits `generator`. + """ + default = inspect.signature(StreamDiffusion.prepare).parameters["generator"].default + assert default is None, ( + f"expected generator default to be None (constructed fresh in the method body), " + f"got a pre-built {type(default).__name__} instance -- this is a shared mutable default" + ) + + +# --------------------------------------------------------------------------- +# StreamDiffusionWrapper.prepare() -- seed threaded across prompt changes +# --------------------------------------------------------------------------- + + +class _FakeStream: + """Stand-in for StreamDiffusion that only records prepare() kwargs.""" + + def __init__(self, current_seed): + self.current_seed = current_seed + self.captured_kwargs = None + + def prepare(self, *args, **kwargs): + self.captured_kwargs = kwargs + + +def _make_wrapper_shell_for_prepare(current_seed) -> StreamDiffusionWrapper: + """Construct a minimal StreamDiffusionWrapper for exercising prepare().""" + w = object.__new__(StreamDiffusionWrapper) + w.stream = _FakeStream(current_seed) + w._reload_text_encoders = lambda: None + w._offload_text_encoders = lambda: None + return w + + +class TestPrepareSeedThreading: + def test_single_prompt_path_forwards_current_seed(self): + """ + Single-prompt runtime prepare() must forward the stream's active + `current_seed`, not silently fall back to StreamDiffusion.prepare()'s + own default (seed=2) and reset the RNG. + """ + w = _make_wrapper_shell_for_prepare(current_seed=12345) + w.prepare("a test prompt") + assert w.stream.captured_kwargs is not None, "stream.prepare() was not called" + assert w.stream.captured_kwargs.get("seed") == 12345 + + def test_prompt_blending_path_forwards_current_seed(self): + """Prompt-blending runtime prepare() must also forward current_seed.""" + w = _make_wrapper_shell_for_prepare(current_seed=54321) + w.update_stream_params = lambda **kwargs: None # blending step, not under test here + w.prepare([("cat", 0.7), ("dog", 0.3)]) + assert w.stream.captured_kwargs is not None, "stream.prepare() was not called" + assert w.stream.captured_kwargs.get("seed") == 54321 + + def test_falls_back_to_default_seed_when_stream_has_no_current_seed(self): + """If the stream has no `current_seed` attribute yet, fall back to 2.""" + w = object.__new__(StreamDiffusionWrapper) + fake_stream = _FakeStream(current_seed=0) + del fake_stream.current_seed # simulate an attribute that was never set + w.stream = fake_stream + w._reload_text_encoders = lambda: None + w._offload_text_encoders = lambda: None + w.prepare("a test prompt") + assert w.stream.captured_kwargs.get("seed") == 2 + + +# --------------------------------------------------------------------------- +# StreamDiffusionWrapper.postprocess_image() -- "latent" clones at boundary +# --------------------------------------------------------------------------- + + +def _make_wrapper_shell_for_postprocess() -> StreamDiffusionWrapper: + """Construct a minimal StreamDiffusionWrapper for exercising postprocess_image().""" + w = object.__new__(StreamDiffusionWrapper) + w.use_cuda_ipc_output = False + w._cuda_ipc_shm_name = None + return w + + +class TestPostprocessImageLatentClone: + def test_latent_output_is_an_independent_tensor(self): + """ + `postprocess_image(..., output_type="latent")` must return a tensor + that does NOT share storage with the input -- the input may be an + internal decode buffer reused across frames (see + StreamDiffusion.__call__/txt2img). A caller retaining the returned + tensor across frames must not see it mutate out from under them. + """ + buf = torch.arange(48, dtype=torch.float32).reshape(1, 3, 4, 4) + w = _make_wrapper_shell_for_postprocess() + out = w.postprocess_image(buf, output_type="latent") + assert out.data_ptr() != buf.data_ptr(), "returned tensor still aliases the input buffer" + assert torch.equal(out, buf), "cloned tensor must have identical values" + + # Mutating the source buffer (simulating buffer reuse on the next frame) + # must NOT be visible in the previously-returned tensor. + buf.fill_(-1.0) + assert not torch.equal(out, buf), "returned tensor was mutated by a later write to the source buffer" From 80c83ce7c99ccba959005b5bca83e2513b32aad7 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 16:20:28 -0400 Subject: [PATCH 06/18] fix: Phase 4 TensorRT engine-write atomicity and graph/refit robustness Guard TRT engine and timing-cache writes against truncation from an interrupted build via a temp-file + os.replace atomic write helper, matching the idiom already used for FP8 calibration data. Add self-healing recovery around cudaGraphLaunch (reset + fallback to plain execution, re-captures next frame) and a stream-sync guard around the defensive, disabled-by-default Engine.refit() path. Verified: 3 new CPU-only unit tests cover the atomic-write helper's happy path and both failure-preservation cases; a live-GPU smoke run against a cached VAE engine exercised capture, happy-path replay, a forced cudaGraphLaunch failure (recovery + safe fallback, no crash), and self-healed re-capture. --- .../acceleration/tensorrt/utilities.py | 61 ++++++++++---- tests/unit/test_trt_atomic_engine_write.py | 80 +++++++++++++++++++ 2 files changed, 127 insertions(+), 14 deletions(-) create mode 100644 tests/unit/test_trt_atomic_engine_write.py diff --git a/src/streamdiffusion/acceleration/tensorrt/utilities.py b/src/streamdiffusion/acceleration/tensorrt/utilities.py index 51e986e7b..12126d9bf 100644 --- a/src/streamdiffusion/acceleration/tensorrt/utilities.py +++ b/src/streamdiffusion/acceleration/tensorrt/utilities.py @@ -443,6 +443,28 @@ def CUASSERT(cuda_ret): return None +def _atomic_write_bytes(path: str, data) -> None: + """Write `data` to `path` atomically via a temp file + os.replace. + + A crash/interrupt mid-write leaves at most a stale ``.tmp`` file; ``path`` is only + ever the fully-written result (os.replace is atomic on a single filesystem). This + guards TRT engine/timing-cache writes against truncation, since the builder's cache + check (builder.py) is a bare os.path.exists with no integrity check. Mirrors the + calibration-data idiom in fp8_quantize.py. + """ + tmp_path = path + ".tmp" + try: + with open(tmp_path, "wb") as f: + f.write(data) + os.replace(tmp_path, path) + except BaseException: + try: + os.remove(tmp_path) + except OSError: + pass + raise + + class TRTProfiler(trt.IProfiler): """ Per-layer TRT timing profiler. @@ -660,6 +682,11 @@ def map_name(name): else: logger.warning(f"No refit weights for layer: {layer_name}") + # Refit reads/writes the engine's weight buffers directly; synchronize first so it + # cannot race in-flight inference that is still reading those buffers. Defensive + # only — refit is unreachable unless the engine was built with enable_refit=True + # (default False), but the guard is cheap and correct either way. + torch.cuda.current_stream().synchronize() if not refitter.refit_cuda_engine(): logger.error("Failed to refit!") raise RuntimeError("TensorRT engine refit failed") @@ -776,8 +803,7 @@ def build( f"notice(s) (createInferBuilder singleton warning) — no impact on engine." ) - with open(self.engine_path, "wb") as f: - f.write(serialized) + _atomic_write_bytes(self.engine_path, serialized) # Save timing cache for next build if timing_cache: @@ -785,8 +811,7 @@ def build( updated_cache = config.get_timing_cache() if updated_cache is not None: os.makedirs(os.path.dirname(timing_cache), exist_ok=True) - with open(timing_cache, "wb") as f: - f.write(updated_cache.serialize()) + _atomic_write_bytes(timing_cache, updated_cache.serialize()) logger.info(f"[TRT Build] Saved timing cache: {timing_cache}") except Exception as e: logger.warning(f"[TRT Build] Could not save timing cache: {e}") @@ -904,8 +929,7 @@ def _build_fp8( f"notice(s) (createInferBuilder singleton warning) — no impact on engine." ) - with open(self.engine_path, "wb") as f: - f.write(serialized) + _atomic_write_bytes(self.engine_path, serialized) # Save timing cache for next build if timing_cache: @@ -913,8 +937,7 @@ def _build_fp8( updated_cache = config.get_timing_cache() if updated_cache is not None: os.makedirs(os.path.dirname(timing_cache), exist_ok=True) - with open(timing_cache, "wb") as f: - f.write(updated_cache.serialize()) + _atomic_write_bytes(timing_cache, updated_cache.serialize()) logger.info(f"[FP8] Saved timing cache: {timing_cache}") except Exception as e: logger.warning(f"[FP8] Could not save timing cache: {e}") @@ -1091,7 +1114,7 @@ def infer(self, feed_dict, stream, use_cuda_graph=False): logger.debug( "TensorRT Engine: filtering unsupported inputs %s (allowed=%s)", missing, - sorted(list(self._allowed_inputs)), + sorted(self._allowed_inputs), ) feed_dict = filtered_feed_dict @@ -1122,10 +1145,20 @@ def infer(self, feed_dict, stream, use_cuda_graph=False): with _gpu_profiler.region("trt_infer"): if use_cuda_graph: if self.cuda_graph_instance is not None: - CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr)) - # No cudaStreamSynchronize — graph replay is async; stream ordering ensures - # downstream GPU ops (copy_, attention) wait for graph completion. - # CPU sync happens only via end.synchronize() in pipeline.__call__. + (launch_status,) = cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr) + if launch_status != cudart.cudaError_t.cudaSuccess: + # Graph replay failed (e.g. stale instance after a context/device + # hiccup). Destroy the graph and fall back to plain execution for + # this frame; the next call re-captures via the branch below. + logger.warning(f"CUDA graph launch failed ({launch_status}); resetting graph and falling back") + self.reset_cuda_graph() + noerror = self.context.execute_async_v3(stream.ptr) + if not noerror: + raise ValueError("ERROR: inference failed.") + # No cudaStreamSynchronize on the success path — graph replay is async; + # stream ordering ensures downstream GPU ops (copy_, attention) wait for + # graph completion. CPU sync happens only via end.synchronize() in + # pipeline.__call__. else: # Warmup passes before graph capture: TRT lazily JIT-compiles tactic # variants on the first few forward calls. Three passes ensure all @@ -1184,7 +1217,7 @@ def decode_images(images: torch.Tensor): def preprocess_image(image: Image.Image): w, h = image.size - w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32 + w, h = (x - x % 32 for x in (w, h)) # resize to integer multiple of 32 image = image.resize((w, h)) init_image = np.array(image).astype(np.float32) / 255.0 init_image = init_image[None].transpose(0, 3, 1, 2) diff --git a/tests/unit/test_trt_atomic_engine_write.py b/tests/unit/test_trt_atomic_engine_write.py new file mode 100644 index 000000000..4b874a2f8 --- /dev/null +++ b/tests/unit/test_trt_atomic_engine_write.py @@ -0,0 +1,80 @@ +""" +Tests for the atomic engine/timing-cache write helper (Phase 4, item 1). + +`_atomic_write_bytes` in `acceleration/tensorrt/utilities.py` replaces the plain +`open(path, "wb")` writes used for TRT engine files and timing caches. A crash or +interrupt mid-write must never leave a truncated file at the final path — the +builder's cache check (`builder.py`) is a bare `os.path.exists`, so a truncated +engine would otherwise be silently treated as valid on the next run. + +These tests exercise only the pure-Python temp-file + os.replace logic; no TRT +context or GPU is required. The module is skipped if `utilities.py` cannot be +imported (e.g. TensorRT/onnx/polygraphy not installed). + +Run with: pytest tests/unit/test_trt_atomic_engine_write.py -v +""" + +import pytest + + +# --------------------------------------------------------------------------- +# Import guard — skip all tests if utilities.py's dependencies are unavailable +# --------------------------------------------------------------------------- + +try: + from streamdiffusion.acceleration.tensorrt import utilities as trt_utilities + + IMPORT_OK = True +except ImportError: + IMPORT_OK = False + +pytestmark = pytest.mark.skipif( + not IMPORT_OK, + reason="acceleration.tensorrt.utilities not importable (TensorRT/onnx/polygraphy missing)", +) + + +class TestAtomicWriteBytes: + def test_atomic_write_creates_file_with_content(self, tmp_path): + """Happy path: final file has exact content, no stray .tmp remains.""" + target = tmp_path / "engine.trt" + payload = b"\x00\x01fake-engine-bytes\xff" * 100 + + trt_utilities._atomic_write_bytes(str(target), payload) + + assert target.read_bytes() == payload + assert not (tmp_path / "engine.trt.tmp").exists() + + def test_failed_write_preserves_existing_file(self, tmp_path, monkeypatch): + """A failed rebuild must not corrupt a previously-good cached engine.""" + target = tmp_path / "engine.trt" + good_bytes = b"known-good-engine-bytes" + target.write_bytes(good_bytes) + + def _boom(*_args, **_kwargs): + raise OSError("simulated interrupt during rename") + + monkeypatch.setattr(trt_utilities.os, "replace", _boom) + + with pytest.raises(OSError, match="simulated interrupt"): + trt_utilities._atomic_write_bytes(str(target), b"new-but-never-committed-bytes") + + # Pre-existing file untouched — os.replace never ran. + assert target.read_bytes() == good_bytes + # Temp file cleaned up, not left dangling next to the real path. + assert not (tmp_path / "engine.trt.tmp").exists() + + def test_failed_write_leaves_no_partial_final(self, tmp_path, monkeypatch): + """No pre-existing file: a failed write must not leave a truncated final file.""" + target = tmp_path / "engine.trt" + + def _boom(*_args, **_kwargs): + raise OSError("simulated interrupt during rename") + + monkeypatch.setattr(trt_utilities.os, "replace", _boom) + + with pytest.raises(OSError, match="simulated interrupt"): + trt_utilities._atomic_write_bytes(str(target), b"partial-write-bytes") + + assert not target.exists() + assert not (tmp_path / "engine.trt.tmp").exists() From d56e8fac41b6a23faebd58e7a94b5687aa96f111 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 17:20:50 -0400 Subject: [PATCH 07/18] perf: instrument unet_step with profiler.region() spans (Sub-phase 5.0) Adds 5 graph-safe profiler.region() spans (prep/sdxl_cond/hooks/engine/post) inside unet_step to attribute the audit's "~66% unattributed" frame time. Baseline capture (RTX 4090, SDXL-turbo FP8) shows that slice is GPU kernel time inside the single TRT UNet engine call (~93% of frame), not host overhead - every host sub-span measures <0.01ms. Instrumentation only, no behavior change; validated by the existing 94/0 unit suite. --- src/streamdiffusion/pipeline.py | 378 ++++++++++++++++---------------- 1 file changed, 193 insertions(+), 185 deletions(-) diff --git a/src/streamdiffusion/pipeline.py b/src/streamdiffusion/pipeline.py index e9b32faa7..cdbd86570 100644 --- a/src/streamdiffusion/pipeline.py +++ b/src/streamdiffusion/pipeline.py @@ -820,204 +820,212 @@ def unet_step( t_list: Union[torch.Tensor, list[int]], idx: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - if self.guidance_scale > 1.0 and (self.cfg_type == "initialize"): - # Pre-allocated buf avoids torch.concat malloc for CFG latent doubling - self._cfg_latent_buf[:1].copy_(x_t_latent[0:1]) - self._cfg_latent_buf[1:].copy_(x_t_latent) - x_t_latent_plus_uc = self._cfg_latent_buf - self._cfg_t_buf[:1].copy_(t_list[0:1]) - self._cfg_t_buf[1:].copy_(t_list) - t_list = self._cfg_t_buf - elif self.guidance_scale > 1.0 and (self.cfg_type == "full"): - self._cfg_latent_buf[: len(x_t_latent)].copy_(x_t_latent) - self._cfg_latent_buf[len(x_t_latent) :].copy_(x_t_latent) - x_t_latent_plus_uc = self._cfg_latent_buf - self._cfg_t_buf[: len(t_list)].copy_(t_list) - self._cfg_t_buf[len(t_list) :].copy_(t_list) - t_list = self._cfg_t_buf - else: - x_t_latent_plus_uc = x_t_latent + with profiler.region("unet_step.prep"): + if self.guidance_scale > 1.0 and (self.cfg_type == "initialize"): + # Pre-allocated buf avoids torch.concat malloc for CFG latent doubling + self._cfg_latent_buf[:1].copy_(x_t_latent[0:1]) + self._cfg_latent_buf[1:].copy_(x_t_latent) + x_t_latent_plus_uc = self._cfg_latent_buf + self._cfg_t_buf[:1].copy_(t_list[0:1]) + self._cfg_t_buf[1:].copy_(t_list) + t_list = self._cfg_t_buf + elif self.guidance_scale > 1.0 and (self.cfg_type == "full"): + self._cfg_latent_buf[: len(x_t_latent)].copy_(x_t_latent) + self._cfg_latent_buf[len(x_t_latent) :].copy_(x_t_latent) + x_t_latent_plus_uc = self._cfg_latent_buf + self._cfg_t_buf[: len(t_list)].copy_(t_list) + self._cfg_t_buf[len(t_list) :].copy_(t_list) + t_list = self._cfg_t_buf + else: + x_t_latent_plus_uc = x_t_latent - # Prepare UNet call arguments (update pre-allocated dict in-place: avoids per-frame dict malloc) - self._unet_kwargs["sample"] = x_t_latent_plus_uc - self._unet_kwargs["timestep"] = t_list - self._unet_kwargs["encoder_hidden_states"] = self.prompt_embeds - unet_kwargs = self._unet_kwargs + # Prepare UNet call arguments (update pre-allocated dict in-place: avoids per-frame dict malloc) + self._unet_kwargs["sample"] = x_t_latent_plus_uc + self._unet_kwargs["timestep"] = t_list + self._unet_kwargs["encoder_hidden_states"] = self.prompt_embeds + unet_kwargs = self._unet_kwargs # Add SDXL-specific conditioning if this is an SDXL model - if self.is_sdxl and hasattr(self, "add_text_embeds") and hasattr(self, "add_time_ids"): - if self.add_text_embeds is not None and self.add_time_ids is not None: - # Handle batching for CFG - replicate conditioning to match batch size - batch_size = x_t_latent_plus_uc.shape[0] - - # Use optimized caching system for SDXL conditioning tensors - cached_conditioning = self._get_cached_sdxl_conditioning( - batch_size, self.cfg_type, self.guidance_scale - ) - if cached_conditioning is not None: - # Cache hit - reuse existing tensors - add_text_embeds = cached_conditioning["text_embeds"] - add_time_ids = cached_conditioning["time_ids"] - else: - # Cache miss - build new tensors using optimized operations - conditioning = self._build_sdxl_conditioning(batch_size) - add_text_embeds = conditioning["text_embeds"] - add_time_ids = conditioning["time_ids"] - # Cache for future use - self._cache_sdxl_conditioning( - batch_size, self.cfg_type, self.guidance_scale, add_text_embeds, add_time_ids + with profiler.region("unet_step.sdxl_cond"): + if self.is_sdxl and hasattr(self, "add_text_embeds") and hasattr(self, "add_time_ids"): + if self.add_text_embeds is not None and self.add_time_ids is not None: + # Handle batching for CFG - replicate conditioning to match batch size + batch_size = x_t_latent_plus_uc.shape[0] + + # Use optimized caching system for SDXL conditioning tensors + cached_conditioning = self._get_cached_sdxl_conditioning( + batch_size, self.cfg_type, self.guidance_scale ) + if cached_conditioning is not None: + # Cache hit - reuse existing tensors + add_text_embeds = cached_conditioning["text_embeds"] + add_time_ids = cached_conditioning["time_ids"] + else: + # Cache miss - build new tensors using optimized operations + conditioning = self._build_sdxl_conditioning(batch_size) + add_text_embeds = conditioning["text_embeds"] + add_time_ids = conditioning["time_ids"] + # Cache for future use + self._cache_sdxl_conditioning( + batch_size, self.cfg_type, self.guidance_scale, add_text_embeds, add_time_ids + ) - unet_kwargs["added_cond_kwargs"] = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} + unet_kwargs["added_cond_kwargs"] = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} # Allow modules to contribute additional UNet kwargs via hooks - if self.unet_hooks: - try: - step_ctx = StepCtx( - x_t_latent=x_t_latent_plus_uc, - t_list=t_list, - step_index=idx if isinstance(idx, int) else (int(idx) if idx is not None else None), - guidance_mode=self.cfg_type if self.guidance_scale > 1.0 else "none", - sdxl_cond=unet_kwargs.get("added_cond_kwargs", None), - ) - extra_from_hooks = {} - for hook in self.unet_hooks: - delta: UnetKwargsDelta = hook(step_ctx) - if delta is None: - continue - if delta.down_block_additional_residuals is not None: - unet_kwargs["down_block_additional_residuals"] = delta.down_block_additional_residuals - if delta.mid_block_additional_residual is not None: - unet_kwargs["mid_block_additional_residual"] = delta.mid_block_additional_residual - if delta.added_cond_kwargs is not None: - # Merge SDXL cond if both exist - base_added = unet_kwargs.get("added_cond_kwargs", {}) - base_added.update(delta.added_cond_kwargs) - unet_kwargs["added_cond_kwargs"] = base_added - if getattr(delta, "extra_unet_kwargs", None): - # Merge extra kwargs from hooks (e.g., ipadapter_scale) - try: - extra_from_hooks.update(delta.extra_unet_kwargs) - except Exception as e: - logger.debug(f"unet_step: failed to merge extra_unet_kwargs from hook: {e}", exc_info=True) - if extra_from_hooks: - unet_kwargs["extra_unet_kwargs"] = extra_from_hooks - except Exception as e: - logger.error(f"unet_step: unet hook failed: {e}") - raise - - # Extract potential ControlNet residual kwargs and generic extra kwargs (e.g., ipadapter_scale) - hook_down_res = unet_kwargs.get("down_block_additional_residuals", None) - hook_mid_res = unet_kwargs.get("mid_block_additional_residual", None) - hook_extra_kwargs = unet_kwargs.get("extra_unet_kwargs", None) if "extra_unet_kwargs" in unet_kwargs else None - - # Call UNet with appropriate conditioning - if self.is_sdxl: - try: - # Detect UNet type and use appropriate calling convention - added_cond_kwargs = unet_kwargs.get("added_cond_kwargs", {}) - - # Check if this is a TensorRT engine or PyTorch UNet - is_tensorrt_engine = self._check_unet_tensorrt() - - if is_tensorrt_engine: - # TensorRT engine expects positional args + kwargs. IP-Adapter scale vector, if any, is provided by hooks via extra_unet_kwargs - extra_kwargs = {} - if isinstance(hook_extra_kwargs, dict): - extra_kwargs.update(hook_extra_kwargs) - - # Include ControlNet residuals if provided by hooks - if hook_down_res is not None: - extra_kwargs["down_block_additional_residuals"] = hook_down_res - if hook_mid_res is not None: - extra_kwargs["mid_block_additional_residual"] = hook_mid_res - - _unet_result = self.unet( - unet_kwargs["sample"], # latent_model_input (positional) - unet_kwargs["timestep"], # timestep (positional) - unet_kwargs["encoder_hidden_states"], # encoder_hidden_states (positional) - kvo_cache=self.kvo_cache, - fio_cache=self.fio_cache, - fi_strength=self._fi_strength_tensor, - fi_threshold=self._fi_threshold_tensor, - **extra_kwargs, - # For TRT engines, ensure SDXL cond shapes match engine builds; if engine expects 81 tokens (77+4), append dummy image tokens when none - **added_cond_kwargs, # SDXL conditioning as kwargs + with profiler.region("unet_step.hooks"): + if self.unet_hooks: + try: + step_ctx = StepCtx( + x_t_latent=x_t_latent_plus_uc, + t_list=t_list, + step_index=idx if isinstance(idx, int) else (int(idx) if idx is not None else None), + guidance_mode=self.cfg_type if self.guidance_scale > 1.0 else "none", + sdxl_cond=unet_kwargs.get("added_cond_kwargs", None), ) - model_pred = _unet_result[0] - kvo_cache_out = _unet_result[1] if len(_unet_result) > 1 else [] - fio_cache_out = _unet_result[2] if len(_unet_result) > 2 else [] - self.update_kvo_cache(kvo_cache_out, fio_cache_out) - else: - # PyTorch UNet expects diffusers-style named arguments. Any processor scaling is handled by IP-Adapter hook - - call_kwargs = { - "sample": unet_kwargs["sample"], - "timestep": unet_kwargs["timestep"], - "encoder_hidden_states": unet_kwargs["encoder_hidden_states"], - "added_cond_kwargs": added_cond_kwargs, - "return_dict": False, - } - # Include ControlNet residuals if present - if hook_down_res is not None: - call_kwargs["down_block_additional_residuals"] = hook_down_res - if hook_mid_res is not None: - call_kwargs["mid_block_additional_residual"] = hook_mid_res - model_pred = self.unet(**call_kwargs)[0] - # No restoration for per-layer scale; next step will set again via updater/time factor - - except Exception as e: - logger.error(f"[PIPELINE] unet_step: *** ERROR: SDXL UNet call failed: {e} ***") - import traceback + extra_from_hooks = {} + for hook in self.unet_hooks: + delta: UnetKwargsDelta = hook(step_ctx) + if delta is None: + continue + if delta.down_block_additional_residuals is not None: + unet_kwargs["down_block_additional_residuals"] = delta.down_block_additional_residuals + if delta.mid_block_additional_residual is not None: + unet_kwargs["mid_block_additional_residual"] = delta.mid_block_additional_residual + if delta.added_cond_kwargs is not None: + # Merge SDXL cond if both exist + base_added = unet_kwargs.get("added_cond_kwargs", {}) + base_added.update(delta.added_cond_kwargs) + unet_kwargs["added_cond_kwargs"] = base_added + if getattr(delta, "extra_unet_kwargs", None): + # Merge extra kwargs from hooks (e.g., ipadapter_scale) + try: + extra_from_hooks.update(delta.extra_unet_kwargs) + except Exception as e: + logger.debug( + f"unet_step: failed to merge extra_unet_kwargs from hook: {e}", exc_info=True + ) + if extra_from_hooks: + unet_kwargs["extra_unet_kwargs"] = extra_from_hooks + except Exception as e: + logger.error(f"unet_step: unet hook failed: {e}") + raise - traceback.print_exc() - raise - else: - # For SD1.5/SD2.1, use the old calling convention for compatibility - # Build kwargs from hooks and include residuals - ip_scale_kw = {} - if isinstance(hook_extra_kwargs, dict): - ip_scale_kw.update(hook_extra_kwargs) - - # PyTorch processor time scaling is handled by the IP-Adapter hook - - # Include ControlNet residuals if present - if hook_down_res is not None: - ip_scale_kw["down_block_additional_residuals"] = hook_down_res - if hook_mid_res is not None: - ip_scale_kw["mid_block_additional_residual"] = hook_mid_res - - _unet_result = self.unet( - x_t_latent_plus_uc, - t_list, - encoder_hidden_states=self.prompt_embeds, - kvo_cache=self.kvo_cache, - fio_cache=self.fio_cache, - fi_strength=self._fi_strength_tensor, - fi_threshold=self._fi_threshold_tensor, - return_dict=False, - **ip_scale_kw, + # Extract potential ControlNet residual kwargs / extra kwargs (e.g., ipadapter_scale), then call UNet + with profiler.region("unet_step.engine"): + hook_down_res = unet_kwargs.get("down_block_additional_residuals", None) + hook_mid_res = unet_kwargs.get("mid_block_additional_residual", None) + hook_extra_kwargs = ( + unet_kwargs.get("extra_unet_kwargs", None) if "extra_unet_kwargs" in unet_kwargs else None ) - model_pred = _unet_result[0] - kvo_cache_out = _unet_result[1] if len(_unet_result) > 1 else [] - fio_cache_out = _unet_result[2] if len(_unet_result) > 2 else [] - self.update_kvo_cache(kvo_cache_out, fio_cache_out) - if self.guidance_scale > 1.0 and (self.cfg_type == "initialize"): - noise_pred_text = model_pred[1:] - self.stock_noise[0:1].copy_(model_pred[0:1]) # in-place: eliminates 1 malloc + copy kernel - elif self.guidance_scale > 1.0 and (self.cfg_type == "full"): - noise_pred_uncond, noise_pred_text = model_pred.chunk(2) - else: - noise_pred_text = model_pred + if self.is_sdxl: + try: + # Detect UNet type and use appropriate calling convention + added_cond_kwargs = unet_kwargs.get("added_cond_kwargs", {}) + + # Check if this is a TensorRT engine or PyTorch UNet + is_tensorrt_engine = self._check_unet_tensorrt() + + if is_tensorrt_engine: + # TensorRT engine expects positional args + kwargs. IP-Adapter scale vector, if any, is provided by hooks via extra_unet_kwargs + extra_kwargs = {} + if isinstance(hook_extra_kwargs, dict): + extra_kwargs.update(hook_extra_kwargs) + + # Include ControlNet residuals if provided by hooks + if hook_down_res is not None: + extra_kwargs["down_block_additional_residuals"] = hook_down_res + if hook_mid_res is not None: + extra_kwargs["mid_block_additional_residual"] = hook_mid_res + + _unet_result = self.unet( + unet_kwargs["sample"], # latent_model_input (positional) + unet_kwargs["timestep"], # timestep (positional) + unet_kwargs["encoder_hidden_states"], # encoder_hidden_states (positional) + kvo_cache=self.kvo_cache, + fio_cache=self.fio_cache, + fi_strength=self._fi_strength_tensor, + fi_threshold=self._fi_threshold_tensor, + **extra_kwargs, + # For TRT engines, ensure SDXL cond shapes match engine builds; if engine expects 81 tokens (77+4), append dummy image tokens when none + **added_cond_kwargs, # SDXL conditioning as kwargs + ) + model_pred = _unet_result[0] + kvo_cache_out = _unet_result[1] if len(_unet_result) > 1 else [] + fio_cache_out = _unet_result[2] if len(_unet_result) > 2 else [] + self.update_kvo_cache(kvo_cache_out, fio_cache_out) + else: + # PyTorch UNet expects diffusers-style named arguments. Any processor scaling is handled by IP-Adapter hook + + call_kwargs = { + "sample": unet_kwargs["sample"], + "timestep": unet_kwargs["timestep"], + "encoder_hidden_states": unet_kwargs["encoder_hidden_states"], + "added_cond_kwargs": added_cond_kwargs, + "return_dict": False, + } + # Include ControlNet residuals if present + if hook_down_res is not None: + call_kwargs["down_block_additional_residuals"] = hook_down_res + if hook_mid_res is not None: + call_kwargs["mid_block_additional_residual"] = hook_mid_res + model_pred = self.unet(**call_kwargs)[0] + # No restoration for per-layer scale; next step will set again via updater/time factor - if self.guidance_scale > 1.0 and (self.cfg_type == "self" or self.cfg_type == "initialize"): - noise_pred_uncond = self.stock_noise * self.delta + except Exception as e: + logger.error(f"[PIPELINE] unet_step: *** ERROR: SDXL UNet call failed: {e} ***") + import traceback - if self.guidance_scale > 1.0 and self.cfg_type != "none": - model_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) - else: - model_pred = noise_pred_text + traceback.print_exc() + raise + else: + # For SD1.5/SD2.1, use the old calling convention for compatibility + # Build kwargs from hooks and include residuals + ip_scale_kw = {} + if isinstance(hook_extra_kwargs, dict): + ip_scale_kw.update(hook_extra_kwargs) + + # PyTorch processor time scaling is handled by the IP-Adapter hook + + # Include ControlNet residuals if present + if hook_down_res is not None: + ip_scale_kw["down_block_additional_residuals"] = hook_down_res + if hook_mid_res is not None: + ip_scale_kw["mid_block_additional_residual"] = hook_mid_res + + _unet_result = self.unet( + x_t_latent_plus_uc, + t_list, + encoder_hidden_states=self.prompt_embeds, + kvo_cache=self.kvo_cache, + fio_cache=self.fio_cache, + fi_strength=self._fi_strength_tensor, + fi_threshold=self._fi_threshold_tensor, + return_dict=False, + **ip_scale_kw, + ) + model_pred = _unet_result[0] + kvo_cache_out = _unet_result[1] if len(_unet_result) > 1 else [] + fio_cache_out = _unet_result[2] if len(_unet_result) > 2 else [] + self.update_kvo_cache(kvo_cache_out, fio_cache_out) + + with profiler.region("unet_step.post"): + if self.guidance_scale > 1.0 and (self.cfg_type == "initialize"): + noise_pred_text = model_pred[1:] + self.stock_noise[0:1].copy_(model_pred[0:1]) # in-place: eliminates 1 malloc + copy kernel + elif self.guidance_scale > 1.0 and (self.cfg_type == "full"): + noise_pred_uncond, noise_pred_text = model_pred.chunk(2) + else: + noise_pred_text = model_pred + + if self.guidance_scale > 1.0 and (self.cfg_type == "self" or self.cfg_type == "initialize"): + noise_pred_uncond = self.stock_noise * self.delta + + if self.guidance_scale > 1.0 and self.cfg_type != "none": + model_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + else: + model_pred = noise_pred_text # compute the previous noisy sample x_t -> x_t-1 if self.use_denoising_batch: From d44f1ed38aef5725f18bc1df7014bfd7973e8dbd Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 17:37:39 -0400 Subject: [PATCH 08/18] perf: skip redundant set_tensor_address rebind in TRT graph replay (Sub-phase 5.1) Engine.infer() unconditionally re-bound every tensor address via set_tensor_address() before each engine call, even in graphed steady state where the addresses are already baked into the captured CUDA graph (self.tensors buffers are persistent and reused via copy_() - see allocate_buffers/_can_reuse_buffers, which resets cuda_graph_instance to None whenever a buffer is actually reallocated). Skip the rebind loop once a graph is captured; only rebind on first capture or after a reset_cuda_graph() recovery. Verified live on the RTX 4090 td_config.yaml runtime path (the UNet engine loads with use_cuda_graph=True unconditionally via engine_manager.py's loader): 24-frame smoke run with real graph capture + replay shows no errors and frame timing (p50 32.10ms, mean 47.33ms) matching a pre-fix A/B run (p50 31.74ms, mean 47.00ms) - confirms correctness with the expected ~0 measured perf delta (the audit's own 5.0 baseline showed host-side region time is <0.01ms; this is a hygiene/false-dependency fix, not a throughput win). Suite 94/0; ruff clean; pyrefly 72 errors (baseline, no new). --- .../acceleration/tensorrt/utilities.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/streamdiffusion/acceleration/tensorrt/utilities.py b/src/streamdiffusion/acceleration/tensorrt/utilities.py index 12126d9bf..6b63c9866 100644 --- a/src/streamdiffusion/acceleration/tensorrt/utilities.py +++ b/src/streamdiffusion/acceleration/tensorrt/utilities.py @@ -1138,9 +1138,15 @@ def infer(self, feed_dict, stream, use_cuda_graph=False): for name, buf in feed_dict.items(): self.tensors[name].copy_(buf) - for name, tensor in self.tensors.items(): - if not self.context.set_tensor_address(name, tensor.data_ptr()): - raise RuntimeError(f"TensorRT: set_tensor_address failed for '{name}'") + # In graphed steady state the tensor addresses are baked into the captured graph + # (self.tensors[name] are persistent buffers reused via copy_() — see + # _can_reuse_buffers/allocate_buffers, which resets cuda_graph_instance to None + # whenever a buffer is actually reallocated). Re-binding every frame is then pure + # host overhead; only rebind on first capture or after a reset (instance is None). + if not (use_cuda_graph and self.cuda_graph_instance is not None): + for name, tensor in self.tensors.items(): + if not self.context.set_tensor_address(name, tensor.data_ptr()): + raise RuntimeError(f"TensorRT: set_tensor_address failed for '{name}'") with _gpu_profiler.region("trt_infer"): if use_cuda_graph: From a1af2b00c9f8881264bd2ccb4620840b039085bc Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 19:18:40 -0400 Subject: [PATCH 09/18] perf: sync-free output/PIL/IPC paths (Sub-phase 5.2) Three output-side sync/allocation fixes, ordered by real-world impact on the TD path: - 5f (only hot-path site): _ipc_pack_rgba / _ipc_pack_unit_rgba now write into a persistent HWC x4 GPU buffer (alpha set once at allocation, B/G/R channels written in place) instead of allocating a fresh torch.full_like alpha + torch.cat every frame. Safe only because the IPC exporters are forced to blocking export (ADR-0001); comments updated at both _lazy_init_ipc_exporter and _lazy_init_cn_ipc_exporter to reflect the buffer is now persistent, not transient. - 5e: _tensor_to_pil_safe (preprocessing_orchestrator.py) reordered CPU-first, following the in-repo base.py:tensor_to_pil template -- the D2H transfer now happens before the tensor.min()/.max() range-check syncs instead of after, collapsing 3 GPU host syncs down to 1. - 5d (hygiene; TD's IPC output path never hits this): _tensor_to_pil_optimized now routes through the shared _output_pin_buf/_d2h_event pinned-buffer + Event machinery instead of a blocking, unpinned .cpu() into pageable memory. Documented the same buffer-reuse caveat as the "np" path (PIL images returned wrap a view of the pinned buffer; callers retaining them across frames must .copy()). Verified via tests/unit/test_sync_free_output_5_2.py: byte-identical parity against frozen pre-5.2 reference formulas for all three sites, plus persistent-buffer reuse/realloc and independent-buffer-identity checks (10 new tests, GPU-gated). Suite 104/0 (94 baseline + 10 new); ruff clean on touched files; pyrefly 72 errors (baseline, no new). --- .../preprocessing_orchestrator.py | 14 +- src/streamdiffusion/wrapper.py | 96 +++++-- tests/unit/test_sync_free_output_5_2.py | 238 ++++++++++++++++++ 3 files changed, 322 insertions(+), 26 deletions(-) create mode 100644 tests/unit/test_sync_free_output_5_2.py diff --git a/src/streamdiffusion/preprocessing/preprocessing_orchestrator.py b/src/streamdiffusion/preprocessing/preprocessing_orchestrator.py index e2b52ae4a..577678bd6 100644 --- a/src/streamdiffusion/preprocessing/preprocessing_orchestrator.py +++ b/src/streamdiffusion/preprocessing/preprocessing_orchestrator.py @@ -596,6 +596,15 @@ def _tensor_to_pil_safe(self, tensor: torch.Tensor) -> Image.Image: # Convert from CHW to HWC tensor = tensor.permute(1, 2, 0) + # 5e: CPU-first — do the D2H transfer before the range-check reductions below, not + # after. tensor.min()/.max() are scalar reductions that force a host sync while the + # tensor is still on GPU; running them on the already-transferred CPU tensor collapses + # 2 GPU syncs into the 1 unavoidable D2H transfer. Same pattern as the in-repo template + # preprocessing/processors/base.py:tensor_to_pil. + tensor = tensor.detach() + if tensor.is_cuda: + tensor = tensor.cpu() + # CRITICAL FIX: Handle VAE output range [-1, 1] -> [0, 1] -> [0, 255] # VAE decode_image() outputs in [-1, 1] range, need to convert to [0, 1] first if tensor.min() < 0: @@ -606,8 +615,8 @@ def _tensor_to_pil_safe(self, tensor: torch.Tensor) -> Image.Image: if tensor.max() <= 1.0: tensor = tensor * 255.0 - # Convert to numpy and then PIL - numpy_image = tensor.detach().cpu().numpy().astype(np.uint8) + # Already on CPU (transferred above) — just cast + build the PIL image. + numpy_image = tensor.numpy().astype(np.uint8) return Image.fromarray(numpy_image) def _pil_to_tensor_safe(self, pil_image: Image.Image, device: str, dtype: torch.dtype) -> torch.Tensor: @@ -789,7 +798,6 @@ def _process_single_preprocessor_group( if getattr(preprocessor, "gpu_native", False): return None - # PIL processing fallback if control_variants["image"] is not None: processed_image = self.prepare_control_image( diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py index 885b47bf9..c981661db 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -349,6 +349,9 @@ def __init__( self._output_pin_buf: Optional[torch.Tensor] = None # pinned CPU buffer for async D2H output self._output_gpu_buf: Optional[torch.Tensor] = None # persistent GPU fp32 staging (avoids per-frame alloc) self._d2h_event: Optional[torch.cuda.Event] = None # event for fine-grained D2H sync + self._ipc_pack_buf: Optional[torch.Tensor] = None # persistent GPU BGRA buffer for _ipc_pack_rgba (5f) + # persistent GPU BGRA buffer for _ipc_pack_unit_rgba (5f) + self._ipc_pack_unit_buf: Optional[torch.Tensor] = None self.use_cuda_ipc_output = use_cuda_ipc_output self._cuda_ipc_shm_name = cuda_ipc_shm_name self._cuda_ipc_num_slots = cuda_ipc_num_slots @@ -1074,15 +1077,35 @@ def postprocess_image( return postprocess_image(image_tensor.cpu(), output_type=output_type)[0] def _ipc_pack_rgba(self, image_tensor: torch.Tensor) -> torch.Tensor: - """Convert pipeline output to HWC uint8 BGRA on GPU for cuda-link wire contract.""" + """Convert pipeline output to HWC uint8 BGRA on GPU for cuda-link wire contract. + + Writes into a persistent HWC×4 GPU buffer (realloc'd only on shape/device change) + instead of allocating a fresh alpha channel + concatenated tensor every frame (5f). + SAFE ONLY because the IPC exporter is forced to blocking export (see + _lazy_init_ipc_exporter / ADR-0001): the frame must be fully copied out via cuda-link + before this buffer is overwritten next call. Async export (CUDALINK_EXPORT_SYNC=0) + would race against this reused buffer and needs double-buffering first — not + implemented. + """ with profiler.region("glue.ipc_pack_rgba"): denorm = self._denormalize_on_gpu(image_tensor) # NCHW [0,1] if denorm.dim() == 4: denorm = denorm[0] # CHW [0,1] rgb_u8 = (denorm * 255).clamp(0, 255).to(torch.uint8) # CHW uint8 rgb_hwc = rgb_u8.permute(1, 2, 0).contiguous() # HWC RGB - alpha = torch.full_like(rgb_hwc[..., :1], 255) - return torch.cat([rgb_hwc[..., 2:3], rgb_hwc[..., 1:2], rgb_hwc[..., 0:1], alpha], dim=-1).contiguous() + h, w = rgb_hwc.shape[0], rgb_hwc.shape[1] + if ( + self._ipc_pack_buf is None + or self._ipc_pack_buf.shape[0] != h + or self._ipc_pack_buf.shape[1] != w + or self._ipc_pack_buf.device != rgb_hwc.device + ): + self._ipc_pack_buf = torch.empty((h, w, 4), dtype=torch.uint8, device=rgb_hwc.device) + self._ipc_pack_buf[..., 3] = 255 # constant alpha, set once at (re)allocation + self._ipc_pack_buf[..., 0] = rgb_hwc[..., 2] # B + self._ipc_pack_buf[..., 1] = rgb_hwc[..., 1] # G + self._ipc_pack_buf[..., 2] = rgb_hwc[..., 0] # R + return self._ipc_pack_buf def _lazy_init_ipc_exporter(self, height: int, width: int): """Initialize Exporter on first frame (lazy to defer CUDA IPC SHM creation).""" @@ -1092,10 +1115,11 @@ def _lazy_init_ipc_exporter(self, height: int, width: int): from cuda_link import Exporter, ExportPolicy, FrameSpec - # SD source buffers (_ipc_pack_rgba output) are transient torch tensors recycled by - # PyTorch's caching allocator the moment the caller returns. Async export would read - # recycled memory → torn frames (ADR-0001 source-buffer lifetime race). Force - # blocking unless the user explicitly opted into async with CUDALINK_EXPORT_SYNC=0. + # SD source buffers (_ipc_pack_rgba output) are a persistent GPU buffer reused every + # frame (5f), not a fresh allocation. Async export would read a buffer that's already + # been overwritten by the next frame → torn frames (ADR-0001 source-buffer lifetime + # race). Force blocking unless the user explicitly opted into async with + # CUDALINK_EXPORT_SYNC=0. policy = ExportPolicy.from_env() if os.environ.get("CUDALINK_EXPORT_SYNC") is None: policy = _dc_replace(policy, export_sync=True) @@ -1118,6 +1142,10 @@ def _ipc_pack_unit_rgba(self, image_tensor: torch.Tensor) -> torch.Tensor: Like _ipc_pack_rgba but skips _denormalize_on_gpu — the ControlNet preprocessor output is already in [0, 1] (not the diffusion [-1, 1] range). + + Writes into a persistent HWC×4 GPU buffer (5f), same rationale/caveat as + _ipc_pack_rgba: SAFE ONLY while the CN-preview exporter stays forced-blocking + (see _lazy_init_cn_ipc_exporter / ADR-0001). """ with profiler.region("glue.ipc_pack_unit_rgba"): t = image_tensor @@ -1125,8 +1153,19 @@ def _ipc_pack_unit_rgba(self, image_tensor: torch.Tensor) -> torch.Tensor: t = t[0] # NCHW → CHW [0,1] rgb_u8 = (t * 255).clamp(0, 255).to(torch.uint8) # CHW uint8 rgb_hwc = rgb_u8.permute(1, 2, 0).contiguous() # HWC RGB - alpha = torch.full_like(rgb_hwc[..., :1], 255) - return torch.cat([rgb_hwc[..., 2:3], rgb_hwc[..., 1:2], rgb_hwc[..., 0:1], alpha], dim=-1).contiguous() + h, w = rgb_hwc.shape[0], rgb_hwc.shape[1] + if ( + self._ipc_pack_unit_buf is None + or self._ipc_pack_unit_buf.shape[0] != h + or self._ipc_pack_unit_buf.shape[1] != w + or self._ipc_pack_unit_buf.device != rgb_hwc.device + ): + self._ipc_pack_unit_buf = torch.empty((h, w, 4), dtype=torch.uint8, device=rgb_hwc.device) + self._ipc_pack_unit_buf[..., 3] = 255 # constant alpha, set once at (re)allocation + self._ipc_pack_unit_buf[..., 0] = rgb_hwc[..., 2] # B + self._ipc_pack_unit_buf[..., 1] = rgb_hwc[..., 1] # G + self._ipc_pack_unit_buf[..., 2] = rgb_hwc[..., 0] # R + return self._ipc_pack_unit_buf def _lazy_init_cn_ipc_exporter(self, height: int, width: int): """Initialize the CN-preview Exporter on first frame (lazy, mirrors _lazy_init_ipc_exporter).""" @@ -1137,7 +1176,7 @@ def _lazy_init_cn_ipc_exporter(self, height: int, width: int): from cuda_link import Exporter, ExportPolicy, FrameSpec # Same source-buffer lifetime race as _lazy_init_ipc_exporter: _ipc_pack_unit_rgba - # returns a transient tensor that the caching allocator recycles on return. + # returns a persistent GPU buffer (5f) reused every frame, overwritten on the next call. # Force blocking unless CUDALINK_EXPORT_SYNC=0 explicitly opts into async. policy = ExportPolicy.from_env() if os.environ.get("CUDALINK_EXPORT_SYNC") is None: @@ -1276,20 +1315,29 @@ def _tensor_to_pil_optimized(self, image_tensor: torch.Tensor) -> List[Image.Ima List[Image.Image] List of PIL RGB images, one for each item in the batch. """ - # Denormalize on GPU first + # 5d: convert to uint8 NHWC on GPU (identical layout to the "np" output path @1045), + # then route through the shared pinned-buffer + Event machinery instead of a blocking, + # unpinned .cpu() into pageable memory. Only one output_type's fast path executes per + # postprocess_image() call, so sharing _output_pin_buf/_d2h_event with the "np" path + # is safe (shape/dtype-guarded realloc below covers either caller). denormalized = self._denormalize_on_gpu(image_tensor) - - # Convert to uint8 on GPU to reduce transfer size - # Scale to [0, 255] and convert to uint8 - # Scale to [0, 255] and convert to uint8 - uint8_tensor = (denormalized * 255).clamp(0, 255).to(torch.uint8) - - # Single efficient CPU transfer - cpu_tensor = uint8_tensor.cpu() - - # Convert to HWC format for PIL - # From BCHW to BHWC - cpu_tensor = cpu_tensor.permute(0, 2, 3, 1) + uint8_nhwc = (denormalized * 255).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).contiguous() + if ( + self._output_pin_buf is None + or self._output_pin_buf.shape != uint8_nhwc.shape + or self._output_pin_buf.dtype != torch.uint8 + ): + self._output_pin_buf = torch.empty(uint8_nhwc.shape, dtype=torch.uint8, pin_memory=True) + self._d2h_event = torch.cuda.Event() + self._output_pin_buf.copy_(uint8_nhwc, non_blocking=True) + with profiler.region("d2h_sync"): + self._d2h_event.record() + self._d2h_event.synchronize() + # NOTE: like the "np" output path, each PIL Image below wraps a view of + # `_output_pin_buf` (Image.fromarray shares the numpy buffer, it does not copy). + # Callers that retain a returned PIL Image across frames must .copy() it themselves; + # this pinned buffer is overwritten in place on the next call. + cpu_tensor = self._output_pin_buf # Convert to PIL images efficiently pil_images = [] @@ -2987,6 +3035,8 @@ def cleanup_gpu_memory(self) -> None: self._output_pin_buf = None self._output_gpu_buf = None self._d2h_event = None + self._ipc_pack_buf = None + self._ipc_pack_unit_buf = None # Force multiple garbage collection cycles for thorough cleanup for i in range(3): diff --git a/tests/unit/test_sync_free_output_5_2.py b/tests/unit/test_sync_free_output_5_2.py new file mode 100644 index 000000000..76780f1fe --- /dev/null +++ b/tests/unit/test_sync_free_output_5_2.py @@ -0,0 +1,238 @@ +""" +Regression tests for Sub-phase 5.2 (output/PIL/IPC sync-free path, +docs/perf_bestpractices_audit_2026-07-10.md follow-ups #2/#... — see the +Phase 5 continuation plan). + +Covers three independent output-side fixes, each converted from a +per-frame-allocation formula to a persistent-buffer / pinned-readback +formula. The reference (`_reference_*`) helpers below are frozen copies of +the pre-5.2 code, used as the numeric-parity oracle: every touched method +must produce byte-identical output to its old formula on the same input. + + - 5f: `_ipc_pack_rgba` / `_ipc_pack_unit_rgba` (wrapper.py) — persistent + HWC x4 GPU buffer instead of a fresh `torch.full_like` alpha + `torch.cat` + every frame. + - 5d: `_tensor_to_pil_optimized` (wrapper.py) — routes through the shared + `_output_pin_buf` / `_d2h_event` pinned-buffer + Event machinery instead + of a blocking, unpinned `.cpu()` into pageable memory. + - 5e: `_tensor_to_pil_safe` (preprocessing_orchestrator.py) — CPU-first + reorder: the D2H transfer now happens before the tensor.min()/.max() + range-check syncs, not after (kills 2 of 3 host syncs; same template as + preprocessing/processors/base.py:tensor_to_pil). + +5d and 5f exercise `torch.cuda.Event()` / `pin_memory=True`, which require a +real CUDA device — the whole module is skipped when unavailable. +""" + +import numpy as np +import pytest +import torch +from PIL import Image + +from streamdiffusion.preprocessing.preprocessing_orchestrator import PreprocessingOrchestrator +from streamdiffusion.wrapper import StreamDiffusionWrapper + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Sub-phase 5.2 output paths require CUDA (pin_memory / cuda.Event)", +) + + +# --------------------------------------------------------------------------- +# helpers: minimal shells + frozen pre-5.2 reference formulas +# --------------------------------------------------------------------------- + + +def _make_wrapper_shell(): + """Minimal StreamDiffusionWrapper shell wired only with the attrs the + touched methods read (object.__new__ pattern, see test_wrapper_exception_hygiene.py).""" + w = object.__new__(StreamDiffusionWrapper) + w._output_pin_buf = None + w._d2h_event = None + w._ipc_pack_buf = None + w._ipc_pack_unit_buf = None + return w + + +def _reference_ipc_pack_rgba(wrapper, image_tensor): + """Pre-5f formula: fresh alpha (full_like) + cat every call.""" + denorm = wrapper._denormalize_on_gpu(image_tensor) + if denorm.dim() == 4: + denorm = denorm[0] + rgb_u8 = (denorm * 255).clamp(0, 255).to(torch.uint8) + rgb_hwc = rgb_u8.permute(1, 2, 0).contiguous() + alpha = torch.full_like(rgb_hwc[..., :1], 255) + return torch.cat([rgb_hwc[..., 2:3], rgb_hwc[..., 1:2], rgb_hwc[..., 0:1], alpha], dim=-1).contiguous() + + +def _reference_ipc_pack_unit_rgba(image_tensor): + """Pre-5f formula for the CN-preview twin (skips denormalize).""" + t = image_tensor + if t.dim() == 4: + t = t[0] + rgb_u8 = (t * 255).clamp(0, 255).to(torch.uint8) + rgb_hwc = rgb_u8.permute(1, 2, 0).contiguous() + alpha = torch.full_like(rgb_hwc[..., :1], 255) + return torch.cat([rgb_hwc[..., 2:3], rgb_hwc[..., 1:2], rgb_hwc[..., 0:1], alpha], dim=-1).contiguous() + + +def _reference_tensor_to_pil_optimized(wrapper, image_tensor): + """Pre-5d formula: blocking, unpinned .cpu().""" + denormalized = wrapper._denormalize_on_gpu(image_tensor) + uint8_tensor = (denormalized * 255).clamp(0, 255).to(torch.uint8) + cpu_tensor = uint8_tensor.cpu() + cpu_tensor = cpu_tensor.permute(0, 2, 3, 1) + pil_images = [] + for i in range(cpu_tensor.shape[0]): + img_array = cpu_tensor[i].numpy() + if img_array.shape[-1] == 1: + pil_images.append(Image.fromarray(img_array.squeeze(-1), mode="L")) + else: + pil_images.append(Image.fromarray(img_array)) + return pil_images + + +def _reference_tensor_to_pil_safe(tensor): + """Pre-5e formula: range-check syncs (tensor.min()/.max()) before the D2H transfer.""" + if tensor.dim() == 4: + tensor = tensor.squeeze(0) + if tensor.dim() == 3 and tensor.shape[0] == 3: + tensor = tensor.permute(1, 2, 0) + if tensor.min() < 0: + tensor = (tensor / 2.0 + 0.5).clamp(0, 1) + if tensor.max() <= 1.0: + tensor = tensor * 255.0 + numpy_image = tensor.detach().cpu().numpy().astype(np.uint8) + return Image.fromarray(numpy_image) + + +# --------------------------------------------------------------------------- +# 5f: _ipc_pack_rgba / _ipc_pack_unit_rgba +# --------------------------------------------------------------------------- + + +class TestIpcPackPersistentBuffer: + def test_ipc_pack_rgba_byte_identical_to_reference(self): + wrapper = _make_wrapper_shell() + torch.manual_seed(0) + image_tensor = torch.rand(1, 3, 64, 96, device="cuda") * 2 - 1 # diffusion range [-1, 1] + + expected = _reference_ipc_pack_rgba(wrapper, image_tensor) + actual = wrapper._ipc_pack_rgba(image_tensor) + + assert torch.equal(actual.cpu(), expected.cpu()) + + def test_ipc_pack_rgba_reuses_same_buffer_across_calls(self): + """5f's entire point: same shape -> same underlying allocation, no per-frame alloc.""" + wrapper = _make_wrapper_shell() + t1 = torch.rand(1, 3, 64, 96, device="cuda") * 2 - 1 + t2 = torch.rand(1, 3, 64, 96, device="cuda") * 2 - 1 + + buf1 = wrapper._ipc_pack_rgba(t1) + ptr1 = buf1.data_ptr() + buf2 = wrapper._ipc_pack_rgba(t2) + + assert buf2.data_ptr() == ptr1, "expected the persistent buffer to be reused, not reallocated" + + def test_ipc_pack_rgba_reallocs_on_shape_change(self): + wrapper = _make_wrapper_shell() + t1 = torch.rand(1, 3, 64, 96, device="cuda") * 2 - 1 + t2 = torch.rand(1, 3, 32, 48, device="cuda") * 2 - 1 + + buf1 = wrapper._ipc_pack_rgba(t1) + assert buf1.shape[:2] == (64, 96) + buf2 = wrapper._ipc_pack_rgba(t2) + assert buf2.shape[:2] == (32, 48) + + def test_ipc_pack_unit_rgba_byte_identical_to_reference(self): + wrapper = _make_wrapper_shell() + torch.manual_seed(1) + image_tensor = torch.rand(1, 3, 48, 64, device="cuda") # already [0, 1], no denorm + + expected = _reference_ipc_pack_unit_rgba(image_tensor) + actual = wrapper._ipc_pack_unit_rgba(image_tensor) + + assert torch.equal(actual.cpu(), expected.cpu()) + + def test_ipc_pack_rgba_and_unit_rgba_use_independent_buffers(self): + """The two packers must not alias the same persistent buffer (separate exporters).""" + wrapper = _make_wrapper_shell() + t = torch.rand(1, 3, 40, 56, device="cuda") + + buf_main = wrapper._ipc_pack_rgba(t * 2 - 1) + buf_unit = wrapper._ipc_pack_unit_rgba(t) + + assert buf_main.data_ptr() != buf_unit.data_ptr() + + +# --------------------------------------------------------------------------- +# 5d: _tensor_to_pil_optimized +# --------------------------------------------------------------------------- + + +class TestTensorToPilOptimizedPinnedReadback: + def test_byte_identical_to_reference(self): + wrapper = _make_wrapper_shell() + torch.manual_seed(2) + image_tensor = torch.rand(2, 3, 40, 56, device="cuda") * 2 - 1 + + expected = _reference_tensor_to_pil_optimized(wrapper, image_tensor) + actual = wrapper._tensor_to_pil_optimized(image_tensor) + + assert len(actual) == len(expected) + for a, e in zip(actual, expected): + assert np.array_equal(np.array(a), np.array(e)) + + def test_reuses_output_pin_buf_across_calls(self): + """5d shares _output_pin_buf/_d2h_event with the 'np' output path (same shape/dtype guard).""" + wrapper = _make_wrapper_shell() + t1 = torch.rand(1, 3, 32, 32, device="cuda") * 2 - 1 + t2 = torch.rand(1, 3, 32, 32, device="cuda") * 2 - 1 + + wrapper._tensor_to_pil_optimized(t1) + buf_ptr_1 = wrapper._output_pin_buf.data_ptr() + wrapper._tensor_to_pil_optimized(t2) + buf_ptr_2 = wrapper._output_pin_buf.data_ptr() + + assert buf_ptr_1 == buf_ptr_2 + + +# --------------------------------------------------------------------------- +# 5e: _tensor_to_pil_safe (CPU-first reorder) +# --------------------------------------------------------------------------- + + +class TestTensorToPilSafeCpuFirst: + def test_byte_identical_vae_range(self): + """VAE-range [-1, 1] input takes both the min()<0 and max()<=1.0 branches.""" + torch.manual_seed(3) + tensor = torch.rand(3, 40, 56, device="cuda") * 2 - 1 + + orchestrator = object.__new__(PreprocessingOrchestrator) + expected = _reference_tensor_to_pil_safe(tensor.clone()) + actual = orchestrator._tensor_to_pil_safe(tensor.clone()) + + assert np.array_equal(np.array(actual), np.array(expected)) + + def test_byte_identical_unit_range(self): + """[0, 1] input only takes the max()<=1.0 branch.""" + torch.manual_seed(4) + tensor = torch.rand(3, 40, 56, device="cuda") + + orchestrator = object.__new__(PreprocessingOrchestrator) + expected = _reference_tensor_to_pil_safe(tensor.clone()) + actual = orchestrator._tensor_to_pil_safe(tensor.clone()) + + assert np.array_equal(np.array(actual), np.array(expected)) + + def test_byte_identical_already_0_255_range(self): + """Values already > 1.0 (pre-scaled to [0, 255]) skip both conditional branches.""" + torch.manual_seed(5) + tensor = torch.rand(3, 40, 56, device="cuda") * 255.0 + + orchestrator = object.__new__(PreprocessingOrchestrator) + expected = _reference_tensor_to_pil_safe(tensor.clone()) + actual = orchestrator._tensor_to_pil_safe(tensor.clone()) + + assert np.array_equal(np.array(actual), np.array(expected)) From cda335bddcb3a8f7d79bb6164625b56b4a5c6d00 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 19:26:13 -0400 Subject: [PATCH 10/18] fix: batch-dynamic engines skip l2tc tiling, silencing benign VALIDATE FAIL build_engine() passed dynamic_shapes=build_dynamic_shape into Engine.build(), which tracks only dynamic resolution. On the default "Flexible" preset the UNet is resolution-static but batch-dynamic (build_static_batch=False), so dynamic_shapes was passed False, and _apply_gpu_profile_to_config's tiling branch ran against a graph that still has a symbolic batch dim -- TRT emitted "[l2tc] VALIDATE FAIL - Graph contains symbolic shape" as a no-op for every applicable layer, for the entire ~11.5 min production TD session analyzed this round. Fix: dynamic_shapes=build_dynamic_shape or not build_static_batch (one line, utilities.py:1372) -- True whenever any dim, resolution or batch, is symbolic, matching l2tc's actual requirement (ALL dims must be concrete). Fully-static engines (e.g. ControlNet: build_static_batch=True, build_dynamic_shape=False) are unaffected -- still get tiling as before. Tightened the _apply_gpu_profile_to_config docstring to match. Verified via new tests/unit/test_l2tc_dynamic_shapes.py: calls build_engine() with GPU detection/memory query/the real TRT build all monkeypatched out (no CUDA/TRT context required) and asserts the dynamic_shapes kwarg it wires into Engine.build() for the previously-broken case (batch-dynamic, resolution-static -> True), the unaffected fully-static case (-> False), and the already-working resolution-dynamic case (-> True). Suite 107/0 (104 baseline + 3 new); ruff clean; pyrefly 72 errors (baseline, no new). Impact is log-spam-only -- zero perf/numerical change, visible on the next engine rebuild. --- .../acceleration/tensorrt/utilities.py | 9 +- tests/unit/test_l2tc_dynamic_shapes.py | 90 +++++++++++++++++++ 2 files changed, 97 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_l2tc_dynamic_shapes.py diff --git a/src/streamdiffusion/acceleration/tensorrt/utilities.py b/src/streamdiffusion/acceleration/tensorrt/utilities.py index 6b63c9866..6eec841d4 100644 --- a/src/streamdiffusion/acceleration/tensorrt/utilities.py +++ b/src/streamdiffusion/acceleration/tensorrt/utilities.py @@ -289,7 +289,7 @@ def _apply_gpu_profile_to_config( Args: config: TRT IBuilderConfig to modify. gpu_profile: Hardware-detected build parameters from detect_gpu_profile(). - dynamic_shapes: Whether this engine uses dynamic input shapes. + dynamic_shapes: Whether this engine has any symbolic dim, incl. batch. - True (default): tiling and l2_limit skipped — TRT confirms these have no effect on symbolic-shape graphs and only produce warning spam. - False (static): tiling and l2_limit applied for full L2 cache benefit. @@ -1369,7 +1369,12 @@ def build_engine( workspace_size=max_workspace_size, fp8=fp8, gpu_profile=gpu_profile, - dynamic_shapes=build_dynamic_shape, + # Any symbolic dim (resolution OR batch) disqualifies l2tc tiling — see + # _apply_gpu_profile_to_config. build_dynamic_shape alone misses the + # batch-dynamic/resolution-static case (e.g. default "Flexible" UNet + # preset), which previously reached the tiling branch and TRT emitted + # "[l2tc] VALIDATE FAIL - Graph contains symbolic shape" as a no-op. + dynamic_shapes=build_dynamic_shape or not build_static_batch, ) return engine diff --git a/tests/unit/test_l2tc_dynamic_shapes.py b/tests/unit/test_l2tc_dynamic_shapes.py new file mode 100644 index 000000000..1d3d3ca79 --- /dev/null +++ b/tests/unit/test_l2tc_dynamic_shapes.py @@ -0,0 +1,90 @@ +""" +Regression test for the l2tc `dynamic_shapes` fix (Phase 5, Commit B — +docs/perf_bestpractices_audit_2026-07-10.md follow-up). + +Root cause: `build_engine()` passed `dynamic_shapes=build_dynamic_shape` into +`Engine.build()`, which tracks only dynamic *resolution*. Batch-dynamic / +resolution-static engines (e.g. the default "Flexible" UNet preset, +`build_static_batch=False, build_dynamic_shape=False`) got `dynamic_shapes=False`, +so `_apply_gpu_profile_to_config`'s tiling branch ran against a graph that still +has a symbolic batch dim, and TRT emitted a benign but noisy +"[l2tc] VALIDATE FAIL - Graph contains symbolic shape" warning for every +applicable layer. + +Fix: `dynamic_shapes=build_dynamic_shape or not build_static_batch` — True +whenever *any* dim (resolution or batch) is symbolic, matching +`_apply_gpu_profile_to_config`'s actual requirement (l2tc needs ALL dims +concrete). Fully-static engines (e.g. ControlNet: `build_static_batch=True, +build_dynamic_shape=False`) are unaffected. + +These tests exercise only `build_engine`'s pure-Python wiring: GPU detection, +the memory query, and the real `Engine.build()` TRT call are all monkeypatched +out, so no CUDA device or TensorRT context is required. +""" + +from unittest.mock import MagicMock + +import pytest + + +try: + from streamdiffusion.acceleration.tensorrt import utilities as trt_utilities + + IMPORT_OK = True +except ImportError: + IMPORT_OK = False + +pytestmark = pytest.mark.skipif( + not IMPORT_OK, + reason="acceleration.tensorrt.utilities not importable (TensorRT/onnx/polygraphy missing)", +) + + +class _FakeModelData: + """Stand-in for a BaseModel subclass — build_engine only calls get_input_profile().""" + + def get_input_profile(self, *args, **kwargs): + return {} + + +def _computed_dynamic_shapes(monkeypatch, *, build_static_batch, build_dynamic_shape): + """Call build_engine() with GPU detection / memory query / the real TRT build + all monkeypatched out, and return the `dynamic_shapes` kwarg it actually + passed to Engine.build().""" + monkeypatch.setattr(trt_utilities.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(trt_utilities, "detect_gpu_profile", lambda device: trt_utilities._fallback_profile()) + monkeypatch.setattr(trt_utilities.cudart, "cudaMemGetInfo", lambda: (0, 8 * 2**30, 16 * 2**30)) + + captured_build = MagicMock(return_value=None) + monkeypatch.setattr(trt_utilities.Engine, "build", captured_build) + + trt_utilities.build_engine( + engine_path="fake.engine", + onnx_opt_path="fake.onnx", + model_data=_FakeModelData(), + opt_image_height=512, + opt_image_width=512, + opt_batch_size=1, + build_static_batch=build_static_batch, + build_dynamic_shape=build_dynamic_shape, + ) + + return captured_build.call_args.kwargs["dynamic_shapes"] + + +class TestL2tcDynamicShapesFix: + def test_batch_dynamic_resolution_static_is_dynamic(self, monkeypatch): + """Default 'Flexible' UNet preset: static resolution, dynamic batch — + this is the bug case; must compute True (previously computed False).""" + result = _computed_dynamic_shapes(monkeypatch, build_static_batch=False, build_dynamic_shape=False) + assert result is True + + def test_fully_static_engine_is_not_dynamic(self, monkeypatch): + """ControlNet-style fully-static engine: unaffected, still computes False.""" + result = _computed_dynamic_shapes(monkeypatch, build_static_batch=True, build_dynamic_shape=False) + assert result is False + + def test_resolution_dynamic_is_dynamic_regardless_of_batch(self, monkeypatch): + """build_dynamic_shape=True alone was already sufficient before this fix.""" + result = _computed_dynamic_shapes(monkeypatch, build_static_batch=True, build_dynamic_shape=True) + assert result is True From d5efd4dbb84f12d88a43a59c092875c1e10246de Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 10 Jul 2026 20:01:03 -0400 Subject: [PATCH 11/18] perf: elide trt.input_staging DtoD copy for kvo/fio UNet cache inputs (Sub-phase 5.6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Engine.infer() copied every feed_dict tensor into its own persistent staging buffer each frame solely to give TensorRT a stable contiguous address to bind. nsys profiling (5.0/5.5) attributed ~9% of frame time (p95 3.1ms, ~161 DtoD copies/frame) to this region on the UNet call, dominated by the kvo/fio StreamV2V cache inputs. Those cache tensors are already persistent, address-stable, TRT-contiguous buffers (create_kvo_cache/create_fi_cache build them with layers_in_bucket as the outermost dim specifically for this) — the only event that moves their address is a batch-size change, which already forces a full buffer realloc+graph reset today. So copying them every frame is pure waste. Add an opt-in zero_copy_names parameter to Engine.infer() and a pure _staging_action() decision helper (copy/bind/bind_and_reset) that lets TensorRT bind directly to the caller's tensor instead of copying into self.tensors[name], guarded by contiguity and dtype-match fallbacks. Wire it in UNet2DConditionModelEngine for the kvo/fio input names only; every other caller (VAE, ControlNet, safety, preprocessing) passes nothing and keeps the original copy path byte-for-byte. --- .../tensorrt/runtime_engines/unet_engine.py | 10 + .../acceleration/tensorrt/utilities.py | 69 ++++++- tests/unit/test_zero_copy_staging_5_6.py | 181 ++++++++++++++++++ 3 files changed, 257 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_zero_copy_staging_5_6.py diff --git a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py index da7da67ff..462d0290a 100644 --- a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py +++ b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py @@ -54,6 +54,13 @@ def __init__(self, filepath: str, stream: "cuda.Stream", use_cuda_graph: bool = self._fio_out_names: List[str] = [] self._fio_cache_len: int = -1 # -1 = not yet initialized + # Sub-phase 5.6: kvo/fio cache inputs are persistent, address-stable, + # TRT-contiguous tensors (see models/utils.py create_kvo_cache/create_fi_cache), + # so Engine.infer() can bind TensorRT directly to them instead of copying into + # its own staging buffer every frame. Rebuilt only when the kvo/fio lazy-init + # branches below fire (cache length change), never per-frame. + self._zero_copy_names: FrozenSet[str] = frozenset() + self._shape_dict: Dict[str, Any] = {} self._input_dict: Dict[str, Any] = {} @@ -90,6 +97,7 @@ def __call__( self._kvo_in_names = [f"kvo_cache_in_{i}" for i in range(n_kvo)] self._kvo_out_names = [f"kvo_cache_out_{i}" for i in range(n_kvo)] self._kvo_cache_len = n_kvo + self._zero_copy_names = frozenset(self._kvo_in_names + self._fio_in_names) # Lazy-init FIO key name lists — same pattern as KVO. n_fio = len(fio_cache) @@ -97,6 +105,7 @@ def __call__( self._fio_in_names = [f"fio_cache_in_{i}" for i in range(n_fio)] self._fio_out_names = [f"fio_cache_out_{i}" for i in range(n_fio)] self._fio_cache_len = n_fio + self._zero_copy_names = frozenset(self._kvo_in_names + self._fio_in_names) # Update pre-allocated dicts in-place — no new dict objects created per call. shape_dict = self._shape_dict @@ -198,6 +207,7 @@ def __call__( input_dict, self.stream, use_cuda_graph=self.use_cuda_graph, + zero_copy_names=self._zero_copy_names, ) except Exception as e: logger.exception(f"UNet2DConditionModelEngine.__call__: Engine.infer failed: {e}") diff --git a/src/streamdiffusion/acceleration/tensorrt/utilities.py b/src/streamdiffusion/acceleration/tensorrt/utilities.py index 6eec841d4..67be6ae74 100644 --- a/src/streamdiffusion/acceleration/tensorrt/utilities.py +++ b/src/streamdiffusion/acceleration/tensorrt/utilities.py @@ -541,6 +541,35 @@ def _median(v): return "\n".join(lines) +def _staging_action( + name: str, + zero_copy_names: frozenset, + is_contiguous: bool, + dtype_match: bool, + prev_ptr: Optional[int], + cur_ptr: int, + graph_exists: bool, +) -> str: + """Decide how Engine.infer() should stage one feed_dict input. + + Pure function (no CUDA calls) so the zero-copy decision table is unit-testable + without a real CUDA graph — see Sub-phase 5.6, + docs/perf_bestpractices_audit_2026-07-10.md. + + Returns one of: + "copy" - copy into self.tensors[name] (today's behavior; always safe) + "bind" - skip the copy; bind TensorRT directly to the caller's tensor + "bind_and_reset" - skip the copy and bind directly, but reset the CUDA graph + first because the caller's address changed while a graph + built from the old address was still live + """ + if name not in zero_copy_names or not is_contiguous or not dtype_match: + return "copy" + if graph_exists and cur_ptr != prev_ptr: + return "bind_and_reset" + return "bind" + + class Engine: def __init__( self, @@ -561,6 +590,11 @@ def __init__( # Cached ExternalStream wrapping the engine's polygraphy stream; allocated on # first infer() call so we avoid constructing a new Python wrapper every frame. self._engine_ext_stream = None + # Sub-phase 5.6: last-bound caller tensor address per zero-copy input name, + # so infer() can detect a pointer change and force a graph reset instead of + # silently replaying stale addresses. Only populated for names actually + # passed via infer()'s zero_copy_names. + self._bound_ptrs: dict[str, int] = {} def __del__(self): # Check if AttributeError: 'Engine' object has no attribute 'buffers' @@ -1086,7 +1120,7 @@ def reset_cuda_graph(self): CUASSERT(cudart.cudaGraphDestroy(self.graph)) self.graph = None - def infer(self, feed_dict, stream, use_cuda_graph=False): + def infer(self, feed_dict, stream, use_cuda_graph=False, zero_copy_names: frozenset = frozenset()): # IProfiler cannot report per-layer times through CUDA graph replay — disable graphs # when profiler is attached. This is automatically set when STREAMDIFFUSION_PROFILE_TRT # is set in activate(), so callers do not need to change anything. @@ -1133,10 +1167,36 @@ def infer(self, feed_dict, stream, use_cuda_graph=False): pt_stream, stream.ptr, ) + # Sub-phase 5.6: for opt-in zero-copy names (kvo/fio UNet cache inputs — already + # persistent, address-stable, TRT-contiguous tensors), skip the DtoD copy and + # point TensorRT directly at the caller's tensor instead. Every other input + # (zero_copy_names defaults to frozenset()) takes the original copy_() path, + # so non-UNet engines (VAE/ControlNet/safety) are byte-for-byte unchanged. + bind_ptrs: dict[str, int] = {} + needs_reset = False + graph_exists = self.cuda_graph_instance is not None with torch.cuda.stream(self._engine_ext_stream): with _gpu_profiler.region("trt.input_staging"): for name, buf in feed_dict.items(): - self.tensors[name].copy_(buf) + cur_ptr = buf.data_ptr() + action = _staging_action( + name, + zero_copy_names, + buf.is_contiguous(), + buf.dtype == self.tensors[name].dtype, + self._bound_ptrs.get(name), + cur_ptr, + graph_exists, + ) + if action == "copy": + self.tensors[name].copy_(buf) + else: + bind_ptrs[name] = cur_ptr + if action == "bind_and_reset": + needs_reset = True + + if needs_reset and self.cuda_graph_instance is not None: + self.reset_cuda_graph() # In graphed steady state the tensor addresses are baked into the captured graph # (self.tensors[name] are persistent buffers reused via copy_() — see @@ -1145,8 +1205,11 @@ def infer(self, feed_dict, stream, use_cuda_graph=False): # host overhead; only rebind on first capture or after a reset (instance is None). if not (use_cuda_graph and self.cuda_graph_instance is not None): for name, tensor in self.tensors.items(): - if not self.context.set_tensor_address(name, tensor.data_ptr()): + address = bind_ptrs.get(name, tensor.data_ptr()) + if not self.context.set_tensor_address(name, address): raise RuntimeError(f"TensorRT: set_tensor_address failed for '{name}'") + if name in bind_ptrs: + self._bound_ptrs[name] = address with _gpu_profiler.region("trt_infer"): if use_cuda_graph: diff --git a/tests/unit/test_zero_copy_staging_5_6.py b/tests/unit/test_zero_copy_staging_5_6.py new file mode 100644 index 000000000..56bbe341f --- /dev/null +++ b/tests/unit/test_zero_copy_staging_5_6.py @@ -0,0 +1,181 @@ +""" +Unit tests for Sub-phase 5.6 (`trt.input_staging` DtoD copy elision, +docs/perf_bestpractices_audit_2026-07-10.md follow-up — see the Phase 5 +continuation plan). + +`Engine.infer()` copies every feed_dict input into its own persistent staging +buffer every frame, purely to give TensorRT a stable contiguous address to +bind. The kvo/fio UNet cache inputs are already persistent, address-stable, +TRT-contiguous tensors (see models/utils.py create_kvo_cache/create_fi_cache), +so that copy is pure waste for them — 5.6 lets `Engine.infer()` bind directly +to an opt-in set of caller tensors instead. + +`_staging_action()` is the pure decision function extracted from `Engine.infer()` +so the per-input zero-copy/copy/reset decision table is testable without a real +CUDA graph (a `cuda_graph_instance` can't be meaningfully MagicMock'd — the same +constraint Sub-phase 5.1's rebind guard hit). This module tests only that pure +function: no CUDA device, no TensorRT context required beyond importing the +module (guarded below, same pattern as test_l2tc_dynamic_shapes.py). +""" + +import pytest + + +try: + from streamdiffusion.acceleration.tensorrt.utilities import _staging_action + + IMPORT_OK = True +except ImportError: + IMPORT_OK = False + +pytestmark = pytest.mark.skipif( + not IMPORT_OK, + reason="acceleration.tensorrt.utilities not importable (TensorRT/onnx/polygraphy missing)", +) + + +PTR_A = 0xDEAD_BEEF +PTR_B = 0xFEED_FACE + + +class TestStagingActionNotZeroCopy: + """Anything not opted in always falls back to the original copy path.""" + + def test_name_not_in_zero_copy_names_copies(self): + result = _staging_action( + "sample", + frozenset({"kvo_cache_in_0"}), + True, + True, + None, + PTR_A, + False, + ) + assert result == "copy" + + def test_empty_zero_copy_names_always_copies(self): + """Default frozenset() → today's behavior exactly, for every name.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset(), + True, + True, + PTR_A, + PTR_A, + True, + ) + assert result == "copy" + + +class TestStagingActionSafetyGuards: + """A zero-copy candidate that fails a safety guard falls back to copy, + never silently binding a non-contiguous tensor or mismatched dtype.""" + + def test_non_contiguous_falls_back_to_copy(self): + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + False, # is_contiguous + True, + None, + PTR_A, + False, + ) + assert result == "copy" + + def test_dtype_mismatch_falls_back_to_copy(self): + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + False, # dtype_match + None, + PTR_A, + False, + ) + assert result == "copy" + + def test_both_guards_fail_falls_back_to_copy(self): + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + False, + False, + None, + PTR_A, + True, + ) + assert result == "copy" + + +class TestStagingActionBind: + """Eligible names bind directly; a reset is forced only when a live graph's + baked address would otherwise go stale.""" + + def test_no_graph_yet_binds_without_reset(self): + """First-ever call: no graph exists, so a fresh bind is always safe.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + True, + None, + PTR_A, + False, + ) + assert result == "bind" + + def test_graph_exists_pointer_unchanged_binds_without_reset(self): + """Steady state: same persistent tensor, same address — no reset needed.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + True, + PTR_A, + PTR_A, + True, + ) + assert result == "bind" + + def test_graph_exists_pointer_changed_binds_and_resets(self): + """Belt-and-suspenders: an unforeseen pointer change while a graph is + live must force a re-capture, not silently rebind a stale graph.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + True, + PTR_A, + PTR_B, + True, + ) + assert result == "bind_and_reset" + + def test_no_graph_pointer_changed_still_just_binds(self): + """No live graph to invalidate → no reset needed even if the pointer moved + (e.g. right after allocate_buffers already reset the graph for a shape change).""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + True, + PTR_A, + PTR_B, + False, + ) + assert result == "bind" + + def test_first_bind_with_no_prior_ptr_recorded(self): + """prev_ptr=None (nothing bound yet) with no live graph must not crash and + must bind cleanly — exercises the None-vs-int comparison short-circuit.""" + result = _staging_action( + "fio_cache_in_0", + frozenset({"fio_cache_in_0"}), + True, + True, + None, + PTR_A, + False, + ) + assert result == "bind" From a9b1669a354122b1808f271f7007a2747b238a5d Mon Sep 17 00:00:00 2001 From: Alex Date: Sat, 11 Jul 2026 08:27:29 -0400 Subject: [PATCH 12/18] fix: lower ControlNet dynamic-shape floor 384->256 via shared min_image_shape (bug #6) ControlNet's dynamic-shape get_input_profile() hardcoded a 384px floor independent of the UNet's 256px floor (BaseModel.min_image_shape), so any resolution in [256, 384) hard-failed at Engine.allocate_buffers the moment ControlNet was enabled. Derive the ControlNet floor/ceiling from the same BaseModel.min_image_shape/max_image_shape the UNet already uses (SDXL inherits via super()), lower the matching engine_manager.py build-option default, and relabel the ControlNet dynamic-engine cache directory (--dyn-384-1024 -> --dyn-256-1024) so a re-floored build never collides with a stale 384-floored cache. Trade-off: widening the dynamic profile can slightly reduce TRT tactic specialization at the optimum resolution -- acceptable given the alternative is a hard failure. Co-Authored-By: Claude Sonnet 5 --- .../acceleration/tensorrt/engine_manager.py | 8 ++++++-- .../tensorrt/models/controlnet_models.py | 14 ++++++++------ 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/streamdiffusion/acceleration/tensorrt/engine_manager.py b/src/streamdiffusion/acceleration/tensorrt/engine_manager.py index a7a301c1f..6ed4ad645 100644 --- a/src/streamdiffusion/acceleration/tensorrt/engine_manager.py +++ b/src/streamdiffusion/acceleration/tensorrt/engine_manager.py @@ -138,7 +138,7 @@ def get_engine_path( if resolution is not None: prefix = f"controlnet_{model_dir_name}--min_batch-{min_batch_size}--max_batch-{max_batch_size}--res-{resolution[0]}x{resolution[1]}" else: - prefix = f"controlnet_{model_dir_name}--min_batch-{min_batch_size}--max_batch-{max_batch_size}--dyn-384-1024" + prefix = f"controlnet_{model_dir_name}--min_batch-{min_batch_size}--max_batch-{max_batch_size}--dyn-256-1024" fp8_suffix = "--fp8" if fp8 else "" return self.engine_dir / (prefix + optlvl_suffix + fp8_suffix) / filename else: @@ -277,7 +277,11 @@ def _get_default_controlnet_build_options( "build_static_batch": True, } if build_dynamic_shape: - opts["min_image_resolution"] = 384 + # Match BaseModel/UNet's 256 floor (was 384) so [256, 384) resolutions + # don't hard-fail with ControlNet active — see get_input_profile in + # controlnet_models.py, which now derives its own floor from the same + # BaseModel.min_image_shape instead of a separate hardcoded literal. + opts["min_image_resolution"] = 256 opts["max_image_resolution"] = 1024 if builder_optimization_level is not None: opts["builder_optimization_level"] = builder_optimization_level diff --git a/src/streamdiffusion/acceleration/tensorrt/models/controlnet_models.py b/src/streamdiffusion/acceleration/tensorrt/models/controlnet_models.py index 125c8797a..cba02b6b6 100644 --- a/src/streamdiffusion/acceleration/tensorrt/models/controlnet_models.py +++ b/src/streamdiffusion/acceleration/tensorrt/models/controlnet_models.py @@ -65,12 +65,14 @@ def get_input_profile(self, batch_size, image_height, image_width, static_batch, min_ctrl_h = max_ctrl_h = opt_ctrl_h = image_height min_ctrl_w = max_ctrl_w = opt_ctrl_w = image_width else: - min_ctrl_h = 384 - max_ctrl_h = 1024 - opt_ctrl_h = 704 - min_ctrl_w = 384 - max_ctrl_w = 1024 - opt_ctrl_w = 704 + # Share the floor/ceiling with BaseModel (min_image_shape/max_image_shape) + # instead of a hardcoded 384 — keeps ControlNet's dynamic range aligned with + # the UNet's (256-1024), so [256, 384) resolutions no longer hard-fail with + # ControlNet active. opt stays 704, clamped into range defensively. + min_ctrl_h, max_ctrl_h = self.min_image_shape, self.max_image_shape + min_ctrl_w, max_ctrl_w = self.min_image_shape, self.max_image_shape + opt_ctrl_h = min(max(704, min_ctrl_h), max_ctrl_h) + opt_ctrl_w = min(max(704, min_ctrl_w), max_ctrl_w) min_latent_h = min_ctrl_h // 8 max_latent_h = max_ctrl_h // 8 From 75ee38d6e397f65efc97d55fd3d3e7a0e58ff002 Mon Sep 17 00:00:00 2001 From: Alex Date: Sat, 11 Jul 2026 08:43:21 -0400 Subject: [PATCH 13/18] perf: normalize blend weights on CPU + persistent ControlNet conditioning_scale scalar (Sub-phase 5.4) Two per-frame GPU creation syncs removed: - _normalize_weights built a CUDA tensor from a Python list every call, but every consumer (.item()/.tolist() in slerp/multi_slerp/cosine_weighted paths, float*cuda_tensor in the linear paths) immediately reads it back or treats it as a scalar operand. Build it on CPU float32 instead -- same consumers work unchanged (a 0-dim CPU tensor is treated as a wrapped scalar when multiplied against a CUDA tensor), the sync and readback are both gone, and float32 is more precise than the model's fp16 for the normalization divide. - ControlNetModelEngine rebuilt a conditioning_scale tensor via torch.tensor(...) on every __call__. Replace with one lazily-allocated persistent scalar updated via fill_() (an async kernel launch, no sync), mirroring the existing _fi_strength_tensor.fill_() idiom. Each active ControlNet has its own engine instance, so a per-instance buffer is correct; the binding rank (_cs_rank1) is fixed at load time so the shape never changes underneath it. Also drops three dead variable assignments and fixes one out-of-order import in stream_parameter_updater.py (pre-existing lint debt in a file this change already touches; the repo's commit gate lints the whole staged file, not just the diff). Co-Authored-By: Claude Sonnet 5 --- .../runtime_engines/controlnet_engine.py | 16 ++++++++---- .../stream_parameter_updater.py | 26 ++++++++++--------- 2 files changed, 25 insertions(+), 17 deletions(-) diff --git a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/controlnet_engine.py b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/controlnet_engine.py index 35569c4f8..7d1ce77c2 100644 --- a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/controlnet_engine.py +++ b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/controlnet_engine.py @@ -42,6 +42,12 @@ def __init__( except Exception: self._cs_rank1 = False # sd15 engines / engines lacking this input + # Persistent conditioning_scale scalar, allocated lazily on first __call__ + # (device isn't known until then). fill_() is an async kernel launch, so + # reusing this buffer every call avoids the per-frame creation sync that + # torch.tensor(...) incurs. + self._cs: Optional[torch.Tensor] = None + # Pre-compute model-specific values to eliminate runtime branching if self.model_type in ["sdxl", "sdxl_turbo"]: self.max_blocks = 9 @@ -119,16 +125,16 @@ def __call__( if timestep.dtype != torch.float32: timestep = timestep.float() + if self._cs is None: + self._cs = torch.zeros((1,) if self._cs_rank1 else (), dtype=torch.float32, device=sample.device) + self._cs.fill_(conditioning_scale) + input_dict = { "sample": sample, "timestep": timestep, "encoder_hidden_states": encoder_hidden_states, "controlnet_cond": controlnet_cond, - "conditioning_scale": ( - torch.tensor([conditioning_scale], dtype=torch.float32, device=sample.device) - if self._cs_rank1 - else torch.tensor(conditioning_scale, dtype=torch.float32, device=sample.device) - ), + "conditioning_scale": self._cs, } if text_embeds is not None: diff --git a/src/streamdiffusion/stream_parameter_updater.py b/src/streamdiffusion/stream_parameter_updater.py index 233845858..13ca8cd1a 100644 --- a/src/streamdiffusion/stream_parameter_updater.py +++ b/src/streamdiffusion/stream_parameter_updater.py @@ -5,9 +5,10 @@ import torch import torch.nn.functional as F +from .preprocessing.orchestrator_user import OrchestratorUser + logger = logging.getLogger(__name__) -from .preprocessing.orchestrator_user import OrchestratorUser class CacheStats: @@ -129,11 +130,9 @@ def register_embedding_preprocessor(self, preprocessor: Any, style_image_key: st def unregister_embedding_preprocessor(self, style_image_key: str) -> None: """Unregister an embedding preprocessor by style image key.""" - original_count = len(self._embedding_preprocessors) self._embedding_preprocessors = [ (preprocessor, key) for preprocessor, key in self._embedding_preprocessors if key != style_image_key ] - removed_count = original_count - len(self._embedding_preprocessors) # Clear cached embeddings for this key if style_image_key in self._embedding_cache: @@ -213,8 +212,15 @@ def get_cached_embeddings(self, style_image_key: str) -> Optional[Tuple[torch.Te return cached_result def _normalize_weights(self, weights: List[float], normalize: bool) -> torch.Tensor: - """Generic weight normalization helper""" - weights_tensor = torch.tensor(weights, device=self.stream.device, dtype=self.stream.dtype) + """Generic weight normalization helper. + + Built on CPU float32 rather than self.stream.device/dtype: every caller either + reads scalars back out (.item()/.tolist()) or multiplies into a CUDA tensor + (a 0-dim CPU operand is treated as a wrapped scalar there), so building on + device was a pure creation-sync + readback with no benefit — and float32 is + more precise than the model's fp16 for the normalization divide. + """ + weights_tensor = torch.tensor(weights, dtype=torch.float32) if normalize: weights_tensor = weights_tensor / weights_tensor.sum() return weights_tensor @@ -992,11 +998,7 @@ def _update_timestep_calculations(self) -> None: # mis-interprets the clean buffer, causing ghost bleed from previous frames. # Threshold 0.75 matches the empirically observed perceptual onset (~t_index 30 in # a 50-step LCM schedule where beta_sqrt crosses 0.78). - if ( - self.stream.use_denoising_batch - and not self.stream.do_add_noise - and len(self.stream.t_list) > 1 - ): + if self.stream.use_denoising_batch and not self.stream.do_add_noise and len(self.stream.t_list) > 1: inter_step_betas = beta_prod_t_sqrt[1:, 0, 0, 0] # per-step, before repeat_interleave max_beta = inter_step_betas.max().item() _BLEED_THRESHOLD = 0.75 @@ -1201,7 +1203,7 @@ def remove_prompt_at_index( return # Remove from current list - removed_prompt = self._current_prompt_list.pop(index) + self._current_prompt_list.pop(index) # Remove from cache and reindex if index in self._prompt_cache: @@ -1302,7 +1304,7 @@ def remove_seed_at_index(self, index: int, interpolation_method: Literal["linear return # Remove from current list - removed_seed = self._current_seed_list.pop(index) + self._current_seed_list.pop(index) # Remove from cache and reindex if index in self._seed_cache: From fcb2ae25de5decd4496799132a89d88d02d87c2d Mon Sep 17 00:00:00 2001 From: Alex Date: Sat, 11 Jul 2026 08:56:28 -0400 Subject: [PATCH 14/18] perf: 1-frame-delayed async NSFW safety-checker readback (Sub-phase 5.3) Removes the cudaStreamSynchronize that NSFWDetectorEngine.__call__ forced on the hot path via probs[0,0].item() whenever use_safety_checker=True. The classifier now writes its GPU probability into a pinned host scalar via a non_blocking copy; _apply_safety_checker reads the *previous* frame's pinned verdict before launching this frame's classification, mirroring the existing SimilarImageFilter 1-frame-delayed readback pattern (image_filter.py). Edge cases: first frame passes through unconditionally (no prior verdict to act on); a 1-frame detection-leak window is possible before flagging lands. Both are acceptable for an opt-in, off-by-default, TensorRT-only feature. Updated tests/unit/test_safety_checker.py for the new (tensor, prob_pin) -> None checker contract and the delayed-readback semantics (10/10 passing). Co-Authored-By: Claude Sonnet 5 --- .../tensorrt/runtime_engines/unet_engine.py | 12 ++- src/streamdiffusion/wrapper.py | 37 +++++++- tests/unit/test_safety_checker.py | 93 +++++++++++++------ 3 files changed, 111 insertions(+), 31 deletions(-) diff --git a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py index 462d0290a..de317065d 100644 --- a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py +++ b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py @@ -506,7 +506,13 @@ def __init__(self, filepath: str, stream: "cuda.Stream", use_cuda_graph: bool = self.engine.load() self.engine.activate() - def __call__(self, image_tensor: torch.Tensor, threshold: float): + def __call__(self, image_tensor: torch.Tensor, prob_pin: torch.Tensor) -> None: + """Launch classification and stash the NSFW probability in a pinned host buffer. + + Does not read the result back (no .item()/host sync here) — callers read + `prob_pin` on a *later* call, mirroring SimilarImageFilter's 1-frame-delayed + pinned readback so this never forces a cudaStreamSynchronize on the hot path. + """ pixel_values = self.image_transforms(image_tensor) self.engine.allocate_buffers( shape_dict={ @@ -522,8 +528,8 @@ def __call__(self, image_tensor: torch.Tensor, threshold: float): )["logits"] probs = F.softmax(logits, dim=-1) - nsfw_prob = 1 - probs[0, 0].item() - return nsfw_prob >= threshold + nsfw_prob = 1 - probs[0, 0] + prob_pin.copy_(nsfw_prob.view(1), non_blocking=True) def to(self, *args, **kwargs): pass diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py index c981661db..17055f54c 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -376,6 +376,10 @@ def __init__( # Caches the last clean (non-flagged) pipeline tensor for the "previous" fallback strategy. # Operates in diffusion range [-1, 1]; set by _apply_safety_checker(). self._prev_clean_tensor: Optional[torch.Tensor] = None + # Pinned CPU scalar for the 1-frame-delayed async NSFW readback (lazy; see + # _apply_safety_checker). Distinct from _output_pin_buf, which is uint8 + # image-shaped and reallocated by the output-postprocessing paths. + self._nsfw_prob_pin: Optional[torch.Tensor] = None self.fp8 = fp8 self.static_shapes = static_shapes self.fp8_allow_fp16_fallback = fp8_allow_fp16_fallback @@ -1360,6 +1364,19 @@ def _apply_safety_checker(self, image_tensor: torch.Tensor) -> torch.Tensor: CUDA-IPC export path: postprocess_image() exports the frame inside its own body and returns None, making any post-hoc substitution unreachable and unsafe. + 1-frame-delayed async readback (mirrors SimilarImageFilter, image_filter.py): the + classifier result for THIS frame is read on the NEXT call, so the substitution below + acts on the PREVIOUS frame's verdict while this frame's classification + pinned + `non_blocking` copy is merely launched. This avoids forcing a cudaStreamSynchronize + on the hot path. No explicit sync guards the pinned read — the same trade-off + SimilarImageFilter already makes: stream ordering means the async copy is enqueued + before all of this frame's remaining GPU work, and by the time Python reaches this + call again next frame (after a full diffusion step + decode, and — for "np"/"pil" + output — a hard _d2h_event.synchronize() in postprocess_image) the tiny 4-byte D2H + copy has long since landed in practice. Worst case on other output types is a stale + read of an even-older value, i.e. one extra frame of detection delay, not corrupted + data. Accepted trade-off for an opt-in, off-by-default, TensorRT-only feature. + Parameters ---------- image_tensor : torch.Tensor @@ -1377,7 +1394,24 @@ def _apply_safety_checker(self, image_tensor: torch.Tensor) -> torch.Tensor: # Denormalize to [0, 1] NCHW for the classifier; stays on GPU. denormalized = self._denormalize_on_gpu(image_tensor) - if self.safety_checker(denormalized, self.safety_checker_threshold): + if self._nsfw_prob_pin is None: + # First frame: no prior async result exists yet. Launch this frame's + # classification so the next frame has a real verdict, but pass this one + # through unflagged (documented 1-frame pass-through edge case). + self._nsfw_prob_pin = torch.zeros(1, dtype=torch.float32, device="cpu").pin_memory() + self.safety_checker(denormalized, self._nsfw_prob_pin) + if self.safety_checker_fallback_type == "previous": + self._prev_clean_tensor = image_tensor.clone() + return image_tensor + + # Step 1: read the PREVIOUS frame's async result (pinned CPU read, no GPU sync). + flagged = self._nsfw_prob_pin.item() >= self.safety_checker_threshold + + # Step 2: launch THIS frame's classification + async pinned copy (no sync). + self.safety_checker(denormalized, self._nsfw_prob_pin) + + # Step 3: branch on the PREVIOUS frame's verdict (1-frame-delayed, no stall). + if flagged: logger.info("NSFW content detected, applying safety fallback frame") if self.safety_checker_fallback_type == "previous" and self._prev_clean_tensor is not None: return self._prev_clean_tensor @@ -3037,6 +3071,7 @@ def cleanup_gpu_memory(self) -> None: self._d2h_event = None self._ipc_pack_buf = None self._ipc_pack_unit_buf = None + self._nsfw_prob_pin = None # Force multiple garbage collection cycles for thorough cleanup for i in range(3): diff --git a/tests/unit/test_safety_checker.py b/tests/unit/test_safety_checker.py index ef5a3d133..1e2b1cbbb 100644 --- a/tests/unit/test_safety_checker.py +++ b/tests/unit/test_safety_checker.py @@ -15,6 +15,14 @@ Fix: _apply_safety_checker() runs *before* postprocess_image, operating on the raw diffusion-range [-1, 1] pipeline tensor so it is always a real tensor regardless of output path. + +Sub-phase 5.3: the checker's verdict is now read 1-frame-delayed from a pinned +async buffer (mirrors SimilarImageFilter) instead of a synchronous .item() call. +`safety_checker` is now `(tensor, prob_pin) -> None`: it writes into `prob_pin` +rather than returning a bool, and `_apply_safety_checker` reads the *previous* +call's pinned value before launching (and writing) this frame's result. Tests +below seed `w._nsfw_prob_pin` directly where they need to bypass the +first-frame pass-through and exercise a specific verdict. """ import torch @@ -39,6 +47,7 @@ def _make_wrapper( w.safety_checker_threshold = threshold w.safety_checker_fallback_type = fallback_type w._prev_clean_tensor = None + w._nsfw_prob_pin = None # safety_checker is set per-test w.safety_checker = None return w @@ -65,9 +74,9 @@ def test_checker_never_receives_none(self): """ received: list = [] - def capturing_checker(tensor, thr): + def capturing_checker(tensor, prob_pin): received.append(tensor) - return True # always flag + prob_pin.fill_(1.0) w = _make_wrapper() w.safety_checker = capturing_checker @@ -83,12 +92,15 @@ def capturing_checker(tensor, thr): # ── Case 2: NSFW → black frame (blank fallback) ────────────────────────── def test_nsfw_blank_fallback_is_black(self): w = _make_wrapper(fallback_type="blank") - w.safety_checker = lambda t, thr: True # always flag + w.safety_checker = lambda t, pin: pin.fill_(1.0) # always "flag" - dummy = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(dummy) + frame1 = torch.randn(1, 3, 64, 64) + _ = w._apply_safety_checker(frame1) # first frame: pass-through, primes the pin + + frame2 = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(frame2) # reads frame1's flagged verdict - assert result is not dummy, "flagged frame should be replaced" + assert result is not frame2, "flagged frame should be replaced" # Denormalize: (x/2+0.5).clamp(0,1); -1.0 → 0.0 = black denorm = (result / 2 + 0.5).clamp(0, 1) assert _black_denorm(denorm), "NSFW blank fallback should produce a black frame" @@ -98,37 +110,42 @@ def test_nsfw_previous_fallback_returns_cached_clean_frame(self): w = _make_wrapper(fallback_type="previous") call_count = [0] - def checker(t, thr): + def checker(t, pin): call_count[0] += 1 - return call_count[0] > 1 # first call: clean; second call: flagged + # Call #1 (launched during frame1) writes a flagged prob; frame2's + # read consumes that async result (1-frame delay). + pin.fill_(1.0 if call_count[0] == 1 else 0.0) w.safety_checker = checker clean = torch.randn(1, 3, 64, 64) - _ = w._apply_safety_checker(clean) # primes _prev_clean_tensor + _ = w._apply_safety_checker(clean) # first frame: pass-through, primes _prev_clean_tensor - flagged = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(flagged) + flagged_frame = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(flagged_frame) assert result is w._prev_clean_tensor, "previous-fallback should return the cached clean tensor" + assert torch.equal(w._prev_clean_tensor, clean) # ── Case 4: Clean frame passthrough ────────────────────────────────────── def test_clean_frame_returned_unchanged(self): w = _make_wrapper() - w.safety_checker = lambda t, thr: False # never flag + w.safety_checker = lambda t, pin: pin.fill_(0.0) # never flag - dummy = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(dummy) + frame1 = torch.randn(1, 3, 64, 64) + _ = w._apply_safety_checker(frame1) # first frame: pass-through, primes the pin - assert torch.equal(result, dummy), "clean frame should be returned unchanged" + frame2 = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(frame2) + + assert torch.equal(result, frame2), "clean frame should be returned unchanged" # ── Case 5: use_safety_checker=False bypasses entirely ─────────────────── def test_disabled_bypasses_checker(self): called = [False] - def should_not_be_called(t, thr): + def should_not_be_called(t, pin): called[0] = True - return False w = _make_wrapper(use_safety_checker=False) w.safety_checker = should_not_be_called @@ -139,14 +156,17 @@ def should_not_be_called(t, thr): assert not called[0], "safety checker must not be called when disabled" assert torch.equal(result, dummy), "disabled checker should return input unchanged" - # ── Case 6: previous fallback with no cached frame → black ─────────────── + # ── Case 6: flagged verdict with no cached frame → black ───────────────── def test_nsfw_previous_fallback_no_cache_falls_back_to_black(self): """ - If the very first frame is flagged and _prev_clean_tensor is None, - the fallback should still produce a black frame rather than raise. + If a frame is flagged (per the previous frame's pinned verdict) and + _prev_clean_tensor is None, the fallback should still produce a black + frame rather than raise. Seeds _nsfw_prob_pin directly to bypass the + first-frame pass-through and exercise this state combination. """ w = _make_wrapper(fallback_type="previous") - w.safety_checker = lambda t, thr: True # flag immediately + w.safety_checker = lambda t, pin: pin.fill_(0.0) + w._nsfw_prob_pin = torch.tensor([1.0]) # a prior flagged verdict already landed assert w._prev_clean_tensor is None # no cache yet dummy = torch.randn(1, 3, 64, 64) @@ -158,7 +178,8 @@ def test_nsfw_previous_fallback_no_cache_falls_back_to_black(self): # ── Case 7: clean frame caches for previous strategy ───────────────────── def test_clean_frame_cached_for_previous_strategy(self): w = _make_wrapper(fallback_type="previous") - w.safety_checker = lambda t, thr: False + w.safety_checker = lambda t, pin: pin.fill_(0.0) + w._nsfw_prob_pin = torch.tensor([0.0]) # a prior clean verdict already landed assert w._prev_clean_tensor is None dummy = torch.randn(1, 3, 64, 64) @@ -167,18 +188,35 @@ def test_clean_frame_cached_for_previous_strategy(self): assert w._prev_clean_tensor is not None, ( "clean frame with previous strategy should populate _prev_clean_tensor" ) + assert torch.equal(w._prev_clean_tensor, dummy) # ── Case 8: clean frame NOT cached for blank strategy ──────────────────── def test_clean_frame_not_cached_for_blank_strategy(self): w = _make_wrapper(fallback_type="blank") - w.safety_checker = lambda t, thr: False + w.safety_checker = lambda t, pin: pin.fill_(0.0) + w._nsfw_prob_pin = torch.tensor([0.0]) # a prior clean verdict already landed dummy = torch.randn(1, 3, 64, 64) w._apply_safety_checker(dummy) assert w._prev_clean_tensor is None, "_prev_clean_tensor should stay None when fallback_type='blank'" - # ── Case 9: _process_skip_diffusion wiring — checker is called, result is black ── + # ── Case 9: first frame always passes through ──────────────────────────── + def test_first_frame_always_passes_through_regardless_of_verdict(self): + """ + The very first frame has no prior async result to read, so it must + pass through unscreened even if the classifier would flag it — the + documented 1-frame pass-through edge case of the delayed-readback design. + """ + w = _make_wrapper(fallback_type="blank") + w.safety_checker = lambda t, pin: pin.fill_(1.0) # would flag if read this frame + + dummy = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(dummy) + + assert torch.equal(result, dummy), "first frame must pass through (no prior verdict to act on)" + + # ── Case 10: _process_skip_diffusion wiring — checker is called, result is black ── def test_skip_diffusion_routes_through_safety_checker(self): """ Verify that _process_skip_diffusion actually feeds its tensor through @@ -191,12 +229,13 @@ def test_skip_diffusion_routes_through_safety_checker(self): """ received: list = [] - def capturing_checker(tensor, thr): + def capturing_checker(tensor, pin): received.append(tensor) - return True # always flag + pin.fill_(0.0) # this frame's own async result — irrelevant to this frame's own read w = _make_wrapper(fallback_type="blank") w.safety_checker = capturing_checker + w._nsfw_prob_pin = torch.tensor([1.0]) # a prior frame's flagged verdict already landed w.mode = "img2img" w.device = torch.device("cpu") w.dtype = torch.float32 @@ -226,7 +265,7 @@ def _apply_image_postprocessing_hooks(self, t): assert len(received) == 1, "safety checker should be called exactly once" assert isinstance(received[0], torch.Tensor), "checker arg must be a Tensor, not None" - # Flagged frame must produce the black substitution (all -1.0 → denorm 0.0) + # Flagged frame (per the pre-seeded pinned verdict) must produce the black substitution denorm = (result / 2 + 0.5).clamp(0, 1) assert torch.allclose(denorm, torch.zeros_like(denorm)), ( "flagged skip-diffusion frame should produce a black output" From 658dcfa21b99c4dbf946ae106ed564761f96a3bb Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 13 Jul 2026 19:29:24 -0400 Subject: [PATCH 15/18] chore: add Claude review workflow to PR2 branch for head-ref validation --- .github/workflows/claude-code-review.yml | 47 ++++++++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 .github/workflows/claude-code-review.yml diff --git a/.github/workflows/claude-code-review.yml b/.github/workflows/claude-code-review.yml new file mode 100644 index 000000000..bef9532a1 --- /dev/null +++ b/.github/workflows/claude-code-review.yml @@ -0,0 +1,47 @@ +name: Claude Code Review + +on: + pull_request: + types: [opened, synchronize, reopened, ready_for_review] + # base-branch filter: PRs *into* SDTD_032_dev, or into any branch in the + # chained remediation-PR stack (2-5), so Claude also fires on chained PRs. + branches: + [ + SDTD_032_dev, + refactor/remediation-phases-1-4, + perf/unet-step-subphase-5x, + perf/trt-zero-copy-audit-hardening, + refactor/param-schema-fp8-preflight, + ] + +jobs: + review: + if: github.event.pull_request.draft == false + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: write + issues: write + id-token: write + steps: + - uses: actions/checkout@v4 + with: { fetch-depth: 0 } + + - uses: anthropics/claude-code-action@v1 + with: + claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }} + track_progress: true # forces tag mode -> sticky comment updates in place + use_sticky_comment: true # one review comment across pushes, not a new one each time + show_full_output: true # surfaces full tool output/denials in the run log for debugging + plugin_marketplaces: "https://github.com/anthropics/claude-code.git" + plugins: "code-review@claude-code-plugins" + prompt: "/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}" + claude_args: >- + --max-turns 40 + --allowedTools "Bash(git diff:*),Bash(git status:*),Bash(git log:*),Bash(git fetch:*),Bash(git merge-base:*),Bash(git rev-parse:*),Bash(git remote:*),Bash(git branch:*),Bash(git show:*),Bash(echo:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh release view:*)" + --append-system-prompt "Repo-specific focus: real-time diffusion hot-path discipline (flag + .item()/.cpu()/.numpy()/print(tensor) in step/frame loops), no hard-coded .cuda()/device, + best-effort error paths must never raise into the render loop, use raise ... from e, TensorRT + zero-copy buffer safety (dtype/shape/lifetime). Honor intentional style: keep + typing.Dict/Optional/Union (no PEP585/604 rewrites), line-length 119, single-letter math vars + OK, torch.device valid as a default arg. Skip nits ruff/pyrefly already cover." From ed0521befed4fd17c10498462354b152c1862e55 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 13 Jul 2026 19:41:14 -0400 Subject: [PATCH 16/18] chore: widen Claude review allowedTools (Skill, python/pytest, /tmp Write) to unblock permission-denied max-turns failures --- .github/workflows/claude-code-review.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/claude-code-review.yml b/.github/workflows/claude-code-review.yml index bef9532a1..b83bd4db1 100644 --- a/.github/workflows/claude-code-review.yml +++ b/.github/workflows/claude-code-review.yml @@ -38,7 +38,7 @@ jobs: prompt: "/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}" claude_args: >- --max-turns 40 - --allowedTools "Bash(git diff:*),Bash(git status:*),Bash(git log:*),Bash(git fetch:*),Bash(git merge-base:*),Bash(git rev-parse:*),Bash(git remote:*),Bash(git branch:*),Bash(git show:*),Bash(echo:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh release view:*)" + --allowedTools "Bash(git diff:*),Bash(git status:*),Bash(git log:*),Bash(git fetch:*),Bash(git merge-base:*),Bash(git rev-parse:*),Bash(git remote:*),Bash(git branch:*),Bash(git show:*),Bash(echo:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh release view:*),Bash(python:*),Bash(python3:*),Bash(pytest:*),Skill,Write(/tmp/*)" --append-system-prompt "Repo-specific focus: real-time diffusion hot-path discipline (flag .item()/.cpu()/.numpy()/print(tensor) in step/frame loops), no hard-coded .cuda()/device, best-effort error paths must never raise into the render loop, use raise ... from e, TensorRT From a6e00fc05ec18a58488ad6743bff59e3a4286249 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 13 Jul 2026 19:51:55 -0400 Subject: [PATCH 17/18] chore: raise Claude review max-turns to 60, discourage filesystem-wide find and output redirection --- .github/workflows/claude-code-review.yml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/workflows/claude-code-review.yml b/.github/workflows/claude-code-review.yml index b83bd4db1..7505cecef 100644 --- a/.github/workflows/claude-code-review.yml +++ b/.github/workflows/claude-code-review.yml @@ -37,11 +37,14 @@ jobs: plugins: "code-review@claude-code-plugins" prompt: "/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}" claude_args: >- - --max-turns 40 + --max-turns 60 --allowedTools "Bash(git diff:*),Bash(git status:*),Bash(git log:*),Bash(git fetch:*),Bash(git merge-base:*),Bash(git rev-parse:*),Bash(git remote:*),Bash(git branch:*),Bash(git show:*),Bash(echo:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh release view:*),Bash(python:*),Bash(python3:*),Bash(pytest:*),Skill,Write(/tmp/*)" --append-system-prompt "Repo-specific focus: real-time diffusion hot-path discipline (flag .item()/.cpu()/.numpy()/print(tensor) in step/frame loops), no hard-coded .cuda()/device, best-effort error paths must never raise into the render loop, use raise ... from e, TensorRT zero-copy buffer safety (dtype/shape/lifetime). Honor intentional style: keep typing.Dict/Optional/Union (no PEP585/604 rewrites), line-length 119, single-letter math vars - OK, torch.device valid as a default arg. Skip nits ruff/pyrefly already cover." + OK, torch.device valid as a default arg. Skip nits ruff/pyrefly already cover. + Tool usage: never run filesystem-wide search (find /, grep -r /), stay inside the repo + checkout. Never redirect Bash output to a file with > or >>, it is blocked; pipe to head/tail + or just read the direct stdout instead." From 47a87eb1f35fac747329bb1672bfa4d718b02130 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 13 Jul 2026 20:32:55 -0400 Subject: [PATCH 18/18] fix: correct NSFW verdict-to-frame attribution and guard pin_memory on CPU-only builds _apply_safety_checker previously read the prior frame's async verdict but applied it to the current frame's pixels, letting an isolated NSFW frame leak through and wrongly blanking the next clean frame. Buffer the raw frame in _pending_frame and emit it once its own verdict lands (+1 frame latency, first call now emits a black startup frame instead of raw pixels). Also guard .pin_memory() behind torch.cuda.is_available() so the CPU-only test contract in test_safety_checker.py holds. --- src/streamdiffusion/wrapper.py | 72 +++++++----- tests/unit/test_safety_checker.py | 182 +++++++++++++++++++++--------- 2 files changed, 170 insertions(+), 84 deletions(-) diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py index 17055f54c..873d33ef7 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -380,6 +380,9 @@ def __init__( # _apply_safety_checker). Distinct from _output_pin_buf, which is uint8 # image-shaped and reallocated by the output-postprocessing paths. self._nsfw_prob_pin: Optional[torch.Tensor] = None + # Raw frame awaiting its own verdict; buffered one call for correct-attribution + # delayed emission (see _apply_safety_checker). + self._pending_frame: Optional[torch.Tensor] = None self.fp8 = fp8 self.static_shapes = static_shapes self.fp8_allow_fp16_fallback = fp8_allow_fp16_fallback @@ -1364,18 +1367,24 @@ def _apply_safety_checker(self, image_tensor: torch.Tensor) -> torch.Tensor: CUDA-IPC export path: postprocess_image() exports the frame inside its own body and returns None, making any post-hoc substitution unreachable and unsafe. - 1-frame-delayed async readback (mirrors SimilarImageFilter, image_filter.py): the - classifier result for THIS frame is read on the NEXT call, so the substitution below - acts on the PREVIOUS frame's verdict while this frame's classification + pinned - `non_blocking` copy is merely launched. This avoids forcing a cudaStreamSynchronize - on the hot path. No explicit sync guards the pinned read — the same trade-off - SimilarImageFilter already makes: stream ordering means the async copy is enqueued - before all of this frame's remaining GPU work, and by the time Python reaches this - call again next frame (after a full diffusion step + decode, and — for "np"/"pil" - output — a hard _d2h_event.synchronize() in postprocess_image) the tiny 4-byte D2H - copy has long since landed in practice. Worst case on other output types is a stale - read of an even-older value, i.e. one extra frame of detection delay, not corrupted - data. Accepted trade-off for an opt-in, off-by-default, TensorRT-only feature. + 1-frame-delayed async classification with delayed EMISSION (mirrors the async-launch + idea in SimilarImageFilter, image_filter.py, but — unlike that filter — buffers the raw + frame so each frame is gated on its OWN verdict, never a neighbor's). Each call: + reads the pinned verdict for the frame buffered on the PREVIOUS call (now landed), + emits that buffered frame gated on its own verdict, launches classification for THIS + frame, and buffers this frame for the next call. No explicit sync guards the pinned + read — the same trade-off SimilarImageFilter already makes: stream ordering means the + async copy is enqueued before all of this frame's remaining GPU work, and by the time + Python reaches this call again next frame (after a full diffusion step + decode, and — + for "np"/"pil" output — a hard _d2h_event.synchronize() in postprocess_image) the tiny + 4-byte D2H copy has long since landed in practice. This avoids forcing a + cudaStreamSynchronize on the hot path. Accepted trade-off for an opt-in, off-by-default, + TensorRT-only feature. + + Two consequences are inherent to gating each frame on its own async verdict without a + sync: (1) output is delayed by exactly one frame; (2) the very first call has no buffered + frame yet, so it emits a black startup frame instead of passing raw pixels through + unscreened, and the final buffered frame of a stream is never emitted at shutdown. Parameters ---------- @@ -1385,8 +1394,9 @@ def _apply_safety_checker(self, image_tensor: torch.Tensor) -> torch.Tensor: Returns ------- torch.Tensor - The original tensor when content is clean, or a fallback tensor (previous clean - frame, or all-black encoded as -1.0 in diffusion range) when content is flagged. + The previously-buffered frame, unchanged when clean, or a fallback tensor (previous + clean frame, or all-black encoded as -1.0 in diffusion range) when flagged. On the + very first call, an all-black startup frame (no buffered frame exists yet). """ if not self.use_safety_checker: return image_tensor @@ -1394,35 +1404,40 @@ def _apply_safety_checker(self, image_tensor: torch.Tensor) -> torch.Tensor: # Denormalize to [0, 1] NCHW for the classifier; stays on GPU. denormalized = self._denormalize_on_gpu(image_tensor) - if self._nsfw_prob_pin is None: - # First frame: no prior async result exists yet. Launch this frame's - # classification so the next frame has a real verdict, but pass this one - # through unflagged (documented 1-frame pass-through edge case). - self._nsfw_prob_pin = torch.zeros(1, dtype=torch.float32, device="cpu").pin_memory() + if self._pending_frame is None: + # First call: nothing buffered yet, so no frame can be gated on its own verdict. + # Prime the pipeline (launch this frame's classification, buffer it) and emit a + # black safety frame rather than ungated pixels. + pin = torch.zeros(1, dtype=torch.float32, device="cpu") + if torch.cuda.is_available(): + pin = pin.pin_memory() + self._nsfw_prob_pin = pin self.safety_checker(denormalized, self._nsfw_prob_pin) - if self.safety_checker_fallback_type == "previous": - self._prev_clean_tensor = image_tensor.clone() - return image_tensor + self._pending_frame = image_tensor.clone() + return torch.full_like(image_tensor, -1.0) - # Step 1: read the PREVIOUS frame's async result (pinned CPU read, no GPU sync). + # Step 1: read the PENDING frame's own async result (pinned CPU read, no GPU sync). flagged = self._nsfw_prob_pin.item() >= self.safety_checker_threshold # Step 2: launch THIS frame's classification + async pinned copy (no sync). self.safety_checker(denormalized, self._nsfw_prob_pin) - # Step 3: branch on the PREVIOUS frame's verdict (1-frame-delayed, no stall). + # Step 3: rotate the pending frame out and this frame in. + pending = self._pending_frame + self._pending_frame = image_tensor.clone() + if flagged: logger.info("NSFW content detected, applying safety fallback frame") if self.safety_checker_fallback_type == "previous" and self._prev_clean_tensor is not None: return self._prev_clean_tensor # -1.0 in diffusion range → 0.0 after denormalization → true black on every output # path (pt, np, pil, CUDA-IPC). - return torch.full_like(image_tensor, -1.0) + return torch.full_like(pending, -1.0) - # Content is clean — cache it for the "previous" fallback strategy. + # Pending frame is clean — cache it for the "previous" fallback strategy. if self.safety_checker_fallback_type == "previous": - self._prev_clean_tensor = image_tensor.clone() - return image_tensor + self._prev_clean_tensor = pending + return pending def _load_model( self, @@ -3072,6 +3087,7 @@ def cleanup_gpu_memory(self) -> None: self._ipc_pack_buf = None self._ipc_pack_unit_buf = None self._nsfw_prob_pin = None + self._pending_frame = None # Force multiple garbage collection cycles for thorough cleanup for i in range(3): diff --git a/tests/unit/test_safety_checker.py b/tests/unit/test_safety_checker.py index 1e2b1cbbb..b4c1d91a7 100644 --- a/tests/unit/test_safety_checker.py +++ b/tests/unit/test_safety_checker.py @@ -16,13 +16,23 @@ raw diffusion-range [-1, 1] pipeline tensor so it is always a real tensor regardless of output path. -Sub-phase 5.3: the checker's verdict is now read 1-frame-delayed from a pinned -async buffer (mirrors SimilarImageFilter) instead of a synchronous .item() call. -`safety_checker` is now `(tensor, prob_pin) -> None`: it writes into `prob_pin` -rather than returning a bool, and `_apply_safety_checker` reads the *previous* -call's pinned value before launching (and writing) this frame's result. Tests -below seed `w._nsfw_prob_pin` directly where they need to bypass the -first-frame pass-through and exercise a specific verdict. +Sub-phase 5.3 (delayed-emission model): the checker's verdict is read +1-frame-delayed from a pinned async buffer (mirrors the async-launch idea in +SimilarImageFilter), but unlike a plain delayed *readback*, the raw frame +itself is buffered in `_pending_frame` so each frame is gated on its OWN +verdict rather than a neighbor's. `safety_checker` is `(tensor, prob_pin) -> +None`: it writes into `prob_pin` rather than returning a bool. Each call to +`_apply_safety_checker`: + 1. reads the verdict for the frame buffered on the PREVIOUS call (now + landed in the pin), + 2. launches classification for the CURRENT frame, + 3. emits the PENDING (previously buffered) frame, gated on its own verdict, + 4. buffers the CURRENT frame for the next call. +This delays output by exactly one frame. The very first call has no buffered +frame yet, so it emits a black startup frame and primes the pipeline instead +of passing raw pixels through unscreened. Tests below seed `w._nsfw_prob_pin` +and `w._pending_frame` directly where they need to bypass the first-call +priming and exercise a specific verdict against a specific frame. """ import torch @@ -48,6 +58,7 @@ def _make_wrapper( w.safety_checker_fallback_type = fallback_type w._prev_clean_tensor = None w._nsfw_prob_pin = None + w._pending_frame = None # safety_checker is set per-test w.safety_checker = None return w @@ -92,15 +103,16 @@ def capturing_checker(tensor, prob_pin): # ── Case 2: NSFW → black frame (blank fallback) ────────────────────────── def test_nsfw_blank_fallback_is_black(self): w = _make_wrapper(fallback_type="blank") - w.safety_checker = lambda t, pin: pin.fill_(1.0) # always "flag" + w.safety_checker = lambda t, pin: None # verdict is pre-seeded below - frame1 = torch.randn(1, 3, 64, 64) - _ = w._apply_safety_checker(frame1) # first frame: pass-through, primes the pin + pending = torch.randn(1, 3, 64, 64) + w._pending_frame = pending + w._nsfw_prob_pin = torch.tensor([1.0]) # pending frame's own verdict: flagged - frame2 = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(frame2) # reads frame1's flagged verdict + current = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(current) - assert result is not frame2, "flagged frame should be replaced" + assert result is not pending, "flagged pending frame should be replaced" # Denormalize: (x/2+0.5).clamp(0,1); -1.0 → 0.0 = black denorm = (result / 2 + 0.5).clamp(0, 1) assert _black_denorm(denorm), "NSFW blank fallback should produce a black frame" @@ -108,37 +120,33 @@ def test_nsfw_blank_fallback_is_black(self): # ── Case 3: NSFW → previous frame ──────────────────────────────────────── def test_nsfw_previous_fallback_returns_cached_clean_frame(self): w = _make_wrapper(fallback_type="previous") - call_count = [0] + w.safety_checker = lambda t, pin: None # verdict is pre-seeded below - def checker(t, pin): - call_count[0] += 1 - # Call #1 (launched during frame1) writes a flagged prob; frame2's - # read consumes that async result (1-frame delay). - pin.fill_(1.0 if call_count[0] == 1 else 0.0) - - w.safety_checker = checker + cached_clean = torch.randn(1, 3, 64, 64) + w._prev_clean_tensor = cached_clean - clean = torch.randn(1, 3, 64, 64) - _ = w._apply_safety_checker(clean) # first frame: pass-through, primes _prev_clean_tensor + pending = torch.randn(1, 3, 64, 64) + w._pending_frame = pending + w._nsfw_prob_pin = torch.tensor([1.0]) # pending frame's own verdict: flagged - flagged_frame = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(flagged_frame) + current = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(current) - assert result is w._prev_clean_tensor, "previous-fallback should return the cached clean tensor" - assert torch.equal(w._prev_clean_tensor, clean) + assert result is cached_clean, "previous-fallback should return the cached clean tensor" # ── Case 4: Clean frame passthrough ────────────────────────────────────── def test_clean_frame_returned_unchanged(self): w = _make_wrapper() - w.safety_checker = lambda t, pin: pin.fill_(0.0) # never flag + w.safety_checker = lambda t, pin: None # verdict is pre-seeded below - frame1 = torch.randn(1, 3, 64, 64) - _ = w._apply_safety_checker(frame1) # first frame: pass-through, primes the pin + pending = torch.randn(1, 3, 64, 64) + w._pending_frame = pending + w._nsfw_prob_pin = torch.tensor([0.0]) # pending frame's own verdict: clean - frame2 = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(frame2) + current = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(current) - assert torch.equal(result, frame2), "clean frame should be returned unchanged" + assert torch.equal(result, pending), "clean pending frame should be emitted unchanged" # ── Case 5: use_safety_checker=False bypasses entirely ─────────────────── def test_disabled_bypasses_checker(self): @@ -159,14 +167,14 @@ def should_not_be_called(t, pin): # ── Case 6: flagged verdict with no cached frame → black ───────────────── def test_nsfw_previous_fallback_no_cache_falls_back_to_black(self): """ - If a frame is flagged (per the previous frame's pinned verdict) and + If the pending frame is flagged (per its own landed verdict) and _prev_clean_tensor is None, the fallback should still produce a black - frame rather than raise. Seeds _nsfw_prob_pin directly to bypass the - first-frame pass-through and exercise this state combination. + frame rather than raise. """ w = _make_wrapper(fallback_type="previous") w.safety_checker = lambda t, pin: pin.fill_(0.0) - w._nsfw_prob_pin = torch.tensor([1.0]) # a prior flagged verdict already landed + w._pending_frame = torch.randn(1, 3, 64, 64) + w._nsfw_prob_pin = torch.tensor([1.0]) # pending frame's own verdict: flagged assert w._prev_clean_tensor is None # no cache yet dummy = torch.randn(1, 3, 64, 64) @@ -179,49 +187,57 @@ def test_nsfw_previous_fallback_no_cache_falls_back_to_black(self): def test_clean_frame_cached_for_previous_strategy(self): w = _make_wrapper(fallback_type="previous") w.safety_checker = lambda t, pin: pin.fill_(0.0) - w._nsfw_prob_pin = torch.tensor([0.0]) # a prior clean verdict already landed + pending = torch.randn(1, 3, 64, 64) + w._pending_frame = pending + w._nsfw_prob_pin = torch.tensor([0.0]) # pending frame's own verdict: clean assert w._prev_clean_tensor is None dummy = torch.randn(1, 3, 64, 64) w._apply_safety_checker(dummy) assert w._prev_clean_tensor is not None, ( - "clean frame with previous strategy should populate _prev_clean_tensor" + "clean pending frame with previous strategy should populate _prev_clean_tensor" ) - assert torch.equal(w._prev_clean_tensor, dummy) + assert torch.equal(w._prev_clean_tensor, pending) # ── Case 8: clean frame NOT cached for blank strategy ──────────────────── def test_clean_frame_not_cached_for_blank_strategy(self): w = _make_wrapper(fallback_type="blank") w.safety_checker = lambda t, pin: pin.fill_(0.0) - w._nsfw_prob_pin = torch.tensor([0.0]) # a prior clean verdict already landed + w._pending_frame = torch.randn(1, 3, 64, 64) + w._nsfw_prob_pin = torch.tensor([0.0]) # pending frame's own verdict: clean dummy = torch.randn(1, 3, 64, 64) w._apply_safety_checker(dummy) assert w._prev_clean_tensor is None, "_prev_clean_tensor should stay None when fallback_type='blank'" - # ── Case 9: first frame always passes through ──────────────────────────── - def test_first_frame_always_passes_through_regardless_of_verdict(self): + # ── Case 9: first call emits black and buffers ──────────────────────────── + def test_first_frame_emits_black_and_buffers(self): """ - The very first frame has no prior async result to read, so it must - pass through unscreened even if the classifier would flag it — the - documented 1-frame pass-through edge case of the delayed-readback design. + The very first call has no buffered frame yet, so it cannot gate + anything on its own verdict. It must emit a black startup frame + (never ungated pixels) and buffer the raw frame for the next call. """ w = _make_wrapper(fallback_type="blank") - w.safety_checker = lambda t, pin: pin.fill_(1.0) # would flag if read this frame + w.safety_checker = lambda t, pin: pin.fill_(1.0) # would flag if read this call dummy = torch.randn(1, 3, 64, 64) result = w._apply_safety_checker(dummy) - assert torch.equal(result, dummy), "first frame must pass through (no prior verdict to act on)" + denorm = (result / 2 + 0.5).clamp(0, 1) + assert _black_denorm(denorm), "first call must emit a black startup frame" + assert w._pending_frame is not None and torch.equal(w._pending_frame, dummy), ( + "first call must buffer the raw frame for next call's emission" + ) # ── Case 10: _process_skip_diffusion wiring — checker is called, result is black ── def test_skip_diffusion_routes_through_safety_checker(self): """ Verify that _process_skip_diffusion actually feeds its tensor through - _apply_safety_checker before postprocess_image, and that a flagged frame - produces the black substitution rather than passing through unscreened. + _apply_safety_checker before postprocess_image, and that a flagged + pending frame produces the black substitution rather than passing + through unscreened. This is a wiring test — the behaviour of _apply_safety_checker itself is covered by the earlier cases. We stub all model-dependent calls with @@ -231,11 +247,12 @@ def test_skip_diffusion_routes_through_safety_checker(self): def capturing_checker(tensor, pin): received.append(tensor) - pin.fill_(0.0) # this frame's own async result — irrelevant to this frame's own read + pin.fill_(0.0) # this frame's own async result — irrelevant to this call's read w = _make_wrapper(fallback_type="blank") w.safety_checker = capturing_checker - w._nsfw_prob_pin = torch.tensor([1.0]) # a prior frame's flagged verdict already landed + w._pending_frame = torch.randn(1, 3, 64, 64) # a prior frame awaiting emission + w._nsfw_prob_pin = torch.tensor([1.0]) # that pending frame's own verdict: flagged w.mode = "img2img" w.device = torch.device("cpu") w.dtype = torch.float32 @@ -265,8 +282,61 @@ def _apply_image_postprocessing_hooks(self, t): assert len(received) == 1, "safety checker should be called exactly once" assert isinstance(received[0], torch.Tensor), "checker arg must be a Tensor, not None" - # Flagged frame (per the pre-seeded pinned verdict) must produce the black substitution + # Flagged pending frame (per the pre-seeded pinned verdict) must produce the black substitution denorm = (result / 2 + 0.5).clamp(0, 1) - assert torch.allclose(denorm, torch.zeros_like(denorm)), ( - "flagged skip-diffusion frame should produce a black output" - ) + assert torch.allclose(denorm, torch.zeros_like(denorm)), "flagged pending frame should produce a black output" + + # ── Case 11: pin_memory is guarded on CPU-only builds ──────────────────── + def test_pin_memory_guarded_on_cpu(self): + """ + Regression test: the first call used to unconditionally call + .pin_memory(), which raises on CPU-only / no-CUDA-driver PyTorch + builds. It must be guarded and still produce a usable CPU tensor. + """ + w = _make_wrapper(fallback_type="blank") + w.safety_checker = lambda t, pin: pin.fill_(0.0) + + dummy = torch.randn(1, 3, 64, 64) + w._apply_safety_checker(dummy) # must not raise + + assert isinstance(w._nsfw_prob_pin, torch.Tensor) + assert w._nsfw_prob_pin.device.type == "cpu" + + # ── Case 12: isolated NSFW frame is blanked, not leaked ────────────────── + def test_isolated_nsfw_frame_is_blanked_not_leaked(self): + """ + Drives [clean, NSFW, clean] through the checker and asserts the NSFW + frame is the one substituted (never emitted raw) and the surrounding + clean frames pass through unscreened — proving the old bug (verdict + misattributed to the wrong frame, causing an isolated NSFW frame to + leak and the following clean frame to be wrongly blanked) is fixed. + """ + verdicts = {"clean1": 0.0, "nsfw": 1.0, "clean2": 0.0} + order = ["clean1", "nsfw", "clean2"] + call_index = [0] + + def checker(t, pin): + key = order[call_index[0]] + call_index[0] += 1 + pin.fill_(verdicts[key]) + + w = _make_wrapper(fallback_type="blank", threshold=0.5) + w.safety_checker = checker + + clean1 = torch.randn(1, 3, 64, 64) + nsfw = torch.randn(1, 3, 64, 64) + clean2 = torch.randn(1, 3, 64, 64) + + r0 = w._apply_safety_checker(clean1) # priming call: black startup frame + r1 = w._apply_safety_checker(nsfw) # emits clean1 (clean) + r2 = w._apply_safety_checker(clean2) # emits nsfw (flagged -> black) + # A 4th call would be needed to emit clean2; verify what we can without it. + + denorm0 = (r0 / 2 + 0.5).clamp(0, 1) + assert _black_denorm(denorm0), "priming call must emit black, not raw pixels" + + assert torch.equal(r1, clean1), "clean1 must be emitted unchanged, gated on its own verdict" + + denorm2 = (r2 / 2 + 0.5).clamp(0, 1) + assert _black_denorm(denorm2), "the NSFW frame itself must be the one blanked, not a neighbor" + assert not torch.equal(r2, clean2), "clean2 must never be substituted for the NSFW frame's verdict"