Merge image-generation bug fixes (#6872)

Fold PR #6872's image-generation fixes into the branch, deduped against the
round-12 dataset-upload and gallery integrity work already on image-generation.

Fixes carried forward from #6872:
- fp8 single-file transformer memory estimate: an fp8 checkpoint loads with no
  quantization_config and diffusers upcasts it to bf16 (~2x resident), so budget
  it accordingly in _plan_memory and estimate_safetensors_dense_mib.
- dense-quant OOM-evict preflight: when the GGUF fits resident but the dense bf16
  transformer this path materializes does not, skip the fast path up front rather
  than evict the current pipeline and OOM in finalization. Combined with the
  existing offload->resident candidate re-plan so both the family-table estimate
  and the on-disk shard measurement gate engagement (unified on the
  transformer_resident_override_mib plan override).
- ControlNet: evict the previous module and its from_pipe wrapper before loading a
  new one so swapping ControlNets within a base-model load cannot accumulate to OOM.
- ControlNet union_control_mode: raise on an unknown control type instead of
  silently defaulting to canny.
- edit-family mask rejection: raise instead of silently dropping a mask on an
  image-editing model that has no inpaint pipeline.
- companion cache: walk the snapshot dir and exclude transformer/ so the
  dense-quant prefetch's cached shards do not inflate the companion total and
  wrongly force offload.
- training: drop piecewise_constant from the LR scheduler enum and force bf16 for
  fp16-incompatible families.
- dataset upload: batch-atomic staging with the same-stem duplicate guard.
- images page: guard negative-prompt restore on guidance>0, clear stale ControlNet
  selection on restore, and revert an optimistic quant label when a pipeline load
  never starts.
- uninstall (sh + ps1): keep the owner-marker guard on sd.cpp removal.

Conflicts resolved in favour of image-generation's evolved memory system,
loadSpecFor catalog, and stop-and-save (lora_path) run detection; #6872's fp8 and
dense-preflight fixes carried forward on top. All affected backend tests pass
(test_diffusion_backend, test_diffusion_training, test_diffusion_lora_trainer,
test_video_gallery, test_diffusion_controlnet).

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

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

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 {

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
@@ -1010,10 +1011,14 @@ class DiffusionBackend:
         return total, base_files
 
     @staticmethod
-    def _cache_bytes(repo_id: str) -> int:
+    def _hub_cache_repo_dir(repo_id: str) -> Path:
+        """Local HF hub cache dir for ``repo_id``."""
         from huggingface_hub import constants
+        return Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}"
 
-        blobs = Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" / "blobs"
+    @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():
@@ -1049,19 +1054,72 @@ 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.
 
-        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."""
-        local = Path(base).expanduser()
-        if local.is_dir():
-            return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True)
-        return DiffusionBackend._cache_bytes(base)
+        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."""
+        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 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")
+                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 ───────────────────────────────
 
@@ -1181,9 +1239,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,
@@ -1227,48 +1284,96 @@ class DiffusionBackend:
                 pipe = None
                 transformer_quant_engaged = None
                 quant_plan = None
+                # The GGUF-size `plan` can mis-budget the dense-quant fast path two ways, so
+                # preflight the real footprint BEFORE evicting the current pipeline. Both
+                # branches need the base repo + a resolved scheme, so gate on the dense-path
+                # preconditions first.
+                dense_declined = False
                 if (
                     kind == "gguf"
                     and normalize_transformer_quant(transformer_quant) is not None
                     and dense_transformer_supported(target)
-                    and plan.offload_policy != OFFLOAD_NONE
                 ):
-                    # The GGUF-size plan picked offload, but the dense-quant artifact has
-                    # a DIFFERENT footprint: int8/fp8 weights are ~half the bf16 bytes, and
-                    # a pre-quantized checkpoint never materialises dense bf16 at all. Ask
-                    # the auto-policy for the candidate's estimate and re-plan against it:
-                    # a resident quantised build beats an offloaded GGUF on speed AND
-                    # quality, so it must be attempted before settling for offload.
-                    candidate = resolve_dense_quant_candidate(
-                        fam = fam,
-                        target = target,
-                        requested = transformer_quant,
-                        base_repo = base,
-                        prequant_path = transformer_prequant_path,
-                        logger = logger,
-                    )
-                    if candidate is not None:
-                        replanned = self._plan_memory(
-                            target,
-                            single_file_path,
-                            base,
-                            fam,
-                            memory_mode,
-                            cpu_offload,
-                            kind = kind,
-                            repo_id = repo_id,
-                            transformer_resident_override_mib = (candidate.transient_transformer_mib),
-                            # The dense path prefetches the base transformer/ shards into the
-                            # cache _companion_cache_bytes reads; pass the auto-policy's own
-                            # companion estimate so the re-plan does not double-count them.
-                            companion_override_mib = candidate.companions_mib,
+                    if plan.offload_policy != OFFLOAD_NONE:
+                        # The GGUF-size plan picked offload, but the dense-quant artifact has
+                        # a DIFFERENT footprint: int8/fp8 weights are ~half the bf16 bytes, and
+                        # a pre-quantized checkpoint never materialises dense bf16 at all. Ask
+                        # the auto-policy for the candidate's estimate and re-plan against it:
+                        # a resident quantised build beats an offloaded GGUF on speed AND
+                        # quality, so it must be attempted before settling for offload.
+                        candidate = resolve_dense_quant_candidate(
+                            fam = fam,
+                            target = target,
+                            requested = transformer_quant,
+                            base_repo = base,
+                            prequant_path = transformer_prequant_path,
+                            logger = logger,
                         )
-                        if replanned.offload_policy == OFFLOAD_NONE:
-                            quant_plan = replanned
+                        if candidate is not None:
+                            replanned = self._plan_memory(
+                                target,
+                                single_file_path,
+                                base,
+                                fam,
+                                memory_mode,
+                                cpu_offload,
+                                kind = kind,
+                                repo_id = repo_id,
+                                transformer_resident_override_mib = (
+                                    candidate.transient_transformer_mib
+                                ),
+                                # The dense path prefetches the base transformer/ shards into the
+                                # cache _companion_cache_bytes reads; pass the auto-policy's own
+                                # companion estimate so the re-plan does not double-count them.
+                                companion_override_mib = candidate.companions_mib,
+                            )
+                            if replanned.offload_policy == OFFLOAD_NONE:
+                                quant_plan = replanned
+                    else:
+                        # The GGUF fits resident, but this path first materialises the base
+                        # repo's dense bf16 transformer -- bigger than the quantised GGUF -- so
+                        # re-check the fit against THAT. A card that fits the GGUF but not the
+                        # dense transformer must skip the fast path up front, not evict the
+                        # current pipeline then OOM in finalization. A prequant checkpoint loads
+                        # a small quantised file (no dense bf16), so skip the re-check there
+                        # (mirrors the prefetch guard). _dense_transformer_resident_bytes reads
+                        # the on-disk shard headers, so it also covers families the size table
+                        # (resolve_dense_quant_candidate) does not list; it returns 0 when the
+                        # shards are absent, in which case the fast path keeps today's behaviour.
+                        scheme = select_transformer_quant_scheme(
+                            target,
+                            transformer_quant,  # normalized above
+                            family = getattr(fam, "name", None),
+                        )
+                        prequant = (
+                            resolve_prequant_source(
+                                fam, scheme, path_override = transformer_prequant_path
+                            )
+                            if scheme is not None
+                            else None
+                        )
+                        if prequant is None:
+                            dense_mib = int(
+                                self._dense_transformer_resident_bytes(base) // (1024 * 1024)
+                            )
+                            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_resident_override_mib = dense_mib,
+                                )
+                                dense_declined = dense_plan.offload_policy != OFFLOAD_NONE
                 if (
                     kind == "gguf"
                     and normalize_transformer_quant(transformer_quant) is not None
                     and dense_transformer_supported(target)
+                    and not dense_declined
                     and (plan.offload_policy == OFFLOAD_NONE or quant_plan is not None)
                 ):
                     try:
@@ -1837,7 +1942,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.
         ``transformer_resident_override_mib`` replaces the file-size transformer estimate
         when the loader is planning for a DIFFERENT artifact than the file on disk (the
@@ -1900,9 +2006,15 @@ class DiffusionBackend:
                 # file-size derivation; companions below stay measured from the cache.
                 transformer_resident = transformer_resident_override_mib
             elif kind == "single_file":
-                # Safetensors single-file: no dequant expansion (it carries its dtype).
+                # 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
+                )
                 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))
