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 01/15] 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 02/15] 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 03/15] 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 04/15] 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 05/15] 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 06/15] 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 f6f198fd5f3ceca45b3de1bbcd9b09ad7e391439 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:07:45 +0000
Subject: [PATCH 07/15] [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 102fa83fa3..6baa0a8a0a 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -1406,9 +1406,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 62cae5fe71cacff56dd2ec72b881a6697b028cce Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 04:41:43 +0000
Subject: [PATCH 08/15] Load image pipelines from the prefetched snapshot
 instead of re-sweeping the hub

The prefetch already scopes the file list (no packaged root singles, no
dtype-variant twins, no ONNX/Flax exports), but from_pretrained was then
called with the hub id, and its own snapshot sweep re-downloaded the
skipped files anyway: 24 GB per FLUX.1 repo and 65 GB on FLUX.2-dev, as
found in the blob cache. Return the snapshot dir from the prefetch (keyed
on the pipeline manifest) and hand it to every pipeline-assembly
from_pretrained site; any prefetch failure keeps the hub id and the old
behavior.
---
 studio/backend/core/inference/diffusion.py    | 41 +++++++++++++++----
 .../backend/tests/test_diffusion_backend.py   | 32 +++++++++++++++
 2 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 09b2fa1b42..2d69c9bb56 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -439,11 +439,17 @@ class DiffusionBackend:
         base: str,
         base_files: list[str],
         hf_token: Optional[str],
-    ) -> None:
+    ) -> Optional[str]:
         """Pre-download the GGUF + the given ``base_files`` into the HF cache,
         WITHOUT the lock and honoring ``_cancel_event``, so load_pipeline's
         from_single_file / from_pretrained hit the cache and the heavy download can
-        be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``."""
+        be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``.
+
+        Returns the base repo's local snapshot dir when the prefetched set includes
+        the pipeline manifest, so from_pretrained can load from disk instead of
+        re-sweeping the hub (its own sweep also pulls files the scoped list skips,
+        e.g. the 24 GB packaged root singles in each FLUX.1 repo); None otherwise
+        (estimate failure, config-only base, local repo) -> hub id as before."""
         from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
 
         # GGUF transformer (hub repos only; a local path is already on disk).
@@ -452,12 +458,16 @@ class DiffusionBackend:
                 repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event
             )
         # Base repo (VAE / text-encoder / scheduler); list comes from the estimate.
+        snapshot_root: Optional[str] = None
         for rfilename in base_files:
             if self._cancel_event.is_set():
                 raise RuntimeError("Cancelled")
-            hf_hub_download_with_xet_fallback(
+            local = hf_hub_download_with_xet_fallback(
                 base, rfilename, hf_token, cancel_event = self._cancel_event
             )
+            if rfilename == "model_index.json":
+                snapshot_root = str(Path(local).parent)
+        return snapshot_root
 
     def validate_load_request(
         self,
@@ -674,7 +684,7 @@ class DiffusionBackend:
                     self._loading.expected_bytes = expected
             # Download outside the lock so unload()/an eviction can preempt the
             # multi-GB pull; load_pipeline below then assembles from the cache.
-            self._prefetch_files(
+            kwargs["_base_local_dir"] = self._prefetch_files(
                 kwargs["repo_id"],
                 kwargs.get("gguf_filename"),
                 base,
@@ -887,6 +897,7 @@ class DiffusionBackend:
         transformer_cache_threshold: Optional[float] = None,
         model_kind: Optional[str] = None,
         _load_token: Optional[int] = None,
+        _base_local_dir: Optional[str] = None,
     ) -> dict[str, Any]:
         # A blank / whitespace-only token must degrade to anonymous access, not be passed
         # as an explicit credential (from_single_file / from_pretrained / the Hub client
@@ -1003,6 +1014,7 @@ class DiffusionBackend:
                             transformer_quant,
                             transformer_quant_fast_accum,
                             fam = fam,
+                            base_local_dir = _base_local_dir,
                             prequant_path = transformer_prequant_path,
                         )
                     except Exception as exc:  # noqa: BLE001 — fall back to the GGUF build
@@ -1033,7 +1045,12 @@ class DiffusionBackend:
                         pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype}
                         if hf_token:
                             pipe_kwargs["token"] = hf_token
-                        pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
+                        # The prefetched snapshot dir keeps from_pretrained off the hub:
+                        # its own snapshot sweep re-downloads files the scoped prefetch
+                        # skipped (root packaged singles, e.g. 24 GB per FLUX.1 repo).
+                        pipe = pipeline_cls.from_pretrained(
+                            _base_local_dir or repo_id, **pipe_kwargs
+                        )
                     elif kind == "single_file" and fam.single_file_is_pipeline:
                         # A single-file SDXL-style checkpoint is the WHOLE pipeline
                         # (U-Net + VAE + both text encoders), not a transformer-only file,
@@ -1069,7 +1086,9 @@ class DiffusionBackend:
                         pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
                         if hf_token:
                             pipe_kwargs["token"] = hf_token
-                        pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
+                        pipe = pipeline_cls.from_pretrained(
+                            _base_local_dir or base, **pipe_kwargs
+                        )
 
                 # Resolve the effective speed mode: GGUF models default to the
                 # near-lossless `default` profile (compile is ~2.2x and sits below
@@ -1301,6 +1320,7 @@ class DiffusionBackend:
         *,
         fam: Optional[DiffusionFamily] = None,
         prequant_path: Optional[str] = None,
+        base_local_dir: Optional[str] = None,
     ) -> tuple[Any, str]:
         """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
 
