From 33c0e031ce7ef374d61475c02c3f9b7fe079dd67 Mon Sep 17 00:00:00 2001 From: Alex Korneev <49823108+forkni@users.noreply.github.com> Date: Sat, 18 Jul 2026 20:19:13 -0400 Subject: [PATCH] perf: UNet step sub-phase optimizations (5.0-5.6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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. * 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. * 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. * 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. * 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). * 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. * 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. * 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). * 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). * 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. * perf: elide trt.input_staging DtoD copy for kvo/fio UNet cache inputs (Sub-phase 5.6) 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. * 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 * 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 * 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 * chore: add Claude review workflow to PR2 branch for head-ref validation * chore: widen Claude review allowedTools (Skill, python/pytest, /tmp Write) to unblock permission-denied max-turns failures * chore: raise Claude review max-turns to 60, discourage filesystem-wide find and output redirection * 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. --------- Co-authored-by: Claude Sonnet 5 --- 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 8910d888..4e6c7fef 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -389,6 +389,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 @@ -1399,18 +1402,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 ---------- @@ -1420,8 +1429,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 @@ -1429,35 +1439,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, @@ -3143,6 +3158,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 bcb657f9..e4e288cd 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 @@ -47,6 +57,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 @@ -91,15 +102,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" @@ -107,37 +119,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): @@ -158,14 +166,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) @@ -178,49 +186,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 @@ -230,11 +246,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 @@ -264,8 +281,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"