From 26a81c5e17542413da86e72ca81425378d35166f Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sat, 4 Jul 2026 02:35:51 -0300
Subject: [PATCH 01/13] Fix diffusion quant gating, MPS load OOM, and native
 sd-server discovery

---
 studio/backend/core/inference/diffusion.py    | 37 +++++++++++++++----
 .../core/inference/diffusion_controlnet.py    | 14 ++++---
 .../core/inference/diffusion_device.py        | 11 ++++++
 .../backend/core/inference/sd_cpp_engine.py   | 18 ++++++++-
 .../backend/tests/test_diffusion_backend.py   |  4 +-
 .../tests/test_diffusion_controlnet.py        |  7 ++--
 .../images/diffusion-train-dialog.tsx         | 17 ++++++---
 .../images/train/diffusion-train-panel.tsx    | 19 +++++++---
 8 files changed, 97 insertions(+), 30 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 28eafe6829..c426d3b5d2 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -842,15 +842,32 @@ class DiffusionBackend:
     def _companion_cache_bytes(base: str) -> int:
         """Resident companion (VAE + text-encoder) size for the memory plan.
 
-        For a hub base repo this is the cached blob total (``_cache_bytes``). For a
-        LOCAL diffusers base directory the blob cache is empty, so sum the on-disk
-        component weights instead, excluding ``transformer/`` (the GGUF supplies the
-        transformer). Without this a local base folds its multi-GB VAE / text-encoder
-        weights to zero and auto planning can pick a resident placement that OOMs."""
+        Sums the cached VAE + text-encoder weights while EXCLUDING ``transformer/`` (the
+        GGUF / single file supplies the transformer, so its ``transformer/`` shards are
+        not resident here). This matters for the dense ``transformer_quant`` path: it
+        prefetches the base repo's ``transformer/`` shards into the cache, and folding
+        those multi-GB shards into the companion size would inflate the plan and wrongly
+        force offload -- gating off the very quant path that fetched them. For a LOCAL
+        diffusers base the blob cache is empty, so walk the on-disk weights; for a hub
+        base, walk the snapshot (whose ``transformer/`` subfolder we can skip) instead of
+        the flat, content-addressed ``blobs/`` dir, which carries no subfolder split."""
         local = Path(base).expanduser()
         if local.is_dir():
             return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True)
-        return DiffusionBackend._cache_bytes(base)
+        from huggingface_hub import constants
+
+        snapshots = Path(constants.HF_HUB_CACHE) / f"models--{base.replace('/', '--')}" / "snapshots"
+        if not snapshots.is_dir():
+            return 0
+        # Multiple revisions may be cached; the active one is the fullest, so take the max.
+        return max(
+            (
+                DiffusionBackend._local_dir_weight_bytes(rev, exclude_transformer = True)
+                for rev in snapshots.iterdir()
+                if rev.is_dir()
+            ),
+            default = 0,
+        )
 
     # ── Synchronous load / generate / unload ───────────────────────────────
 
@@ -1653,7 +1670,13 @@ class DiffusionBackend:
         resolution/batch changed, or a stale-cache reuse otherwise. Best-effort: a
         transformer without the hook (uncached load) is a silent no-op."""
         transformer = getattr(pipe, "transformer", None)
-        reset = getattr(transformer, "reset_stateful_hooks", None)
+        # A diffusers CacheMixin transformer clears FBCache via ``_reset_stateful_cache``
+        # (which drives its HookRegistry.reset_stateful_hooks internally). The public
+        # ``reset_stateful_hooks`` name lives only on the HookRegistry, not on the
+        # transformer, so keep it only as a version fallback.
+        reset = getattr(transformer, "_reset_stateful_cache", None) or getattr(
+            transformer, "reset_stateful_hooks", None
+        )
         if callable(reset):
             try:
                 reset()
diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py
index b318b040c8..9f42350c73 100644
--- a/studio/backend/core/inference/diffusion_controlnet.py
+++ b/studio/backend/core/inference/diffusion_controlnet.py
@@ -205,7 +205,8 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve
 # Union ControlNet mode indices. A single "union" model covers several control modes and
 # selects the active one via an integer ``control_mode`` argument; these are the standard
 # indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" (an already-made