@@ -1348,7 +1368,7 @@ class DiffusionBackend:
                 )
                 if transformer is not None:
                     pipe = self._assemble_pipe(
-                        pipeline_cls, base, transformer, dtype, hf_token, device
+                        pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
                     )
                     return pipe, scheme
 
@@ -1356,7 +1376,9 @@ class DiffusionBackend:
         transformer = transformer_cls.from_pretrained(
             base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
         )
-        pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device)
+        pipe = self._assemble_pipe(
+            pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
+        )
         scheme = quantize_transformer(
             pipe,
             target,
@@ -1377,13 +1399,14 @@ class DiffusionBackend:
         dtype: Any,
         hf_token: Optional[str],
         device: str,
+        base_local_dir: Optional[str] = None,
     ) -> Any:
         """Assemble the diffusers pipeline around ``transformer`` and place it on ``device``
         (a no-op for an already-placed pre-quantized transformer; it moves the companions)."""
         pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer}
         if hf_token:
             pipe_kwargs["token"] = hf_token
-        pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
+        pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs)
         pipe.to(device)
         return pipe
 
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index 6baa0a8a0a..1dccc2d0ac 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -2121,3 +2121,35 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
     backend.generate(prompt = "a sloth")
     backend.generate(prompt = "another sloth")
     assert resets == [True, True]
+
+
+def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch):
+    # The prefetched pipeline manifest's directory is the local snapshot root; a
+    # config-only base list (no manifest) returns None so the hub id stays in use.
+    backend = DiffusionBackend()
+    monkeypatch.setattr(
+        "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback",
+        lambda repo, fn, tok, **k: f"/cache/snap/{fn}",
+    )
+    root = backend._prefetch_files(
+        "base/repo", None, "base/repo", ["model_index.json", "vae/x.safetensors"], None
+    )
+    assert root == "/cache/snap"
+    assert (
+        backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None)
+        is None
+    )
+
+
+def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path):
+    # With a prefetched snapshot, from_pretrained must receive the local dir --
+    # its own hub sweep would re-download the root packaged singles the scoped
+    # prefetch skips (24 GB per FLUX.1 repo).
+    backend = DiffusionBackend()
+    backend.load_pipeline(
+        "unsloth/Qwen-Image-2512-bnb-4bit",
+        model_kind = "pipeline",
+        _base_local_dir = str(tmp_path),
+    )
+    assert _FakePipeline.last["base"] == str(tmp_path)
+    backend.unload()

From 919661ddf084f63b4df31ae0f598f86f9f24b27b 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 04:42:14 +0000
Subject: [PATCH 09/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

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

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 2d69c9bb56..3b47287722 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1086,9 +1086,7 @@ class DiffusionBackend:
                         pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
                         if hf_token:
                             pipe_kwargs["token"] = hf_token
-                        pipe = pipeline_cls.from_pretrained(
-                            _base_local_dir or base, **pipe_kwargs
-                        )
+                        pipe = pipeline_cls.from_pretrained(_base_local_dir or base, **pipe_kwargs)
 
                 # Resolve the effective speed mode: GGUF models default to the
                 # near-lossless `default` profile (compile is ~2.2x and sits below
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index 1dccc2d0ac..65516fb5e9 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -2136,8 +2136,7 @@ def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch):
     )
     assert root == "/cache/snap"
     assert (
-        backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None)
-        is None
+        backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) is None
     )
 
 

From c53a6cb65e05e231a8bb007d7f5a6463dd6d19e4 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 04:53:37 +0000
Subject: [PATCH 10/15] Address review findings: dataset preflight, sd.cpp
 unload barrier, caption tombstone, ControlNet cache race

- Run the trainer's caption discovery in the start route BEFORE freeing GPU
  residents, so a missing or uncaptionable dataset 400s without evicting the
  loaded chat/Images model.