@@ -2011,6 +2123,13 @@ class DiffusionBackend:
                 _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None)
                 if _cn_fs.blocked:
                     raise ValueError(_cn_fs.reason)
+            # 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()
+                clear_gpu_cache()
             import torch
 
             # state.dtype is the display string saved at load ("bfloat16"), NOT a
@@ -2415,6 +2534,12 @@ 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 mask would be silently dropped.
+                        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:

studio/backend/core/inference/diffusion_controlnet.py

@@ -204,8 +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 it maps to nothing (the caller omits control_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,
@@ -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.
 
-    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 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())
+    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:

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,11 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
         mps_available = False
 
     if mps_available:
+        # 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
         # far outside float16's finite range (~6.5e4) -- Z-Image's MLP
@@ -254,7 +260,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",

studio/backend/core/inference/diffusion_lora.py

@@ -293,6 +293,13 @@ 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."""
+    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]],
     *,
@@ -301,20 +308,31 @@ 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."""
+    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,
+        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

studio/backend/core/inference/diffusion_memory.py

@@ -236,14 +236,19 @@ 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).
-    """
+    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:
+        return storage_mib * 2
     return storage_mib
 
 

studio/backend/core/inference/sd_cpp_backend.py

@@ -768,6 +768,10 @@ 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, 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:
             raise ValueError(
                 "ControlNet is not yet supported on the native sd.cpp engine; run on a GPU "

studio/backend/core/inference/sd_cpp_engine.py

@@ -125,7 +125,9 @@ 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 lands in its own versioned subdir rather than flattening into ``root``."""
     name = _binary_name(stem)
     cands = [
         root / "build" / "bin" / name,
@@ -133,6 +135,15 @@ def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
         root / "bin" / name,
         root / name,
     ]
