From 7f59cd6c1e8e6e9d1b7881f2a4ac703640c9c4f2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 02:06:08 +0000 Subject: [PATCH] Review pass over the merged diffusion phases: seven correctness fixes Re-reviewed each merged phase PR against this branch's tip and fixed what is still real: - A superseded background load no longer cancels the current model's in-flight generation: the load-token check now runs BEFORE the cancel signal, with a re-check under the generate lock (Phase 1 review). - enable_model_cpu_offload / enable_sequential_cpu_offload now forward the resolved target device; diffusers defaults to CUDA, which broke offloaded loads on non-CUDA accelerators such as Intel XPU (Phase 2 review). - build_sd_cpp_command rejects a None prompt (str(None) previously slipped into argv as the literal "None") and a mask without an init image, which is an invalid sd-cli inpaint invocation (Phase 4/6 review). - The dense-quant OOM fallback drops the caught exception before clear_gpu_cache(): the traceback pinned the partially built dense transformer, so the VRAM this cleanup exists to reclaim stayed allocated through the GGUF rebuild (Phase 8 review). - Pre-quantized transformers (built via from_config) are eval()'d to match the from_pretrained paths, so train-mode layers cannot make prequant inference nondeterministic (Phase 9 review). - FBCache state is reset before each generation when a step cache is engaged: diffusers never clears the stateful first-block residuals on the resident transformer, so a resolution or batch change on the next request hit a shape mismatch, and an unchanged request could reuse stale residuals (Phase 12 review). Each fix carries a regression test; the full diffusion battery passes. --- studio/backend/core/inference/diffusion.py | 43 +++++++++++- .../core/inference/diffusion_memory.py | 6 +- .../core/inference/diffusion_prequant.py | 9 +++ studio/backend/core/inference/sd_cpp_args.py | 9 ++- .../backend/tests/test_diffusion_backend.py | 67 ++++++++++++++++++- studio/backend/tests/test_diffusion_memory.py | 22 ++++-- .../backend/tests/test_diffusion_prequant.py | 13 ++++ studio/backend/tests/test_sd_cpp_args.py | 19 ++++++ 8 files changed, 176 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 894e3ce692..78c2e242b6 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -499,12 +499,18 @@ class DiffusionBackend: # The cancel makes that wait ~one step (or the rest of the denoise for a # pipeline that ignores the step callback). with self._lock: + # Bail BEFORE signalling any cancel if this load was already superseded (an + # unload/eviction or a newer load bumped the token while we were resolving / + # downloading). Otherwise a stale worker would abort an unrelated, still-live + # generation from the CURRENT model and only then discover it has nothing to do. + if _load_token is not None and _load_token != self._load_token: + raise RuntimeError("Diffusion load was cancelled.") if self._active_generate_cancel is not None: self._active_generate_cancel.set() with self._generate_lock: with self._lock: - # Bail before the (slow, VRAM-heavy) build if an unload/eviction or a - # newer load superseded this one while we were resolving/downloading. + # Re-check under the generate lock: a newer load/unload may have superseded + # this one while we waited for the in-flight denoise to exit. if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Diffusion load was cancelled.") @@ -554,6 +560,12 @@ class DiffusionBackend: ) pipe = None transformer_quant_engaged = None + # Drop the exception (and its traceback) BEFORE clearing the cache: + # exc.__traceback__ keeps _load_dense_quant_pipeline's frame -- and + # thus its partially-built dense bf16 transformer/pipe -- alive, so + # clear_gpu_cache() could not otherwise reclaim that VRAM before the + # GGUF build (the OOM-fallback path this cleanup exists for). + del exc clear_gpu_cache() if pipe is None: @@ -826,6 +838,26 @@ class DiffusionBackend: explicit_offload = cpu_offload, ) + @staticmethod + def _reset_step_cache(pipe: Any) -> None: + """Clear the transformer's stateful step cache (FBCache) before a generation. + + diffusers keys FBCache residuals by cache context ("cond"/"uncond") on the + long-lived transformer, and neither the pipeline nor the context exit resets + them (``StateManager`` only clears via ``reset_stateful_hooks``, which no + pipeline calls). This backend reuses one resident pipe across generations, so + without a reset the next generation's first step compares its first-block + residual against the PREVIOUS request's -- a tensor-shape mismatch when the + resolution/batch changed, or a stale-cache reuse otherwise. Best-effort: a + transformer without the hook (uncached load) is a silent no-op.""" + transformer = getattr(pipe, "transformer", None) + reset = getattr(transformer, "reset_stateful_hooks", None) + if callable(reset): + try: + reset() + except Exception: # noqa: BLE001 — reset is best-effort, never fail a generation + pass + def generate( self, *, @@ -910,6 +942,13 @@ class DiffusionBackend: if "callback_on_step_end" in call_params: kwargs["callback_on_step_end"] = _on_step + # Start each generation from a clean step cache: FBCache residuals from + # a prior request on this resident pipe would otherwise be compared + # against this generation's first step (shape mismatch on a resolution/ + # batch change, or stale reuse). No-op when no cache is engaged. + if state.transformer_cache: + self._reset_step_cache(state.pipe) + self._gen = gen try: # inference_mode is strictly faster than the no_grad diffusers diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index b2c1e25f11..d3bfb4c175 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -423,20 +423,20 @@ def apply_memory_plan( # is in the low-VRAM situation where the decode-time spike can OOM, so turn VAE # tiling on now (if not already engaged) to cap it. nonlocal tiling_engaged - pipe.enable_model_cpu_offload() + pipe.enable_model_cpu_offload(device = device) if not tiling_engaged: tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger) policy = plan.offload_policy if policy == OFFLOAD_MODEL: - pipe.enable_model_cpu_offload() + pipe.enable_model_cpu_offload(device = device) elif policy == OFFLOAD_GROUP: if not _apply_group_offload(pipe, device, logger): _fallback_to_model_offload() policy = OFFLOAD_MODEL elif policy == OFFLOAD_SEQUENTIAL: try: - pipe.enable_sequential_cpu_offload() + pipe.enable_sequential_cpu_offload(device = device) except Exception as exc: # noqa: BLE001 — keep the model loadable if logger is not None: logger.warning( diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 207f305220..2232e726c8 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -193,6 +193,15 @@ def load_prequantized_transformer( transformer.load_state_dict(state_dict, strict = True, assign = True) transformer = transformer.to(device) + # Built via from_config (not from_pretrained), so it starts in TRAIN mode; the + # dense and GGUF paths load through from_pretrained, which diffusers documents as + # returning an eval()'d module. Match that here so any train/eval-sensitive layer + # (e.g. dropout) can't make prequant inference nondeterministic or diverge from + # the other load paths. + try: + transformer.eval() + except Exception: # noqa: BLE001 — eval() is best-effort + pass try: # diagnostic marker, mirrors the runtime-quant path transformer._unsloth_runtime_quant = scheme except Exception: # noqa: BLE001 — marker is best-effort diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 8069e8f610..5e0db4f0ce 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -206,8 +206,15 @@ def build_sd_cpp_command( """ if not files.diffusion_model: raise ValueError("diffusion_model path is required") - if not str(params.prompt).strip(): + # ``(prompt or "")`` so a None prompt is rejected here rather than slipping past + # ``str(None)`` == "None" (truthy) and landing in argv as a literal "None". + if not (params.prompt or "").strip(): raise ValueError("prompt is required") + # sd-cli inpaint needs the source image too: a --mask with no --init-img is an + # invalid invocation (sd-cli has nothing to inpaint into), so reject it here with a + # clear error instead of emitting a command that fails deep in sd-cli. + if params.mask and not params.init_img: + raise ValueError("init_img is required when mask is set (inpaint needs a source image)") cmd: list[str] = [binary, "--mode", DEFAULT_MODE, "--diffusion-model", files.diffusion_model] for flag, value in ( diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 5ec8bb3954..bc9b3d8e69 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -139,11 +139,13 @@ class _FakePipe: self.moved_to = device return self - def enable_model_cpu_offload(self) -> None: + def enable_model_cpu_offload(self, device = None) -> None: self.offloaded = True + self.offload_device = device - def enable_sequential_cpu_offload(self) -> None: + def enable_sequential_cpu_offload(self, device = None) -> None: self.sequential_offloaded = True + self.offload_device = device def enable_vae_tiling(self) -> None: self.vae_tiled = True @@ -499,6 +501,29 @@ def test_unload_cancels_in_flight_load(fake_runtime): ) +def test_superseded_load_does_not_cancel_live_generation(fake_runtime): + # A superseded background load (its token was bumped by a newer load/unload) that + # finally reaches load_pipeline must bail WITHOUT signalling the current model's + # in-flight generation: the token check has to run before the cancel is set, or a + # stale worker aborts an unrelated, still-live denoise. + import threading as _threading + + backend = DiffusionBackend() + fam = detect_family("unsloth/Z-Image-Turbo-GGUF") + live_cancel = _threading.Event() + backend._active_generate_cancel = live_cancel # a generation from the CURRENT model + token = 11 + backend._load_token = token + 1 # this load has already been superseded + with pytest.raises(RuntimeError, match = "cancelled"): + backend.load_pipeline( + "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z-image-turbo-Q4_K_S.gguf", + base_repo = fam.base_repo, + _load_token = token, + ) + assert not live_cancel.is_set() # the live generation was left untouched + + def test_pick_dtype_bf16_only_on_ampere(fake_runtime, monkeypatch): # BF16 only on Ampere+ (cc >= 8); pre-Ampere cards must fall back to FP16. torch = sys.modules["torch"] @@ -1134,3 +1159,41 @@ def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, mo assert status["transformer_quant"] is None assert status["offload_policy"] == "model" assert _FakeTransformer.last["path"] # GGUF path used + + +def test_reset_step_cache_helper_is_best_effort(): + # Calls the transformer's reset hook when present. + calls = [] + pipe = types.SimpleNamespace( + transformer = types.SimpleNamespace(reset_stateful_hooks = lambda: calls.append(True)) + ) + DiffusionBackend._reset_step_cache(pipe) + assert calls == [True] + # No transformer, or a transformer without the hook -> silent no-op (never raises). + DiffusionBackend._reset_step_cache(types.SimpleNamespace()) + DiffusionBackend._reset_step_cache(types.SimpleNamespace(transformer = object())) + + +def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path): + # FBCache residuals live on the resident transformer across generations, so each + # generate() must reset the stateful cache first -- but only when a cache is engaged. + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + ) + resets = [] + backend._state.pipe.transformer = types.SimpleNamespace( + reset_stateful_hooks = lambda: resets.append(True) + ) + # No cache engaged (transformer_cache is None) -> reset must NOT run. + backend.generate(prompt = "a sloth") + assert resets == [] + # Engage a cache; every subsequent generation resets the stateful cache first. + object.__setattr__(backend._state, "transformer_cache", "fbcache") + backend.generate(prompt = "a sloth") + backend.generate(prompt = "another sloth") + assert resets == [True, True] diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index cfdddb051e..bc9ba592ce 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -327,16 +327,19 @@ def test_snapshot_never_raises_on_probe_failure(monkeypatch): class _RecordingPipe: def __init__(self) -> None: self.calls: list[str] = [] + self.offload_device = None def to(self, device): self.calls.append(f"to:{device}") return self - def enable_model_cpu_offload(self): + def enable_model_cpu_offload(self, device = None): self.calls.append("model_offload") + self.offload_device = device - def enable_sequential_cpu_offload(self): + def enable_sequential_cpu_offload(self, device = None): self.calls.append("sequential_offload") + self.offload_device = device def enable_vae_tiling(self): self.calls.append("vae_tiling") @@ -385,6 +388,16 @@ def test_apply_model_offload_engages_offload_and_tiling(): assert "to:cuda" not in pipe.calls # offload owns placement; never both assert "vae_tiling" in pipe.calls and "vae_slicing" in pipe.calls assert effective == OFFLOAD_MODEL and tiled is True + assert pipe.offload_device == "cuda" # device threaded to enable_model_cpu_offload + + +def test_apply_model_offload_passes_target_device(): + # enable_model_cpu_offload defaults to CUDA in diffusers; on a non-CUDA accelerator + # (e.g. Intel XPU, which this backend supports) the target device must be forwarded + # or diffusers offloads to the wrong backend and the load fails. + pipe = _RecordingPipe() + apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = False), device = "xpu") + assert pipe.offload_device == "xpu" def test_apply_vae_tiling_falls_back_to_vae_submodule(): @@ -404,7 +417,7 @@ def test_apply_vae_tiling_falls_back_to_vae_submodule(): def _slice(self): self.vae.sliced = True - def enable_model_cpu_offload(self): + def enable_model_cpu_offload(self, device = None): self.offloaded = True pipe = _VaeOnly() @@ -439,13 +452,14 @@ def test_apply_sequential_offload(): ) assert "sequential_offload" in pipe.calls and "to:cuda" not in pipe.calls assert effective == OFFLOAD_SEQUENTIAL + assert pipe.offload_device == "cuda" # device threaded to sequential offload too def test_apply_sequential_falls_back_to_model_offload_when_unsupported(): # Sequential offload is unreliable for GGUF on some diffusers versions; the # applier must fall back to whole-module offload and report what actually ran. class _NoSeqPipe(_RecordingPipe): - def enable_sequential_cpu_offload(self): + def enable_sequential_cpu_offload(self, device = None): raise RuntimeError("sequential offload not supported for this transformer") pipe = _NoSeqPipe() diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index a233f6225f..2f1a37f8d6 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -66,6 +66,7 @@ class _FakeTransformer: def __init__(self): self.assigned = None self.moved = None + self.eval_called = False @classmethod def load_config(cls, base, **kw): @@ -101,6 +102,10 @@ class _FakeTransformer: self.moved = device return self + def eval(self): + self.eval_called = True + return self + def _stub_torch_accelerate( monkeypatch, @@ -182,6 +187,14 @@ def test_load_meta_init_and_assign(monkeypatch, tmp_path): assert t._unsloth_runtime_quant == "fp8" +def test_load_puts_transformer_in_eval_mode(monkeypatch, tmp_path): + # Built via from_config (not from_pretrained), so the loader must eval() it to match + # the dense/GGUF paths; otherwise train-mode dropout makes inference nondeterministic. + t = _load(monkeypatch, tmp_path, _good_ckpt()) + assert t is not None + assert t.eval_called is True + + def test_load_missing_file_is_none(monkeypatch, tmp_path): assert _load(monkeypatch, tmp_path, _good_ckpt(), exists = False) is None diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index 200672aa26..d3133b08d2 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -222,6 +222,25 @@ def test_build_inpaint_adds_mask(): assert _pair(cmd, "--init-img") == "/in/src.png" +def test_build_inpaint_mask_without_init_img_rejected(): + # sd-cli inpaint needs a source image; a --mask with no --init-img is invalid argv, + # so the builder must reject it up front instead of emitting a doomed command. + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + params = SdCppGenParams(prompt = "x", mask = "/in/mask.png") + with pytest.raises(ValueError, match = "init_img is required"): + build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + + +def test_build_rejects_none_prompt(): + # A None prompt must be rejected, not coerced to the literal string "None" and + # forwarded into argv. + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + with pytest.raises(ValueError, match = "prompt is required"): + build_sd_cpp_command( + "/bin/sd-cli", files, SdCppGenParams(prompt = None), output_path = "/o.png" + ) + + def test_build_edit_repeats_ref_image(): files = SdCppModelFiles(diffusion_model = "/m/flux.gguf") params = SdCppGenParams(prompt = "add a hat", ref_images = ("/r/a.png", "/r/b.png"))