- sd.cpp unload now waits out a cancelled one-shot generation on the generate
  lock before reporting the device free, matching the diffusers backend.
- Clearing a caption that came from metadata.jsonl writes an empty sidecar
  tombstone instead of unlinking (both readers treat an existing sidecar as
  authoritative), so the cleared label cannot resurface.
- The ControlNet wrapper pipe is only cached while its load is still current,
  closing the unload race the model cache already handled.
---
 studio/backend/core/inference/diffusion.py    |  9 ++++-
 .../backend/core/inference/sd_cpp_backend.py  |  7 ++++
 studio/backend/routes/training.py             | 37 ++++++++++++++++---
 .../tests/test_diffusion_controlnet.py        | 18 +++++++++
 .../backend/tests/test_diffusion_training.py  |  7 ++++
 5 files changed, 72 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 2d69c9bb56..ddf31db97a 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1582,7 +1582,14 @@ class DiffusionBackend:
             pipe = getattr(diffusers, pipe_cls_name).from_pipe(
                 state.pipe, controlnet = cn_model, torch_dtype = None
             )
-            self._cn_pipes[key] = pipe
+            with self._lock:
+                # Same race as the model cache above: an unload/superseding load may
+                # have cleared _cn_pipes while from_pipe ran; caching now would pin a
+                # pipeline built around the UNLOADED base and hand it to the next load.
+                if cancel.is_set() or self._state is not state:
+                    del pipe
+                    raise RuntimeError(DIFFUSION_CANCELLED_MSG)
+                self._cn_pipes[key] = pipe
         return pipe
 
     @staticmethod
diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 75572b500a..baee6605e0 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -1096,6 +1096,13 @@ class SdCppDiffusionBackend:
             state.server.stop()
         if pending is not None and pending is not (state.server if state else None):
             pending.stop()
+        # Wait for a signalled one-shot generation to actually exit before reporting
+        # unloaded: callers (the GPU arbiter, training cleanup) treat this return as
+        # "the device is free", but a one-shot sd-cli child killed by the cancel above
+        # unwinds under _generate_lock. A bare acquire is the exit barrier (never taken
+        # while holding _lock; same pattern as DiffusionBackend.unload).
+        with self._generate_lock:
+            pass
         return self.status()
 
     def status(self) -> dict[str, Any]:
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index a89a8d01c2..56514d2272 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1234,6 +1234,21 @@ async def start_diffusion_training(
     # user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
     _preflight_gated_base(config.get("base_model", ""), config.get("hf_token"))
 
+    # Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise
+    # fails inside the spawned trainer AFTER the user's chat/Images model was
+    # evicted. Same discovery the trainer runs, so the two cannot disagree.
+    from core.training import diffusion_train_common as _dtc
+
+    try:
+        await asyncio.to_thread(
+            _dtc.discover_image_caption_pairs,
+            config["data_dir"],
+            instance_prompt = config.get("instance_prompt") or None,
+            caption_column = config.get("caption_column") or "text",
+        )
+    except (FileNotFoundError, ValueError) as e:
+        raise HTTPException(status_code = 400, detail = str(e))
+
     # Free resident GPU workloads (export / Images pipeline / chat) before the trainer
     # loads its own pipeline.
     _free_gpu_for_diffusion_training()
@@ -1636,12 +1651,24 @@ async def set_diffusion_dataset_caption(
         sidecar = image_path.with_suffix(".txt")
         if caption:
             sidecar.write_text(caption, encoding = "utf-8")
-        else:
-            # Blank clears the sidecar; also drop a stale .caption so the image reads as
-            # uncaptioned afterwards.
-            sidecar.unlink(missing_ok = True)
             image_path.with_suffix(".caption").unlink(missing_ok = True)
-        return _image_record(folder, image_path, _load_metadata_captions(folder))
+            return _image_record(folder, image_path, _load_metadata_captions(folder))
+        # Blank must actually clear. Unlinking alone would resurface this image's
+        # metadata.jsonl / captions.jsonl caption (the fallback source), so when one
+        # exists write an EMPTY sidecar instead: both the record reader and the
+        # trainer's discovery treat an existing sidecar as authoritative even when
+        # empty, which makes it a tombstone. No metadata caption -> plain cleanup.
+        meta = _load_metadata_captions(folder)
+        try:
+            rel = image_path.relative_to(folder).as_posix()
+        except ValueError:
+            rel = image_path.name
+        if image_path.name in meta or rel in meta:
+            sidecar.write_text("", encoding = "utf-8")
+        else:
+            sidecar.unlink(missing_ok = True)
+        image_path.with_suffix(".caption").unlink(missing_ok = True)
+        return _image_record(folder, image_path, meta)
 
     return await asyncio.to_thread(write)
 
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index 5708c3a3ae..3d4f7c98cf 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -239,6 +239,9 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
     monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
     b = DiffusionBackend()
     st = _state()
+    # The pipe cache only commits while ``st`` is the CURRENT load (an unload racing
+    # from_pipe must not repopulate the cache), so mirror the loaded invariant.
+    b._state = st
     resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
     p1 = b._controlnet_pipe(st, resolved, threading.Event())
     assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel)
@@ -259,3 +262,18 @@ def test_controlnet_pipe_rejects_family_without_classes():
     st.family.controlnet_pipeline_class = None
     with pytest.raises(ValueError, match = "not supported"):
         b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event())