+    # 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)
+        for sub in subdirs:
+            cands.append(sub / name)
+            cands.append(sub / "bin" / name)
+    except OSError:
+        pass
     return cands
 
 

studio/backend/core/training/diffusion_train_common.py

@@ -37,6 +37,28 @@ 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). 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",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+        "constant",
+        "constant_with_warmup",
+    }
+)
+
+# 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"}
 _CAPTION_EXTS = (".txt", ".caption")
 # diffusers' canonical single-file LoRA name, so load_lora_weights(dir) finds it.
@@ -85,7 +107,11 @@ 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:
+    # 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):
         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()}"
@@ -471,6 +497,17 @@ 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")
+        # 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 "
+                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}"
+            )
         if not 1 <= int(self.cache_variants) <= 16:
             raise ValueError("cache_variants must be between 1 and 16")
         compile_transformer = str(self.compile_transformer or "auto").strip().lower()
@@ -870,8 +907,19 @@ 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: pick the next
+        # free numeric suffix instead.
+        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)

studio/backend/models/training.py

@@ -722,7 +722,14 @@ class DiffusionTrainingStartRequest(BaseModel):
         5.0, gt = 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",
+    ] = Field("constant")
     lr_warmup_steps: int = Field(0, ge = 0)
     center_crop: bool = Field(False)
     random_flip: bool = Field(True)

studio/backend/routes/training.py

@@ -1581,6 +1581,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:
         # Normalise to a safe basename. Path.name does not split on a backslash on POSIX, so a
         # Windows client that sends a backslash path in the multipart filename would otherwise be
@@ -1597,40 +1600,61 @@ async def upload_diffusion_dataset(
                 status_code = 400,
                 detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
             )
-        dest = folder / filename
         # Reject a second IMAGE that shares this one's stem but differs by extension (sample.png
         # vs sample.jpg): both resolve to the same <stem>.txt caption sidecar (the kohya/diffusers
         # convention the reader, editor, and delete paths all use), so keeping both would silently
-        # make them share -- and corrupt -- one caption during training. Scan the whole folder, not
-        # just this batch, because uploads accumulate (earlier-in-batch files are already on disk).
-        # Re-uploading the exact same name (same stem AND extension) stays allowed as an overwrite;
-        # caption/text files are exempt (attaching sample.txt to sample.png is the intended flow).
+        # share -- and corrupt -- one caption during training. Check both files already on disk
+        # (uploads accumulate) and earlier images validated in THIS batch (nothing is on disk yet
+        # in this up-front pass). Re-uploading the exact same name (same stem AND extension) stays
+        # an overwrite; caption/text files are exempt (sample.txt for sample.png is intended).
         if ext in _DIFFUSION_DATASET_IMAGE_EXTS:
             stem = Path(filename).stem
