From 6a8b0b47e7a1e32413d2af9dd51e3aef7da18802 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 01:00:47 +0000
Subject: [PATCH 1/8] Fix review findings on image generation: failed-load
 VRAM, API defaults, preflights

- Free reserved VRAM in the diffusion load worker's failure path: a load-time OOM
  never commits _state and the next load's _unload_locked early-returns, so nothing
  else reclaimed the half-built pipeline's memory
- Use a monotonic clock for the denoise ETA rate
- Sync _GENERATION_DEFAULTS with the UI table: kontext, flux.2-dev, sdxl-turbo and
  SDXL base rows so /v1/images/generations stops falling back to 9 steps / CFG 0
- 400 (not sanitized 500) when /v1/images/generations hits an edit-only model
- Fail fast on pre-Ampere CUDA in the DiT trainer instead of dying in model load
- Run the trainer trust gate in the diffusion training route before freeing GPU
  residents so an untrusted base cannot tear down loaded chat/Images models
- Protect native sd.cpp companion VAE/text-encoder repos from cache deletion while
  a load is downloading them
- Exempt the task-scoped Images picker from the chat-only GGUF/MLX format gate so
  local diffusers pipelines stay selectable on no-GPU hosts
---
 studio/backend/core/inference/diffusion.py    |  7 +++++-
 .../core/inference/diffusion_families.py      |  9 ++++++++
 .../backend/core/inference/sd_cpp_backend.py  | 23 ++++++++++++++++---
 .../core/training/diffusion_dit_trainer.py    |  7 ++++++
 studio/backend/routes/inference.py            | 15 ++++++++++++
 studio/backend/routes/training.py             | 10 ++++++++
 .../assistant-ui/model-selector/pickers.tsx   |  5 +++-
 7 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index fda2a27a4b..0531b39aef 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -691,6 +691,10 @@ class DiffusionBackend:
             if self._load_token != token:
                 return
             logger.error("diffusion.load_failed: %s", exc)
+            # Free the debris of a failed construction (e.g. a load-time OOM): _state was
+            # never committed, and the next load's _unload_locked early-returns on a None
+            # state, so nothing else releases the reserved VRAM.
+            clear_gpu_cache()
             # Redact native paths: this error is surfaced verbatim via the
             # load-progress poll, and Studio can run as a shared server.
             from utils.native_path_leases import redact_native_paths
@@ -1971,7 +1975,8 @@ class DiffusionBackend:
                 gen = _GenState(total_steps = steps)
 
                 def _on_step(pipe, step_index, timestep, callback_kwargs):
-                    now = time.time()
+                    # Monotonic: a wall-clock adjustment (NTP) mid-denoise would skew the ETA.
+                    now = time.monotonic()
                     gen.step = step_index + 1
                     if gen.first_step_at == 0.0:
                         gen.first_step_at = now
diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index f49843ed3d..88ce4007dc 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -430,10 +430,19 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
 _GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
     ("z-image-turbo", 9, 0.0),
     ("flux.1-schnell", 4, 0.0),
+    # Kontext (editing) before the generic flux.1: ~28 steps, lower guidance (~2.5).
+    ("kontext", 28, 2.5),
     ("flux.1", 28, 3.5),
     ("flux.2-klein", 4, 0.0),
+    # FLUX.2-dev is the full (non-distilled) model: more steps + real guidance.
+    ("flux.2-dev", 28, 4.0),
     ("qwen-image", 20, 4.0),
     ("z-image", 20, 4.0),
+    # SDXL: Turbo is distilled (few steps, no CFG); base/full SDXL wants ~30 steps and
+    # real CFG (~7). "sdxl-turbo" must precede the generic "sdxl" substring match.
+    ("sdxl-turbo", 3, 0.0),
+    ("stable-diffusion-xl", 30, 7.0),
+    ("sdxl", 30, 7.0),
 )
 # Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback.
 _GENERATION_DEFAULT_FALLBACK = (9, 0.0)
diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index d8cb21d409..78425833d6 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -241,6 +241,9 @@ class _SdLoading:
 
     repo_id: str
     base_repo: str
+    # Companion asset repos (VAE / text encoders) this load fetches, so the
+    # delete-cached guard protects them for the whole download/finalize window.
+    asset_repos: tuple[str, ...] = ()
     expected_bytes: int = 0
     downloaded_bytes: int = 0
     error: Optional[str] = None
@@ -407,7 +410,17 @@ class SdCppDiffusionBackend:
             self._load_token += 1
             token = self._load_token
             self._cancel_event.clear()