+
+def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
+    # An unload that lands while from_pipe is assembling must not let the wrapper
+    # repopulate the cache around the torn-down base pipe.
+    import threading
+
+    from core.inference.diffusion import DiffusionBackend
+
+    monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
+    b = DiffusionBackend()
+    st = _state()  # never committed to b._state: the load is already gone
+    resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
+    with pytest.raises(RuntimeError, match = "cancelled"):
+        b._controlnet_pipe(st, resolved, threading.Event())
+    assert b._cn_pipes == {}
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index 1f85b08ccb..b5acaa7bb9 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -265,6 +265,13 @@ def client(monkeypatch):
 
     monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False))
     monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
+    # The dataset preflight runs the trainer's discovery against _BODY's fake
+    # data_dir; stub it here so wiring tests pass, and let the dedicated preflight
+    # tests below re-point it at a real tmp dataset.
+    monkeypatch.setattr(
+        "core.training.diffusion_train_common.discover_image_caption_pairs",
+        lambda data_dir, **kw: [("img.png", "caption")],
+    )
     app = FastAPI()
     app.include_router(training_router, prefix = "/api/train")
     app.dependency_overrides[get_current_subject] = lambda: "test-user"

From bcf25ca569d4a98b28d629915ba6097216a8a049 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 04:56:25 +0000
Subject: [PATCH 11/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/tests/test_diffusion_controlnet.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index 3d4f7c98cf..50bf7159b0 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -263,6 +263,7 @@ def test_controlnet_pipe_rejects_family_without_classes():
     with pytest.raises(ValueError, match = "not supported"):
         b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event())
 
+
 def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
     # An unload that lands while from_pipe is assembling must not let the wrapper
     # repopulate the cache around the torn-down base pipe.

From 7a05d8655b8b2e29676b7e7fdbbb62802f8a36f5 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 05:30:58 +0000
Subject: [PATCH 12/15] Stream every DiT through group offload, not just the
 primary transformer

