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"