-            self._loading = _SdLoading(repo_id = repo_id, base_repo = base)
+            self._loading = _SdLoading(
+                repo_id = repo_id,
+                base_repo = base,
+                asset_repos = tuple(
+                    dict.fromkeys(
+                        r
+                        for r, _f, kind in self._asset_specs(repo_id, gguf_filename, fam)
+                        if kind != "diffusion_model"
+                    )
+                ),
+            )
 
         threading.Thread(
             target = self._run_load,
@@ -678,12 +691,16 @@ class SdCppDiffusionBackend:
     def loading_repo_ids(self) -> tuple[str, ...]:
         """Repo ids an in-flight background load is downloading (empty when idle).
         Mirrors the diffusers backend so the delete-cached guard can query whichever
-        engine is active without caring which one it got."""
+        engine is active without caring which one it got. Includes the companion
+        VAE / text-encoder repos: deleting one of those mid-load would remove files
+        the committed SdCppModelFiles paths need."""
         with self._lock:
             loading = self._loading
             if loading is None or loading.error is not None:
                 return ()
-            return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
+            return tuple(
+                r for r in (loading.repo_id, loading.base_repo, *loading.asset_repos) if r
+            )
 
     # ── Generate ───────────────────────────────────────────────────────────
 
diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index 50bb0f5d3b..d29a85a140 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -528,6 +528,13 @@ def run_dit_lora_training(
     device = "cuda" if torch.cuda.is_available() else "cpu"
     # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is
     # unsupported for real runs but keeps import/unit tests architecture-agnostic).
+    # Fail fast on pre-Ampere CUDA (T4/V100/RTX 20xx): bf16 compute is required and the
+    # run would otherwise die deep in model load with an opaque dtype error.
+    if device == "cuda" and not torch.cuda.is_bf16_supported():
+        raise ValueError(
+            "This trainer requires a bfloat16-capable GPU (Ampere or newer); "
+            "this CUDA device does not support bf16."
+        )
     weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32
     use_lora_targets = _select_lora_targets(cfg.lora_target_modules, spec.lora_targets)
 
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 1cc1c83ad4..32f7b48f58 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -12076,6 +12076,21 @@ async def openai_image_generations(
         # isn't loaded; the global handler turns this into the OpenAI envelope.
         raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
 
+    # An edit-only model (Qwen-Image-Edit, FLUX Kontext) needs an input image this API
+    # cannot supply; refuse up front with a 400 instead of letting the backend's
+    # ValueError surface as a sanitized 500.
+    workflows = status.get("workflows") or []
+    if workflows and "txt2img" not in workflows:
+        raise HTTPException(
+            status_code = 400,
+            detail = openai_error_body(
+                "The loaded image model is edit-only (it requires an input image); "
+                "load a text-to-image model to use this endpoint.",
+                status = 400,
+                param = "model",
+            ),
+        )
+
     # Fall back to the resolved base repo so a local-path load (whose repo_id is a
     # filesystem path) still gets the right per-model steps/guidance.
     steps, guidance = default_generation_params(status.get("repo_id"), status.get("base_repo"))
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 54324dd544..a89a8d01c2 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1219,6 +1219,16 @@ async def start_diffusion_training(
     except ValueError as e:
         raise HTTPException(status_code = 400, detail = str(e))
 
+    # Run the trainers' trust gate here too (both assert the same predicate before
+    # from_pretrained), so an untrusted/typoed base 400s BEFORE freeing GPU residents
+    # instead of tearing down the user's chat/Images model and failing in the child.
+    from core.training.diffusion_train_common import _assert_trusted_base_model
+
+    try:
+        _assert_trusted_base_model(config.get("base_model", ""))
+    except ValueError as e:
+        raise HTTPException(status_code = 400, detail = str(e))
+
     # Preflight access to a gated base repo with the user's token BEFORE freeing GPU
     # residents, so a missing/insufficient token fails fast (400) without tearing down the
     # user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index 336dfece69..c380e04244 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -1982,7 +1982,9 @@ export function HubModelPicker({
   );
   // Local ./models entries. Chat-only Studio runs GGUF (any host) and MLX (Mac
   // only), so raw checkpoints there are hidden (mirrors the cached non-GGUF
-  // rule). An MLX build a Mac user dropped in ./models stays selectable.
+  // rule). An MLX build a Mac user dropped in ./models stays selectable. A
+  // task-scoped picker (Images) is exempt: the image backend loads local
+  // diffusers/safetensors pipelines even on chat-only (no-GPU, native) hosts.
   const sortedLocalDir = useMemo(
     () =>
       sortLocalModels(
@@ -1990,6 +1992,7 @@ export function HubModelPicker({
           (m) =>
             passesTaskGate(m.task, m.model_id ?? m.id, task) &&
             (!chatOnly ||
+              task != null ||
               localModelIsGguf(m) ||
               (isMac && localModelIsMlx(m))) &&
             localModelMatchesFormat(m, formatFilter) &&

From f1d9c88606046669b82518e60cf8b4676af54491 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 01:47:02 +0000
Subject: [PATCH 2/8] Validate load modes before eviction and wait out a
 cancelled denoise on unload

---
 studio/backend/core/inference/diffusion.py    | 47 ++++++++++++++----
 .../backend/tests/test_diffusion_backend.py   | 48 ++++++++++++++++---
 2 files changed, 79 insertions(+), 16 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 0531b39aef..09b2fa1b42 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -58,18 +58,20 @@ from .diffusion_speed import (
     SPEED_OFF,
     apply_speed_optims,
     compile_eligible,
+    normalize_speed_mode,
     resolve_speed_mode,
     restore_backend_flags,
     snapshot_backend_flags,
 )
 from .diffusion_attention import (
     apply_attention_backend,
+    normalize_attention_backend,
     select_attention_backend,
 )
 from . import diffusion_compile_cache as compile_cache
 from . import diffusion_gguf_compile as gguf_compile
-from .diffusion_cache import apply_step_cache
-from .diffusion_precision import quantize_text_encoders
+from .diffusion_cache import apply_step_cache, normalize_transformer_cache
+from .diffusion_precision import normalize_te_quant, quantize_text_encoders
 from .diffusion_prequant import (
     load_prequantized_transformer,
     resolve_prequant_source,
@@ -693,8 +695,13 @@ class DiffusionBackend:
             logger.error("diffusion.load_failed: %s", exc)
             # Free the debris of a failed construction (e.g. a load-time OOM): _state was
             # never committed, and the next load's _unload_locked early-returns on a None
-            # state, so nothing else releases the reserved VRAM.
-            clear_gpu_cache()
+            # state, so nothing else releases the reserved VRAM. Guarded: a sticky CUDA
+            # error makes synchronize() raise, which would skip stamping the REAL error
+            # below and leave the client polling forever.
+            try:
+                clear_gpu_cache()
+            except Exception:  # noqa: BLE001
+                pass
             # Redact native paths: this error is surfaced verbatim via the
             # load-progress poll, and Studio can run as a shared server.
             from utils.native_path_leases import redact_native_paths
@@ -900,6 +907,14 @@ class DiffusionBackend:
             model_kind = model_kind,
         )
         kind = resolve_model_kind(gguf_filename, model_kind)
+        # Validate every mode string that can raise NOW, before this load evicts the
+        # previous pipeline below: their first in-line uses all sit past _unload_locked,
+        # where a bad request would cost the user their working model.
+        transformer_quant = normalize_transformer_quant(transformer_quant)
+        normalize_speed_mode(speed_mode)
+        normalize_attention_backend(attention_backend)
+        normalize_transformer_cache(transformer_cache)
+        normalize_te_quant(text_encoder_quant)
         # For a full pipeline the repo itself supplies every component, so it is its
         # own base; the single-file kinds resolve the companion base diffusers repo.
         base = (
@@ -972,7 +987,7 @@ class DiffusionBackend:
                 transformer_quant_engaged = None
                 if (
                     kind == "gguf"
-                    and normalize_transformer_quant(transformer_quant) is not None
+                    and transformer_quant is not None  # normalized above, pre-eviction
                     and dense_transformer_supported(target)
                     and plan.offload_policy == OFFLOAD_NONE
                 ):
@@ -1002,7 +1017,12 @@ class DiffusionBackend:
                         # 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()
+                        # Guarded: after an OOM/sticky CUDA error synchronize() can
+                        # raise, and this fallback path must still reach the GGUF build.
+                        try:
+                            clear_gpu_cache()
+                        except Exception:  # noqa: BLE001
+                            pass
 
                 if pipe is None:
                     if kind == "pipeline":
@@ -2053,9 +2073,8 @@ class DiffusionBackend:
         self._cancel_event.set()
         with self._lock:
             # Abort an in-flight denoise too by setting ITS cancel event, so the step
-            # callback stops it. unload does NOT take _generate_lock — it must return
-            # promptly; the running generate keeps its own pipe reference, so freeing
-            # _state here can't crash it, and its VRAM is reclaimed when it returns
+            # callback stops it. The running generate keeps its own pipe reference, so
+            # freeing _state here can't crash it; its VRAM is reclaimed when it exits
             # (within ~one step thanks to the cancel).
             if self._active_generate_cancel is not None:
                 self._active_generate_cancel.set()
@@ -2064,6 +2083,14 @@ class DiffusionBackend:
             # committing) and drop the marker so the next load starts clean.
             self._load_token += 1
             self._loading = None
+        # Wait for the signalled denoise to actually exit before reporting unloaded:
+        # callers treat this return as "VRAM is free" (the GPU arbiter hands the GPU
+        # to chat next; the training routes size their run against it), and the
+        # denoise holds its pipe until the next step callback. generate() holds
+        # _generate_lock for its full body, so a bare acquire is the exit barrier
+        # (never while holding _lock -- generate takes _lock inside _generate_lock).
+        with self._generate_lock:
+            pass
         return self.status()
 
     def _unload_locked(self) -> None:
@@ -2088,7 +2115,7 @@ class DiffusionBackend:
             uninstall_patches()
             uninstall_arch_patches()
         # NOTE: we deliberately do NOT call state.pipe.unload_lora_weights() here. unload()
-        # sets the cancel event but does not take _generate_lock, so a LoRA-backed denoise
+        # only acquires _generate_lock AFTER this teardown, so a LoRA-backed denoise
         # can still be running on this same pipe for up to one more callback; mutating its
         # adapter layers now would race that in-flight generation. The whole pipe is dropped
         # just below (self._state = None; del state; clear_gpu_cache()), so the adapter
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index cd25103a24..102fa83fa3 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -1384,6 +1384,34 @@ def test_load_promotes_fp16_to_fp32_for_zimage_only(fake_runtime, monkeypatch, t
     assert q["dtype"] == "float16"  # fp16-compatible family keeps fp16 on pre-Ampere
 
 
+def test_bad_mode_strings_fail_before_eviction(fake_runtime):
+    # Every mode normalizer that can raise runs BEFORE the load evicts the previous
+    # pipeline, so a bad request never costs the user their working model.
+    backend = DiffusionBackend()
+    fam = detect_family("unsloth/Z-Image-GGUF")
+    backend._state = _LoadState(
+        pipe = object(),
+        family = fam,
+        repo_id = "r",
+        base_repo = "b",
+        device = "cpu",
+        dtype = "float32",
+        cpu_offload = False,
+    )
+    for kwargs in (
+        {"transformer_quant": "int7"},
+        {"speed_mode": "warp"},
+        {"attention_backend": "bogus"},
+        {"transformer_cache": "bogus"},
+        {"text_encoder_quant": "fp3"},
+    ):
+        with pytest.raises(ValueError):
+            backend.load_pipeline(
+                "unsloth/Z-Image-GGUF", gguf_filename = "m.gguf", **kwargs
+            )
+        assert backend._state is not None
+
+
 # Lock split + mid-denoise cancellation
 
 
@@ -1427,17 +1455,25 @@ def test_generate_lock_split_keeps_status_and_unload_responsive(fake_runtime):
     assert backend.status()["loaded"] is True
     assert backend.generate_progress()["active"] is True
 
-    # unload() must return promptly (it does not wait on _generate_lock) and signal
-    # THIS in-flight generation's cancel event.
+    cancel_ref = backend._active_generate_cancel
+    assert cancel_ref is not None
+
+    # unload() signals THIS generation's cancel event, then waits for the denoise to
+    # actually exit before returning: callers treat its return as "VRAM is free" (the
+    # GPU arbiter hands the GPU to chat on it). Release the pipe once the cancel
+    # lands, standing in for the step callback of a real pipeline.
+    releaser = threading.Thread(target = lambda: (cancel_ref.wait(5), release.set()))
+    releaser.start()
     backend.unload()
-    assert backend._active_generate_cancel is not None
-    assert backend._active_generate_cancel.is_set()
+    releaser.join(5)
+    assert cancel_ref.is_set()
     assert backend.status()["loaded"] is False
 
-    release.set()
     t.join(5)
-    # The cancelled generation raised rather than returning a now-evicted image.
+    # The cancelled generation raised rather than returning a now-evicted image, and
+    # it had already exited (deregistering its cancel) before unload() returned.
     assert "exc" in out and "cancelled" in str(out["exc"]).lower()
+    assert backend._active_generate_cancel is None
 
 
 def test_callback_cancellation_interrupts_denoise(fake_runtime):

From 76eee534eabd8df6d4b05e3dacc0c9083f1d09a5 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 01:48:12 +0000
Subject: [PATCH 3/8] Gate explicit attention kernels on NVIDIA CUDA and roll
 back partial FBCache hooks

---
 .../core/inference/diffusion_attention.py     |  5 ++++
 .../backend/core/inference/diffusion_cache.py | 11 ++++++++-
 .../backend/tests/test_diffusion_attention.py | 11 ++++++++-
 studio/backend/tests/test_diffusion_cache.py  | 23 +++++++++++++++++++
 4 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py
index e652b0068c..8b1ce29ec8 100644
--- a/studio/backend/core/inference/diffusion_attention.py
+++ b/studio/backend/core/inference/diffusion_attention.py
@@ -129,6 +129,11 @@ def select_attention_backend(
         backend = _ALIASES[alias]
         if backend == "native":
             return None
+        # Every explicit kernel here (cuDNN / flash* / sage) is CUDA+NVIDIA-only; on
+        # ROCm / MPS / CPU diffusers accepts the name at set time and the first
+        # generation crashes, so drop to the native default up front.
+        if not _is_cuda_nvidia(target):
+            return None
         # An arch-gated kernel (flash3/flash4) on a card that can't run it would set fine
         # then crash mid-generation, so drop it to the native default up front.
         if not _backend_arch_supported(backend):
diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py
index b7a2b7ac45..3cb74d7f80 100644
--- a/studio/backend/core/inference/diffusion_cache.py
+++ b/studio/backend/core/inference/diffusion_cache.py
@@ -89,7 +89,10 @@ def apply_step_cache(
         _warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)"))
         return None
     try:
-        from diffusers import FirstBlockCacheConfig
+        try:
+            from diffusers import FirstBlockCacheConfig
+        except ImportError:  # older diffusers exports it only from diffusers.hooks
+            from diffusers.hooks import FirstBlockCacheConfig
 
         config = FirstBlockCacheConfig(threshold = thr)
         enable_cache(config)
@@ -101,6 +104,12 @@ def apply_step_cache(
             logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr)
         return mode
     except Exception as exc:  # noqa: BLE001 — incompatible model -> run uncached
+        # enable_cache can fail after hooking some blocks; drop any partial hooks so
+        # the reported-uncached model doesn't actually run half-cached.
+        try:
+            transformer.disable_cache()
+        except Exception:  # noqa: BLE001
+            pass
         _warn(logger, mode, exc)
         return None
 
diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py
index 3e43aab8a9..99c5e966f2 100644
--- a/studio/backend/tests/test_diffusion_attention.py
+++ b/studio/backend/tests/test_diffusion_attention.py
@@ -75,7 +75,7 @@ def test_auto_stays_native_off_nvidia(monkeypatch):
 
 
 def test_explicit_backend_honored_regardless_of_speed(monkeypatch):
-    monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False)
+    monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True)
     # Pin a high capability so the arch-gated flash4 isn't dropped by the runtime check.
     monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0))
     assert select_attention_backend(_target(), "sage", speed_active = False) == "sage"
@@ -83,6 +83,15 @@ def test_explicit_backend_honored_regardless_of_speed(monkeypatch):
     assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn"
 
 
+def test_explicit_backend_dropped_off_nvidia_cuda(monkeypatch):
+    # Explicit cuDNN/flash/sage on ROCm / MPS / CPU passes diffusers' set-time check
+    # and crashes at the first generation, so selection drops to the native default.
+    monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False)
+    monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0))
+    for alias in ("sage", "flash", "flash4", "cudnn"):
+        assert select_attention_backend(_target(device = "mps"), alias, speed_active = True) is None
+
+
 def test_explicit_native_returns_none():
     # native is the default -> nothing to set.
     assert select_attention_backend(_target(), "native", speed_active = True) is None
diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py
index 62071d9aa6..fe2f331281 100644
--- a/studio/backend/tests/test_diffusion_cache.py
+++ b/studio/backend/tests/test_diffusion_cache.py
@@ -136,6 +136,29 @@ def test_incompatible_model_runs_uncached(monkeypatch):
     assert apply_step_cache(_pipe(t), mode = "fbcache") is None
 
 
+def test_enable_cache_failure_rolls_back_partial_hooks(monkeypatch):
+    # enable_cache can raise after hooking some blocks; the reported-uncached model
+    # must not actually run half-cached, so the failure path calls disable_cache.
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer(fail = True)
+    t.disabled = False
+    t.disable_cache = lambda: setattr(t, "disabled", True)
+    assert apply_step_cache(_pipe(t), mode = "fbcache") is None
+    assert t.disabled is True
+
+
+def test_config_import_falls_back_to_hooks_module(monkeypatch):
+    # Older diffusers exports FirstBlockCacheConfig only from diffusers.hooks.
+    diffusers = types.ModuleType("diffusers")  # no FirstBlockCacheConfig attribute
+    monkeypatch.setitem(sys.modules, "diffusers", diffusers)
+    hooks = types.ModuleType("diffusers.hooks")
+    hooks.FirstBlockCacheConfig = _Config
+    monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
+    t = _MixinTransformer()
+    assert apply_step_cache(_pipe(t), mode = "fbcache") == TC_FBCACHE
+    assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD
+
+
 def test_missing_transformer_is_none(monkeypatch):
     _stub_diffusers(monkeypatch)
     pipe = types.SimpleNamespace(transformer = None)

From 098809d2fe721dd16668a3b596d65d35a531951b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 01:49:19 +0000
Subject: [PATCH 4/8] Reject sd-cli batch runs and clear stale output targets
 before a run

---
 studio/backend/core/inference/sd_cpp_args.py   |  7 ++++++-
 studio/backend/core/inference/sd_cpp_engine.py |  3 +++
 studio/backend/tests/test_sd_cpp_args.py       | 12 ++++++++++--
 studio/backend/tests/test_sd_cpp_engine.py     | 16 ++++++++++++++++
 4 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py
index c370777bab..09e3faa778 100644
--- a/studio/backend/core/inference/sd_cpp_args.py
+++ b/studio/backend/core/inference/sd_cpp_args.py
@@ -269,7 +269,12 @@ def build_sd_cpp_command(
     if params.seed is not None:
         cmd += ["--seed", str(int(params.seed))]
     if params.batch_count and params.batch_count != 1:
-        cmd += ["--batch-count", str(int(params.batch_count))]
+        # sd-cli names the extra batch images itself (output_2.png, ...) and the runner
+        # collects only the literal --output path, so a CLI batch would silently drop
+        # every image after the first. Batches go through the sdcpp server API instead.
+        raise ValueError(
+            "sd-cli runs are single-image; use the sdcpp server API for batch generation."
+        )
 
     cmd += ["--output", output_path]
     if threads is not None:
diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py
index 52e4d6d693..c6582b0cb8 100644
--- a/studio/backend/core/inference/sd_cpp_engine.py
+++ b/studio/backend/core/inference/sd_cpp_engine.py
@@ -364,6 +364,9 @@ class SdCppEngine:
     def _prepare_out(output_path: str) -> Path:
         out = Path(output_path)
         out.parent.mkdir(parents = True, exist_ok = True)
+        # Drop a stale file at the target so the post-run is_file() check proves THIS
+        # run produced the image, not a leftover from an earlier run at the same path.
+        out.unlink(missing_ok = True)
         return out
 
     def _run(
diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py
index 4157d0e01f..fdc5aec863 100644
--- a/studio/backend/tests/test_sd_cpp_args.py
+++ b/studio/backend/tests/test_sd_cpp_args.py
@@ -166,10 +166,18 @@ def test_build_appends_offload_and_extra_args_last():
 
 def test_build_negative_prompt_and_batch():
     files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
-    params = SdCppGenParams(prompt = "x", negative_prompt = "blurry", batch_count = 3)
+    params = SdCppGenParams(prompt = "x", negative_prompt = "blurry")
     cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
     assert _pair(cmd, "--negative-prompt") == "blurry"
-    assert _pair(cmd, "--batch-count") == "3"
+    # A CLI batch would silently drop every image after the first (the runner only
+    # collects the literal --output path), so the builder rejects it outright.
+    with pytest.raises(ValueError, match = "single-image"):
+        build_sd_cpp_command(
+            "/bin/sd-cli",
+            files,
+            SdCppGenParams(prompt = "x", batch_count = 3),
+            output_path = "/o.png",
+        )
 
 
 def test_build_omits_unset_optional_params():
diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py
index daaf5223a8..8fee117b3a 100644
--- a/studio/backend/tests/test_sd_cpp_engine.py
+++ b/studio/backend/tests/test_sd_cpp_engine.py
@@ -323,6 +323,22 @@ def test_generate_raises_when_no_output_despite_success(tmp_path, monkeypatch):
         )
 
 
+def test_generate_does_not_return_stale_preexisting_output(tmp_path, monkeypatch):
+    # A leftover file at the target path must not satisfy the post-run output check
+    # when the run itself produced nothing: the target is cleared before the run.
+    e = _engine(tmp_path)
+    out = tmp_path / "img.png"
+    out.write_bytes(b"stale")
+    _patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out, write = False)
+    with pytest.raises(RuntimeError, match = "no image"):
+        e.generate(
+            SdCppModelFiles(diffusion_model = "/m/z.gguf"),
+            SdCppGenParams(prompt = "x"),
+            output_path = str(out),
+        )
+    assert not out.exists()
+
+
 def test_generate_raises_when_binary_missing():
     e = SdCppEngine(binary = None)
     with pytest.raises(RuntimeError, match = "not found"):

From e605075508a412dd512bf2845858491ddd12be95 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 01:51:34 +0000
Subject: [PATCH 5/8] Fix diffusion training validation and honor lr_scheduler
 and batch size in the DiT trainer

---
 .../core/training/diffusion_dit_trainer.py    | 24 +++++++++++++++----
 .../core/training/diffusion_train_common.py   |  4 ++++
 studio/backend/models/training.py             |  8 +++++--
 .../tests/test_diffusion_lora_trainer.py      | 10 ++++++++
 .../backend/tests/test_diffusion_training.py  | 14 +++++++++++
 5 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index d29a85a140..54ff06d2ce 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -592,6 +592,16 @@ def run_dit_lora_training(
     lora_params = [p for p in transformer.parameters() if p.requires_grad]
 
     optimizer = _make_optimizer(lora_params, cfg.learning_rate)
+    # One lr_sched.step() per optimizer update (cfg.train_steps total), matching the
+    # SDXL trainer: counting micro-steps instead would stretch warmup past the run.
+    from diffusers.optimization import get_scheduler
+
+    lr_sched = get_scheduler(
+        cfg.lr_scheduler,
+        optimizer = optimizer,
+        num_warmup_steps = cfg.lr_warmup_steps,
+        num_training_steps = cfg.train_steps,
+    )
     scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
         cfg.base_model, subfolder = "scheduler", token = cfg.hf_token
     )
@@ -604,10 +614,15 @@ def run_dit_lora_training(
     peak_gb = 0.0
     t_start = time.time()
     done = 0
+    # Honor train_batch_size by folding it into the micro-step count: averaging the
+    # gradient over batch * accum single-image passes is mathematically identical to
+    # true batching with a mean loss, and keeps the QLoRA memory profile flat (one
+    # image's activations at a time). Previously batch_size > 1 silently trained at 1.
+    micro_steps = cfg.gradient_accumulation_steps * cfg.train_batch_size
     for opt_step in range(cfg.train_steps):
         optimizer.zero_grad(set_to_none = True)
         step_loss = 0.0
-        for _ in range(cfg.gradient_accumulation_steps):
+        for _ in range(micro_steps):
             i = rng.randrange(len(image_paths))
             px = (
                 _load_pixel_tensor(
@@ -645,8 +660,8 @@ def run_dit_lora_training(
                 )
                 target = noise - latents
                 loss = F.mse_loss(model_pred.float(), target.float(), reduction = "mean")
-            (loss / cfg.gradient_accumulation_steps).backward()
-            step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
+            (loss / micro_steps).backward()
+            step_loss += float(loss.detach()) / micro_steps
 
         grad_norm: Optional[float] = None
         if cfg.max_grad_norm and cfg.max_grad_norm > 0:
@@ -654,6 +669,7 @@ def run_dit_lora_training(
             # chart wants (spikes stay visible even when clipping flattens the update).
             grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
         optimizer.step()
+        lr_sched.step()
 
         running_loss += step_loss
         done = opt_step + 1
@@ -672,7 +688,7 @@ def run_dit_lora_training(
                 total_steps = cfg.train_steps,
                 loss = round(step_loss, 5),
                 avg_loss = round(running_loss / done, 5),
-                learning_rate = cfg.learning_rate,
+                learning_rate = lr_sched.get_last_lr()[0],
                 grad_norm = round(grad_norm, 5) if grad_norm is not None else None,
                 samples_per_second = sps,
                 peak_memory_gb = peak_gb or None,
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 62a16eef8c..d70f9bce00 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -259,6 +259,10 @@ class DiffusionLoraConfig:
             raise ValueError("resolution must be a multiple of 8 and >= 64")
         if self.mixed_precision not in ("bf16", "fp16", "no"):
             raise ValueError("mixed_precision must be one of bf16 / fp16 / no")
+        # A zero/negative gamma would zero out (or invert) the min-SNR weight and
+        # silently train on a degenerate loss; None is the documented disable.
+        if self.snr_gamma is not None and float(self.snr_gamma) <= 0:
+            raise ValueError("snr_gamma must be > 0, or null to disable min-SNR weighting")
         # learning_rate can arrive as a string ("1e-4") from the Studio config path, which
         # preserves it as a string after validation; coerce so AdamW receives a float.
         try:
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 965a854042..3256dfa8e7 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -704,10 +704,14 @@ class DiffusionTrainingStartRequest(BaseModel):
         default_factory = lambda: ["to_k", "to_q", "to_v", "to_out.0"],
         description = "U-Net modules to attach LoRA to",
     )
-    max_grad_norm: float = Field(1.0, gt = 0, description = "Gradient clipping max-norm")
+    max_grad_norm: float = Field(
+        1.0, ge = 0, description = "Gradient clipping max-norm; 0 disables clipping"
+    )
     seed: int = Field(42)
     mixed_precision: Literal["bf16", "fp16", "no"] = Field("bf16")
-    snr_gamma: Optional[float] = Field(5.0, description = "Min-SNR loss weighting; null disables")
+    snr_gamma: Optional[float] = Field(
+        5.0, gt = 0, description = "Min-SNR loss weighting; null disables"
+    )
     gradient_checkpointing: bool = Field(True)
     lr_scheduler: str = Field("constant")
     lr_warmup_steps: int = Field(0, ge = 0)
diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py
index 75046165d0..388a2379e8 100644
--- a/studio/backend/tests/test_diffusion_lora_trainer.py
+++ b/studio/backend/tests/test_diffusion_lora_trainer.py
@@ -138,6 +138,16 @@ def test_config_rejects_zero_lora_alpha():
         DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", lora_alpha = 0).normalized()
 
 
+def test_config_rejects_nonpositive_snr_gamma():
+    # gamma <= 0 zeroes/inverts the min-SNR weight; None is the documented disable.
+    with pytest.raises(ValueError, match = "snr_gamma"):
+        DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = 0).normalized()
+    cfg = DiffusionLoraConfig(
+        base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = None
+    ).normalized()
+    assert cfg.snr_gamma is None
+
+
 def test_config_coerces_string_learning_rate():
     # The Studio config path preserves learning_rate as a string; normalize to float.
     cfg = DiffusionLoraConfig(
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index a7ef2240f4..1f85b08ccb 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -303,6 +303,20 @@ def test_route_start_forwards_extra_training_knobs(client):
     assert client._fake.started_with["lora_target_modules"] == ["to_q", "to_v"]
 
 
+def test_route_start_accepts_zero_max_grad_norm(client):
+    # 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_);
+    # the request model must not reject it.
+    r = client.post("/api/train/diffusion/start", json = {**_BODY, "max_grad_norm": 0.0})
+    assert r.status_code == 200, r.text
+    assert client._fake.started_with["max_grad_norm"] == 0.0
+
+
+def test_route_start_rejects_nonpositive_snr_gamma(client):
+    # gamma <= 0 zeroes/inverts the min-SNR loss weight; null is the disable value.
+    r = client.post("/api/train/diffusion/start", json = {**_BODY, "snr_gamma": 0})
+    assert r.status_code == 422
+
+
 def test_route_start_rejects_uncontained_paths(client):
     # An absolute path outside the Studio dataset roots is a 400, not silently accepted.
     r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"})

From 25d9cf960409b739a340c81d90f5b2e7b9749a6e Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 02:04:16 +0000
Subject: [PATCH 6/8] Stub diffusers.hooks too in the no-diffusers cache test

---
 studio/backend/tests/test_diffusion_cache.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py
index fe2f331281..befab46a6c 100644
--- a/studio/backend/tests/test_diffusion_cache.py
+++ b/studio/backend/tests/test_diffusion_cache.py
@@ -166,7 +166,10 @@ def test_missing_transformer_is_none(monkeypatch):
 
 
 def test_diffusers_unavailable_runs_uncached(monkeypatch):
-    # no diffusers import -> best-effort returns None, load proceeds uncached.
+    # no diffusers import -> best-effort returns None, load proceeds uncached. Block the
+    # hooks module too: the config import falls back to diffusers.hooks, which a REAL
+    # earlier import in the test session may have left cached in sys.modules.
     monkeypatch.setitem(sys.modules, "diffusers", None)
+    monkeypatch.setitem(sys.modules, "diffusers.hooks", None)
     t = _MixinTransformer()
     assert apply_step_cache(_pipe(t), mode = "fbcache") is None

From 78abe263873a179eb311b8f1067ed7ee2c9ae440 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 5 Jul 2026 02:13:55 +0000
Subject: [PATCH 7/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/sd_cpp_backend.py | 4 +---
 studio/backend/tests/test_diffusion_backend.py  | 4 +---
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 78425833d6..75572b500a 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -698,9 +698,7 @@ class SdCppDiffusionBackend:
             loading = self._loading
             if loading is None or loading.error is not None:
                 return ()
-            return tuple(
-                r for r in (loading.repo_id, loading.base_repo, *loading.asset_repos) if r
-            )
+            return tuple(r for r in (loading.repo_id, loading.base_repo, *loading.asset_repos) if r)
 
     # ── Generate ───────────────────────────────────────────────────────────
 
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index 2131385e43..0507987fbb 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -1414,9 +1414,7 @@ def test_bad_mode_strings_fail_before_eviction(fake_runtime):
         {"text_encoder_quant": "fp3"},
     ):
         with pytest.raises(ValueError):
-            backend.load_pipeline(
-                "unsloth/Z-Image-GGUF", gguf_filename = "m.gguf", **kwargs
-            )
+            backend.load_pipeline("unsloth/Z-Image-GGUF", gguf_filename = "m.gguf", **kwargs)
         assert backend._state is not None
 
 

From 68735819cda89a7c788dca3bf53d19bc26bb9d7b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 02:14:31 +0000
Subject: [PATCH 8/8] Keep transformer_quant tri-state through pre-eviction
 validation

---
 studio/backend/core/inference/diffusion.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 6db556bf7a..951598a6c3 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -927,8 +927,9 @@ class DiffusionBackend:
         kind = resolve_model_kind(gguf_filename, model_kind)
         # Validate every mode string that can raise NOW, before this load evicts the
         # previous pipeline below: their first in-line uses all sit past _unload_locked,
-        # where a bad request would cost the user their working model.
-        transformer_quant = normalize_transformer_quant(transformer_quant)
+        # where a bad request would cost the user their working model. Validate-only for
+        # transformer_quant: the raw value keeps the unset/auto vs explicit-off tri-state.
+        normalize_transformer_quant(transformer_quant)
         normalize_speed_mode(speed_mode)
         normalize_attention_backend(attention_backend)
         normalize_transformer_cache(transformer_cache)
@@ -1018,7 +1019,7 @@ class DiffusionBackend:
                 quant_plan = None
                 if (
                     kind == "gguf"
-                    and transformer_quant is not None  # normalized above, pre-eviction
+                    and normalize_transformer_quant(transformer_quant) is not None
                     and dense_transformer_supported(target)
                     and plan.offload_policy != OFFLOAD_NONE
                 ):