A dual-DiT pipeline (Ideogram 4's unconditional tower) placed its second
denoiser resident under the group tier, which defeats the tier since the
pair rarely fits where one alone did not. Stream transformer_2 and
unconditional_transformer alongside the transformer and keep only the
smaller companions resident.
---
 studio/backend/core/inference/diffusion_memory.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py
index c56d1592ce..4c4a3cca2c 100644
--- a/studio/backend/core/inference/diffusion_memory.py
+++ b/studio/backend/core/inference/diffusion_memory.py
@@ -502,6 +502,16 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
         import torch
         from diffusers.hooks import apply_group_offloading
 
+        # A dual-DiT pipeline (e.g. Ideogram 4's unconditional tower) carries a second
+        # denoiser as large as the first; leaving it resident would defeat this tier
+        # (the pair rarely fits where one alone did not). Stream every DiT and keep
+        # only the genuinely smaller companions resident.
+        streamed: dict[str, Any] = {"transformer": transformer}
+        for extra in ("transformer_2", "unconditional_transformer"):
+            module = getattr(pipe, extra, None)
+            if isinstance(module, torch.nn.Module):
+                streamed[extra] = module
+
         onload = torch.device(device)
         use_stream = onload.type == "cuda"  # overlap H2D copies with compute on CUDA
         gkwargs: dict[str, Any] = {
@@ -531,11 +541,12 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
         # load-time crash. The streamed transformer manages its own placement via the
         # offloading hooks applied next.
         for name, comp in getattr(pipe, "components", {}).items():
-            if name == "transformer":
+            if name in streamed:
                 continue
             if isinstance(comp, torch.nn.Module):
                 comp.to(onload)
-        apply_group_offloading(transformer, **gkwargs)
+        for module in streamed.values():
+            apply_group_offloading(module, **gkwargs)
         return True
     except Exception as exc:  # noqa: BLE001 — fall back to whole-module offload
         if logger is not None:

From f313cfd7e50121ef08104a227f21520c42e8f393 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 07:40:45 +0000
Subject: [PATCH 13/15] Security audit: baseline the new huggingface-hub
 Sandboxes findings

The latest huggingface-hub release added the Sandboxes feature. Its
bootstrap (_sandbox.py) fetches the static sbx-server binary into /tmp with
an Authorization header and marks it executable, which is exactly the
staged-dropper pattern the scanner hunts, and three while True polling loops
in _sandbox.py / hf_api.py / utils/_http.py match the beaconing heuristic.
All four verified against the official huggingface/huggingface_hub
repository: the snippet is the documented sandbox server injection and the
loops are deadline-style job and sandbox polling. Entries generated with
--write-baseline and reviewed line by line; scan_packages.py huggingface-hub
now exits 0 with the four findings suppressed.
---
 scripts/scan_packages_baseline.json | 32 +++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json
index 8e7cd4ea5a..3b33a7ce84 100644
--- a/scripts/scan_packages_baseline.json
+++ b/scripts/scan_packages_baseline.json
@@ -226,6 +226,22 @@
    "evidence": "FS: L528:         history: list[Response] | None = None, sha256:f56272dccd651b2644aa41ef6e688e211462427aad07fef5150240ec7347446e\nNetwork: L9: import urllib.request | L1243:     class _CookieCompatRequest(urllib.request.Request):",
    "evidence_hash": "b32f79e58c938680d89efa74113eeba76c9fc5aedf5de18086f93bef274c4bda"
   },
+  {
+   "package": "huggingface-hub",
+   "file": "huggingface_hub/_sandbox.py",
+   "check": "C2 polling/beaconing loop detected",
+   "severity": "CRITICAL",
+   "evidence": "L1179:             while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38",
+   "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8"
+  },
+  {
+   "package": "huggingface-hub",
+   "file": "huggingface_hub/_sandbox.py",
+   "check": "Writes to /tmp and executes (staged dropper)",
+   "severity": "CRITICAL",
+   "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"",
+   "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22"
+  },
   {
    "package": "huggingface-hub",
    "file": "huggingface_hub/hf_api.py",
@@ -242,6 +258,14 @@
    "evidence": "L4600:         while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272",
    "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602"
   },
+  {
+   "package": "huggingface-hub",
+   "file": "huggingface_hub/hf_api.py",
+   "check": "C2 polling/beaconing loop detected",
+   "severity": "CRITICAL",
+   "evidence": "L4613:         while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6",
+   "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4"
+  },
   {
    "package": "huggingface-hub",
    "file": "huggingface_hub/hf_api.py",
@@ -266,6 +290,14 @@
    "evidence": "L298:     while True: sha256:6b8e5e569594caf7c4eca6137646dae471a7c3aae7294096cf876f30b5f90306",
    "evidence_hash": "c066cc27bce31ee7b6ce07411ee7a7d9ecfbf3aafc8848f6641fabfe522a7703"
   },
+  {
+   "package": "huggingface-hub",
+   "file": "huggingface_hub/utils/_http.py",
+   "check": "C2 polling/beaconing loop detected",
+   "severity": "CRITICAL",
+   "evidence": "L462:     while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958",
+   "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e"
+  },
   {
    "package": "ipython",
    "file": "IPython/core/debugger.py",

From fc1e0991244f07045ac06152a5a82f4cc79685d3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 11:36:58 +0000
Subject: [PATCH 14/15] Harden ControlNet loads, thumbnail cache keys, API
 training guard, and picker roving keys

Review follow-ups on the image-generation PR:

- ControlNet: resolve_controlnet accepts a bare owner/name repo without the
  non-GGUF base trust gate, and _controlnet_pipe hands it straight to
  from_pretrained. A malicious pickle .bin would deserialize on load, so run
  the same Hugging Face malware preflight (evaluate_file_security) the chat and
  export loaders use before any remote ControlNet load; local dirs are exempt.
- Dataset thumbnails: key the cache on the full filename instead of the stem so
  sample.png and sample.jpg no longer collide on one .thumbs file (which could
  serve or delete the wrong image); the delete cleanup globs the same key.
- Diffusion training start: mirror start_training's API-key guard so an API
  client cannot start training (which frees VRAM by unloading chat) while an
  inference request is streaming; it now returns 409 before any GPU is freed.
- Model picker: include the curated safetensors row keys in the recommended
  roving key list so arrow-key navigation reaches those rows instead of hitting
  the duplicate option-missing id.

Tests: ControlNet malware gate (remote blocked before from_pretrained, local
skipped), thumbnail same-stem cache separation, API-key diffusion-start 409
before GPU free. Full diffusion suites green.
---
 studio/backend/core/inference/diffusion.py    | 11 +++
 studio/backend/routes/training.py             | 31 +++++++-
 .../tests/test_diffusion_controlnet.py        | 70 +++++++++++++++++++
 .../tests/test_diffusion_dataset_api.py       | 20 +++++-
 .../backend/tests/test_diffusion_training.py  | 37 +++++++++-
 .../assistant-ui/model-selector/pickers.tsx   |  9 +++
 6 files changed, 171 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 007d6d349d..603e4ca637 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1537,6 +1537,17 @@ class DiffusionBackend:
         if cn_model is None:
             if cancel.is_set():
                 raise RuntimeError(DIFFUSION_CANCELLED_MSG)
+            # resolve_controlnet accepts a bare owner/name repo without the non-GGUF base
+            # trust gate, and from_pretrained below downloads and deserializes it. A
+            # malicious pickle .bin would execute on load, so run the same Hub malware
+            # preflight the chat/export loaders use before any remote ControlNet load. A
+            # local dir the user picked has no Hub scan and is exempt (fail-open there).
+            if not getattr(resolved_cn, "is_local", False):
+                from utils.security import evaluate_file_security
+
+                _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None)
+                if _cn_fs.blocked:
+                    raise ValueError(_cn_fs.reason)
             import torch
 
             # state.dtype is the display string saved at load ("bfloat16"), NOT a
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 56514d2272..07dbd5c812 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1177,11 +1177,29 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
 
 @router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse)
 async def start_diffusion_training(
-    body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject)
+    body: DiffusionTrainingStartRequest,
+    current_subject: str = Depends(get_current_subject),
+    via_api_key: bool = Depends(authenticated_via_api_key),
 ):
     """Start an SDXL LoRA training job from an image + caption dataset."""
     from core.training.diffusion_training_service import get_diffusion_training_service
 