-# map) carries no intrinsic mode, so it maps to nothing (the caller omits control_mode).
+# map) carries no intrinsic mode, so union_control_mode() defaults it to 0 (a union model
+# still requires a concrete mode).
 _UNION_CONTROL_MODES: dict[str, int] = {
     "canny": 0,
     "tile": 1,
@@ -220,13 +221,16 @@ _UNION_CONTROL_MODES: dict[str, int] = {
 def union_control_mode(spec_id: str, control_type: str) -> Optional[int]:
     """The integer ``control_mode`` for a union ControlNet, or None.
 
-    Returns a mode only for a curated *union* catalog entry AND a control type that maps to a
-    known index; otherwise None so the caller omits the kwarg (a non-union ControlNet has a
-    single fixed mode, and 'passthrough' does not name one). Pure lookup, no network."""
+    A union model REQUIRES a concrete ``control_mode`` (diffusers raises when it is None),
+    so for a curated union entry always return an index: the mapped mode, or a default
+    (0 / canny) for a type that carries no intrinsic mode such as 'passthrough' (an
+    already-made control map, which is also the UI's default for these models). For a
+    non-union entry return None so the caller omits the kwarg (it has a single fixed
+    mode). Pure lookup, no network."""
     entry = _catalog_by_id().get(spec_id)
     if entry is None or not entry.is_union:
         return None
-    return _UNION_CONTROL_MODES.get((control_type or "").strip().lower())
+    return _UNION_CONTROL_MODES.get((control_type or "").strip().lower(), 0)
 
 
 def preprocess_control(image: Any, control_type: str) -> Any:
diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py
index 2d2cc11343..b58b1f84e2 100644
--- a/studio/backend/core/inference/diffusion_device.py
+++ b/studio/backend/core/inference/diffusion_device.py
@@ -14,6 +14,7 @@ dtype choice and the capability flags the backend keys optimisation paths off.
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import Any, Optional
 
@@ -234,6 +235,16 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
         mps_available = False
 
     if mps_available:
+        # Relax the MPS memory watermark BEFORE the first MPS allocation (the bfloat16
+        # probe just below). torch reads PYTORCH_MPS_HIGH_WATERMARK_RATIO exactly once,
+        # when the MPS allocator first initializes, so setting it any later is a no-op.
+        # The allocator otherwise caps a process at ~1.7x recommendedMaxWorkingSetSize,
+        # and a model that fits in unified system RAM but exceeds that cap OOMs at
+        # pipe.to("mps") (observed on an 8GB M1 mac mini: "MPS allocated 9.06 GiB, max
+        # allowed 9.07 GiB"). CPU offload can't help on unified memory (it frees no
+        # device bytes). Lifting the cap lets MPS spill into system RAM; a model larger
+        # than RAM would fail either way. setdefault respects a user-provided override.
+        os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
         # Prefer bfloat16; otherwise fall back to float32, NEVER silent float16.
         # Modern diffusion transformers (Z-Image, FLUX.2, ...) produce activations
         # far outside float16's finite range (~6.5e4) -- Z-Image's MLP
diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py
index 52e4d6d693..af896007b5 100644
--- a/studio/backend/core/inference/sd_cpp_engine.py
+++ b/studio/backend/core/inference/sd_cpp_engine.py
@@ -125,7 +125,12 @@ def runtime_env(binary: str, base_env: Optional[dict[str, str]] = None) -> dict[
 def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
     """``stem`` locations under a stable-diffusion.cpp checkout/install ``root``,
     highest priority first: the cmake ``build/bin`` tree, then a Windows Release
-    subdir, then the root itself."""
+    subdir, then the root itself, then the prebuilt archive's versioned subdir.
+
+    The prebuilt archive extracts into a top-level versioned dir
+    (``sd-master-<tag>-bin-<host>/``) rather than flattening into ``root``, so without
+    the ``root/*/`` glob a fresh prebuilt install is invisible here -- which silently
+    demotes the persistent sd-server to one-shot mode and re-downloads on every start."""
     name = _binary_name(stem)
     cands = [
         root / "build" / "bin" / name,
@@ -133,6 +138,17 @@ def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
         root / "bin" / name,
         root / name,
     ]
+    # Prebuilt archive layout: root/sd-master-<tag>-bin-<host>/<name> (+ its own bin/).
+    # Newest install first (by mtime -- tag strings don't sort numerically, so a lexical
+    # sort would rank build 99 above build 100).
+    try:
+        subdirs = [p for p in root.iterdir() if p.is_dir()]
+        subdirs.sort(key = lambda p: p.stat().st_mtime, reverse = True)
+        for sub in subdirs:
+            cands.append(sub / name)
+            cands.append(sub / "bin" / name)
+    except OSError:
+        pass
     return cands
 
 
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index d36e62f3cf..cf5f8ac07b 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -2043,7 +2043,7 @@ def test_reset_step_cache_helper_is_best_effort():
     # Calls the transformer's reset hook when present.
     calls = []
     pipe = types.SimpleNamespace(
-        transformer = types.SimpleNamespace(reset_stateful_hooks = lambda: calls.append(True))
+        transformer = types.SimpleNamespace(_reset_stateful_cache = lambda: calls.append(True))
     )
     DiffusionBackend._reset_step_cache(pipe)
     assert calls == [True]
@@ -2065,7 +2065,7 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
     )
     resets = []
     backend._state.pipe.transformer = types.SimpleNamespace(
-        reset_stateful_hooks = lambda: resets.append(True)
+        _reset_stateful_cache = lambda: resets.append(True)
     )
     # No cache engaged (transformer_cache is None) -> reset must NOT run.
     backend.generate(prompt = "a sloth")
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index 5708c3a3ae..ced1b752ac 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -56,12 +56,13 @@ def test_resolve_controlnet_enforces_family_match():
 
 
 def test_union_control_mode_maps_only_union_entries():
-    # Union entries map a known control type to its integer mode; passthrough / unknown
-    # types and non-union ids return None so the caller omits control_mode.
+    # Union entries map a known control type to its integer mode; a union model always
+    # needs a concrete mode, so an unmapped type (passthrough) defaults to 0. A non-union
+    # id returns None so the caller omits control_mode.
     assert dc.union_control_mode("flux-union-pro", "canny") == 0
     assert dc.union_control_mode("flux-union-pro", "depth") == 2
     assert dc.union_control_mode("flux-union-pro", "pose") == 4
-    assert dc.union_control_mode("flux-union-pro", "passthrough") is None
+    assert dc.union_control_mode("flux-union-pro", "passthrough") == 0
     assert dc.union_control_mode("some/bare-repo", "canny") is None
 
 
diff --git a/studio/frontend/src/features/images/diffusion-train-dialog.tsx b/studio/frontend/src/features/images/diffusion-train-dialog.tsx
index 25b260a03e..016087bd8d 100644
--- a/studio/frontend/src/features/images/diffusion-train-dialog.tsx
+++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx
@@ -120,24 +120,29 @@ export function DiffusionTrainDialog({
   }, [open, poll]);
 
   const active = Boolean(status?.active) || status?.status === "running";
-  const completed = status?.status === "completed";
+  // A stopped run still saves + publishes a deployable adapter (catalog_path set), so
+  // treat it as finished-with-adapter too; a save=False cancel has no catalog_path.
+  const hasSavedAdapter =
+    status?.status === "completed" ||
+    (status?.status === "stopped" && Boolean(status?.catalog_path));
+  const completed = hasSavedAdapter;
   const pct =
     status && status.total_steps > 0
       ? Math.min(100, Math.round((status.step / status.total_steps) * 100))
       : 0;
 
-  // Notify the parent exactly once when a run reaches "completed", so it can rescan the
-  // LoRA picker (a LoRA trained while a model is loaded is otherwise invisible until a
-  // model swap re-runs the discovery effect).
+  // Notify the parent exactly once when a run finishes with a saved adapter, so it can
+  // rescan the LoRA picker (a LoRA trained while a model is loaded is otherwise invisible
+  // until a model swap re-runs the discovery effect).
   const [notifiedComplete, setNotifiedComplete] = useState(false);
   useEffect(() => {
-    if (status?.status === "completed" && !notifiedComplete) {
+    if (hasSavedAdapter && !notifiedComplete) {
       setNotifiedComplete(true);
       onTrainingComplete?.();
     } else if (status?.status === "running" && notifiedComplete) {
       setNotifiedComplete(false); // arm again for the next run
     }
-  }, [status?.status, notifiedComplete, onTrainingComplete]);
+  }, [hasSavedAdapter, status?.status, notifiedComplete, onTrainingComplete]);
 
   const selectedDataset =
     dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined;
diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
index 45df1d50a6..75b1d00ad1 100644
--- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
+++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
@@ -318,23 +318,30 @@ export function DiffusionTrainPanel({
   // terminal "completed" status until the next start, so we can't rely on it clearing).
   const [dismissedJobId, setDismissedJobId] = useState<string | null>(null);
   const running = Boolean(status?.active) || status?.status === "running";
-  const completed =
-    status?.status === "completed" && status.job_id !== dismissedJobId;
+  // A stopped run still saves + catalog-publishes a real, deployable adapter (status
+  // carries catalog_path), so treat "stopped with an adapter" as terminal-with-adapter
+  // too -- otherwise the normal "stop once the loss looks good" flow leaves the trained
+  // adapter with no Deploy button and no picker refresh. A save=False cancel has no
+  // catalog_path, so it correctly still shows nothing.
+  const hasSavedAdapter =
+    status?.status === "completed" ||
+    (status?.status === "stopped" && Boolean(status?.catalog_path));
+  const completed = hasSavedAdapter && status?.job_id !== dismissedJobId;
   const pct =
     status && status.total_steps > 0
       ? Math.min(100, Math.round((status.step / status.total_steps) * 100))
       : 0;
 
-  // Notify the parent exactly once per completed run so it rescans the LoRA picker.
+  // Notify the parent exactly once per finished run so it rescans the LoRA picker.
   const notifiedComplete = useRef(false);
   useEffect(() => {
-    if (status?.status === "completed" && !notifiedComplete.current) {
+    if (hasSavedAdapter && !notifiedComplete.current) {
       notifiedComplete.current = true;
       onTrainingComplete?.();
     } else if (status?.status === "running" && notifiedComplete.current) {
       notifiedComplete.current = false;
     }
-  }, [status?.status, onTrainingComplete]);
+  }, [hasSavedAdapter, status?.status, onTrainingComplete]);
 
   const selectedDataset =
     dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined;
@@ -780,7 +787,7 @@ export function DiffusionTrainPanel({
       <div className="flex min-w-0 flex-1 flex-col gap-4 overflow-y-auto">
         {status &&
         status.status !== "idle" &&
-        !(status.status === "completed" && status.job_id === dismissedJobId) ? (
+        !(hasSavedAdapter && status.job_id === dismissedJobId) ? (
           <>
             <div className="bg-card corner-squircle flex flex-col gap-3 rounded-3xl p-5 ring-1 ring-foreground/10">
               <div className="flex items-center justify-between">

From e5c63cdfffbbb8cec3c0457f01216a07cb6f4615 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sat, 4 Jul 2026 03:17:12 -0300
Subject: [PATCH 02/13] Fix diffusion training validation, dataset upload
 atomicity, and LoRA error mapping

---
 scripts/nvfp4_t211_probe.py                   |  2 +-
 scripts/sparse_accum_probe.py                 |  2 +-
 scripts/uninstall.ps1                         |  7 +-
 scripts/uninstall.sh                          |  7 +-
 studio/backend/core/inference/diffusion.py    | 58 +++++++++++--
 .../core/inference/diffusion_device.py        |  4 +-
 .../backend/core/inference/diffusion_lora.py  | 48 +++++++++--
 .../core/inference/diffusion_memory.py        | 21 ++++-
 .../backend/core/inference/sd_cpp_backend.py  |  7 ++
 .../core/training/diffusion_dit_trainer.py    | 11 ++-
 .../core/training/diffusion_lora_trainer.py   |  4 +-
 .../core/training/diffusion_train_common.py   | 72 ++++++++++++++--
 studio/backend/models/training.py             | 19 +++--
 studio/backend/routes/training.py             | 38 ++++++---
 .../tests/test_diffusion_dataset_api.py       |  6 +-
 .../tests/test_diffusion_dit_trainer.py       | 20 +++--
 studio/backend/tests/test_diffusion_lora.py   | 20 +++++
 .../tests/test_diffusion_lora_trainer.py      | 85 ++++++++++++++++++-
 .../backend/tests/test_diffusion_training.py  | 18 ++++
 studio/backend/tests/test_sd_cpp_backend.py   | 16 ++++
 .../src/features/images/images-page.tsx       | 34 +++++++-
 21 files changed, 428 insertions(+), 71 deletions(-)

diff --git a/scripts/nvfp4_t211_probe.py b/scripts/nvfp4_t211_probe.py
index e69e01b2de..de697165ea 100644
--- a/scripts/nvfp4_t211_probe.py
+++ b/scripts/nvfp4_t211_probe.py
@@ -24,7 +24,7 @@ import numpy as np
 
 BASE = "Tongyi-MAI/Z-Image-Turbo"
 PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
-OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_t211_images")
+OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_t211_images"
 
 
 # ----------------------------------------------------------------------------- diag
diff --git a/scripts/sparse_accum_probe.py b/scripts/sparse_accum_probe.py
index 4f91f905b6..74006334f3 100644
--- a/scripts/sparse_accum_probe.py
+++ b/scripts/sparse_accum_probe.py
@@ -27,7 +27,7 @@ import numpy as np
 
 BASE = "Tongyi-MAI/Z-Image-Turbo"
 PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
-OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/sparse_images")
+OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "sparse_images"
 
 
 def _psnr(a, b):
diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1
index 5d60bacf90..2de89637b6 100644
--- a/scripts/uninstall.ps1
+++ b/scripts/uninstall.ps1
@@ -315,7 +315,7 @@ function Uninstall-UnslothStudio {
         # 2. A loaded module under a target root (orphaned mp-fork python holding a
         #    venv DLL). Scoped to names that load our DLLs to keep the scan fast.
         try {
-            $cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli -ErrorAction SilentlyContinue
+            $cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli, sd-cli, sd-server -ErrorAction SilentlyContinue
             foreach ($proc in $cands) {
                 $hit = $false
                 try {
@@ -366,7 +366,7 @@ function Uninstall-UnslothStudio {
     _StopStudioProcesses -KnownRoots $knownRoots
     # Also stop anything holding a handle on the exact paths we delete (llama-server,
     # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused.
-    _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode))
+    _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultSdCpp, $defaultCache, $defaultNode))
 
     # ── Remove custom-root install trees ──
     _Step "Removing data and install directories..."
@@ -380,6 +380,9 @@ function Uninstall-UnslothStudio {
             continue
         }
         _RemovePath $r
+        # The native diffusion sibling (<custom root>.parent\stable-diffusion.cpp) is
+        # intentionally NOT removed: sd.cpp writes no owner marker and sits in the user's
+        # own parent dir, so auto-deleting it could destroy a user-managed clone.
     }
     # Default install dir (always at %USERPROFILE%\.unsloth\studio when present).
     if ($defaultStudioHome) { _RemovePath $defaultStudioHome }
diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh
index f68fd37131..36d974844d 100755
--- a/scripts/uninstall.sh
+++ b/scripts/uninstall.sh
@@ -217,8 +217,11 @@ _remove_path "$HOME/.unsloth/studio"
 # when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept.
 _remove_path "$HOME/.unsloth/llama.cpp"
 # Default-mode native diffusion (stable-diffusion.cpp / sd-cli) build, a sibling of
-# studio like llama.cpp (install_sd_cpp_prebuilt.default_install_dir()). No-op in
-# env/custom mode and when absent. A user-set UNSLOTH_SD_CPP_PATH is kept.
+# studio like llama.cpp (install_sd_cpp_prebuilt.default_install_dir()). Only the default
+# location is removed. In env/custom mode the install is <custom root>.parent/
+# stable-diffusion.cpp, which is intentionally left in place: sd.cpp writes no owner
+# marker and sits in the user's own parent dir, so auto-deleting it could destroy a
+# user-managed stable-diffusion.cpp clone. A user-set UNSLOTH_SD_CPP_PATH is kept.
 _remove_path "$HOME/.unsloth/stable-diffusion.cpp"
 _remove_path "$HOME/.unsloth/.cache"
 # Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in
diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index c426d3b5d2..a986b31bdd 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -408,9 +408,14 @@ class DiffusionBackend:
         the transformer), but ``_load_dense_quant_pipeline`` fetches them with
         ``from_pretrained(subfolder = "transformer")`` under the load lock during
         "finalizing", after the previous pipeline was already evicted, where
-        unload/cancellation cannot preempt the download. Mirrors the dense-path
-        gates in ``load_pipeline``: quant requested and supported for this device,
-        and no pre-quantized checkpoint that would shortcut the dense build."""
+        unload/cancellation cannot preempt the download. Checks the dense-path gates
+        in ``load_pipeline`` that are knowable pre-download: quant requested and
+        supported for this device, and no pre-quantized checkpoint that would shortcut
+        the dense build. It deliberately does NOT mirror the ``plan.offload_policy ==
+        OFFLOAD_NONE`` gate: the memory plan needs the GGUF's on-disk size, which
+        isn't known until the GGUF is cached (after this prefetch runs). So the
+        transformer/ shards can be prefetched for a load that the plan then routes to
+        offload -- they stay cached for a later resident load rather than being wasted."""
         mode = normalize_transformer_quant(kwargs.get("transformer_quant"))
         if mode is None:
             return False
@@ -1014,6 +1019,21 @@ class DiffusionBackend:
                         # GGUF build (the OOM-fallback path this cleanup exists for).
                         del exc
                         clear_gpu_cache()
+                elif (
+                    kind == "gguf"
+                    and normalize_transformer_quant(transformer_quant) is not None
+                    and dense_transformer_supported(target)
+                    and plan.offload_policy != OFFLOAD_NONE
+                ):
+                    # The dense fast path needs the transformer resident, so a memory_mode
+                    # (balanced / low_vram) that forces offload silently drops the requested
+                    # quant. Warn so the disengage is diagnosable rather than a null status.
+                    logger.warning(
+                        "diffusion.transformer_quant: %s requested but memory_mode forces "
+                        "offload (%s); loading GGUF without dense quant",
+                        normalize_transformer_quant(transformer_quant),
+                        plan.offload_policy,
+                    )
 
                 if pipe is None:
                     if kind == "pipeline":
@@ -1390,7 +1410,8 @@ class DiffusionBackend:
 
         The size estimate is per-kind: diffusers keeps GGUF weights packed (per-matmul
         transient dequant), so a GGUF loads near its on-disk size; a safetensors
-        single-file loads near its on-disk size (it carries its dtype); and a full
+        single-file loads near its on-disk size (it carries its dtype), except an fp8
+        transformer file that gets upcast to bf16 on load (~2x resident); and a full
         pipeline is one cached download (transformer + companions), already compressed."""
         device_memory = snapshot_device_memory(target)
         if kind == "pipeline":
@@ -1408,9 +1429,19 @@ class DiffusionBackend:
             companion_mib = None
         else:
             if kind == "single_file":
-                # Safetensors single-file: no dequant expansion (it carries its dtype).
+                # Safetensors single-file. A dense bf16 file loads near its on-disk size,
+                # but a transformer-only fp8 checkpoint is loaded via from_single_file with
+                # a bf16 compute dtype and NO quantization_config, so diffusers upcasts it
+                # fp8 -> bf16 (~2x resident). Detect fp8 from the basename and budget the
+                # expansion. The single-file-is-pipeline (SDXL) path is a full bf16 pipeline
+                # checkpoint, not this fp8 transformer path, so it stays at on-disk size.
+                fp8_upcast = not getattr(fam, "single_file_is_pipeline", False) and (
+                    "fp8" in Path(single_file_path).name.lower()
+                    if single_file_path
+                    else False
+                )
                 transformer_resident = estimate_safetensors_dense_mib(
-                    file_size_mib(single_file_path)
+                    file_size_mib(single_file_path), fp8_upcast = fp8_upcast
                 )
             else:
                 transformer_resident = estimate_gguf_resident_mib(file_size_mib(single_file_path))
@@ -1500,6 +1531,14 @@ class DiffusionBackend:
         if cn_model is None:
             if cancel.is_set():
                 raise RuntimeError(DIFFUSION_CANCELLED_MSG)
+            # A single generation uses exactly one ControlNet, so keep at most one resident:
+            # on a miss for a new id, drop the previously-cached module + its from_pipe wrapper
+            # (both dicts, kept consistent) and free the VRAM before loading the new one, or
+            # swapping distinct ControlNets within a base-model load accumulates until OOM.
+            if self._cn_models or self._cn_pipes:
+                self._cn_models.clear()
+                self._cn_pipes.clear()
+                clear_gpu_cache()
             import torch
 
             # state.dtype is the display string saved at load ("bfloat16"), NOT a
@@ -1784,6 +1823,13 @@ class DiffusionBackend:
                         raise ValueError(
                             f"{state.family.name} is an image-editing model: provide an input image."
                         )
+                    if mask_image is not None:
+                        # The edit family has no inpaint pipeline; a supplied mask would be
+                        # silently dropped (this branch wins over the inpaint branch below).
+                        raise ValueError(
+                            f"{state.family.name} is an image-editing model and does not "
+                            "support masks (mask_image)."
+                        )
                     workflow = "edit"
                     init_pil = _decode_b64_image(init_image, mode = "RGB")
                 elif mask_image is not None and init_image is not None:
diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py
index b58b1f84e2..1c5dff9580 100644
--- a/studio/backend/core/inference/diffusion_device.py
+++ b/studio/backend/core/inference/diffusion_device.py
@@ -265,7 +265,9 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
 
 
 def _cpu_target(torch: Any, dtype: Any = None) -> DiffusionDeviceTarget:
-    if dtype is None:
+    # torch is None on the no-torch CPU fallback; leave dtype=None then (matching the
+    # no-torch DiffusionDeviceTarget elsewhere) rather than crashing on torch.float32.
+    if dtype is None and torch is not None:
         dtype = torch.float32
     return DiffusionDeviceTarget(
         device = "cpu",
diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py
index cad6d3eae7..bd3e2d3784 100644
--- a/studio/backend/core/inference/diffusion_lora.py
+++ b/studio/backend/core/inference/diffusion_lora.py
@@ -266,6 +266,17 @@ def _pick_repo_weight_file(repo_id: str, hf_token: Optional[str]) -> str:
     raise FileNotFoundError(f"no .safetensors/.gguf LoRA file found in '{repo_id}'")
 
 
+def _scrub_hub_url(msg: str) -> str:
+    """Strip embedded http(s) URLs from a Hub error message before it hits a 400 body.
+
+    huggingface_hub errors interpolate the request URL (and a request id) into their
+    message; a raw endpoint URL is noise in a client-facing 400, so drop it.
+    """
+    cleaned = re.sub(r"https?://\S+", "", msg)
+    # Collapse the whitespace / stray separators the URL removal leaves behind.
+    return re.sub(r"\s{2,}", " ", cleaned).strip()
+
+
 def resolve_specs(
     specs: list[tuple[str, float]],
     *,
@@ -274,20 +285,41 @@ def resolve_specs(
 ) -> list[ResolvedLora]:
     """Resolve request (id, weight) pairs, dropping zero-weight entries.
 
-    A stale / unknown id raises FileNotFoundError inside resolve_one; convert it to
-    ValueError so the route (which maps only ValueError to a 400) reports bad client
-    input instead of a generic 500. A Hub download can also raise
-    ``RuntimeError("Cancelled")`` when the user unloads / starts a superseding load
-    mid-download; convert that to the diffusion cancellation sentinel so the route
-    maps it to a 409 instead of a generic server error toast."""
+    A stale / unknown id raises FileNotFoundError inside resolve_one; a mistyped Hub
+    repo id makes the Hub resolution raise a huggingface_hub client error (a missing
+    repo -> RepositoryNotFoundError, a bad revision -> RevisionNotFoundError, a missing
+    weight file -> EntryNotFoundError, a gated model -> GatedRepoError). Convert those
+    NAMED not-found/gated errors to ValueError so the route (which maps only ValueError
+    to a 400) reports bad client input instead of a generic 500 -- Hub error messages
+    embed the request URL, so scrub it out before it reaches the 400 body. Catch them by
+    name rather than their common HfHubHTTPError base on purpose: a Hub-side 5xx / 429
+    (an outage, not bad input) is a bare HfHubHTTPError and must stay a 500. A Hub
+    download can also raise ``RuntimeError("Cancelled")`` when the user unloads / starts a
+    superseding load mid-download; convert that to the diffusion cancellation sentinel so
+    the route maps it to a 409 instead of a generic server error toast. A non-cancellation
+    RuntimeError (e.g. a stalled download, disk full) stays a 500 -- it is not bad
+    client input."""
+    from huggingface_hub.errors import (
+        EntryNotFoundError,
+        GatedRepoError,
+        RepositoryNotFoundError,
+        RevisionNotFoundError,
+    )
+
     out: list[ResolvedLora] = []
     try:
         for spec_id, weight in specs:
             if weight == 0:
                 continue
             out.append(resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event))
-    except FileNotFoundError as exc:
-        raise ValueError(str(exc)) from exc
+    except (
+        FileNotFoundError,
+        RepositoryNotFoundError,
+        RevisionNotFoundError,
+        EntryNotFoundError,
+        GatedRepoError,
+    ) as exc:
+        raise ValueError(_scrub_hub_url(str(exc))) from exc
     except RuntimeError as exc:
         if str(exc) == "Cancelled":
             raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc
diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py
index c56d1592ce..090e295a2e 100644
--- a/studio/backend/core/inference/diffusion_memory.py
+++ b/studio/backend/core/inference/diffusion_memory.py
@@ -236,14 +236,27 @@ def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]:
     return int(storage_mib * 1.05)  # small margin for allocator + bf16 norms/biases
 
 
-def estimate_safetensors_dense_mib(storage_mib: Optional[int]) -> Optional[int]:
+def estimate_safetensors_dense_mib(
+    storage_mib: Optional[int], *, fp8_upcast: bool = False
+) -> Optional[int]:
     """Resident size of a safetensors checkpoint, in MiB.
 
     Unlike a GGUF (which is dequantised to bf16/fp16 on load, so a 4-bit file
-    expands ~4x), a safetensors checkpoint loads near its on-disk size: a dense
-    bf16 file is already bf16, and a bnb-4bit / fp8 file stays compressed in VRAM.
-    So the on-disk size is the estimate, returned unchanged (None passes through).
+    expands ~4x), a safetensors checkpoint usually loads near its on-disk size: a
+    dense bf16 file is already bf16, and a bnb-4bit file stays compressed in VRAM
+    (it carries its own quantization_config). So the on-disk size is the estimate,
+    returned unchanged (None passes through).
+
+    The exception is ``fp8_upcast``: the fp8 single-file transformer path loads via
+    ``from_single_file`` with a bf16 compute dtype and NO quantization_config, so
+    diffusers upcasts the fp8 weights (1 byte/param) to bf16 (2 bytes/param) --
+    roughly 2x the on-disk bytes resident. Budget that, or the plan under-reserves
+    and OOMs.
     """
+    if storage_mib is None:
+        return None
+    if fp8_upcast:
+        return storage_mib * 2
     return storage_mib
 
 
diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index d8cb21d409..6c54af5e87 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -737,6 +737,13 @@ class SdCppDiffusionBackend:
                 "img2img / inpaint / reference / upscale are not yet supported on the native "
                 "sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows."
             )
+        # strength 0 (or None) disables ControlNet -- documented on the request model, and
+        # the diffusers path treats it as plain txt2img -- so a strength-0 spec must be a
+        # no-op here too, not a hard 400. Only a genuinely active (strength > 0) ControlNet
+        # is rejected. Strength is element 3 of the tuple
+        # (id, image, type, strength, guidance_start, guidance_end).
+        if controlnet is not None and controlnet[3] in (None, 0, 0.0):
+            controlnet = None
         if controlnet is not None:
             raise ValueError(
                 "ControlNet is not yet supported on the native sd.cpp engine; run on a GPU "
diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index d4ddf9d6d4..03f886155d 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -32,7 +32,6 @@ from typing import Any, Callable, Optional
 
 from core.training.diffusion_train_common import (
     DEFAULT_LORA_FILENAME,
-    DEFAULT_LORA_TARGETS,
     DiffusionLoraConfig,
     EventCb,
     StopCb,
@@ -64,11 +63,11 @@ def _select_lora_targets(
 ) -> tuple[str, ...]:
     """Pick the LoRA target modules for a DiT run.
 
-    ``normalized()`` always fills ``lora_target_modules`` with the generic
-    ``DEFAULT_LORA_TARGETS`` when a caller does not set it, so that value means "unset"
-    here: prefer the family's ``spec.lora_targets`` (which add the DiT-specific
-    projections). Any OTHER explicit tuple is a deliberate override and still wins."""
-    if tuple(cfg_targets) == DEFAULT_LORA_TARGETS:
+    ``normalized()`` leaves ``lora_target_modules`` empty when a caller does not set it, so
+    an empty tuple means "unset" here: use the family's ``spec.lora_targets`` (which add the
+    DiT-specific joint-attention projections). Any explicit tuple is a deliberate override
+    and still wins."""
+    if not tuple(cfg_targets):
         return tuple(spec_targets)
     return tuple(cfg_targets)
 
diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py
index deef4d03ce..be006061a7 100644
--- a/studio/backend/core/training/diffusion_lora_trainer.py
+++ b/studio/backend/core/training/diffusion_lora_trainer.py
@@ -211,13 +211,15 @@ def run_diffusion_lora_training(
     for m in (unet, *text_encoders):
         m.to(device, dtype = weight_dtype)
 
+    # An empty (unset) config means "use the family default": the SDXL attention projections.
+    unet_targets = list(cfg.lora_target_modules) or list(DEFAULT_LORA_TARGETS)
     unet.add_adapter(
         LoraConfig(
             r = cfg.lora_rank,
             lora_alpha = cfg.lora_alpha,
             lora_dropout = cfg.lora_dropout,
             init_lora_weights = "gaussian",
-            target_modules = list(cfg.lora_target_modules),
+            target_modules = unet_targets,
         )
     )
     if cfg.gradient_checkpointing:
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index fe077020ac..d854cc4173 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -30,11 +30,32 @@ from core.inference.diffusion_families import (
     trainable_family_names,
 )
 
-# Default LoRA target modules: the attention projections common to the SDXL U-Net and the
-# DiT transformers (the diffusers/kohya convention). A family whose trainer wants a wider
-# set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback.
+# Default LoRA target modules: the attention projections of the SDXL U-Net (the
+# diffusers/kohya convention). Used by the SDXL trainer as its fallback when the config
+# leaves ``lora_target_modules`` empty; the DiT trainers supply their own wider set. Kept
+# here so the SDXL trainer has a named default even for an empty (unset) config.
 DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
 
+# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). Validated in
+# normalized() so a typo fails fast at request time, not minutes later in the subprocess.
+_LR_SCHEDULERS: frozenset[str] = frozenset(
+    {
+        "linear",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+        "constant",
+        "constant_with_warmup",
+        "piecewise_constant",
+    }
+)
+
+# DiT families that overflow fp16 (their RoPE / embedder run in fp32), so they train in bf16
+# only. Encoded here -- keyed by resolved family -- so normalized() can reject an fp16 request
+# before spawn without importing the DiT trainer's _SPECS (which would create an import
+# cycle). The DiT trainer keeps a matching guard as defense in depth.
+_FORCE_BF16_FAMILIES: frozenset[str] = frozenset({"qwen-image", "z-image"})
+
 _IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
 _CAPTION_EXTS = (".txt", ".caption")
 # diffusers' canonical single-file LoRA name, so load_lora_weights(dir) finds it.
@@ -83,7 +104,15 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
     # GGUF weights (a ``.gguf`` file or a ``*-GGUF`` repo) are inference-only: training needs
     # the full diffusers pipeline (transformer + VAE + text encoders), which a GGUF repo does
     # not provide. Reject by name even when the family itself is trainable.
-    if name.endswith(".gguf") or "gguf" in name:
+    # A ``.gguf`` file always rejects. The broad ``"gguf" in name`` catch (for ``*-GGUF``
+    # repos) must NOT reject a real local diffusers directory that merely has "gguf" in its
+    # path, so it is skipped for a local diffusers checkout -- identified by its
+    # ``model_index.json`` marker (the same marker the loader uses), NOT a bare ``is_dir()``:
+    # a GGUF-only folder must still reject here and fail fast, rather than pass and fail late
+    # in the subprocess after the resident chat/Images models were already evicted.
+    local = Path(base_model).expanduser() if base_model else None
+    is_local_diffusers = bool(local and (local / "model_index.json").is_file())
+    if name.endswith(".gguf") or ("gguf" in name and not is_local_diffusers):
         raise ValueError(
             f"'{base_model}' is a GGUF checkpoint/repo, which can't be a training base "
             f"(training needs the full diffusers model). {_trainable_hint()}"
@@ -215,7 +244,9 @@ class DiffusionLoraConfig:
     lora_rank: int = 16
     lora_alpha: Optional[int] = None  # defaults to lora_rank
     lora_dropout: float = 0.0
-    lora_target_modules: tuple[str, ...] = DEFAULT_LORA_TARGETS
+    # Empty = "unset": each trainer supplies its family default (SDXL DEFAULT_LORA_TARGETS,
+    # or the DiT family's wider joint-attention set). A non-empty tuple is an explicit override.
+    lora_target_modules: tuple[str, ...] = ()
     seed: int = 42
     mixed_precision: str = "bf16"  # "bf16" | "fp16" | "no"
     snr_gamma: Optional[float] = 5.0  # min-SNR loss weighting; None disables
@@ -259,6 +290,19 @@ 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 bf16-only DiT family (Qwen-Image / Z-Image) must refuse fp16 up front rather than
+        # accepting the request, evicting resident models, and only then failing in the
+        # subprocess. The DiT trainer keeps a matching guard as defense in depth.
+        if self.mixed_precision == "fp16" and resolved_family in _FORCE_BF16_FAMILIES:
+            raise ValueError(
+                f"'{resolved_family}' LoRA training requires bf16: fp16 overflows its fp32 "
+                f"RoPE / embedder internals. Set mixed precision to bf16."
+            )
+        if str(self.lr_scheduler) not in _LR_SCHEDULERS:
+            raise ValueError(
+                f"lr_scheduler must be one of {', '.join(sorted(_LR_SCHEDULERS))}; "
+                f"got {self.lr_scheduler!r}"
+            )
         # 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:
@@ -268,7 +312,9 @@ class DiffusionLoraConfig:
         if learning_rate <= 0:
             raise ValueError("learning_rate must be > 0")
         alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank
-        targets = tuple(self.lora_target_modules) or DEFAULT_LORA_TARGETS
+        # Leave an unset (empty) target list empty: the trainer fills the family default
+        # (SDXL DEFAULT_LORA_TARGETS, or the DiT family's wider set) so the family spec wins.
+        targets = tuple(self.lora_target_modules)
         # A blank Hub token (the Studio default when none is configured) must load
         # anonymously, not as an explicit empty credential.
         token = self.hf_token.strip() if isinstance(self.hf_token, str) else self.hf_token
@@ -390,8 +436,20 @@ def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Option
             else Path(cfg.output_dir).name
         )
         alias = sanitize_alias(base)
+        src_resolved = Path(lora_path).resolve()
         dest = loras_dir() / f"{alias}.safetensors"
-        if Path(lora_path).resolve() != dest.resolve():
+        # A retrain with the same adapter name must not clobber a prior mirror: if the
+        # destination already exists and is a different file, pick the next free numeric
+        # suffix (<alias>-2, <alias>-3, ...) for both the weights and their .json sidecar.
+        if dest.exists() and dest.resolve() != src_resolved:
+            n = 2
+            while True:
+                candidate = loras_dir() / f"{alias}-{n}.safetensors"
+                if not candidate.exists() or candidate.resolve() == src_resolved:
+                    dest = candidate
+                    break
+                n += 1
+        if src_resolved != dest.resolve():
             shutil.copy2(lora_path, dest)
         _write_lora_sidecar(dest.with_suffix(".json"), cfg)
         return str(dest)
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 2479571357..c9606eb5b1 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -698,18 +698,27 @@ class DiffusionTrainingStartRequest(BaseModel):
     lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank")
     lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0)
     # Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that
-    # sets them is not silently trained with defaults. Default the target list to the SDXL
-    # attention projections (the trainer's DEFAULT_LORA_TARGETS) so it is never None.
+    # sets them is not silently trained with defaults. Empty (the default) means "unset": each
+    # trainer supplies its own family targets -- the SDXL DEFAULT_LORA_TARGETS for SDXL, the
+    # wider joint-attention set for the DiT families. A non-empty list is an explicit override.
     lora_target_modules: List[str] = Field(
-        default_factory = lambda: ["to_k", "to_q", "to_v", "to_out.0"],
-        description = "U-Net modules to attach LoRA to",
+        default_factory = list,
+        description = "Modules to attach LoRA to; empty = the trainer's family default",
     )
     max_grad_norm: float = Field(1.0, gt = 0, description = "Gradient clipping max-norm")
     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")
     gradient_checkpointing: bool = Field(True)
-    lr_scheduler: str = Field("constant")
+    lr_scheduler: Literal[
+        "linear",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+        "constant",
+        "constant_with_warmup",
+        "piecewise_constant",
+    ] = Field("constant")
     lr_warmup_steps: int = Field(0, ge = 0)
     center_crop: bool = Field(False)
     random_flip: bool = Field(True)
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index e61042f1f6..8868eb1e2a 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1375,6 +1375,9 @@ async def upload_diffusion_dataset(
     total_bytes = 0
     uploaded = 0
     allowed = _DIFFUSION_DATASET_IMAGE_EXTS | _DIFFUSION_DATASET_TEXT_EXTS
+    # Validate every filename up front so a valid image ahead of a bad one is not left
+    # written on disk when the 400 fires -- make the upload all-or-nothing.
+    names: list[str] = []
     for f in files:
         filename = Path(f.filename or "").name.strip().replace("\x00", "")
         ext = Path(filename).suffix.lower()
@@ -1384,9 +1387,17 @@ async def upload_diffusion_dataset(
                 status_code = 400,
                 detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
             )
-        dest = folder / filename
-        complete = False
-        try:
+        names.append(filename)
+    # Roll back every file written this request if the batch does not fully commit, so a
+    # mid-batch 413 (or a disk error / client disconnect) leaves the dataset unchanged
+    # rather than partially populated -- the upload is all-or-nothing, not just for the
+    # extension check above but for the size limit too.
+    written: list[Path] = []
+    committed = False
+    try:
+        for f, filename in zip(files, names):
+            dest = folder / filename
+            written.append(dest)
             with open(dest, "wb") as out:
                 while chunk := await f.read(1024 * 1024):
                     total_bytes += len(chunk)
@@ -1400,14 +1411,15 @@ async def upload_diffusion_dataset(
                             ),
                         )
                     out.write(chunk)
-            complete = True
-        finally:
-            if not complete:
+            uploaded += 1
+        committed = True
+    finally:
+        if not committed:
+            for p in written:
                 try:
-                    dest.unlink(missing_ok = True)
+                    p.unlink(missing_ok = True)
                 except OSError:
                     pass
-        uploaded += 1
 
     summary = _diffusion_dataset_summary(folder)
     return DiffusionDatasetUploadResponse(
@@ -1572,7 +1584,7 @@ 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"
+        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
@@ -1645,7 +1657,7 @@ 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"):
+            for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"):
                 t.unlink(missing_ok = True)
         return {"deleted": image_path.name}
 
@@ -1852,7 +1864,7 @@ def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
     )
     # Map basename -> caption from every jsonl carrying file_name + caption column.
     captions: dict[str, str] = {}
-    for jf in snap.rglob("*.jsonl"):
+    for jf in sorted(snap.rglob("*.jsonl")):
         for line in jf.read_text(encoding = "utf-8").splitlines():
             line = line.strip()
             if not line:
@@ -1863,7 +1875,9 @@ def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
                 continue
             fn = row.get("file_name") or row.get("image") or row.get("file")
             if fn and caption_col in row:
-                captions[Path(str(fn)).name] = str(row[caption_col])
+                # First writer wins over sorted manifests, so the plain manifest
+                # (e.g. output_file.jsonl) is deterministic rather than OS-visit order.
+                captions.setdefault(Path(str(fn)).name, str(row[caption_col]))
     # Copy images (those with a caption first, so a cap keeps captioned pairs).
     images = sorted(
         p
diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py
index b10dd04f4c..154b671215 100644
--- a/studio/backend/tests/test_diffusion_dataset_api.py
+++ b/studio/backend/tests/test_diffusion_dataset_api.py
@@ -179,13 +179,15 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
     (folder / "x.txt").write_text("cap", encoding = "utf-8")
     # Generate a thumbnail so we can assert it is cleaned up too.
     client.get("/api/train/diffusion/dataset/d/image/x.png?thumb=32")
-    assert list((folder / ".thumbs").glob("x_*.jpg"))
+    # Thumb cache key includes the extension (x.png_32.jpg), so png and jpg
+    # siblings can't collide.
+    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"))
 
 
 # ── traversal / validation ───────────────────────────────────────────────────
diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py
index 63cc01922c..79a3777113 100644
--- a/studio/backend/tests/test_diffusion_dit_trainer.py
+++ b/studio/backend/tests/test_diffusion_dit_trainer.py
@@ -41,23 +41,25 @@ def test_specs_cover_the_three_dit_families():
 
 
 def test_select_lora_targets_uses_family_default_for_generic_config():
-    # normalized() fills lora_target_modules with the generic DEFAULT_LORA_TARGETS when a
-    # caller doesn't set it, so that value must resolve to the family's targets (which add
-    # the DiT-specific projections), not stay stuck on the generic SDXL list.
-    assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == _FLUX_TARGETS
-    assert _select_lora_targets(DEFAULT_LORA_TARGETS, _QWEN_TARGETS) == _QWEN_TARGETS
-    assert _select_lora_targets(DEFAULT_LORA_TARGETS, _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS
+    # normalized() leaves lora_target_modules empty when a caller doesn't set it, so an empty
+    # tuple must resolve to the family's targets (which add the DiT-specific joint-attention
+    # projections), not the generic SDXL list.
+    assert _select_lora_targets((), _FLUX_TARGETS) == _FLUX_TARGETS
+    assert _select_lora_targets((), _QWEN_TARGETS) == _QWEN_TARGETS
+    assert _select_lora_targets((), _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS
+    # The generic SDXL default is NOT treated as unset here: an explicit list wins.
+    assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == DEFAULT_LORA_TARGETS
 
 
 def test_select_lora_targets_explicit_override_wins():
-    # Any OTHER explicit tuple is a deliberate override and must win over the family spec.
+    # Any explicit tuple is a deliberate override and must win over the family spec.
     override = ("to_q", "to_k")
     assert _select_lora_targets(override, _FLUX_TARGETS) == override
-    # The default request path (config carrying the generic default) reaches the spec.
+    # The default request path (config leaving targets unset) reaches the spec.
     cfg = DiffusionLoraConfig(
         base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o"
     ).normalized()
-    assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
+    assert cfg.lora_target_modules == ()
     assert (
         _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets)
         == _FLUX_TARGETS
diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py
index 2743546255..a40807a3a2 100644
--- a/studio/backend/tests/test_diffusion_lora.py
+++ b/studio/backend/tests/test_diffusion_lora.py
@@ -189,6 +189,26 @@ def test_resolve_specs_maps_unknown_id_to_valueerror(tmp_path, monkeypatch):
         dl.resolve_specs([("nope", 1.0)])
 
 
+def test_resolve_specs_maps_hub_error_to_valueerror(tmp_path, monkeypatch):
+    # A mistyped Hub repo id makes the Hub resolution raise a huggingface_hub client error
+    # (RepositoryNotFoundError, an HfHubHTTPError). resolve_specs must surface it as
+    # ValueError so the route returns 400, not a generic 500. The Hub message embeds the
+    # request URL, which must be scrubbed out of the client-facing 400.
+    from huggingface_hub.errors import RepositoryNotFoundError
+
+    def _boom(spec_id, weight, **kw):
+        raise RepositoryNotFoundError(
+            "404 Client Error. Repository Not Found for url: "
+            "https://huggingface.co/api/models/nope/nope (Request ID: abc)"
+        )
+
+    monkeypatch.setattr(dl, "resolve_one", _boom)
+    with pytest.raises(ValueError) as ei:
+        dl.resolve_specs([("nope/nope", 1.0)])
+    assert "http" not in str(ei.value)  # request URL scrubbed
+    assert "Repository Not Found" in str(ei.value)
+
+
 def test_scan_local_disambiguates_identical_stems(tmp_path, monkeypatch):
     # foo.safetensors and foo.gguf must get distinct ids so each is addressable; a
     # unique stem keeps its clean stem id.
diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py
index 75046165d0..dc84b70f17 100644
--- a/studio/backend/tests/test_diffusion_lora_trainer.py
+++ b/studio/backend/tests/test_diffusion_lora_trainer.py
@@ -94,7 +94,8 @@ def test_discover_missing_dir_raises(tmp_path):
 def test_config_normalized_defaults():
     cfg = DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o").normalized()
     assert cfg.lora_alpha == cfg.lora_rank  # alpha defaults to rank
-    assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
+    # Targets stay empty (unset) after normalize; each trainer fills its own family default.
+    assert cfg.lora_target_modules == ()
 
 
 @pytest.mark.parametrize(
@@ -335,3 +336,85 @@ def test_publish_writes_metadata_sidecar(tmp_path, monkeypatch):
     assert meta["lora_rank"] == 8
     assert meta["trigger_prompt"] == "a photo in sks style"
     assert meta["source"] == "studio-trained"
+
+
+def test_publish_does_not_clobber_same_name_adapter(tmp_path, monkeypatch):
+    # A retrain with the same adapter name must not overwrite a prior mirror: the second
+    # publish lands under a numeric suffix (my-style -> my-style-2), sidecar alongside it.
+    from pathlib import Path
+
+    from core.inference import diffusion_lora
+    from core.training.diffusion_lora_trainer import _publish_to_lora_catalog
+
+    loras = tmp_path / "loras"
+    loras.mkdir()
+    monkeypatch.setattr(diffusion_lora, "loras_dir", lambda: loras)
+
+    def _publish(payload: bytes) -> str:
+        src = tmp_path / "run" / "pytorch_lora_weights.safetensors"
+        src.parent.mkdir(parents = True, exist_ok = True)
+        src.write_bytes(payload)
+        cfg = DiffusionLoraConfig(
+            base_model = "stabilityai/sdxl-turbo",
+            data_dir = "d",
+            output_dir = str(tmp_path / "run"),
+            adapter_name = "my-style",
+        ).normalized()
+        return _publish_to_lora_catalog(str(src), cfg)
+
+    first = _publish(b"adapter-v1")
+    second = _publish(b"adapter-v2")
+    assert Path(first).name == "my-style.safetensors"
+    assert Path(second).name == "my-style-2.safetensors"
+    # The first mirror is intact (not clobbered) and the second is the new content.
+    assert Path(first).read_bytes() == b"adapter-v1"
+    assert Path(second).read_bytes() == b"adapter-v2"
+    assert Path(second).with_suffix(".json").is_file()
+
+
+def test_config_rejects_bad_lr_scheduler():
+    # A typo'd scheduler ('constnat') must fail at normalize time, not later in the subprocess.
+    with pytest.raises(ValueError, match = "lr_scheduler"):
+        DiffusionLoraConfig(
+            base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "constnat"
+        ).normalized()
+    # A valid diffusers scheduler passes.
+    cfg = DiffusionLoraConfig(
+        base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "cosine"
+    ).normalized()
+    assert cfg.lr_scheduler == "cosine"
+
+
+def test_config_rejects_fp16_on_bf16_only_family():
+    # qwen-image / z-image are bf16-only: an fp16 request must be rejected before spawn,
+    # in normalized(), not only by the subprocess-side guard.
+    for base in ("Tongyi-MAI/Z-Image-Turbo", "unsloth/Qwen-Image-2512-unsloth-bnb-4bit"):
+        with pytest.raises(ValueError, match = "bf16"):
+            DiffusionLoraConfig(
+                base_model = base, data_dir = "d", output_dir = "o", mixed_precision = "fp16"
+            ).normalized()
+    # FLUX (not force-bf16) still accepts fp16.
+    cfg = DiffusionLoraConfig(
+        base_model = "black-forest-labs/FLUX.1-dev",
+        data_dir = "d",
+        output_dir = "o",
+        mixed_precision = "fp16",
+    ).normalized()
+    assert cfg.mixed_precision == "fp16"
+
+
+def test_gguf_substring_does_not_reject_local_diffusers_dir(tmp_path):
+    # A local diffusers directory whose path merely contains 'gguf' is a valid training base
+    # (it carries model_index.json, not GGUF weights); the broad substring must not reject it.
+    from core.training.diffusion_train_common import resolve_trainable_family
+
+    local = tmp_path / "my-gguf-experiments" / "sdxl-finetune"
+    local.mkdir(parents = True)
+    (local / "model_index.json").write_text("{}", encoding = "utf-8")
+    assert resolve_trainable_family(str(local)) == "sdxl"
+    # A real .gguf file still rejects even inside such a dir.
+    with pytest.raises(ValueError, match = "GGUF"):
+        resolve_trainable_family(str(local / "weights.gguf"))
+    # A *-GGUF repo id (not a local dir) still rejects.
+    with pytest.raises(ValueError, match = "GGUF"):
+        resolve_trainable_family("unsloth/FLUX.1-dev-GGUF")
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index a7ef2240f4..92429d60f8 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -492,6 +492,24 @@ def test_diffusion_dataset_upload_rejects_unsupported_files(client, dataset_root
     assert "Unsupported file" in r.json()["detail"]
 
 
+def test_diffusion_dataset_upload_over_cap_rolls_back_whole_batch(client, dataset_roots, monkeypatch):
+    # All-or-nothing: a valid image ahead of the one that trips the size cap must NOT be
+    # left on disk (the 413 mid-batch rolls back every file written this request).
+    import utils.upload_limits as ul
+
+    monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 100)
+    ds_root, _ = dataset_roots
+    files = [
+        ("files", ("small.png", b"x" * 50, "image/png")),
+        ("files", ("big.png", b"y" * 200, "image/png")),
+    ]
+    r = client.post("/api/train/diffusion/dataset", data = {"name": "rollback"}, files = files)
+    assert r.status_code == 413, r.text
+    folder = ds_root / "rollback"
+    assert not (folder / "small.png").exists()  # the earlier valid file was rolled back
+    assert not (folder / "big.png").exists()
+
+
 def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypatch):
     # A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed,
     # so a bad pick never unloads the user's working chat/Images model.
diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py
index be2fc22a4b..68caeb58dd 100644
--- a/studio/backend/tests/test_sd_cpp_backend.py
+++ b/studio/backend/tests/test_sd_cpp_backend.py
@@ -771,6 +771,22 @@ def test_generate_rejects_controlnet_on_native_engine():
         b.generate(prompt = "x", steps = 4, seed = 1, controlnet = ("id", "img", "canny", 1.0, 0.0, 1.0))
 
 
+@pytest.mark.parametrize("cn_strength", [0, 0.0, None])
+def test_generate_treats_zero_strength_controlnet_as_disabled(cn_strength):
+    # strength 0 (or None) disables ControlNet -- the diffusers path treats it as plain
+    # txt2img and the request model documents it -- so a strength-0 spec must succeed on the
+    # native engine too, not 400. Only a genuinely active (strength > 0) ControlNet is rejected.
+    eng = _FakeEngine()
+    b = _loaded_backend(engine = eng)
+    out = b.generate(
+        prompt = "x",
+        steps = 4,
+        seed = 1,
+        controlnet = ("id", "img", "canny", cn_strength, 0.0, 1.0),
+    )
+    assert len(out["images"]) == 1
+
+
 def test_generate_rejects_image_conditioned_on_native_engine():
     # img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call
     # with an init image on the native engine gets a clean ValueError, not a silent txt2img.
diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx
index e712f24c21..0bf4179a25 100644
--- a/studio/frontend/src/features/images/images-page.tsx
+++ b/studio/frontend/src/features/images/images-page.tsx
@@ -1237,7 +1237,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
   // Load an image's recipe back into the form inputs.
   const restoreSettings = useCallback((image: GalleryImage) => {
     setPrompt(image.prompt);
-    setNegativePrompt(image.negative_prompt ?? "");
+    // The Negative-prompt field only renders (and submits) when guidance>0, so a
+    // guidance=0 recipe must not restore a hidden negative prompt that would
+    // resurface if guidance is later raised. Mirror the submit-path gating.
+    setNegativePrompt(image.guidance > 0 ? (image.negative_prompt ?? "") : "");
     setSteps(image.steps);
     setGuidance(image.guidance);
     setSeed(String(image.seed));
@@ -1272,6 +1275,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
       if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight });
     }
     setLoras(restoredLoras);
+    // The recipe carries no control image (it isn't persisted), so a faithful
+    // restore can't reproduce a ControlNet run -- clear any stale form selection
+    // rather than leaking it into the restored recipe, mirroring the LoRA clear.
+    setControlnetId("");
+    setControlImage(null);
     toast.success("Settings restored to inputs");
   }, []);
 
@@ -1520,11 +1528,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
       // Curated non-GGUF model: load as a full pipeline or single-file safetensors.
       const spec = SAFETENSORS_MODELS[id];
       if (spec) {
+        // Optimistically drop the quant label, but revert if the load never starts
+        // so a failed load doesn't leave a stale 'GGUF · variant' label (mirrors the
+        // GGUF branches below; the poll owns the after-start revert via quantRevert).
+        const prevQuant = quant;
+        quantRevert.current = { prev: prevQuant };
         setQuant(null);
         const d = defaultsFor(id);
         setSteps(d.steps);
         setGuidance(d.guidance);
-        void handleLoad(id, { kind: spec.kind, filename: spec.filename });
+        void handleLoad(id, { kind: spec.kind, filename: spec.filename }).then((started) => {
+          if (!started) {
+            setQuant(prevQuant);
+            quantRevert.current = null;
+          }
+        });
         return;
       }
       // GGUF quant pick from the variant expander. Optimistic for instant picker
@@ -1582,11 +1600,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         toast.error("Only unsloth or on-device image models can be loaded here");
         return;
       }
+      // Optimistically drop the quant label, but revert if the load never starts so
+      // a failed load doesn't leave a stale 'GGUF · variant' label (mirrors the GGUF
+      // branches above; the poll owns the after-start revert via quantRevert).
+      const prevQuant = quant;
+      quantRevert.current = { prev: prevQuant };
       setQuant(null);
       const d = defaultsFor(id);
       setSteps(d.steps);
       setGuidance(d.guidance);
-      void handleLoad(id, { kind: "pipeline" });
+      void handleLoad(id, { kind: "pipeline" }).then((started) => {
+        if (!started) {
+          setQuant(prevQuant);
+          quantRevert.current = null;
+        }
+      });
     },
     [busy, handleLoad, quant],
   );

From 4161d5d87a7f063c48a86bcfd94d2bfa94c6ae5d Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sat, 4 Jul 2026 11:09:37 -0300
Subject: [PATCH 03/13] Dedup the HF hub cache-dir path construction into a
 helper

---
 studio/backend/core/inference/diffusion.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 8120f544ad..16683b4b6e 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -806,10 +806,15 @@ class DiffusionBackend:
         return total, base_files
 
     @staticmethod
-    def _cache_bytes(repo_id: str) -> int:
+    def _hub_cache_repo_dir(repo_id: str) -> Path:
+        """The local HF hub cache dir for ``repo_id`` (``.../models--org--name``)."""
         from huggingface_hub import constants
 
-        blobs = Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" / "blobs"
+        return Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}"
+
+    @staticmethod
+    def _cache_bytes(repo_id: str) -> int:
+        blobs = DiffusionBackend._hub_cache_repo_dir(repo_id) / "blobs"
         total = 0
         try:
             for entry in blobs.iterdir():
@@ -861,9 +866,7 @@ class DiffusionBackend:
         local = Path(base).expanduser()
         if local.is_dir():
             return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True)
-        from huggingface_hub import constants
-
-        snapshots = Path(constants.HF_HUB_CACHE) / f"models--{base.replace('/', '--')}" / "snapshots"
+        snapshots = DiffusionBackend._hub_cache_repo_dir(base) / "snapshots"
         if not snapshots.is_dir():
             return 0
         # Multiple revisions may be cached; the active one is the fullest, so take the max.

From 3bb3734879bb95fa17744f64a29c35edfb641563 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 4 Jul 2026 14:10:56 +0000
Subject: [PATCH 04/13] [pre-commit.ci] auto fixes from pre-commit.com hooks

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

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 16683b4b6e..fa3e168c48 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -809,7 +809,6 @@ class DiffusionBackend:
     def _hub_cache_repo_dir(repo_id: str) -> Path:
         """The local HF hub cache dir for ``repo_id`` (``.../models--org--name``)."""
         from huggingface_hub import constants
-
         return Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}"
 
     @staticmethod
@@ -1448,9 +1447,7 @@ class DiffusionBackend:
                 # expansion. The single-file-is-pipeline (SDXL) path is a full bf16 pipeline
                 # checkpoint, not this fp8 transformer path, so it stays at on-disk size.
                 fp8_upcast = not getattr(fam, "single_file_is_pipeline", False) and (
-                    "fp8" in Path(single_file_path).name.lower()
-                    if single_file_path
-                    else False
+                    "fp8" in Path(single_file_path).name.lower() if single_file_path else False
                 )
                 transformer_resident = estimate_safetensors_dense_mib(
                     file_size_mib(single_file_path), fp8_upcast = fp8_upcast
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index 92429d60f8..2cff7f65c3 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -492,7 +492,9 @@ def test_diffusion_dataset_upload_rejects_unsupported_files(client, dataset_root
     assert "Unsupported file" in r.json()["detail"]
 
 
-def test_diffusion_dataset_upload_over_cap_rolls_back_whole_batch(client, dataset_roots, monkeypatch):
+def test_diffusion_dataset_upload_over_cap_rolls_back_whole_batch(
+    client, dataset_roots, monkeypatch
+):
     # All-or-nothing: a valid image ahead of the one that trips the size cap must NOT be
     # left on disk (the 413 mid-batch rolls back every file written this request).
     import utils.upload_limits as ul

From 756c7157218ea2fb90f61b893923fe5fc3bb2790 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sat, 4 Jul 2026 12:54:47 -0300
Subject: [PATCH 05/13] Trim the verbose comments added by the fixes

---
 scripts/uninstall.ps1                         |  5 +-
 scripts/uninstall.sh                          |  8 +--
 studio/backend/core/inference/diffusion.py    | 62 +++++++------------
 .../core/inference/diffusion_controlnet.py    | 14 ++---
 .../core/inference/diffusion_device.py        | 13 ++--
 .../backend/core/inference/diffusion_lora.py  | 24 ++-----
 .../core/inference/diffusion_memory.py        | 16 ++---
 .../backend/core/inference/sd_cpp_backend.py  |  7 +--
 .../backend/core/inference/sd_cpp_engine.py   |  9 +--
 .../core/training/diffusion_dit_trainer.py    |  6 +-
 .../core/training/diffusion_lora_trainer.py   |  2 +-
 .../core/training/diffusion_train_common.py   | 38 ++++--------
 studio/backend/models/training.py             |  5 +-
 studio/backend/routes/training.py             |  9 +--
 .../images/diffusion-train-dialog.tsx         |  4 +-
 .../src/features/images/images-page.tsx       | 16 ++---
 .../images/train/diffusion-train-panel.tsx    |  7 +--
 17 files changed, 75 insertions(+), 170 deletions(-)

diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1
index 2de89637b6..0f2aca669b 100644
--- a/scripts/uninstall.ps1
+++ b/scripts/uninstall.ps1
@@ -380,9 +380,8 @@ function Uninstall-UnslothStudio {
             continue
         }
         _RemovePath $r
-        # The native diffusion sibling (<custom root>.parent\stable-diffusion.cpp) is
-        # intentionally NOT removed: sd.cpp writes no owner marker and sits in the user's
-        # own parent dir, so auto-deleting it could destroy a user-managed clone.
+        # The native diffusion sibling (<custom root>.parent\stable-diffusion.cpp) is left
+        # in place: sd.cpp writes no owner marker, so auto-deleting it could destroy a clone.
     }
     # Default install dir (always at %USERPROFILE%\.unsloth\studio when present).
     if ($defaultStudioHome) { _RemovePath $defaultStudioHome }
diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh
index 36d974844d..320eb29c07 100755
--- a/scripts/uninstall.sh
+++ b/scripts/uninstall.sh
@@ -217,11 +217,9 @@ _remove_path "$HOME/.unsloth/studio"
 # when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept.
 _remove_path "$HOME/.unsloth/llama.cpp"
 # Default-mode native diffusion (stable-diffusion.cpp / sd-cli) build, a sibling of
-# studio like llama.cpp (install_sd_cpp_prebuilt.default_install_dir()). Only the default
-# location is removed. In env/custom mode the install is <custom root>.parent/
-# stable-diffusion.cpp, which is intentionally left in place: sd.cpp writes no owner
-# marker and sits in the user's own parent dir, so auto-deleting it could destroy a
-# user-managed stable-diffusion.cpp clone. A user-set UNSLOTH_SD_CPP_PATH is kept.
+# studio like llama.cpp. Only the default location is removed; the env/custom-mode
+# sibling (<custom root>.parent/stable-diffusion.cpp) is left in place since sd.cpp
+# writes no owner marker and auto-deleting it could destroy a user-managed clone.
 _remove_path "$HOME/.unsloth/stable-diffusion.cpp"
 _remove_path "$HOME/.unsloth/.cache"
 # Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in
diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index fa3e168c48..cdac37c2b1 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -404,18 +404,11 @@ class DiffusionBackend:
         """True when ``load_pipeline`` may take the dense transformer-quant path, so
         the prefetch should also pull the base repo's ``transformer/`` shards.
 
-        Those shards are excluded from the prefetch by default (the GGUF supplies
-        the transformer), but ``_load_dense_quant_pipeline`` fetches them with
-        ``from_pretrained(subfolder = "transformer")`` under the load lock during
-        "finalizing", after the previous pipeline was already evicted, where
-        unload/cancellation cannot preempt the download. Checks the dense-path gates
-        in ``load_pipeline`` that are knowable pre-download: quant requested and
-        supported for this device, and no pre-quantized checkpoint that would shortcut
-        the dense build. It deliberately does NOT mirror the ``plan.offload_policy ==
-        OFFLOAD_NONE`` gate: the memory plan needs the GGUF's on-disk size, which
-        isn't known until the GGUF is cached (after this prefetch runs). So the
-        transformer/ shards can be prefetched for a load that the plan then routes to
-        offload -- they stay cached for a later resident load rather than being wasted."""
+        Those shards are excluded from the prefetch by default (the GGUF supplies the
+        transformer), but ``_load_dense_quant_pipeline`` fetches them later under the
+        load lock, where unload/cancellation cannot preempt the download. Checks only
+        the dense-path gates knowable pre-download; skips the ``offload_policy`` gate
+        since that needs the GGUF's on-disk size, not known until after this runs."""
         mode = normalize_transformer_quant(kwargs.get("transformer_quant"))
         if mode is None:
             return False
@@ -807,7 +800,7 @@ class DiffusionBackend:
 
     @staticmethod
     def _hub_cache_repo_dir(repo_id: str) -> Path:
-        """The local HF hub cache dir for ``repo_id`` (``.../models--org--name``)."""
+        """Local HF hub cache dir for ``repo_id``."""
         from huggingface_hub import constants
         return Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}"
 
@@ -853,15 +846,10 @@ class DiffusionBackend:
     def _companion_cache_bytes(base: str) -> int:
         """Resident companion (VAE + text-encoder) size for the memory plan.
 
-        Sums the cached VAE + text-encoder weights while EXCLUDING ``transformer/`` (the
-        GGUF / single file supplies the transformer, so its ``transformer/`` shards are
-        not resident here). This matters for the dense ``transformer_quant`` path: it
-        prefetches the base repo's ``transformer/`` shards into the cache, and folding
-        those multi-GB shards into the companion size would inflate the plan and wrongly
-        force offload -- gating off the very quant path that fetched them. For a LOCAL
-        diffusers base the blob cache is empty, so walk the on-disk weights; for a hub
-        base, walk the snapshot (whose ``transformer/`` subfolder we can skip) instead of
-        the flat, content-addressed ``blobs/`` dir, which carries no subfolder split."""
+        Excludes ``transformer/`` (supplied by the GGUF/single file, not resident here) --
+        otherwise the dense-quant prefetch's cached transformer shards would inflate this
+        and wrongly force offload. Walks the snapshot dir, not the flat ``blobs/`` cache,
+        since only the snapshot preserves the subfolder split needed to exclude it."""
         local = Path(base).expanduser()
         if local.is_dir():
             return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True)
@@ -1029,9 +1017,7 @@ class DiffusionBackend:
                     and dense_transformer_supported(target)
                     and plan.offload_policy != OFFLOAD_NONE
                 ):
-                    # The dense fast path needs the transformer resident, so a memory_mode
-                    # (balanced / low_vram) that forces offload silently drops the requested
-                    # quant. Warn so the disengage is diagnosable rather than a null status.
+                    # memory_mode forcing offload silently drops the requested quant; warn so it's diagnosable.
                     logger.warning(
                         "diffusion.transformer_quant: %s requested but memory_mode forces "
                         "offload (%s); loading GGUF without dense quant",
@@ -1440,12 +1426,10 @@ class DiffusionBackend:
             companion_mib = None
         else:
             if kind == "single_file":
-                # Safetensors single-file. A dense bf16 file loads near its on-disk size,
-                # but a transformer-only fp8 checkpoint is loaded via from_single_file with
-                # a bf16 compute dtype and NO quantization_config, so diffusers upcasts it
-                # fp8 -> bf16 (~2x resident). Detect fp8 from the basename and budget the
-                # expansion. The single-file-is-pipeline (SDXL) path is a full bf16 pipeline
-                # checkpoint, not this fp8 transformer path, so it stays at on-disk size.
+                # An fp8 transformer checkpoint loads via from_single_file with a bf16
+                # compute dtype and no quantization_config, so diffusers upcasts it to
+                # bf16 (~2x resident); detect it from the basename. Excludes the
+                # single-file-is-pipeline (SDXL) case, which is already a bf16 pipeline.
                 fp8_upcast = not getattr(fam, "single_file_is_pipeline", False) and (
                     "fp8" in Path(single_file_path).name.lower() if single_file_path else False
                 )
@@ -1540,10 +1524,9 @@ class DiffusionBackend:
         if cn_model is None:
             if cancel.is_set():
                 raise RuntimeError(DIFFUSION_CANCELLED_MSG)
-            # A single generation uses exactly one ControlNet, so keep at most one resident:
-            # on a miss for a new id, drop the previously-cached module + its from_pipe wrapper
-            # (both dicts, kept consistent) and free the VRAM before loading the new one, or
-            # swapping distinct ControlNets within a base-model load accumulates until OOM.
+            # Keep at most one ControlNet resident: evict the previous module + its
+            # from_pipe wrapper before loading the new one, or swapping ControlNets
+            # within a base-model load accumulates until OOM.
             if self._cn_models or self._cn_pipes:
                 self._cn_models.clear()
                 self._cn_pipes.clear()
@@ -1718,10 +1701,8 @@ class DiffusionBackend:
         resolution/batch changed, or a stale-cache reuse otherwise. Best-effort: a
         transformer without the hook (uncached load) is a silent no-op."""
         transformer = getattr(pipe, "transformer", None)
-        # A diffusers CacheMixin transformer clears FBCache via ``_reset_stateful_cache``
-        # (which drives its HookRegistry.reset_stateful_hooks internally). The public
-        # ``reset_stateful_hooks`` name lives only on the HookRegistry, not on the
-        # transformer, so keep it only as a version fallback.
+        # ``_reset_stateful_cache`` is the transformer-level entry point; the public
+        # ``reset_stateful_hooks`` lives only on the HookRegistry, kept as a fallback.
         reset = getattr(transformer, "_reset_stateful_cache", None) or getattr(
             transformer, "reset_stateful_hooks", None
         )
@@ -1833,8 +1814,7 @@ class DiffusionBackend:
                             f"{state.family.name} is an image-editing model: provide an input image."
                         )
                     if mask_image is not None:
-                        # The edit family has no inpaint pipeline; a supplied mask would be
-                        # silently dropped (this branch wins over the inpaint branch below).
+                        # The edit family has no inpaint pipeline; a mask would be silently dropped.
                         raise ValueError(
                             f"{state.family.name} is an image-editing model and does not "
                             "support masks (mask_image)."
diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py
index 9f42350c73..170b48bd0f 100644
--- a/studio/backend/core/inference/diffusion_controlnet.py
+++ b/studio/backend/core/inference/diffusion_controlnet.py
@@ -204,9 +204,8 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve
 
 # Union ControlNet mode indices. A single "union" model covers several control modes and
 # selects the active one via an integer ``control_mode`` argument; these are the standard
-# indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" (an already-made
-# map) carries no intrinsic mode, so union_control_mode() defaults it to 0 (a union model
-# still requires a concrete mode).
+# indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" carries no
+# intrinsic mode, so union_control_mode() defaults it to 0.
 _UNION_CONTROL_MODES: dict[str, int] = {
     "canny": 0,
     "tile": 1,
@@ -221,12 +220,9 @@ _UNION_CONTROL_MODES: dict[str, int] = {
 def union_control_mode(spec_id: str, control_type: str) -> Optional[int]:
     """The integer ``control_mode`` for a union ControlNet, or None.
 
-    A union model REQUIRES a concrete ``control_mode`` (diffusers raises when it is None),
-    so for a curated union entry always return an index: the mapped mode, or a default
-    (0 / canny) for a type that carries no intrinsic mode such as 'passthrough' (an
-    already-made control map, which is also the UI's default for these models). For a
-    non-union entry return None so the caller omits the kwarg (it has a single fixed
-    mode). Pure lookup, no network."""
+    A union model requires a concrete mode (diffusers raises on None), so a curated union
+    entry always gets an index, defaulting to 0 for types like 'passthrough' that carry
+    none. A non-union entry returns None so the caller omits the kwarg."""
     entry = _catalog_by_id().get(spec_id)
     if entry is None or not entry.is_union:
         return None
diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py
index 1c5dff9580..f0dadcf312 100644
--- a/studio/backend/core/inference/diffusion_device.py
+++ b/studio/backend/core/inference/diffusion_device.py
@@ -235,15 +235,10 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
         mps_available = False
 
     if mps_available:
-        # Relax the MPS memory watermark BEFORE the first MPS allocation (the bfloat16
-        # probe just below). torch reads PYTORCH_MPS_HIGH_WATERMARK_RATIO exactly once,
-        # when the MPS allocator first initializes, so setting it any later is a no-op.
-        # The allocator otherwise caps a process at ~1.7x recommendedMaxWorkingSetSize,
-        # and a model that fits in unified system RAM but exceeds that cap OOMs at
-        # pipe.to("mps") (observed on an 8GB M1 mac mini: "MPS allocated 9.06 GiB, max
-        # allowed 9.07 GiB"). CPU offload can't help on unified memory (it frees no
-        # device bytes). Lifting the cap lets MPS spill into system RAM; a model larger
-        # than RAM would fail either way. setdefault respects a user-provided override.
+        # torch reads PYTORCH_MPS_HIGH_WATERMARK_RATIO once, at the first MPS allocation
+        # (the bfloat16 probe below), so it must be relaxed before that. Otherwise the
+        # allocator caps the process at ~1.7x recommendedMaxWorkingSetSize and can OOM a
+        # model that would otherwise fit in unified RAM. setdefault respects an override.
         os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
         # Prefer bfloat16; otherwise fall back to float32, NEVER silent float16.
         # Modern diffusion transformers (Z-Image, FLUX.2, ...) produce activations
diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py
index bd3e2d3784..cee8d51242 100644
--- a/studio/backend/core/inference/diffusion_lora.py
+++ b/studio/backend/core/inference/diffusion_lora.py
@@ -267,11 +267,7 @@ def _pick_repo_weight_file(repo_id: str, hf_token: Optional[str]) -> str:
 
 
 def _scrub_hub_url(msg: str) -> str:
-    """Strip embedded http(s) URLs from a Hub error message before it hits a 400 body.
-
-    huggingface_hub errors interpolate the request URL (and a request id) into their
-    message; a raw endpoint URL is noise in a client-facing 400, so drop it.
-    """
+    """Strip embedded http(s) URLs from a Hub error message before it hits a 400 body."""
     cleaned = re.sub(r"https?://\S+", "", msg)
     # Collapse the whitespace / stray separators the URL removal leaves behind.
     return re.sub(r"\s{2,}", " ", cleaned).strip()
@@ -285,20 +281,10 @@ def resolve_specs(
 ) -> list[ResolvedLora]:
     """Resolve request (id, weight) pairs, dropping zero-weight entries.
 
-    A stale / unknown id raises FileNotFoundError inside resolve_one; a mistyped Hub
-    repo id makes the Hub resolution raise a huggingface_hub client error (a missing
-    repo -> RepositoryNotFoundError, a bad revision -> RevisionNotFoundError, a missing
-    weight file -> EntryNotFoundError, a gated model -> GatedRepoError). Convert those
-    NAMED not-found/gated errors to ValueError so the route (which maps only ValueError
-    to a 400) reports bad client input instead of a generic 500 -- Hub error messages
-    embed the request URL, so scrub it out before it reaches the 400 body. Catch them by
-    name rather than their common HfHubHTTPError base on purpose: a Hub-side 5xx / 429
-    (an outage, not bad input) is a bare HfHubHTTPError and must stay a 500. A Hub
-    download can also raise ``RuntimeError("Cancelled")`` when the user unloads / starts a
-    superseding load mid-download; convert that to the diffusion cancellation sentinel so
-    the route maps it to a 409 instead of a generic server error toast. A non-cancellation
-    RuntimeError (e.g. a stalled download, disk full) stays a 500 -- it is not bad
-    client input."""
+    Maps the named not-found/gated Hub errors (bad repo/revision/file/gating) to a 400
+    and scrubs the URL from the message; deliberately does NOT catch the base
+    HfHubHTTPError so a Hub 5xx stays a 500. A mid-download cancel also maps to a 409
+    instead of a generic 500."""
     from huggingface_hub.errors import (
         EntryNotFoundError,
         GatedRepoError,
diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py
index 090e295a2e..065fd4eb11 100644
--- a/studio/backend/core/inference/diffusion_memory.py
+++ b/studio/backend/core/inference/diffusion_memory.py
@@ -241,18 +241,10 @@ def estimate_safetensors_dense_mib(
 ) -> Optional[int]:
     """Resident size of a safetensors checkpoint, in MiB.
 
-    Unlike a GGUF (which is dequantised to bf16/fp16 on load, so a 4-bit file
-    expands ~4x), a safetensors checkpoint usually loads near its on-disk size: a
-    dense bf16 file is already bf16, and a bnb-4bit file stays compressed in VRAM
-    (it carries its own quantization_config). So the on-disk size is the estimate,
-    returned unchanged (None passes through).
-
-    The exception is ``fp8_upcast``: the fp8 single-file transformer path loads via
-    ``from_single_file`` with a bf16 compute dtype and NO quantization_config, so
-    diffusers upcasts the fp8 weights (1 byte/param) to bf16 (2 bytes/param) --
-    roughly 2x the on-disk bytes resident. Budget that, or the plan under-reserves
-    and OOMs.
-    """
+    Unlike a GGUF (dequantised on load, so a 4-bit file expands ~4x), a safetensors
+    checkpoint usually loads near its on-disk size (None passes through unchanged).
+    Exception: ``fp8_upcast`` -- an fp8 single-file transformer loads with no
+    quantization_config, so diffusers upcasts it to bf16 (~2x on-disk resident)."""
     if storage_mib is None:
         return None
     if fp8_upcast:
diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 6c54af5e87..577370468b 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -737,11 +737,8 @@ class SdCppDiffusionBackend:
                 "img2img / inpaint / reference / upscale are not yet supported on the native "
                 "sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows."
             )
-        # strength 0 (or None) disables ControlNet -- documented on the request model, and
-        # the diffusers path treats it as plain txt2img -- so a strength-0 spec must be a
-        # no-op here too, not a hard 400. Only a genuinely active (strength > 0) ControlNet
-        # is rejected. Strength is element 3 of the tuple
-        # (id, image, type, strength, guidance_start, guidance_end).
+        # strength 0 (or None) disables ControlNet (documented on the request model, matches
+        # the diffusers path), so it must be a no-op here too, not a hard 400.
         if controlnet is not None and controlnet[3] in (None, 0, 0.0):
             controlnet = None
         if controlnet is not None:
diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py
index af896007b5..fe1e7ebc92 100644
--- a/studio/backend/core/inference/sd_cpp_engine.py
+++ b/studio/backend/core/inference/sd_cpp_engine.py
@@ -127,10 +127,7 @@ def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
     highest priority first: the cmake ``build/bin`` tree, then a Windows Release
     subdir, then the root itself, then the prebuilt archive's versioned subdir.
 
-    The prebuilt archive extracts into a top-level versioned dir
-    (``sd-master-<tag>-bin-<host>/``) rather than flattening into ``root``, so without
-    the ``root/*/`` glob a fresh prebuilt install is invisible here -- which silently
-    demotes the persistent sd-server to one-shot mode and re-downloads on every start."""
+    The prebuilt lands in its own versioned subdir rather than flattening into ``root``."""
     name = _binary_name(stem)
     cands = [
         root / "build" / "bin" / name,
@@ -138,9 +135,7 @@ def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
         root / "bin" / name,
         root / name,
     ]
-    # Prebuilt archive layout: root/sd-master-<tag>-bin-<host>/<name> (+ its own bin/).
-    # Newest install first (by mtime -- tag strings don't sort numerically, so a lexical
-    # sort would rank build 99 above build 100).
+    # Newest install first, by mtime -- tag strings don't sort numerically.
     try:
         subdirs = [p for p in root.iterdir() if p.is_dir()]
         subdirs.sort(key = lambda p: p.stat().st_mtime, reverse = True)
diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index 80e40abcaa..226b011320 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -63,10 +63,8 @@ def _select_lora_targets(
 ) -> tuple[str, ...]:
     """Pick the LoRA target modules for a DiT run.
 
-    ``normalized()`` leaves ``lora_target_modules`` empty when a caller does not set it, so
-    an empty tuple means "unset" here: use the family's ``spec.lora_targets`` (which add the
-    DiT-specific joint-attention projections). Any explicit tuple is a deliberate override
-    and still wins."""
+    An empty ``cfg_targets`` means "unset": use the family's ``spec.lora_targets``. Any
+    explicit tuple is a deliberate override and still wins."""
     if not tuple(cfg_targets):
         return tuple(spec_targets)
     return tuple(cfg_targets)
diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py
index b96dc1fa07..3424df7dc3 100644
--- a/studio/backend/core/training/diffusion_lora_trainer.py
+++ b/studio/backend/core/training/diffusion_lora_trainer.py
@@ -211,7 +211,7 @@ def run_diffusion_lora_training(
     for m in (unet, *text_encoders):
         m.to(device, dtype = weight_dtype)
 
-    # An empty (unset) config means "use the family default": the SDXL attention projections.
+    # Empty (unset) config means use the family default.
     unet_targets = list(cfg.lora_target_modules) or list(DEFAULT_LORA_TARGETS)
     unet.add_adapter(
         LoraConfig(
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 0040650f0c..e4a83b8c50 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -30,14 +30,11 @@ from core.inference.diffusion_families import (
     trainable_family_names,
 )
 
-# Default LoRA target modules: the attention projections of the SDXL U-Net (the
-# diffusers/kohya convention). Used by the SDXL trainer as its fallback when the config
-# leaves ``lora_target_modules`` empty; the DiT trainers supply their own wider set. Kept
-# here so the SDXL trainer has a named default even for an empty (unset) config.
+# Default LoRA target modules: the SDXL U-Net attention projections. DiT trainers supply
+# their own wider set instead.
 DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
 
-# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). Validated in
-# normalized() so a typo fails fast at request time, not minutes later in the subprocess.
+# diffusers' SchedulerType names (diffusers.optimization.get_scheduler).
 _LR_SCHEDULERS: frozenset[str] = frozenset(
     {
         "linear",
@@ -50,10 +47,8 @@ _LR_SCHEDULERS: frozenset[str] = frozenset(
     }
 )
 
-# DiT families that overflow fp16 (their RoPE / embedder run in fp32), so they train in bf16
-# only. Encoded here -- keyed by resolved family -- so normalized() can reject an fp16 request
-# before spawn without importing the DiT trainer's _SPECS (which would create an import
-# cycle). The DiT trainer keeps a matching guard as defense in depth.
+# DiT families whose fp32 RoPE/embedder overflow fp16, so they train in bf16 only. Must stay
+# in sync with the DiT trainer's own specs (kept separate to avoid an import cycle).
 _FORCE_BF16_FAMILIES: frozenset[str] = frozenset({"qwen-image", "z-image"})
 
 _IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
@@ -104,12 +99,8 @@ def resolve_trainable_family(base_model: str, model_family: Optional[str] = None
     # GGUF weights (a ``.gguf`` file or a ``*-GGUF`` repo) are inference-only: training needs
     # the full diffusers pipeline (transformer + VAE + text encoders), which a GGUF repo does
     # not provide. Reject by name even when the family itself is trainable.
-    # A ``.gguf`` file always rejects. The broad ``"gguf" in name`` catch (for ``*-GGUF``
-    # repos) must NOT reject a real local diffusers directory that merely has "gguf" in its
-    # path, so it is skipped for a local diffusers checkout -- identified by its
-    # ``model_index.json`` marker (the same marker the loader uses), NOT a bare ``is_dir()``:
-    # a GGUF-only folder must still reject here and fail fast, rather than pass and fail late
-    # in the subprocess after the resident chat/Images models were already evicted.
+    # Exempt a local diffusers checkout that merely has "gguf" in its path, identified by its
+    # ``model_index.json`` marker (same marker the loader uses), not a bare ``is_dir()``.
     local = Path(base_model).expanduser() if base_model else None
     is_local_diffusers = bool(local and (local / "model_index.json").is_file())
     if name.endswith(".gguf") or ("gguf" in name and not is_local_diffusers):
@@ -244,8 +235,7 @@ class DiffusionLoraConfig:
     lora_rank: int = 16
     lora_alpha: Optional[int] = None  # defaults to lora_rank
     lora_dropout: float = 0.0
-    # Empty = "unset": each trainer supplies its family default (SDXL DEFAULT_LORA_TARGETS,
-    # or the DiT family's wider joint-attention set). A non-empty tuple is an explicit override.
+    # Empty = "unset": each trainer supplies its own family default.
     lora_target_modules: tuple[str, ...] = ()
     seed: int = 42
     mixed_precision: str = "bf16"  # "bf16" | "fp16" | "no"
@@ -290,9 +280,7 @@ 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 bf16-only DiT family (Qwen-Image / Z-Image) must refuse fp16 up front rather than
-        # accepting the request, evicting resident models, and only then failing in the
-        # subprocess. The DiT trainer keeps a matching guard as defense in depth.
+        # Refuse fp16 for a bf16-only DiT family up front, before evicting resident models.
         if self.mixed_precision == "fp16" and resolved_family in _FORCE_BF16_FAMILIES:
             raise ValueError(
                 f"'{resolved_family}' LoRA training requires bf16: fp16 overflows its fp32 "
@@ -312,8 +300,7 @@ class DiffusionLoraConfig:
         if learning_rate <= 0:
             raise ValueError("learning_rate must be > 0")
         alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank
-        # Leave an unset (empty) target list empty: the trainer fills the family default
-        # (SDXL DEFAULT_LORA_TARGETS, or the DiT family's wider set) so the family spec wins.
+        # Leave an unset (empty) target list empty so the trainer fills the family default.
         targets = tuple(self.lora_target_modules)
         # A blank Hub token (the Studio default when none is configured) must load
         # anonymously, not as an explicit empty credential.
@@ -441,9 +428,8 @@ def _publish_to_lora_catalog(lora_path: str, cfg: DiffusionLoraConfig) -> Option
         alias = sanitize_alias(base)
         src_resolved = Path(lora_path).resolve()
         dest = loras_dir() / f"{alias}.safetensors"
-        # A retrain with the same adapter name must not clobber a prior mirror: if the
-        # destination already exists and is a different file, pick the next free numeric
-        # suffix (<alias>-2, <alias>-3, ...) for both the weights and their .json sidecar.
+        # A retrain with the same adapter name must not clobber a prior mirror: pick the next
+        # free numeric suffix instead.
         if dest.exists() and dest.resolve() != src_resolved:
             n = 2
             while True:
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 522d5d3521..224caad35d 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -698,9 +698,8 @@ class DiffusionTrainingStartRequest(BaseModel):
     lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank")
     lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0)
     # Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that
-    # sets them is not silently trained with defaults. Empty (the default) means "unset": each
-    # trainer supplies its own family targets -- the SDXL DEFAULT_LORA_TARGETS for SDXL, the
-    # wider joint-attention set for the DiT families. A non-empty list is an explicit override.
+    # sets them is not silently trained with defaults. Empty = "unset": the trainer fills in
+    # its family default.
     lora_target_modules: List[str] = Field(
         default_factory = list,
         description = "Modules to attach LoRA to; empty = the trainer's family default",
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index a82491c04f..134fdbf8c2 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1389,10 +1389,8 @@ async def upload_diffusion_dataset(
                 detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
             )
         names.append(filename)
-    # Roll back every file written this request if the batch does not fully commit, so a
-    # mid-batch 413 (or a disk error / client disconnect) leaves the dataset unchanged
-    # rather than partially populated -- the upload is all-or-nothing, not just for the
-    # extension check above but for the size limit too.
+    # Roll back all files written this request on a mid-batch failure (size limit,
+    # disk error, disconnect) so the dataset is never left partially populated.
     written: list[Path] = []
     committed = False
     try:
@@ -1884,8 +1882,7 @@ def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
                 continue
             fn = row.get("file_name") or row.get("image") or row.get("file")
             if fn and caption_col in row:
-                # First writer wins over sorted manifests, so the plain manifest
-                # (e.g. output_file.jsonl) is deterministic rather than OS-visit order.
+                # First writer wins over sorted manifests, for deterministic results.
                 captions.setdefault(Path(str(fn)).name, str(row[caption_col]))
     # Copy images (those with a caption first, so a cap keeps captioned pairs).
     images = sorted(
diff --git a/studio/frontend/src/features/images/diffusion-train-dialog.tsx b/studio/frontend/src/features/images/diffusion-train-dialog.tsx
index 016087bd8d..17d6dc8193 100644
--- a/studio/frontend/src/features/images/diffusion-train-dialog.tsx
+++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx
@@ -131,9 +131,7 @@ export function DiffusionTrainDialog({
       ? Math.min(100, Math.round((status.step / status.total_steps) * 100))
       : 0;
 
-  // Notify the parent exactly once when a run finishes with a saved adapter, so it can
-  // rescan the LoRA picker (a LoRA trained while a model is loaded is otherwise invisible
-  // until a model swap re-runs the discovery effect).
+  // Notify the parent exactly once per finished run so it rescans the LoRA picker.
   const [notifiedComplete, setNotifiedComplete] = useState(false);
   useEffect(() => {
     if (hasSavedAdapter && !notifiedComplete) {
diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx
index b518ad2cb9..bb4f95b8e3 100644
--- a/studio/frontend/src/features/images/images-page.tsx
+++ b/studio/frontend/src/features/images/images-page.tsx
@@ -1237,9 +1237,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
   // Load an image's recipe back into the form inputs.
   const restoreSettings = useCallback((image: GalleryImage) => {
     setPrompt(image.prompt);
-    // The Negative-prompt field only renders (and submits) when guidance>0, so a
-    // guidance=0 recipe must not restore a hidden negative prompt that would
-    // resurface if guidance is later raised. Mirror the submit-path gating.
+    // Negative prompt only applies when guidance>0; don't restore a hidden value.
     setNegativePrompt(image.guidance > 0 ? (image.negative_prompt ?? "") : "");
     setSteps(image.steps);
     setGuidance(image.guidance);
@@ -1275,9 +1273,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
       if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight });
     }
     setLoras(restoredLoras);
-    // The recipe carries no control image (it isn't persisted), so a faithful
-    // restore can't reproduce a ControlNet run -- clear any stale form selection
-    // rather than leaking it into the restored recipe, mirroring the LoRA clear.
+    // The control image isn't persisted, so clear any stale ControlNet selection.
     setControlnetId("");
     setControlImage(null);
     toast.success("Settings restored to inputs");
@@ -1528,9 +1524,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
       // Curated non-GGUF model: load as a full pipeline or single-file safetensors.
       const spec = SAFETENSORS_MODELS[id];
       if (spec) {
-        // Optimistically drop the quant label, but revert if the load never starts
-        // so a failed load doesn't leave a stale 'GGUF · variant' label (mirrors the
-        // GGUF branches below; the poll owns the after-start revert via quantRevert).
+        // Optimistically clear the quant label, revert it if the load never starts.
         const prevQuant = quant;
         quantRevert.current = { prev: prevQuant };
         setQuant(null);
@@ -1600,9 +1594,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         toast.error("Only unsloth or on-device image models can be loaded here");
         return;
       }
-      // Optimistically drop the quant label, but revert if the load never starts so
-      // a failed load doesn't leave a stale 'GGUF · variant' label (mirrors the GGUF
-      // branches above; the poll owns the after-start revert via quantRevert).
+      // Optimistically clear the quant label, revert it if the load never starts.
       const prevQuant = quant;
       quantRevert.current = { prev: prevQuant };
       setQuant(null);
diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
index 04fdeded5c..7111307907 100644
--- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
+++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
@@ -318,11 +318,8 @@ export function DiffusionTrainPanel({
   // terminal "completed" status until the next start, so we can't rely on it clearing).
   const [dismissedJobId, setDismissedJobId] = useState<string | null>(null);
   const running = Boolean(status?.active) || status?.status === "running";
-  // A stopped run still saves + catalog-publishes a real, deployable adapter (status
-  // carries catalog_path), so treat "stopped with an adapter" as terminal-with-adapter
-  // too -- otherwise the normal "stop once the loss looks good" flow leaves the trained
-  // adapter with no Deploy button and no picker refresh. A save=False cancel has no
-  // catalog_path, so it correctly still shows nothing.
+  // A stopped run still saves + publishes a deployable adapter (catalog_path set), so
+  // treat it as finished-with-adapter too; a save=False cancel has no catalog_path.
   const hasSavedAdapter =
     status?.status === "completed" ||
     (status?.status === "stopped" && Boolean(status?.catalog_path));

From 639d4061afaa86ded54bdd19afa7e5f8010b1f79 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sat, 4 Jul 2026 20:26:27 -0300
Subject: [PATCH 06/13] Revert the LoRA target-module default change (already
 fixed upstream)

---
 .../core/training/diffusion_dit_trainer.py    |  9 ++++++---
 .../core/training/diffusion_lora_trainer.py   |  4 +---
 .../core/training/diffusion_train_common.py   | 11 +++++-----
 studio/backend/models/training.py             |  8 ++++----
 .../tests/test_diffusion_dit_trainer.py       | 20 +++++++++----------
 .../tests/test_diffusion_lora_trainer.py      |  3 +--
 6 files changed, 26 insertions(+), 29 deletions(-)

diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index 226b011320..50bb0f5d3b 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -32,6 +32,7 @@ from typing import Any, Callable, Optional
 
 from core.training.diffusion_train_common import (
     DEFAULT_LORA_FILENAME,
+    DEFAULT_LORA_TARGETS,
     DiffusionLoraConfig,
     EventCb,
     StopCb,
@@ -63,9 +64,11 @@ def _select_lora_targets(
 ) -> tuple[str, ...]:
     """Pick the LoRA target modules for a DiT run.
 
-    An empty ``cfg_targets`` means "unset": use the family's ``spec.lora_targets``. Any
-    explicit tuple is a deliberate override and still wins."""
-    if not tuple(cfg_targets):
+    ``normalized()`` always fills ``lora_target_modules`` with the generic
+    ``DEFAULT_LORA_TARGETS`` when a caller does not set it, so that value means "unset"
+    here: prefer the family's ``spec.lora_targets`` (which add the DiT-specific
+    projections). Any OTHER explicit tuple is a deliberate override and still wins."""
+    if tuple(cfg_targets) == DEFAULT_LORA_TARGETS:
         return tuple(spec_targets)
     return tuple(cfg_targets)
 
diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py
index 3424df7dc3..db2eca1011 100644
--- a/studio/backend/core/training/diffusion_lora_trainer.py
+++ b/studio/backend/core/training/diffusion_lora_trainer.py
@@ -211,15 +211,13 @@ def run_diffusion_lora_training(
     for m in (unet, *text_encoders):
         m.to(device, dtype = weight_dtype)
 
-    # Empty (unset) config means use the family default.
-    unet_targets = list(cfg.lora_target_modules) or list(DEFAULT_LORA_TARGETS)
     unet.add_adapter(
         LoraConfig(
             r = cfg.lora_rank,
             lora_alpha = cfg.lora_alpha,
             lora_dropout = cfg.lora_dropout,
             init_lora_weights = "gaussian",
-            target_modules = unet_targets,
+            target_modules = list(cfg.lora_target_modules),
         )
     )
     if cfg.gradient_checkpointing:
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index e4a83b8c50..cdf78995b3 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -30,8 +30,9 @@ from core.inference.diffusion_families import (
     trainable_family_names,
 )
 
-# Default LoRA target modules: the SDXL U-Net attention projections. DiT trainers supply
-# their own wider set instead.
+# Default LoRA target modules: the attention projections common to the SDXL U-Net and the
+# DiT transformers (the diffusers/kohya convention). A family whose trainer wants a wider
+# set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback.
 DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
 
 # diffusers' SchedulerType names (diffusers.optimization.get_scheduler).
@@ -235,8 +236,7 @@ class DiffusionLoraConfig:
     lora_rank: int = 16
     lora_alpha: Optional[int] = None  # defaults to lora_rank
     lora_dropout: float = 0.0
-    # Empty = "unset": each trainer supplies its own family default.
-    lora_target_modules: tuple[str, ...] = ()
+    lora_target_modules: tuple[str, ...] = DEFAULT_LORA_TARGETS
     seed: int = 42
     mixed_precision: str = "bf16"  # "bf16" | "fp16" | "no"
     snr_gamma: Optional[float] = 5.0  # min-SNR loss weighting; None disables
@@ -300,8 +300,7 @@ class DiffusionLoraConfig:
         if learning_rate <= 0:
             raise ValueError("learning_rate must be > 0")
         alpha = self.lora_alpha if self.lora_alpha is not None else self.lora_rank
-        # Leave an unset (empty) target list empty so the trainer fills the family default.
-        targets = tuple(self.lora_target_modules)
+        targets = tuple(self.lora_target_modules) or DEFAULT_LORA_TARGETS
         # A blank Hub token (the Studio default when none is configured) must load
         # anonymously, not as an explicit empty credential.
         token = self.hf_token.strip() if isinstance(self.hf_token, str) else self.hf_token
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 224caad35d..3883b9847f 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -698,11 +698,11 @@ class DiffusionTrainingStartRequest(BaseModel):
     lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank")
     lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0)
     # Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that
-    # sets them is not silently trained with defaults. Empty = "unset": the trainer fills in
-    # its family default.
+    # sets them is not silently trained with defaults. Default the target list to the SDXL
+    # attention projections (the trainer's DEFAULT_LORA_TARGETS) so it is never None.
     lora_target_modules: List[str] = Field(
-        default_factory = list,
-        description = "Modules to attach LoRA to; empty = the trainer's family default",
+        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")
     seed: int = Field(42)
diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py
index 79a3777113..63cc01922c 100644
--- a/studio/backend/tests/test_diffusion_dit_trainer.py
+++ b/studio/backend/tests/test_diffusion_dit_trainer.py
@@ -41,25 +41,23 @@ def test_specs_cover_the_three_dit_families():
 
 
 def test_select_lora_targets_uses_family_default_for_generic_config():
-    # normalized() leaves lora_target_modules empty when a caller doesn't set it, so an empty
-    # tuple must resolve to the family's targets (which add the DiT-specific joint-attention
-    # projections), not the generic SDXL list.
-    assert _select_lora_targets((), _FLUX_TARGETS) == _FLUX_TARGETS
-    assert _select_lora_targets((), _QWEN_TARGETS) == _QWEN_TARGETS
-    assert _select_lora_targets((), _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS
-    # The generic SDXL default is NOT treated as unset here: an explicit list wins.
-    assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == DEFAULT_LORA_TARGETS
+    # normalized() fills lora_target_modules with the generic DEFAULT_LORA_TARGETS when a
+    # caller doesn't set it, so that value must resolve to the family's targets (which add
+    # the DiT-specific projections), not stay stuck on the generic SDXL list.
+    assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == _FLUX_TARGETS
+    assert _select_lora_targets(DEFAULT_LORA_TARGETS, _QWEN_TARGETS) == _QWEN_TARGETS
+    assert _select_lora_targets(DEFAULT_LORA_TARGETS, _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS
 
 
 def test_select_lora_targets_explicit_override_wins():
-    # Any explicit tuple is a deliberate override and must win over the family spec.
+    # Any OTHER explicit tuple is a deliberate override and must win over the family spec.
     override = ("to_q", "to_k")
     assert _select_lora_targets(override, _FLUX_TARGETS) == override
-    # The default request path (config leaving targets unset) reaches the spec.
+    # The default request path (config carrying the generic default) reaches the spec.
     cfg = DiffusionLoraConfig(
         base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o"
     ).normalized()
-    assert cfg.lora_target_modules == ()
+    assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
     assert (
         _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets)
         == _FLUX_TARGETS
diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py
index dc84b70f17..d9078bf0e7 100644
--- a/studio/backend/tests/test_diffusion_lora_trainer.py
+++ b/studio/backend/tests/test_diffusion_lora_trainer.py
@@ -94,8 +94,7 @@ def test_discover_missing_dir_raises(tmp_path):
 def test_config_normalized_defaults():
     cfg = DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o").normalized()
     assert cfg.lora_alpha == cfg.lora_rank  # alpha defaults to rank
-    # Targets stay empty (unset) after normalize; each trainer fills its own family default.
-    assert cfg.lora_target_modules == ()
+    assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
 
 
 @pytest.mark.parametrize(

From ee61c7b650a08c1afce69f03b8d9d43f2b54a9b0 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 4 Jul 2026 23:36:40 +0000
Subject: [PATCH 07/13] Fix hub error test construction for huggingface_hub 1.x
 for PR #6872

huggingface_hub 1.x makes HfHubHTTPError's response argument required and
reads only .headers and .request from it, while 0.x defaults it to None.
Pass a small stub response so the test constructs the error on both,
verified against 0.36.2 and the 1.22.0 wheel.
---
 studio/backend/tests/test_diffusion_lora.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py
index a40807a3a2..668226e27a 100644
--- a/studio/backend/tests/test_diffusion_lora.py
+++ b/studio/backend/tests/test_diffusion_lora.py
@@ -197,9 +197,12 @@ def test_resolve_specs_maps_hub_error_to_valueerror(tmp_path, monkeypatch):
     from huggingface_hub.errors import RepositoryNotFoundError
 
     def _boom(spec_id, weight, **kw):
+        # response is optional in huggingface_hub 0.x but required in 1.x; both only
+        # read .headers / .request, so a stub keeps the test working on either.
         raise RepositoryNotFoundError(
             "404 Client Error. Repository Not Found for url: "
-            "https://huggingface.co/api/models/nope/nope (Request ID: abc)"
+            "https://huggingface.co/api/models/nope/nope (Request ID: abc)",
+            response = types.SimpleNamespace(headers = {}, request = None),
         )
 
     monkeypatch.setattr(dl, "resolve_one", _boom)

From 01d0aa1a83460d1b40b92d157d1f814e4267eee8 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Mon, 6 Jul 2026 03:07:57 -0300
Subject: [PATCH 08/13] Preserve pre-existing files on upload rollback and
 budget the dense transformer in the quant preflight

---
 studio/backend/core/inference/diffusion.py    | 150 ++++++++++++++----
 studio/backend/routes/training.py             |  21 ++-
 .../backend/tests/test_diffusion_backend.py   |  96 +++++++++++
 .../backend/tests/test_diffusion_training.py  |  24 +++
 4 files changed, 250 insertions(+), 41 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index cdac37c2b1..de55a062e4 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -19,11 +19,12 @@ bar. GPU-handoff policy lives in the arbiter the routes call, not here.
 from __future__ import annotations
 
 import inspect
+import json
 import threading
 import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 from loggers import get_logger
 from utils.hardware import clear_gpu_cache
@@ -842,6 +843,19 @@ class DiffusionBackend:
                 continue
         return total
 
+    @staticmethod
+    def _max_over_cached_revs(base: str, fn: Callable[[Path], int]) -> int:
+        """Apply ``fn`` to a LOCAL diffusers dir, or to the fullest cached hub snapshot
+        revision (the active one is the fullest), returning that count. 0 when nothing is
+        cached. Multiple revisions may be cached, so take the max."""
+        local = Path(base).expanduser()
+        if local.is_dir():
+            return fn(local)
+        snapshots = DiffusionBackend._hub_cache_repo_dir(base) / "snapshots"
+        if not snapshots.is_dir():
+            return 0
+        return max((fn(rev) for rev in snapshots.iterdir() if rev.is_dir()), default = 0)
+
     @staticmethod
     def _companion_cache_bytes(base: str) -> int:
         """Resident companion (VAE + text-encoder) size for the memory plan.
@@ -850,22 +864,51 @@ class DiffusionBackend:
         otherwise the dense-quant prefetch's cached transformer shards would inflate this
         and wrongly force offload. Walks the snapshot dir, not the flat ``blobs/`` cache,
         since only the snapshot preserves the subfolder split needed to exclude it."""
-        local = Path(base).expanduser()
-        if local.is_dir():
-            return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True)
-        snapshots = DiffusionBackend._hub_cache_repo_dir(base) / "snapshots"
-        if not snapshots.is_dir():
-            return 0
-        # Multiple revisions may be cached; the active one is the fullest, so take the max.
-        return max(
-            (
-                DiffusionBackend._local_dir_weight_bytes(rev, exclude_transformer = True)
-                for rev in snapshots.iterdir()
-                if rev.is_dir()
-            ),
-            default = 0,
+        return DiffusionBackend._max_over_cached_revs(
+            base, lambda d: DiffusionBackend._local_dir_weight_bytes(d, exclude_transformer = True)
         )
 
+    @staticmethod
+    def _safetensors_param_count(path: Path) -> int:
+        """Total tensor elements in a safetensors file, read from its JSON header (an
+        8-byte little-endian length prefix then the header) without touching tensor data."""
+        try:
+            with open(path, "rb") as fh:
+                header_len = int.from_bytes(fh.read(8), "little")
+                header = json.loads(fh.read(header_len))
+            total = 0
+            for name, meta in header.items():
+                if name == "__metadata__" or not isinstance(meta, dict):
+                    continue
+                numel = 1
+                for dim in meta.get("shape", []):
+                    numel *= dim
+                total += numel
+            return total
+        except Exception:  # noqa: BLE001 — best-effort estimate; a corrupt/crafted shard
+            # (bad header length, non-dict header, odd shape) must degrade to 0 so the caller
+            # gates on the plain plan, never crash the load.
+            return 0
+
+    @staticmethod
+    def _dense_transformer_resident_bytes(base: str) -> int:
+        """Resident bf16 size of the base repo's dense ``transformer/`` for the dense-quant
+        preflight. That fast path loads the transformer at the compute dtype (bf16, 2
+        bytes/param) before quantizing, so budget num_params * 2 -- NOT the on-disk bytes,
+        which for an F32 base (e.g. Z-Image) are ~2x the resident size. Read from the
+        safetensors shard headers. Returns 0 when no ``transformer/*.safetensors`` shards
+        are present (an uncached base, or a .bin-only transformer); the caller then gates
+        the fast path on the plain plan."""
+        def _params(d: Path) -> int:
+            tdir = d / "transformer"
+            if not tdir.is_dir():
+                return 0
+            return sum(
+                DiffusionBackend._safetensors_param_count(s) for s in tdir.glob("*.safetensors")
+            )
+
+        return DiffusionBackend._max_over_cached_revs(base, _params) * 2  # bf16: 2 bytes/param
+
     # ── Synchronous load / generate / unload ───────────────────────────────
 
     def load_pipeline(
@@ -954,9 +997,8 @@ class DiffusionBackend:
                 pipeline_cls = getattr(diffusers, fam.pipeline_class)
 
                 # Decide placement up front (the weights are still on CPU, so free VRAM is
-                # the real budget) -- this also doubles as the dense-quant preflight: the
-                # dense bf16 transformer must fit resident, so the fast path is offered only
-                # when the plan is `none`.
+                # the real budget). This plan budgets the GGUF file and places the plain
+                # load; the dense-quant fast path is preflighted separately below.
                 plan = self._plan_memory(
                     target,
                     single_file_path,
@@ -976,14 +1018,52 @@ class DiffusionBackend:
                 # GGUF kind offers it: it materialises the dense bf16 transformer from the
                 # base repo, which the safetensors kinds (a single-file or already-quantized
                 # pipeline) do not have.
-                pipe = None
-                transformer_quant_engaged = None
-                if (
+                dense_quant_requested = (
                     kind == "gguf"
                     and normalize_transformer_quant(transformer_quant) is not None
                     and dense_transformer_supported(target)
-                    and plan.offload_policy == OFFLOAD_NONE
-                ):
+                )
+                # `plan` budgets the GGUF file, but this path materializes the base repo's
+                # dense bf16 transformer -- re-check the fit against THAT so a card that fits
+                # the GGUF but not the dense transformer skips up front instead of evicting +
+                # OOMing in finalization. A prequant checkpoint loads a small quantized file
+                # (no dense bf16), so skip the re-check there (mirrors the prefetch guard).
+                dense_fits = plan.offload_policy == OFFLOAD_NONE
+                if dense_quant_requested and dense_fits:
+                    scheme = select_transformer_quant_scheme(
+                        target,
+                        normalize_transformer_quant(transformer_quant),
+                        family = getattr(fam, "name", None),
+                    )
+                    prequant = (
+                        resolve_prequant_source(
+                            fam, scheme, path_override = transformer_prequant_path
+                        )
+                        if scheme is not None
+                        else None
+                    )
+                    dense_mib = (
+                        int(self._dense_transformer_resident_bytes(base) // (1024 * 1024))
+                        if prequant is None
+                        else 0
+                    )
+                    if dense_mib > 0:
+                        dense_plan = self._plan_memory(
+                            target,
+                            single_file_path,
+                            base,
+                            fam,
+                            memory_mode,
+                            cpu_offload,
+                            kind = kind,
+                            repo_id = repo_id,
+                            transformer_mib_override = dense_mib,
+                        )
+                        dense_fits = dense_plan.offload_policy == OFFLOAD_NONE
+
+                pipe = None
+                transformer_quant_engaged = None
+                if dense_quant_requested and dense_fits:
                     try:
                         pipe, transformer_quant_engaged = self._load_dense_quant_pipeline(
                             transformer_cls,
@@ -1011,18 +1091,15 @@ class DiffusionBackend:
                         # GGUF build (the OOM-fallback path this cleanup exists for).
                         del exc
                         clear_gpu_cache()
-                elif (
-                    kind == "gguf"
-                    and normalize_transformer_quant(transformer_quant) is not None
-                    and dense_transformer_supported(target)
-                    and plan.offload_policy != OFFLOAD_NONE
-                ):
-                    # memory_mode forcing offload silently drops the requested quant; warn so it's diagnosable.
+                elif dense_quant_requested:
+                    # Quant requested but the dense fast path needs a resident load that isn't
+                    # available here (an explicit memory_mode offload, or the dense transformer
+                    # is too large to fit resident). Warn so the disengage is diagnosable, and
+                    # load the GGUF build without it.
                     logger.warning(
-                        "diffusion.transformer_quant: %s requested but memory_mode forces "
-                        "offload (%s); loading GGUF without dense quant",
+                        "diffusion.transformer_quant: %s requested but the dense fast path needs "
+                        "a resident load that doesn't fit here; loading GGUF without dense quant",
                         normalize_transformer_quant(transformer_quant),
-                        plan.offload_policy,
                     )
 
                 if pipe is None:
@@ -1399,6 +1476,7 @@ class DiffusionBackend:
         *,
         kind: str = "gguf",
         repo_id: Optional[str] = None,
+        transformer_mib_override: Optional[int] = None,
     ):
         """Build the memory plan for this load: snapshot free device memory and
         estimate the model's resident footprint, then let the planner pick an
@@ -1425,7 +1503,11 @@ class DiffusionBackend:
             model_dense_mib = estimate_safetensors_dense_mib(cached_mib)
             companion_mib = None
         else:
-            if kind == "single_file":
+            if transformer_mib_override is not None:
+                # Dense-quant preflight: budget the dense bf16 transformer the fast path
+                # materializes, not the (much smaller) single-file/GGUF the plain load uses.
+                transformer_resident = transformer_mib_override
+            elif kind == "single_file":
                 # An fp8 transformer checkpoint loads via from_single_file with a bf16
                 # compute dtype and no quantization_config, so diffusers upcasts it to
                 # bf16 (~2x resident); detect it from the basename. Excludes the
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 134fdbf8c2..d80a8e3b06 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1389,15 +1389,20 @@ async def upload_diffusion_dataset(
                 detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
             )
         names.append(filename)
-    # Roll back all files written this request on a mid-batch failure (size limit,
-    # disk error, disconnect) so the dataset is never left partially populated.
-    written: list[Path] = []
+    # Stage each file to a temp name and only move it into place once the whole batch is
+    # written, so a mid-batch failure (size limit, disk error, disconnect) leaves the
+    # dataset untouched -- including any pre-existing file that shares a name, which a
+    # direct write would have truncated (repeat uploads into the same name accumulate).
+    staged: list[tuple[Path, Path]] = []  # (temp, final)
     committed = False
     try:
         for f, filename in zip(files, names):
             dest = folder / filename
-            written.append(dest)
-            with open(dest, "wb") as out:
+            # A filename-independent temp name so a long (but valid, <= NAME_MAX) filename
+            # can't overflow NAME_MAX once the staging suffix is added.
+            tmp = folder / f".upload-{_uuid.uuid4().hex}.part"
+            staged.append((tmp, dest))
+            with open(tmp, "wb") as out:
                 while chunk := await f.read(1024 * 1024):
                     total_bytes += len(chunk)
                     if total_bytes > limit_bytes:
@@ -1411,12 +1416,14 @@ async def upload_diffusion_dataset(
                         )
                     out.write(chunk)
             uploaded += 1
+        for tmp, dest in staged:
+            tmp.replace(dest)  # atomic on the same filesystem
         committed = True
     finally:
         if not committed:
-            for p in written:
+            for tmp, _ in staged:
                 try:
-                    p.unlink(missing_ok = True)
+                    tmp.unlink(missing_ok = True)
                 except OSError:
                     pass
 
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index a30d3ffadf..1fb6714cf7 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -1950,6 +1950,102 @@ def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, mo
     assert _FakeTransformer.last["path"]  # GGUF path used
 
 
+def test_dense_quant_skipped_when_dense_transformer_does_not_fit(
+    fake_runtime, tmp_path, monkeypatch
+):
+    # The GGUF fits resident (plan `none`), but the DENSE bf16 transformer the fast path
+    # materializes does not. The fast path must be skipped up front (preflighted against
+    # the dense transformer, not the GGUF), and GGUF loads RESIDENT -- not evicted, OOMed
+    # in finalization, then offloaded.
+    from core.inference import diffusion as dmod
+
+    backend = DiffusionBackend()
+    _force_cuda_target(backend, monkeypatch)
+    monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
+    # A scheme resolves and there is no prequant, so the dense bf16 is materialized and the
+    # dense-fit re-check runs against a large (won't-fit) dense transformer.
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
+    )
+    monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
+    monkeypatch.setattr(
+        DiffusionBackend,
+        "_dense_transformer_resident_bytes",
+        staticmethod(lambda base: 40 * 1024 ** 3),
+    )
+    orig_plan = DiffusionBackend._plan_memory
+
+    def plan_wrap(self, *a, transformer_mib_override = None, **k):
+        # GGUF budget fits (real plan -> none); the dense-transformer preflight does not.
+        if transformer_mib_override is not None:
+            return types.SimpleNamespace(offload_policy = "model")
+        return orig_plan(self, *a, **k)
+
+    monkeypatch.setattr(DiffusionBackend, "_plan_memory", plan_wrap)
+
+    @classmethod
+    def _fp_fail(cls, *a, **k):
+        pytest.fail("dense transformer must not load when it won't fit resident")
+
+    monkeypatch.setattr(_FakeTransformer, "from_pretrained", _fp_fail, raising = False)
+    (tmp_path / "m.gguf").write_bytes(b"x")
+    status = backend.load_pipeline(
+        str(tmp_path),
+        gguf_filename = "m.gguf",
+        family_override = "z-image",
+        transformer_quant = "fp8",
+    )
+    assert status["transformer_quant"] is None       # dense quant skipped
+    assert status["offload_policy"] == "none"         # GGUF loaded resident, not offloaded
+    assert _FakeTransformer.last["path"]              # GGUF path used
+
+
+def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypatch):
+    # With a prequant checkpoint, the fast path loads the small quantized file, not the
+    # dense bf16 -- so the dense-transformer re-check must NOT run and must NOT decline the
+    # fast path, even when the base's dense shards happen to be cached and large.
+    from core.inference import diffusion as dmod
+
+    backend = DiffusionBackend()
+    _force_cuda_target(backend, monkeypatch)
+    monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
+    )
+    monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: "prequant/path")
+    # Large dense shards cached: if the re-check ran, it would wrongly decline the fast path.
+    monkeypatch.setattr(
+        DiffusionBackend,
+        "_dense_transformer_resident_bytes",
+        staticmethod(lambda base: 999 * 1024 ** 3),
+    )
+    dense_refit_ran = []
+    orig_plan = DiffusionBackend._plan_memory
+
+    def spy_plan(self, *a, transformer_mib_override = None, **k):
+        if transformer_mib_override is not None:
+            dense_refit_ran.append(True)
+        return orig_plan(self, *a, **k)
+
+    monkeypatch.setattr(DiffusionBackend, "_plan_memory", spy_plan)
+    attempted = []
+
+    def fake_dense_load(self, *a, **k):
+        attempted.append(True)
+        return None, None  # fall through to GGUF; we only assert the path was reached
+
+    monkeypatch.setattr(DiffusionBackend, "_load_dense_quant_pipeline", fake_dense_load)
+    (tmp_path / "m.gguf").write_bytes(b"x")
+    backend.load_pipeline(
+        str(tmp_path),
+        gguf_filename = "m.gguf",
+        family_override = "z-image",
+        transformer_quant = "fp8",
+    )
+    assert dense_refit_ran == []   # prequant -> dense re-check skipped
+    assert attempted == [True]     # fast path still attempted (with the prequant)
+
+
 def test_transformer_quant_unsupported_scheme_skips_dense_download(
     fake_runtime, tmp_path, monkeypatch
 ):
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index 2cff7f65c3..cd271c2f05 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -512,6 +512,30 @@ def test_diffusion_dataset_upload_over_cap_rolls_back_whole_batch(
     assert not (folder / "big.png").exists()
 
 
+def test_diffusion_dataset_upload_over_cap_preserves_existing_file(
+    client, dataset_roots, monkeypatch
+):
+    # A failed batch that reuses an existing filename must NOT delete the user's
+    # pre-existing file (repeat uploads accumulate). Staging to a temp file keeps the
+    # original intact until the whole batch commits.
+    import utils.upload_limits as ul
+
+    monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 100)
+    ds_root, _ = dataset_roots
+    folder = ds_root / "keep"
+    folder.mkdir(parents = True)
+    (folder / "existing.png").write_bytes(b"ORIGINAL")  # from an earlier upload
+    files = [
+        ("files", ("existing.png", b"NEW", "image/png")),  # re-upload, small
+        ("files", ("big.png", b"y" * 200, "image/png")),   # trips the cap
+    ]
+    r = client.post("/api/train/diffusion/dataset", data = {"name": "keep"}, files = files)
+    assert r.status_code == 413, r.text
+    assert (folder / "existing.png").read_bytes() == b"ORIGINAL"  # untouched
+    assert not (folder / "big.png").exists()
+    assert not list(folder.glob(".*.part"))  # no leftover temp files
+
+
 def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypatch):
     # A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed,
     # so a bad pick never unloads the user's working chat/Images model.

From 35dccf2da68fcbe8614a5f1b003a4dbef66d02da Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Mon, 6 Jul 2026 10:08:19 -0300
Subject: [PATCH 09/13] Drop redundant transformer_quant re-normalization in
 the dense-quant gate

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

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 4d4f95541f..a252e39790 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -891,8 +891,8 @@ class DiffusionBackend:
 
     @staticmethod
     def _safetensors_param_count(path: Path) -> int:
-        """Total tensor elements in a safetensors file, read from its JSON header (an
-        8-byte little-endian length prefix then the header) without touching tensor data."""
+        """Total tensor elements in a safetensors file, read from its JSON header without
+        touching the tensor data. 0 on any read/parse failure."""
         try:
             with open(path, "rb") as fh:
                 header_len = int.from_bytes(fh.read(8), "little")
@@ -1062,7 +1062,7 @@ class DiffusionBackend:
                 if dense_quant_requested and dense_fits:
                     scheme = select_transformer_quant_scheme(
                         target,
-                        normalize_transformer_quant(transformer_quant),
+                        transformer_quant,  # normalized above
                         family = getattr(fam, "name", None),
                     )
                     prequant = (
@@ -1135,7 +1135,7 @@ class DiffusionBackend:
                     logger.warning(
                         "diffusion.transformer_quant: %s requested but the dense fast path needs "
                         "a resident load that doesn't fit here; loading GGUF without dense quant",
-                        normalize_transformer_quant(transformer_quant),
+                        transformer_quant,  # normalized above
                     )
 
                 if pipe is None:

From 04ddf5449c11641ce5e119078716a5ac846d27f8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 6 Jul 2026 13:18:19 +0000
Subject: [PATCH 10/13] [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 +
 .../backend/tests/test_diffusion_backend.py   | 28 +++++++++++++------
 .../backend/tests/test_diffusion_training.py  |  2 +-
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index a252e39790..7c632788e4 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -920,6 +920,7 @@ class DiffusionBackend:
         safetensors shard headers. Returns 0 when no ``transformer/*.safetensors`` shards
         are present (an uncached base, or a .bin-only transformer); the caller then gates
         the fast path on the plain plan."""
+
         def _params(d: Path) -> int:
             tdir = d / "transformer"
             if not tdir.is_dir():
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index 57f9e5a24c..3da7b8840c 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -2005,11 +2005,16 @@ def test_dense_quant_skipped_when_dense_transformer_does_not_fit(
     monkeypatch.setattr(
         DiffusionBackend,
         "_dense_transformer_resident_bytes",
-        staticmethod(lambda base: 40 * 1024 ** 3),
+        staticmethod(lambda base: 40 * 1024**3),
     )
     orig_plan = DiffusionBackend._plan_memory
 
-    def plan_wrap(self, *a, transformer_mib_override = None, **k):
+    def plan_wrap(
+        self,
+        *a,
+        transformer_mib_override = None,
+        **k,
+    ):
         # GGUF budget fits (real plan -> none); the dense-transformer preflight does not.
         if transformer_mib_override is not None:
             return types.SimpleNamespace(offload_policy = "model")
@@ -2029,9 +2034,9 @@ def test_dense_quant_skipped_when_dense_transformer_does_not_fit(
         family_override = "z-image",
         transformer_quant = "fp8",
     )
-    assert status["transformer_quant"] is None       # dense quant skipped
-    assert status["offload_policy"] == "none"         # GGUF loaded resident, not offloaded
-    assert _FakeTransformer.last["path"]              # GGUF path used
+    assert status["transformer_quant"] is None  # dense quant skipped
+    assert status["offload_policy"] == "none"  # GGUF loaded resident, not offloaded
+    assert _FakeTransformer.last["path"]  # GGUF path used
 
 
 def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypatch):
@@ -2051,12 +2056,17 @@ def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypa
     monkeypatch.setattr(
         DiffusionBackend,
         "_dense_transformer_resident_bytes",
-        staticmethod(lambda base: 999 * 1024 ** 3),
+        staticmethod(lambda base: 999 * 1024**3),
     )
     dense_refit_ran = []
     orig_plan = DiffusionBackend._plan_memory
 
-    def spy_plan(self, *a, transformer_mib_override = None, **k):
+    def spy_plan(
+        self,
+        *a,
+        transformer_mib_override = None,
+        **k,
+    ):
         if transformer_mib_override is not None:
             dense_refit_ran.append(True)
         return orig_plan(self, *a, **k)
@@ -2076,8 +2086,8 @@ def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypa
         family_override = "z-image",
         transformer_quant = "fp8",
     )
-    assert dense_refit_ran == []   # prequant -> dense re-check skipped
-    assert attempted == [True]     # fast path still attempted (with the prequant)
+    assert dense_refit_ran == []  # prequant -> dense re-check skipped
+    assert attempted == [True]  # fast path still attempted (with the prequant)
 
 
 def test_transformer_quant_unsupported_scheme_skips_dense_download(
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index 9233a54def..ff0b2f12ff 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -583,7 +583,7 @@ def test_diffusion_dataset_upload_over_cap_preserves_existing_file(
     (folder / "existing.png").write_bytes(b"ORIGINAL")  # from an earlier upload
     files = [
         ("files", ("existing.png", b"NEW", "image/png")),  # re-upload, small
-        ("files", ("big.png", b"y" * 200, "image/png")),   # trips the cap
+        ("files", ("big.png", b"y" * 200, "image/png")),  # trips the cap
     ]
     r = client.post("/api/train/diffusion/dataset", data = {"name": "keep"}, files = files)
     assert r.status_code == 413, r.text

From 0d0a6d2b9603ae75532f2e44cffcf972e4db8882 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 7 Jul 2026 13:39:34 +0000
Subject: [PATCH 11/13] Drop piecewise_constant from the trainable lr_scheduler
 allow-list

piecewise_constant is the only diffusers scheduler that needs a step_rules string, and
neither diffusion trainer passes one (get_scheduler is called with only warmup/training
steps, and there is no config field for it). Accepting it let /diffusion/start pass
normalized(), free the resident GPU workloads, spawn the trainer, and only then crash in
the subprocess (get_piecewise_constant_schedule does step_rules.split(",") on None) -- the
exact evict-then-fail the up-front validation exists to prevent. Reject it now with a clear
400. The remaining six schedulers all run with only warmup/training steps.
---
 .../core/training/diffusion_train_common.py   |  9 +++++--
 .../tests/test_diffusion_lora_trainer.py      | 26 +++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index aabfbf42a6..aefe2de3e1 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -35,7 +35,13 @@ from core.inference.diffusion_families import (
 # set overrides this in its own defaults; kept here so DiffusionLoraConfig has a sane fallback.
 DEFAULT_LORA_TARGETS: tuple[str, ...] = ("to_k", "to_q", "to_v", "to_out.0")
 
-# diffusers' SchedulerType names (diffusers.optimization.get_scheduler).
+# diffusers' SchedulerType names (diffusers.optimization.get_scheduler). piecewise_constant is
+# intentionally excluded: it is the only scheduler that needs a `step_rules` string, which the
+# trainers never pass (get_scheduler is called with only warmup/training steps, and there is no
+# config field for it). Accepting it would pass normalized(), free the resident GPU workloads,
+# then crash in the trainer subprocess (get_piecewise_constant_schedule does step_rules.split(",")
+# on None) -- the exact evict-then-fail the up-front validation exists to prevent. The remaining
+# six all run with only warmup/training steps.
 _LR_SCHEDULERS: frozenset[str] = frozenset(
     {
         "linear",
@@ -44,7 +50,6 @@ _LR_SCHEDULERS: frozenset[str] = frozenset(
         "polynomial",
         "constant",
         "constant_with_warmup",
-        "piecewise_constant",
     }
 )
 
diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py
index bd91eae9c8..d89a472bd3 100644
--- a/studio/backend/tests/test_diffusion_lora_trainer.py
+++ b/studio/backend/tests/test_diffusion_lora_trainer.py
@@ -114,6 +114,32 @@ def test_config_normalized_validation(kw):
         DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw).normalized()
 
 
+def test_normalized_rejects_piecewise_constant():
+    # piecewise_constant needs a step_rules string the trainers never supply, so get_scheduler()
+    # would crash in the trainer subprocess AFTER the resident GPU workloads are freed. It must be
+    # rejected up front (a clean ValueError -> 400), not accepted like the other schedulers.
+    with pytest.raises(ValueError, match = "lr_scheduler"):
+        DiffusionLoraConfig(
+            base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = "piecewise_constant"
+        ).normalized()
+
+
+def test_normalized_accepts_supported_schedulers():
+    # Every scheduler in the allow-list runs with only warmup/training steps (no extra required arg).
+    for sched in (
+        "linear",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+        "constant",
+        "constant_with_warmup",
+    ):
+        cfg = DiffusionLoraConfig(
+            base_model = "b", data_dir = "d", output_dir = "o", lr_scheduler = sched
+        ).normalized()
+        assert cfg.lr_scheduler == sched
+
+
 def test_compute_sdxl_add_time_ids():
     assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024)
 

From 46a0a21d53b8dd9b76903c697f882e38c4c0b902 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 7 Jul 2026 14:23:14 +0000
Subject: [PATCH 12/13] Drop piecewise_constant from the diffusion training API
 scheduler enum

Follow-up to removing piecewise_constant from the trainable scheduler allow-list: the
DiffusionTrainingStartRequest.lr_scheduler Literal still advertised it, so a client that
picked it straight from the schema passed request validation and then hit the 400 from
normalized(). Remove it from the enum too so the API only offers schedulers the trainers
can actually run, and add a test asserting the enum never advertises a scheduler outside
the validation allow-list.
---
 studio/backend/models/training.py               |  1 -
 .../tests/test_diffusion_lora_trainer.py        | 17 +++++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 1d57322781..8b32929466 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -720,7 +720,6 @@ class DiffusionTrainingStartRequest(BaseModel):
         "polynomial",
         "constant",
         "constant_with_warmup",
-        "piecewise_constant",
     ] = Field("constant")
     lr_warmup_steps: int = Field(0, ge = 0)
     center_crop: bool = Field(False)
diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py
index d89a472bd3..67a0115ae5 100644
--- a/studio/backend/tests/test_diffusion_lora_trainer.py
+++ b/studio/backend/tests/test_diffusion_lora_trainer.py
@@ -140,6 +140,23 @@ def test_normalized_accepts_supported_schedulers():
         assert cfg.lr_scheduler == sched
 
 
+def test_api_scheduler_enum_never_advertises_a_rejected_scheduler():
+    # The request-model enum must not offer a scheduler that normalized() rejects: a client that
+    # picks it straight from the schema would get a 400. Every option the API advertises must be in
+    # the validation allow-list (this guards against the enum and allow-list drifting apart again,
+    # e.g. piecewise_constant left in one but removed from the other).
+    import typing
+
+    from core.training.diffusion_train_common import _LR_SCHEDULERS
+    from models.training import DiffusionTrainingStartRequest
+
+    api_options = set(
+        typing.get_args(DiffusionTrainingStartRequest.model_fields["lr_scheduler"].annotation)
+    )
+    assert api_options and api_options <= _LR_SCHEDULERS, api_options - _LR_SCHEDULERS
+    assert "piecewise_constant" not in api_options
+
+
 def test_compute_sdxl_add_time_ids():
     assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024)
 

From 41bdc116be7f0ecf258560706ae72fc8f4caa803 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 7 Jul 2026 14:55:37 +0000
Subject: [PATCH 13/13] Reject unknown union ControlNet control types instead
 of defaulting to canny

union_control_mode fell back to control_mode=0 (the canny head) for ANY unmapped control
type, so a typo'd or unsupported value like 'detph' silently conditioned the map as canny
instead of failing. preprocess_control passes non-canny maps through unchanged, so that map
would be interpreted under the wrong mode with no error. Now only 'passthrough' (or an empty
type) keeps the deliberate mode-0 default; any other unknown type raises ValueError, which
the generate route maps to a 400. Known modes are unchanged.
---
 .../core/inference/diffusion_controlnet.py    | 19 +++++++++++++++----
 .../tests/test_diffusion_controlnet.py        | 15 +++++++++++++++
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py
index 170b48bd0f..a4e482121a 100644
--- a/studio/backend/core/inference/diffusion_controlnet.py
+++ b/studio/backend/core/inference/diffusion_controlnet.py
@@ -220,13 +220,24 @@ _UNION_CONTROL_MODES: dict[str, int] = {
 def union_control_mode(spec_id: str, control_type: str) -> Optional[int]:
     """The integer ``control_mode`` for a union ControlNet, or None.
 
-    A union model requires a concrete mode (diffusers raises on None), so a curated union
-    entry always gets an index, defaulting to 0 for types like 'passthrough' that carry
-    none. A non-union entry returns None so the caller omits the kwarg."""
+    A union model requires a concrete mode (diffusers raises on None). A known mode maps to its
+    index; ``passthrough`` (or an empty type) carries no intrinsic mode and defaults to 0 (the
+    canny head). An unknown/typo'd type (e.g. 'detph') raises ValueError so the route rejects it
+    with a 400 instead of silently running the canny head against a map meant for another mode,
+    which would produce wrong conditioning. A non-union entry returns None so the caller omits the
+    kwarg."""
     entry = _catalog_by_id().get(spec_id)
     if entry is None or not entry.is_union:
         return None
-    return _UNION_CONTROL_MODES.get((control_type or "").strip().lower(), 0)
+    ct = (control_type or "").strip().lower()
+    if ct in _UNION_CONTROL_MODES:
+        return _UNION_CONTROL_MODES[ct]
+    if ct in ("", "passthrough"):
+        return 0  # already-preprocessed map with no intrinsic mode; canny is the default head
+    raise ValueError(
+        f"Unknown control type {control_type!r} for a union ControlNet. Use one of: "
+        f"{', '.join(sorted(_UNION_CONTROL_MODES))}, or passthrough."
+    )
 
 
 def preprocess_control(image: Any, control_type: str) -> Any:
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index ab5eb71b22..d5282299d4 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -66,6 +66,21 @@ def test_union_control_mode_maps_only_union_entries():
     assert dc.union_control_mode("some/bare-repo", "canny") is None
 
 
+def test_union_control_mode_rejects_unknown_type():
+    # An unknown / typo'd control type (e.g. 'detph') must NOT silently fall back to the canny
+    # head (0): preprocess_control passes non-canny maps through unchanged, so mode 0 would
+    # condition a map meant for another mode as canny -- silently wrong. Only passthrough (or an
+    # empty type) defaults to 0; anything else raises so the route returns a 400.
+    with pytest.raises(ValueError, match = "Unknown control type"):
+        dc.union_control_mode("flux-union-pro", "detph")
+    with pytest.raises(ValueError, match = "Unknown control type"):
+        dc.union_control_mode("flux-union-pro", "scribble")
+    # passthrough and empty still default to 0 (the intended no-intrinsic-mode case); a non-union
+    # entry is unaffected (returns None, never raises).
+    assert dc.union_control_mode("flux-union-pro", "") == 0
+    assert dc.union_control_mode("some/bare-repo", "detph") is None
+
+
 def test_resolve_controlnet_local(tmp_path, monkeypatch):
     d = tmp_path / "controlnets"
     d.mkdir()