-            for existing in folder.iterdir():
-                if (
-                    existing.is_file()
-                    and existing.name != filename
-                    and existing.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS
-                    and existing.stem == stem
-                ):
-                    raise HTTPException(
-                        status_code = 400,
-                        detail = (
-                            f"Duplicate image name '{stem}'. '{existing.name}' is already in this "
-                            f"dataset; two images sharing a name would share one '{stem}.txt' "
-                            f"caption. Rename one before uploading."
-                        ),
-                    )
-        # Stream into a sibling temp file and only atomically promote it once the whole file
-        # is written and within the limit. A mid-stream 413 (or any abort) then removes the
-        # TEMP file, never dest, so re-uploading a batch that trips the limit can no longer
-        # truncate/delete an example that was already stored under the same name.
-        fd, tmp_name = tempfile.mkstemp(dir = folder, prefix = ".upload-", suffix = ext)
-        tmp = Path(tmp_name)
-        complete = False
-        try:
-            with os.fdopen(fd, "wb") as out:
+            clash = next(
+                (
+                    p.name
+                    for p in folder.iterdir()
+                    if p.is_file()
+                    and p.name != filename
+                    and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS
+                    and p.stem == stem
+                ),
+                None,
+            )
+            if clash is None:
+                clash = next(
+                    (
+                        n
+                        for n in names
+                        if n != filename
+                        and Path(n).suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS
+                        and Path(n).stem == stem
+                    ),
+                    None,
+                )
+            if clash is not None:
+                raise HTTPException(
+                    status_code = 400,
+                    detail = (
+                        f"Duplicate image name '{stem}'. '{clash}' is already in this "
+                        f"dataset; two images sharing a name would share one '{stem}.txt' "
+                        f"caption. Rename one before uploading."
+                    ),
+                )
+        names.append(filename)
+    # 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
+            # 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:
@@ -1643,15 +1667,17 @@ async def upload_diffusion_dataset(
                             ),
                         )
                     out.write(chunk)
-            os.replace(tmp, dest)
-            complete = True
-        finally:
-            if not complete:
+            uploaded += 1
+        for tmp, dest in staged:
+            tmp.replace(dest)  # atomic on the same filesystem
+        committed = True
+    finally:
+        if not committed:
+            for tmp, _ in staged:
                 try:
                     tmp.unlink(missing_ok = True)
                 except OSError:
                     pass
-        uploaded += 1
 
     summary = _diffusion_dataset_summary(folder)
     return DiffusionDatasetUploadResponse(
@@ -2123,7 +2149,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:
@@ -2134,7 +2160,8 @@ 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, 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(
         p

studio/backend/tests/test_diffusion_backend.py

@@ -2396,6 +2396,112 @@ 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_resident_override_mib = None,
+        **k,
+    ):
+        # GGUF budget fits (real plan -> none); the dense-transformer preflight does not.
+        if transformer_resident_override_mib 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_resident_override_mib = None,
+        **k,
+    ):
+        if transformer_resident_override_mib 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
 ):

studio/backend/tests/test_diffusion_controlnet.py

@@ -56,15 +56,31 @@ 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
 
 
+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()

studio/backend/tests/test_diffusion_lora.py

@@ -190,6 +190,29 @@ 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):
+        # 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)",
+            response = types.SimpleNamespace(headers = {}, request = None),
+        )
+
+    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.

studio/backend/tests/test_diffusion_lora_trainer.py

@@ -240,6 +240,49 @@ def test_config_from_dict_threads_num_epochs():
     assert cfg.num_epochs == 12
 
 
+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_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)
 
@@ -471,3 +514,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")

studio/backend/tests/test_diffusion_training.py

@@ -1015,6 +1015,50 @@ def test_import_example_partial_failure_leaves_no_partial_dataset(
     assert not any(p.name.startswith(".dreambooth-dog.import-") for p in ds_root.iterdir())
 
 
+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_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.

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.

studio/frontend/src/features/images/diffusion-train-dialog.tsx

@@ -120,24 +120,27 @@ 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 per finished run so it rescans the LoRA picker.
   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;

studio/frontend/src/features/images/images-page.tsx

@@ -1277,7 +1277,8 @@ 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 ?? "");
+    // 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);
     setSeed(String(image.seed));
@@ -1312,6 +1313,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
       if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight });
     }
     setLoras(restoredLoras);
+    // The control image isn't persisted, so clear any stale ControlNet selection.
+    setControlnetId("");
+    setControlImage(null);
     toast.success("Settings restored to inputs");
   }, []);
 
@@ -1649,11 +1653,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         toast.error("Only unsloth or on-device image models can be loaded here");
         return;
       }
+      // Optimistically clear the quant label, revert it if the load never starts.
+      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],
   );