+    # When Studio is driven as an inference API (API-key auth), refuse to start training
+    # while a request is in flight: _free_gpu_for_diffusion_training() below unloads the
+    # chat backends to reclaim VRAM, which would kill the stream. Mirrors start_training so
+    # a diffusion start cannot silently drop an active API inference request.
+    if via_api_key is True:
+        from core.inference.llama_keepwarm import other_inference_request_count
+        if other_inference_request_count(current_request_counted = False) > 0:
+            raise HTTPException(
+                status_code = 409,
+                detail = (
+                    "Cannot start diffusion (Images) training over the API while an inference "
+                    "request is in progress. Wait for it to finish, or start training from the "
+                    "Studio UI."
+                ),
+            )
+
     # Interlock: refuse while an LLM training run holds the GPU (symmetric with the
     # diffusion check in start_training), so the two trainers never contend for VRAM.
     try:
@@ -1606,7 +1624,11 @@ async def get_diffusion_dataset_image(
 
         thumbs_dir = folder / _THUMBS_DIRNAME
         thumbs_dir.mkdir(exist_ok = True)
-        thumb_path = thumbs_dir / f"{image_path.stem}_{size}.jpg"
+        # Key on the full filename (stem + extension), not the stem: two images that
+        # share a stem but differ by extension (sample.png / sample.jpg) would otherwise
+        # collide on one cache file, and an mtime-newer cache built for the first would
+        # be served for the second, showing the wrong image in the labeling grid.
+        thumb_path = thumbs_dir / f"{image_path.name}_{size}.jpg"
         src_mtime = image_path.stat().st_mtime
         if thumb_path.is_file() and thumb_path.stat().st_mtime >= src_mtime:
             return thumb_path
@@ -1691,7 +1713,10 @@ async def delete_diffusion_dataset_image(
             image_path.with_suffix(ext).unlink(missing_ok = True)
         thumbs_dir = folder / _THUMBS_DIRNAME
         if thumbs_dir.is_dir():
-            for t in thumbs_dir.glob(f"{image_path.stem}_*.jpg"):
+            # Thumbs are keyed on the full filename (stem + extension), so match that
+            # here too; a stem-only glob would leave this image's thumbs behind and
+            # could delete a same-stem sibling's (sample.png vs sample.jpg).
+            for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"):
                 t.unlink(missing_ok = True)
         return {"deleted": image_path.name}
 
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index 50bf7159b0..b7ae037197 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -231,12 +231,24 @@ def _state():
     )
 
 
+def _allow_cn_security(monkeypatch):
+    """Stub the Hub malware preflight to allow the load (hermetic, no network)."""
+    import utils.security
+
+    monkeypatch.setattr(
+        utils.security,
+        "evaluate_file_security",
+        lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = False, reason = ""),
+    )
+
+
 def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
     import threading
 
     from core.inference.diffusion import DiffusionBackend
 
     monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
