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 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/acceleration/tensorrt/runtime_engines/unet_engine.py b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py index da7da67ff..de317065d 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}") @@ -496,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={ @@ -512,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/acceleration/tensorrt/utilities.py b/src/streamdiffusion/acceleration/tensorrt/utilities.py index 218101e3f..7cea0b456 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. @@ -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' @@ -1089,7 +1123,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. @@ -1136,14 +1170,49 @@ 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) - - 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}'") + 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 + # _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(): + 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: @@ -1366,7 +1435,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/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: 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/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: diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py index 9f7276237..c75149704 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -350,6 +350,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 @@ -374,6 +377,13 @@ 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 + # 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 @@ -1075,15 +1085,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).""" @@ -1093,10 +1123,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) @@ -1119,6 +1150,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 @@ -1126,8 +1161,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).""" @@ -1138,7 +1184,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: @@ -1299,20 +1345,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 = [] @@ -1335,6 +1390,25 @@ 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 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 ---------- image_tensor : torch.Tensor @@ -1343,8 +1417,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 @@ -1352,18 +1427,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.safety_checker(denormalized, self.safety_checker_threshold): + 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) + self._pending_frame = image_tensor.clone() + return torch.full_like(image_tensor, -1.0) + + # 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: 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, @@ -3010,6 +3107,10 @@ 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 + 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_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 diff --git a/tests/unit/test_safety_checker.py b/tests/unit/test_safety_checker.py index ef5a3d133..b4c1d91a7 100644 --- a/tests/unit/test_safety_checker.py +++ b/tests/unit/test_safety_checker.py @@ -15,6 +15,24 @@ 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 (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 @@ -39,6 +57,8 @@ 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 + w._pending_frame = None # safety_checker is set per-test w.safety_checker = None return w @@ -65,9 +85,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 +103,16 @@ 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: None # verdict is pre-seeded below - dummy = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(dummy) + 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 - assert result is not dummy, "flagged frame should be replaced" + current = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(current) + + 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" @@ -96,39 +120,40 @@ 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] - - def checker(t, thr): - call_count[0] += 1 - return call_count[0] > 1 # first call: clean; second call: flagged + w.safety_checker = lambda t, pin: None # verdict is pre-seeded below - 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) # 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 = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(flagged) + 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 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, thr: False # never flag + w.safety_checker = lambda t, pin: None # verdict is pre-seeded below - dummy = torch.randn(1, 3, 64, 64) - result = w._apply_safety_checker(dummy) + 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 torch.equal(result, dummy), "clean frame should be returned unchanged" + current = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(current) + + 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): 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 +164,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 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. """ 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._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) @@ -158,32 +186,58 @@ 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) + 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, 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, thr: False + w.safety_checker = lambda t, pin: pin.fill_(0.0) + 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: _process_skip_diffusion wiring — checker is called, result is black ── + # ── Case 9: first call emits black and buffers ──────────────────────────── + def test_first_frame_emits_black_and_buffers(self): + """ + 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 call + + dummy = torch.randn(1, 3, 64, 64) + result = w._apply_safety_checker(dummy) + + 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 @@ -191,12 +245,14 @@ 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 call's read w = _make_wrapper(fallback_type="blank") w.safety_checker = capturing_checker + 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 @@ -226,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 must produce the black substitution (all -1.0 → denorm 0.0) + # 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" 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)) 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"