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

@@ -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

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(

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

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 (

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]

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()

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
 

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"))