+    _allow_cn_security(monkeypatch)
     b = DiffusionBackend()
     st = _state()
     # The pipe cache only commits while ``st`` is the CURRENT load (an unload racing
@@ -252,6 +264,64 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
     assert b._cn_models["flux-union-pro"] is p1.controlnet
 
 
+def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
+    # A bare owner/name ControlNet is accepted by resolve_controlnet without the base
+    # trust gate, so the load path must run the Hub malware preflight: a flagged remote
+    # repo must raise BEFORE from_pretrained downloads/deserializes it.
+    import threading
+
+    import utils.security
+    from core.inference.diffusion import DiffusionBackend
+
+    loaded = {"called": False}
+
+    class _TrapModel(_FakeCNModel):
+        @classmethod
+        def from_pretrained(cls, path, torch_dtype = None, token = None):
+            loaded["called"] = True
+            return super().from_pretrained(path, torch_dtype = torch_dtype, token = token)
+
+    mod = _fake_diffusers()
+    mod.FluxControlNetModel = _TrapModel
+    monkeypatch.setitem(sys.modules, "diffusers", mod)
+    monkeypatch.setattr(
+        utils.security,
+        "evaluate_file_security",
+        lambda name, hf_token = None, **kw: types.SimpleNamespace(
+            blocked = True, reason = "Hugging Face security scan flagged unsafe files: evil.bin"
+        ),
+    )
+    b = DiffusionBackend()
+    st = _state()
+    b._state = st
+    resolved = dc.ResolvedControlNet("evil/cn", "evil/cn", is_local = False)
+    with pytest.raises(ValueError, match = "security scan flagged"):
+        b._controlnet_pipe(st, resolved, threading.Event())
+    assert loaded["called"] is False
+
+
+def test_controlnet_pipe_skips_scan_for_local_dir(monkeypatch, tmp_path):
+    # A local dir the user picked has no Hub scan; the preflight must not block it even
+    # if the (unused) scan stub would say blocked.
+    import threading
+
+    import utils.security
+    from core.inference.diffusion import DiffusionBackend
+
+    monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
+    monkeypatch.setattr(
+        utils.security,
+        "evaluate_file_security",
+        lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = True, reason = "x"),
+    )
+    b = DiffusionBackend()
+    st = _state()
+    b._state = st
+    resolved = dc.ResolvedControlNet("my-cn", str(tmp_path), is_local = True)
+    p = b._controlnet_pipe(st, resolved, threading.Event())
+    assert isinstance(p, _FakeCNPipe)
+
+
 def test_controlnet_pipe_rejects_family_without_classes():
     import threading
 
diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py
index b10dd04f4c..451f594ba1 100644
--- a/studio/backend/tests/test_diffusion_dataset_api.py
+++ b/studio/backend/tests/test_diffusion_dataset_api.py
@@ -177,15 +177,29 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
     folder.mkdir()
     _write_png(folder / "x.png")
     (folder / "x.txt").write_text("cap", encoding = "utf-8")
-    # Generate a thumbnail so we can assert it is cleaned up too.
+    # Generate a thumbnail so we can assert it is cleaned up too. Thumbs are keyed on
+    # the full filename (stem + extension) to avoid same-stem collisions across formats.
     client.get("/api/train/diffusion/dataset/d/image/x.png?thumb=32")
-    assert list((folder / ".thumbs").glob("x_*.jpg"))
+    assert list((folder / ".thumbs").glob("x.png_*.jpg"))
 
     r = client.delete("/api/train/diffusion/dataset/d/image/x.png")
     assert r.status_code == 200, r.text
     assert not (folder / "x.png").exists()
     assert not (folder / "x.txt").exists()
-    assert not list((folder / ".thumbs").glob("x_*.jpg"))
+    assert not list((folder / ".thumbs").glob("x.png_*.jpg"))
+
+
+def test_thumb_cache_key_distinguishes_same_stem_extensions(client, ds_root):
+    # sample.png and sample.jpg share a stem; each must get its OWN thumbnail cache
+    # file, so the labeling grid never serves one image's thumbnail for the other.
+    folder = ds_root / "d"
+    folder.mkdir()
+    Image.new("RGB", (8, 8), (10, 20, 30)).save(folder / "sample.png", format = "PNG")
+    Image.new("RGB", (8, 8), (200, 210, 220)).save(folder / "sample.jpg", format = "JPEG")
+    client.get("/api/train/diffusion/dataset/d/image/sample.png?thumb=32")
+    client.get("/api/train/diffusion/dataset/d/image/sample.jpg?thumb=32")
+    thumbs = sorted(p.name for p in (folder / ".thumbs").glob("*.jpg"))
+    assert thumbs == ["sample.jpg_32.jpg", "sample.png_32.jpg"]
 
 
 # ── traversal / validation ───────────────────────────────────────────────────
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index b5acaa7bb9..b138de3980 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -20,7 +20,7 @@ import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
-from auth.authentication import get_current_subject
+from auth.authentication import authenticated_via_api_key, get_current_subject
 from core.training.diffusion_training_service import DiffusionTrainingService
 from routes.training import router as training_router
 
@@ -275,8 +275,12 @@ def client(monkeypatch):
     app = FastAPI()
     app.include_router(training_router, prefix = "/api/train")
     app.dependency_overrides[get_current_subject] = lambda: "test-user"
+    # Default to session (UI) auth: the API-key inference-in-flight guard is a no-op there,
+    # so the wiring tests below behave as before. The guard test flips this override.
+    app.dependency_overrides[authenticated_via_api_key] = lambda: False
     c = TestClient(app)
     c._fake = fake  # type: ignore[attr-defined]
+    c._app = app  # type: ignore[attr-defined]
     return c
 
 
@@ -365,6 +369,37 @@ def test_route_start_conflict_maps_to_409(client):
     assert r.status_code == 409
 
 
+def test_route_start_over_api_with_inference_in_flight_is_409(client, monkeypatch):
+    # An API-key client must not start diffusion training (which frees VRAM by unloading
+    # chat) while an inference request is streaming; it should 409 instead of killing it.
+    client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
+    monkeypatch.setattr(
+        "core.inference.llama_keepwarm.other_inference_request_count",
+        lambda current_request_counted = False: 1,
+    )
+    freed = {"called": False}
+    import routes.training as tr
+
+    monkeypatch.setattr(
+        tr, "_free_gpu_for_diffusion_training", lambda: freed.__setitem__("called", True)
+    )
+    r = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r.status_code == 409
+    # The guard must run BEFORE any GPU is freed, so the live inference stream survives.
+    assert freed["called"] is False
+
+
+def test_route_start_over_api_without_inference_proceeds(client, monkeypatch):
+    # Same API-key path but no inference in flight: the start proceeds normally.
+    client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
+    monkeypatch.setattr(
+        "core.inference.llama_keepwarm.other_inference_request_count",
+        lambda current_request_counted = False: 0,
+    )
+    r = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r.status_code == 200
+
+
 def test_route_status_and_stop(client):
     client.post("/api/train/diffusion/start", json = _BODY)
     s = client.get("/api/train/diffusion/status")
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 c380e04244..f096af73c8 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -2282,6 +2282,14 @@ export function HubModelPicker({
     }
 
     if (section === "recommended") {
+      // Curated safetensors rows render ABOVE the recommended rows (and call
+      // getOptionProps), so their keys must lead here or they fall back to the
+      // duplicate ...-option-missing id and drop out of arrow-key navigation.
+      keys.push(
+        ...curatedSafetensorsRows.map((m) =>
+          makeModelOptionKey("curated-safetensors", m.id),
+        ),
+      );
       keys.push(
         ...recommendedRows.map((r) => makeModelOptionKey("recommended", r.id)),
       );
@@ -2291,6 +2299,7 @@ export function HubModelPicker({
   }, [
     cachedReady,
     chatOnly,
+    curatedSafetensorsRows,
     sortedCustomFolderModels,
     customFoldersCollapsed,
     downloadedCollapsed,

From e80067512845c19f21ef89b92ced7c52307cc148 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 11:40:31 +0000
Subject: [PATCH 15/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/diffusion.py        | 1 -
 studio/backend/tests/test_diffusion_controlnet.py | 8 ++++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 603e4ca637..b40add9dd2 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1544,7 +1544,6 @@ class DiffusionBackend:
             # local dir the user picked has no Hub scan and is exempt (fail-open there).
             if not getattr(resolved_cn, "is_local", False):
                 from utils.security import evaluate_file_security
-
                 _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None)
                 if _cn_fs.blocked:
                     raise ValueError(_cn_fs.reason)
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index b7ae037197..bcb7393457 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -234,7 +234,6 @@ def _state():
 def _allow_cn_security(monkeypatch):
     """Stub the Hub malware preflight to allow the load (hermetic, no network)."""
     import utils.security
-
     monkeypatch.setattr(
         utils.security,
         "evaluate_file_security",
@@ -277,7 +276,12 @@ def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
 
     class _TrapModel(_FakeCNModel):
         @classmethod
-        def from_pretrained(cls, path, torch_dtype = None, token = None):
+        def from_pretrained(
+            cls,
+            path,
+            torch_dtype = None,
+            token = None,
+        ):
             loaded["called"] = True
             return super().from_pretrained(path, torch_dtype = torch_dtype, token = token)