Merge remote-tracking branch 'origin/diffusion-lora-training-api' into diffusion-lora-training-ui

studio/backend/core/inference/diffusion.py

@@ -68,14 +68,6 @@ from .diffusion_attention import (
 )
 from . import diffusion_compile_cache as compile_cache
 from . import diffusion_gguf_compile as gguf_compile
-from .diffusion_eager_patches import (
-    install_compile_safe_patches,
-    uninstall_patches,
-)
-from .diffusion_arch_patches import (
-    install_arch_patches,
-    uninstall_arch_patches,
-)
 from .diffusion_cache import apply_step_cache
 from .diffusion_precision import quantize_text_encoders
 from .diffusion_prequant import (
@@ -149,11 +141,6 @@ def _decode_b64_image(data: str, *, mode: str = "RGB") -> Any:
         blob = base64.b64decode(raw, validate = False)
     except (binascii.Error, ValueError) as exc:
         raise ValueError(f"Invalid base64 image data: {exc}") from exc
-    try:
-        img = Image.open(io.BytesIO(blob))
-        img.load()
-    except Exception as exc:  # noqa: BLE001 — surfaced as a 400 to the client
-        raise ValueError(f"Could not decode image: {exc}") from exc
     # Bound the decoded size. Every image-conditioned workflow (img2img / inpaint / upscale /
     # reference / edit) decodes through here, so this single guard protects init, mask, and
     # each reference image uniformly. PIL only WARNS in its 89-178MP "decompression bomb" soft
@@ -161,9 +148,19 @@ def _decode_b64_image(data: str, *, mode: str = "RGB") -> Any:
     # well below that. 4096px covers txt2img's 2048 max, upscales, and normal outpaint canvases;
     # anything larger is rejected with a clear 400 instead of risking an OOM.
     max_side = 4096
-    w, h = img.size
-    if w > max_side or h > max_side:
-        raise ValueError(f"Image is too large ({w}x{h}); maximum is {max_side}px per side.")
+    try:
+        img = Image.open(io.BytesIO(blob))
+        # Read the declared dimensions from the header (Image.open is lazy) and reject an
+        # over-limit image BEFORE img.load() decompresses its pixels, so a crafted
+        # small-payload/huge-dimension file can't spike memory before the guard runs.
+        w, h = img.size
+        if w > max_side or h > max_side:
+            raise ValueError(f"Image is too large ({w}x{h}); maximum is {max_side}px per side.")
+        img.load()
+    except ValueError:
+        raise  # the size guard's own message; don't wrap it as a decode error
+    except Exception as exc:  # noqa: BLE001 — surfaced as a 400 to the client
+        raise ValueError(f"Could not decode image: {exc}") from exc
     return img.convert(mode)
 
 
@@ -507,6 +504,23 @@ class DiffusionBackend:
         if kind in ("gguf", "single_file"):
             if not gguf_filename:
                 raise ValueError(f"a single-file checkpoint name is required for a '{kind}' load.")
+            # Fail a kind/extension mismatch here (before the route evicts chat and grabs the
+            # GPU), instead of deep in the background from_single_file: a "gguf" load needs a
+            # .gguf file, and a "single_file" load must not be handed a .gguf.
+            is_gguf_name = gguf_filename.lower().endswith(".gguf")
+            if kind == "gguf" and not is_gguf_name:
+                raise ValueError("a 'gguf' load requires a .gguf checkpoint name.")
+            if kind == "single_file" and is_gguf_name:
+                raise ValueError("a .gguf checkpoint needs model_kind 'gguf', not 'single_file'.")
+            # A single-file load must name an actual checkpoint: an arbitrary repo file
+            # (README.md, config.json) would pass preflight, evict the chat model, and
+            # only fail in the background from_single_file -- the eviction this
+            # validation exists to prevent.
+            if kind == "single_file" and not gguf_filename.lower().endswith(".safetensors"):
+                raise ValueError(
+                    f"'{gguf_filename}' is not a loadable single-file checkpoint "
+                    f"(expected a .safetensors name; use a .gguf name for a GGUF load)."
+                )
             if local_root.exists():
                 resolve_local_gguf_child(local_root, gguf_filename)
             elif path_shaped:
@@ -523,6 +537,16 @@ class DiffusionBackend:
                     )
             elif path_shaped:
                 raise FileNotFoundError(f"Local model path does not exist: {repo_id}")
+            elif repo_id.upper().endswith("-GGUF"):
+                # A remote "*-GGUF" id is a single-file GGUF repo, not a full diffusers
+                # pipeline: loading it as a pipeline passes the trusted-repo check, evicts
+                # chat, then fails in the background when from_pretrained finds no
+                # model_index.json. Reject the certain case here (no network round-trip)
+                # so the bad pick fails before the GPU handoff, as the route expects.
+                raise ValueError(
+                    f"'{repo_id}' is a single-file GGUF repo; load it with model_kind 'gguf' "
+                    f"and a .gguf filename, not as a full pipeline."
+                )
         return fam
 
     # ── Background load + progress ─────────────────────────────────────────
@@ -783,6 +807,30 @@ class DiffusionBackend:
             return 0  # repo not in cache yet
         return total
 
+    @staticmethod
+    def _local_dir_weight_bytes(path: Path, *, exclude_transformer: bool) -> int:
+        """Sum the on-disk weight files under a local diffusers directory. The HF blob
+        cache is empty for a local path, so this is the only size signal for auto memory
+        planning; without it a large local model folds to zero and the planner skips
+        offload and OOMs. ``exclude_transformer`` drops the ``transformer/`` subfolder
+        for GGUF/single-file loads (their transformer is the single file, not resident
+        here); a full pipeline load keeps it (the whole repo is resident)."""
+        total = 0
+        for f in path.rglob("*"):
+            if f.suffix.lower() not in (".safetensors", ".bin", ".pt", ".ckpt"):
+                continue
+            try:
+                rel = f.relative_to(path)
+            except ValueError:
+                continue
+            if exclude_transformer and rel.parts and rel.parts[0] == "transformer":
+                continue
+            try:
+                total += f.stat().st_size
+            except OSError:
+                continue
+        return total
+
     @staticmethod
     def _companion_cache_bytes(base: str) -> int:
         """Resident companion (VAE + text-encoder) size for the memory plan.
@@ -794,21 +842,7 @@ class DiffusionBackend:
         weights to zero and auto planning can pick a resident placement that OOMs."""
         local = Path(base).expanduser()
         if local.is_dir():
-            total = 0
-            for f in local.rglob("*"):
-                if f.suffix.lower() not in (".safetensors", ".bin", ".pt", ".ckpt"):
-                    continue
-                try:
-                    rel = f.relative_to(local)
-                except ValueError:
-                    continue
-                if rel.parts and rel.parts[0] == "transformer":
-                    continue  # supplied by the GGUF single-file; not resident here
-                try:
-                    total += f.stat().st_size
-                except OSError:
-                    continue
-            return total
+            return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True)
         return DiffusionBackend._cache_bytes(base)
 
     # ── Synchronous load / generate / unload ───────────────────────────────
@@ -1076,6 +1110,18 @@ class DiffusionBackend:
                 eager_patched = False
                 compile_ctx = None
                 state_committed = False
+                # Lazy import: these patch modules import torch at module level, so
+                # importing them here (not at module load) keeps diffusion.py torch-free
+                # to import, letting get_diffusion_backend() run on a torchless native install.
+                from .diffusion_eager_patches import (
+                    install_compile_safe_patches,
+                    uninstall_patches,
+                )
+                from .diffusion_arch_patches import (
+                    install_arch_patches,
+                    uninstall_arch_patches,
+                )
+
                 try:
                     if effective_speed != SPEED_OFF:
                         install_compile_safe_patches()
@@ -1326,7 +1372,13 @@ class DiffusionBackend:
         if kind == "pipeline":
             # The whole repo (transformer + companions) is one cached download; the
             # cached bytes are the resident estimate (bnb-4bit / fp8 stay compressed).
-            cached = self._cache_bytes(repo_id) if repo_id else 0
+            # A LOCAL pipeline path isn't in the HF blob cache, so sum its on-disk weights
+            # (transformer included) instead of folding to zero and skipping offload.
+            local_repo = Path(repo_id).expanduser() if repo_id else None
+            if local_repo is not None and local_repo.is_dir():
+                cached = self._local_dir_weight_bytes(local_repo, exclude_transformer = False)
+            else:
+                cached = self._cache_bytes(repo_id) if repo_id else 0
             cached_mib = int(cached // (1024 * 1024)) if cached else None
             model_dense_mib = estimate_safetensors_dense_mib(cached_mib)
             companion_mib = None
@@ -1396,7 +1448,14 @@ class DiffusionBackend:
         # reuse the resident modules AT THEIR LOADED dtype, which is the whole point of
         # from_pipe (component reuse, no reload, no extra VRAM).
         pipe = getattr(diffusers, class_name).from_pipe(state.pipe, torch_dtype = None)
-        self._aux_pipes[class_name] = pipe
+        # Only publish to the shared aux cache if THIS load is still current. from_pipe runs
+        # under _generate_lock but NOT _lock, so an unload()/superseding load can clear
+        # _aux_pipes and null _state while it builds; caching unconditionally would re-insert
+        # a wrapper over now-stale modules that a later same-workflow load would reuse (or
+        # keep the old VRAM pinned). This generation still uses the returned pipe.
+        with self._lock:
+            if self._state is state:
+                self._aux_pipes[class_name] = pipe
         return pipe
 
     def _controlnet_pipe(self, state: _LoadState, resolved_cn: Any, cancel: threading.Event) -> Any:
@@ -1430,6 +1489,13 @@ class DiffusionBackend:
                 # raise on a blank credential instead of falling back, so coerce to None.
                 token = state.hf_token or None,
             )
+            if cancel.is_set():
+                # An unload/eviction raced the blocking download above and may have already
+                # cleared the load. Bail BEFORE any device placement so we don't allocate
+                # several GB onto the GPU after _unload_locked() freed it (which would OOM
+                # or make the unload appear to free memory only to repopulate it).
+                del cn_model
+                raise RuntimeError(DIFFUSION_CANCELLED_MSG)
             # Placement must follow the base model's offload policy. A resident base moves
             # the ControlNet resident too; an offloaded (low-VRAM) base streams it through
             # the device with group offloading instead of forcing the whole module onto the
@@ -1516,11 +1582,13 @@ class DiffusionBackend:
             family = getattr(state.family, "name", None),
             model_kind = state.kind,
             transformer_quant = state.transformer_quant,
+            compiled = "compiled" in (getattr(state, "speed_optims", ()) or ()),
         ):
             raise ValueError(
                 "LoRA is not supported for this model/quantisation on the diffusers engine "
-                "(GGUF-via-diffusers or torchao fp8/int8). Use a bf16 or bnb-4bit load, or the "
-                "native engine for GGUF models."
+                "(GGUF-via-diffusers, torchao fp8/int8, or a torch.compile'd Speed=default/max "
+                "load). Use a bf16 or bnb-4bit load at Speed=off/eager, or the native engine "
+                "for GGUF models."
             )
 
         resolved = diffusion_lora.resolve_specs(specs, hf_token = state.hf_token, cancel_event = cancel)
@@ -1713,6 +1781,15 @@ class DiffusionBackend:
                     fit = min(1.0, max_side / max(tw_f, th_f))
                     tw = max(16, int(round(tw_f * fit / 16.0)) * 16)
                     th = max(16, int(round(th_f * fit / 16.0)) * 16)
+                    # After the absolute cap, the target must still exceed the input, or
+                    # "upscale" would shrink it (e.g. a 3000px source at 2x clamps to 2048).
+                    # Reject rather than silently return a smaller image than uploaded.
+                    if max(tw, th) <= max(iw, ih):
+                        raise ValueError(
+                            f"Upscale would not enlarge this image: its longest side "
+                            f"({max(iw, ih)}px) already meets the {max_side}px output limit. "
+                            f"Use a smaller source image."
+                        )
                     init_pil = init_pil.resize((tw, th), Image.LANCZOS)
                     if strength is None:
                         # Hires-fix default: low enough to preserve content, high enough to
@@ -1983,6 +2060,10 @@ class DiffusionBackend:
         # bit-identical dequant. Idempotent.
         gguf_compile.uninstall_all()
         if state.eager_patched:
+            # Lazy import (torch at module level) to keep diffusion.py torch-free to import.
+            from .diffusion_eager_patches import uninstall_patches
+            from .diffusion_arch_patches import uninstall_arch_patches
+
             uninstall_patches()
             uninstall_arch_patches()
         # NOTE: we deliberately do NOT call state.pipe.unload_lora_weights() here. unload()
@@ -2056,6 +2137,7 @@ class DiffusionBackend:
                 family = state.family.name,
                 model_kind = state.kind,
                 transformer_quant = state.transformer_quant,
+                compiled = "compiled" in (getattr(state, "speed_optims", ()) or ()),
             ),
             "supports_controlnet": diffusion_controlnet.supports_controlnet(
                 engine = "diffusers",

studio/backend/core/inference/diffusion_controlnet.py

@@ -97,6 +97,26 @@ def sanitize_id(raw: str) -> str:
     return stem or "controlnet"
 
 
+def _has_controlnet_weights(p: Path) -> bool:
+    """True when ``p`` holds a loadable diffusers ControlNet weight (or shard index).
+
+    A config-only folder (interrupted copy/download) would otherwise be advertised and
+    then fail deep inside ``from_pretrained`` as a generic 500. Accept the standard
+    single-file weights, a sharded weight index, or any ``.safetensors`` shard."""
+    names = (
+        "diffusion_pytorch_model.safetensors",
+        "diffusion_pytorch_model.bin",
+        "diffusion_pytorch_model.safetensors.index.json",
+        "diffusion_pytorch_model.bin.index.json",
+    )
+    if any((p / n).exists() for n in names):
+        return True
+    try:
+        return any(child.suffix == ".safetensors" for child in p.iterdir())
+    except OSError:
+        return False
+
+
 def _scan_local() -> list[ControlNetCatalogEntry]:
     """A local ControlNet is a directory containing a diffusers config + weights."""
     entries: list[ControlNetCatalogEntry] = []
@@ -108,7 +128,9 @@ def _scan_local() -> list[ControlNetCatalogEntry]:
     for p in children:
         if not p.is_dir():
             continue
-        if not (p / "config.json").exists():
+        # Require BOTH the config and a loadable weight/index: a config-only folder is an
+        # incomplete copy/download, and advertising it would fail later in from_pretrained.
+        if not (p / "config.json").exists() or not _has_controlnet_weights(p):
             continue
         entries.append(
             ControlNetCatalogEntry(

studio/backend/core/inference/diffusion_families.py

@@ -298,14 +298,26 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
 _EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "layered")
 
 
+def _token_in_needle(token: str, needle: str) -> bool:
+    """True when ``token`` appears in ``needle`` as a whole path/name segment, i.e.
+    delimited by a separator (``- _ . / \\``) or a string boundary, not merely as a
+    raw substring. This keeps multi-part tokens matching where they should
+    ('qwen-image-edit' in 'qwen-image-edit-2511') while preventing a short token from
+    matching inside an unrelated word ('kontext' must not match 'kontextual', 'edit'
+    must not match 'edition')."""
+    return re.search(r"(?:^|[-_./\\])" + re.escape(token) + r"(?:$|[-_./\\])", needle) is not None
+
+
 def _best_family_match(needle: str) -> Optional[DiffusionFamily]:
-    """The family whose name/alias is the LONGEST substring of ``needle``. Longest =
-    most specific, so an edit checkpoint ('...qwen-image-edit-2511...') matches the
-    'qwen-image-edit' family rather than the generic 'qwen-image' one."""
+    """The family whose name/alias is the LONGEST whole-segment token of ``needle``.
+    Longest = most specific, so an edit checkpoint ('...qwen-image-edit-2511...')
+    matches the 'qwen-image-edit' family rather than the generic 'qwen-image' one.
+    Segment matching (not raw substring) stops a short alias like 'kontext' from
+    hijacking an unrelated path such as '.../kontextual/z-image-...gguf'."""
     best: Optional[tuple[DiffusionFamily, int]] = None
     for fam in _FAMILIES:
         for token in (fam.name, *fam.aliases):
-            if token in needle and (best is None or len(token) > best[1]):
+            if _token_in_needle(token, needle) and (best is None or len(token) > best[1]):
                 best = (fam, len(token))
     return best[0] if best else None
 
@@ -331,9 +343,16 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff
         # Don't let a generic base family (e.g. qwen-image) swallow a variant it can't run
         # (qwen-image-LAYERED, ...-Inpaint): if the id still carries a reject keyword the
         # matched family does not itself declare, reject so the load fails fast + clearly.
+        # Scope the keyword check to the LAST path component (the model id or
+        # filename), not arbitrary parent directories: a valid file selected as
+        # repo_id `/models/edit` + filename `Z-Image-Turbo-Q4.gguf` must not be
+        # rejected because a parent folder happens to be named `edit`. The
+        # combined `repo_id/gguf_filename` fallback passes the filename last.
+        basename = re.split(r"[/\\]+", needle)[-1]
         matched_tokens = (match.name, *match.aliases)
         if any(
-            kw in needle and not any(kw in tok for tok in matched_tokens) for kw in _EDIT_KEYWORDS
+            _token_in_needle(kw, basename) and not any(kw in tok for tok in matched_tokens)
+            for kw in _EDIT_KEYWORDS
         ):
             return None
         return match

studio/backend/core/inference/diffusion_lora.py

@@ -25,6 +25,8 @@ from typing import Optional
 from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
 from utils.paths.storage_roots import studio_root
 
+from .diffusion_families import DIFFUSION_CANCELLED_MSG
+
 # LoRA file formats we accept. sd-cli probes .safetensors/.gguf/.pt; diffusers loads
 # .safetensors. We expose safetensors + gguf (pt is legacy/pickled -> excluded for safety).
 _NATIVE_EXTS = (".safetensors", ".gguf")
@@ -241,7 +243,10 @@ def resolve_specs(
 
     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."""
+    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."""
     out: list[ResolvedLora] = []
     try:
         for spec_id, weight in specs:
@@ -250,6 +255,10 @@ def resolve_specs(
             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 RuntimeError as exc:
+        if str(exc) == "Cancelled":
+            raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc
+        raise
     return out
 
 
@@ -297,14 +306,15 @@ def inject_prompt_tags(prompt: str, resolved: list[ResolvedLora]) -> str:
     sd-cli strips these tags before they reach the model, so appending them is safe and
     deterministic. A selected adapter's weight is validated (0-2) and recorded in the
     request/gallery, so the injected tag must WIN over any `<lora:ALIAS:...>` the user
-    typed for that same alias: strip a user tag whose alias matches a selected adapter,
-    then append the validated one. Tags for aliases the user typed that are NOT selected
-    are left untouched (free-form use).
+    typed. Strip ALL user-typed tags first: only the selected adapters are materialized in
+    the managed `--lora-model-dir`, so a tag for an unselected alias can never resolve
+    anyway (sd-cli's extract_and_remove_lora silently removes unresolved tags), and a tag
+    for a selected alias must not override the validated weight. Then append the validated
+    tags for the selected adapters.
     """
-    selected = {r.alias for r in resolved}
-    # Drop any user-typed tag whose alias is one of the selected adapters, so the typed
-    # weight can't override the validated weight (or slip outside the 0-2 bounds).
-    cleaned = _TAG_RE.sub(lambda m: "" if m.group(1) in selected else m.group(0), prompt)
+    # Drop every user-typed tag: unselected ones are dead (not in the managed dir) and
+    # selected ones must not override the validated weight / 0-2 bounds.
+    cleaned = _TAG_RE.sub("", prompt)
     # Collapse whitespace left by stripped tags without disturbing the user's text.
     cleaned = re.sub(r"[ \t]{2,}", " ", cleaned).strip()
     tags = [f"<lora:{r.alias}:{_fmt_weight(r.weight)}>" for r in resolved]
@@ -343,12 +353,17 @@ def supports_lora(
     family: Optional[str],
     model_kind: Optional[str],
     transformer_quant: Optional[str],
+    compiled: bool = False,
 ) -> bool:
     """Single gate for whether the current load can apply LoRA (used by status + backends).
 
     Native (sd_cpp): GGUF via sd-cli, for the LoRA-capable families only (Qwen excluded).
     Diffusers: bf16 or bnb-4bit transformers, but NOT the dense torchao fp8/int8 fast path
-    (tensor-subclass weights) and NOT GGUF-via-diffusers.
+    (tensor-subclass weights) and NOT GGUF-via-diffusers. A diffusers transformer that was
+    torch.compile'd at load (Speed=default/max) also can't take a non-hotswap adapter:
+    diffusers requires the adapter to be loaded BEFORE compilation, so applying one to the
+    already-compiled module fails with adapter-key mismatches. ``compiled`` is diffusers-only
+    (the native sd-cli path has no torch compile).
     """
     fam = (family or "").lower()
     if engine == "sd_cpp":
@@ -358,4 +373,6 @@ def supports_lora(
         return False  # GGUF diffusers transformer: use the native engine for LoRA
     if transformer_quant and transformer_quant.lower() in _DIFFUSERS_LORA_BLOCKED_QUANT:
         return False
+    if compiled:
+        return False  # can't load an adapter onto an already-compiled transformer
     return True

studio/backend/core/inference/sd_cpp_backend.py

@@ -181,7 +181,7 @@ def _map_guidance(
     classifier-free ``--cfg-scale``. A distilled 0/1 means CFG off (sd-cli's 1.0); a
     value > 1 is real CFG. Mirrors the engine mapping validated in the CPU benchmark.
     """
-    if fam.name in ("flux.1", "flux.2-klein"):
+    if fam.name in ("flux.1", "flux.2-klein", "flux.2-dev"):
         return None, (float(guidance) if guidance is not None else None)
     cfg = float(guidance) if (guidance is not None and guidance > 1.0) else 1.0
     return cfg, None
@@ -509,10 +509,18 @@ class SdCppDiffusionBackend:
 
         from core.inference import diffusion_lora
 
-        if init_image is not None or mask_image is not None or reference_images:
+        if (
+            init_image is not None
+            or mask_image is not None
+            or reference_images
+            or (upscale is not None and upscale > 1)
+        ):
+            # upscale needs an input image, so a direct API call with upscale > 1 but no
+            # init_image must be rejected too rather than silently returning a plain,
+            # un-upscaled text-to-image result (the diffusers backend rejects the same).
             raise ValueError(
-                "img2img / inpaint / reference are not yet supported on the native sd.cpp "
-                "engine; run on a GPU (diffusers) for image-conditioned workflows."
+                "img2img / inpaint / reference / upscale are not yet supported on the native "
+                "sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows."
             )
         if controlnet is not None:
             raise ValueError(
@@ -536,8 +544,13 @@ class SdCppDiffusionBackend:
                 cfg_scale, flux_guidance = _map_guidance(state.family, guidance)
                 # Resolve any selected LoRA adapters up front (downloads land in the HF
                 # cache; a bad id fails here as a clear 400 before we spawn sd-cli).
+                # Drop weight-0 rows BEFORE the support gate: LoraSpec documents weight 0
+                # as disabling the adapter (the diffusers path treats it as empty), so a
+                # request carrying only disabled rows must stay a no-op even on a family
+                # where native LoRA is unsupported, rather than 400 on a dead selection.
                 lora_resolved: list = []
-                if loras:
+                active_loras = [(i, w) for (i, w) in (loras or []) if w != 0]
+                if active_loras:
                     if not diffusion_lora.supports_lora(
                         engine = "sd_cpp",
                         family = state.family.name,
@@ -549,7 +562,7 @@ class SdCppDiffusionBackend:
                             "sd.cpp engine."
                         )
                     lora_resolved = diffusion_lora.resolve_specs(
-                        loras, hf_token = state.hf_token, cancel_event = cancel
+                        active_loras, hf_token = state.hf_token, cancel_event = cancel
                     )
                 extra_args: list[str] = []
                 if state.vae_format:
@@ -695,6 +708,7 @@ class SdCppDiffusionBackend:
                 "engine": "sd_cpp",
                 "supports_lora": False,
                 "supports_controlnet": False,
+                "workflows": [],
             }
         from core.inference import diffusion_lora
 
@@ -728,6 +742,11 @@ class SdCppDiffusionBackend:
             ),
             # Native ControlNet (sd-cli --control-net) is a follow-up; off for now.
             "supports_controlnet": False,
+            # The native engine supports plain text-to-image only (generate() rejects
+            # img2img / inpaint / reference / upscale), so advertise just txt2img. Without
+            # this the status omits workflows, the UI reads [], and it disables the Create
+            # tab for a loaded native model, stranding the user on an image-only tab.
+            "workflows": ["txt2img"],
         }
 
 

studio/backend/core/inference/sd_cpp_engine.py

@@ -75,6 +75,14 @@ def _terminate(proc: "subprocess.Popen") -> None:
             proc.kill()
         except Exception:  # noqa: BLE001 -- best-effort teardown
             pass
+    # Reap the killed child so it does not linger as a zombie until the next Popen
+    # cleanup / interpreter exit. Callers raise immediately after _terminate (the
+    # cancellation and timeout paths), so without this a burst of image cancellations
+    # leaks process-table entries. SIGKILL is prompt, so a short bounded wait suffices.
+    try:
+        proc.wait(timeout = 5)
+    except Exception:  # noqa: BLE001 -- best-effort reap; never block teardown
+        pass
 
 
 def _binary_name(stem: str) -> str:

studio/backend/core/training/diffusion_lora_trainer.py

@@ -391,11 +391,15 @@ def run_diffusion_lora_training(
 
     lora_params = [p for p in unet.parameters() if p.requires_grad]
     optimizer = torch.optim.AdamW(lora_params, lr = cfg.learning_rate)
+    # The scheduler advances once per optimizer update: lr_sched.step() runs a single
+    # time per outer opt_step (after the accumulation inner loop), for cfg.train_steps
+    # total. Count warmup/decay in those optimizer steps -- multiplying by the
+    # accumulation factor would stretch warmup past the run and never reach the decay.
     lr_sched = get_scheduler(
         cfg.lr_scheduler,
         optimizer = optimizer,
-        num_warmup_steps = cfg.lr_warmup_steps * cfg.gradient_accumulation_steps,
-        num_training_steps = cfg.train_steps * cfg.gradient_accumulation_steps,
+        num_warmup_steps = cfg.lr_warmup_steps,
+        num_training_steps = cfg.train_steps,
     )
 
     vae_scale = vae.config.scaling_factor

studio/backend/core/training/diffusion_training_service.py

@@ -31,12 +31,25 @@ _CTX = mp.get_context("spawn")
 _TERMINAL = ("complete", "error")
 
 
-def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
+def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
     # Imported lazily so this module (and the route layer) stays torch-free at import.
     from .diffusion_lora_trainer import run_diffusion_training_process
     run_diffusion_training_process(event_queue = event_queue, stop_queue = stop_queue, config = config)
 
 
+def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
+    # First thing in the spawned child (before torch is imported): bind to the parent's
+    # death on Linux and scrub the native path lease secret, exactly like the inference /
+    # export / LLM-training workers. multiprocessing children cannot be given a
+    # parent-set preexec_fn, so the child must self-bind; otherwise a Studio crash or
+    # kill leaves this trainer holding the GPU. Tests inject their own target, so this
+    # binding only runs for the real production spawn.
+    from utils.native_path_leases import run_without_native_path_secret
+    run_without_native_path_secret(
+        _run_diffusion_child, event_queue = event_queue, stop_queue = stop_queue, config = config
+    )
+
+
 def _idle_state() -> dict[str, Any]:
     return {
         "active": False,
@@ -225,8 +238,12 @@ class DiffusionTrainingService:
                     message = "Training...",
                 )
             elif etype == "complete":
+                # Reset in_model_load: a stop during model load emits complete without a
+                # preceding model_load_completed, which would otherwise leave a stale
+                # loading indicator after the job ended.
                 s.update(
                     active = False,
+                    in_model_load = False,
                     status = "stopped" if ev.get("stopped") else "completed",
                     output_dir = ev.get("output_dir"),
                     lora_path = ev.get("lora_path"),
@@ -235,7 +252,14 @@ class DiffusionTrainingService:
                     else "Training complete.",
                 )
             elif etype == "error":
-                s.update(active = False, status = "error", message = str(ev.get("message", "error")))
+                # Reset in_model_load too: an error raised during model loading has no
+                # model_load_completed, so the terminal state must clear it explicitly.
+                s.update(
+                    active = False,
+                    in_model_load = False,
+                    status = "error",
+                    message = str(ev.get("message", "error")),
+                )
 
 
 _service: Optional[DiffusionTrainingService] = None

studio/backend/models/inference.py

@@ -1947,6 +1947,23 @@ class DiffusionGenerateRequest(BaseModel):
         "the loaded model or its quantisation can't apply ControlNet.",
     )
 
+    @field_validator("loras")
+    @classmethod
+    def _unique_lora_ids(cls, value: Optional[list[LoraSpec]]) -> Optional[list[LoraSpec]]:
+        # Both apply paths break alias collisions by suffixing the adapter name/file, so a
+        # repeated id would load the SAME adapter as several distinct adapters and stack
+        # its effect past the per-adapter weight bound. The UI already prevents duplicates;
+        # reject them for API clients too so each adapter takes effect at most once.
+        if value:
+            seen: set[str] = set()
+            for spec in value:
+                if spec.id in seen:
+                    raise ValueError(
+                        f"duplicate LoRA id '{spec.id}'; list each adapter at most once"
+                    )
+                seen.add(spec.id)
+        return value
+
     @field_validator("reference_images")
     @classmethod
     def _bounded_reference_items(cls, value: Optional[list[str]]) -> Optional[list[str]]:

studio/backend/routes/inference.py

@@ -2823,12 +2823,26 @@ def _guard_chat_load_against_training(
     from routes.training_vram import can_load_chat_during_training
 
     try:
-        if not get_training_backend().is_training_active():
-            return
+        llm_active = get_training_backend().is_training_active()
     except Exception as e:
         logger.warning("Could not check training state for chat-load guard: %s", e)
         return
 
+    if not llm_active:
+        # An SDXL LoRA trainer runs in its own subprocess and its VRAM can't be cheaply
+        # fit-checked here, so refuse the chat load outright while one is active rather
+        # than risk OOMing the run. Symmetric with the image-load guard.
+        if _diffusion_training_active():
+            raise HTTPException(
+                status_code = 409,
+                detail = (
+                    "Can't load this model while diffusion (Images) training is running: "
+                    "its GPU memory use can't be verified against the trainer, so the load "
+                    "was refused to protect the run. Try again after training finishes."
+                ),
+            )
+        return
+
     is_gguf = bool(getattr(config, "is_gguf", False))
     required_override_gb = (
         _estimate_gguf_required_gb(
@@ -11021,6 +11035,16 @@ async def _openai_passthrough_non_streaming(
 # ──────────────────────────────────────────────────────────────────────────
 
 
+def _diffusion_training_active() -> bool:
+    """Whether a diffusion (SDXL) LoRA job is running. Best-effort so a load is never
+    blocked just because the training service could not be imported/read."""
+    try:
+        from core.training.diffusion_training_service import get_diffusion_training_service
+        return get_diffusion_training_service().is_active()
+    except Exception:  # noqa: BLE001
+        return False
+
+
 def _guard_diffusion_load_against_training() -> None:
     """Refuse loading an image model while a training run is active. Unlike chat,
     a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the
@@ -11029,11 +11053,15 @@ def _guard_diffusion_load_against_training() -> None:
     from core.training import get_training_backend
 
     try:
-        if not get_training_backend().is_training_active():
-            return
+        llm_active = get_training_backend().is_training_active()
     except Exception as e:
         logger.warning("Could not check training state for image-load guard: %s", e)
         return
+    # An SDXL LoRA trainer runs in its own subprocess on the same GPU, so an image
+    # load must be refused while one is active too -- otherwise the resident pipeline
+    # competes with the trainer for VRAM. Symmetric with the diffusion-start interlock.
+    if not llm_active and not _diffusion_training_active():
+        return
     raise HTTPException(
         status_code = 409,
         detail = (
@@ -11051,12 +11079,10 @@ async def load_diffusion_model(
     from core.inference.diffusion import get_diffusion_backend, resolve_model_kind
     from core.inference.diffusion_device import resolve_diffusion_device_target
     from core.inference.diffusion_engine_router import (
-        active_engine_name,
         annotate_status,
         select_and_activate_engine,
     )
     from core.inference.gpu_arbiter import acquire_for, release, DIFFUSION
-    from core.inference.sd_cpp_engine import ENGINE_SD_CPP
     from utils.native_path_leases import redact_native_paths
 
     backend = get_diffusion_backend()
@@ -11084,12 +11110,15 @@ async def load_diffusion_model(
         engine = await asyncio.to_thread(
             select_and_activate_engine, fam, hf_token = request.hf_token, model_kind = kind
         )
-        # Take the GPU from the chat backend only when this load will actually use it.
-        # diffusers always does; a *force-native* sd.cpp load on a CUDA/XPU/MPS box does
-        # too. But a native sd.cpp load on a pure-CPU host never touches the GPU, so
-        # acquiring would evict the resident chat model for nothing -- skip the handoff.
+        # Take the GPU from the chat backend only when this load will actually use it,
+        # which is exactly the resolved device being non-CPU. diffusers on an accelerator
+        # and a force-native sd.cpp load on CUDA/XPU/MPS both resolve to that device; a
+        # native sd.cpp load on a pure-CPU host does not. Crucially, a CPU-only host with
+        # no usable sd-cli falls back to diffusers ON CPU -- that also never touches GPU
+        # memory, so keying off the engine name (not the device) would wrongly evict a
+        # resident chat model for a load that cannot use the GPU. Gate on the device.
         device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device)
-        needs_gpu = active_engine_name() != ENGINE_SD_CPP or device != "cpu"
+        needs_gpu = device != "cpu"
         if needs_gpu:
             # Then kick the (slow) load onto a background thread and return at once --
             # the client polls images/load-progress.
@@ -11212,8 +11241,13 @@ async def generate_diffusion_image(
                     {
                         "prompt": request.prompt,
                         "negative_prompt": request.negative_prompt,
-                        "width": request.width,
-                        "height": request.height,
+                        # Persist the ACTUAL output size, not the request sliders: Transform/
+                        # Inpaint/Edit derive it from the uploaded image, Extend grows the
+                        # canvas, and Upscale resizes it, so request.width/height would record
+                        # (and later restore) the wrong dimensions for those workflows. For
+                        # plain txt2img the image size equals the sliders anyway.
+                        "width": getattr(image, "width", None) or request.width,
+                        "height": getattr(image, "height", None) or request.height,
                         "steps": request.steps,
                         "guidance": request.guidance,
                         "seed": seed,
@@ -11230,7 +11264,10 @@ async def generate_diffusion_image(
                         "controlnet": (
                             f"{request.controlnet.id}:{request.controlnet.control_type}:"
                             f"{request.controlnet.strength:g}"
-                            if request.controlnet
+                            # strength 0 is treated as disabled and skipped before loading /
+                            # conditioning, so the image is unconditioned; don't claim a
+                            # ControlNet was applied in the recipe/metadata.
+                            if request.controlnet and request.controlnet.strength > 0
                             else None
                         ),
                         "created_at": created_at,
@@ -11254,6 +11291,8 @@ async def list_gallery_images(
     offset: int = 0,
     current_subject: str = Depends(get_current_subject),
 ):
+    from pydantic import ValidationError
+
     from core.inference import image_gallery
 
     limit = max(1, min(limit, 200))
@@ -11261,10 +11300,18 @@ async def list_gallery_images(
     # Fetch one extra to learn whether more remain, without a second scan.
     records = await asyncio.to_thread(image_gallery.list_images, limit + 1, offset)
     has_more = len(records) > limit
-    return GalleryListResponse(
-        images = [GalleryImage(**r) for r in records[:limit]],
-        has_more = has_more,
-    )
+    # Build the response per record and drop any that fail schema validation: a PNG
+    # whose recipe chunk has all required keys but a wrong value type (e.g. a
+    # hand-dropped or corrupted file) passes the presence-only read but would raise
+    # inside GalleryImage(**r). Skipping it keeps one bad file from 500-ing the whole
+    # gallery listing.
+    images = []
+    for r in records[:limit]:
+        try:
+            images.append(GalleryImage(**r))
+        except ValidationError:
+            continue
+    return GalleryListResponse(images = images, has_more = has_more)
 
 
 @studio_router.get("/images/gallery/{image_id}/file")

studio/backend/tests/test_diffusion_backend.py

@@ -22,6 +22,14 @@ from core.inference.diffusion import (
     _base_file_downloaded,
     _resolve_diffusion_compute_dtype,
 )
+
+# diffusion.py imports the compile/arch patch modules LAZILY (they pull torch at module
+# level, and diffusion.py must stay importable on a torchless native install). Import them
+# here at collection time -- under the real torch -- so they are cached in sys.modules
+# before the fake-torch fixtures swap it out; otherwise the lazy import inside load_pipeline
+# would try to build them against the incomplete stub torch.
+import core.inference.diffusion_eager_patches  # noqa: E402,F401
+import core.inference.diffusion_arch_patches  # noqa: E402,F401
 from core.inference.diffusion_families import (
     detect_family,
     resolve_base_repo,
@@ -77,6 +85,34 @@ def test_detect_family_from_repo_id():
     assert detect_family("meta-llama/Llama-3-8B") is None
 
 
+def test_detect_family_matches_reject_and_alias_by_segment():
+    # Reject keywords and short aliases must match whole path/name segments, not raw
+    # substrings, so an unrelated word that merely CONTAINS one does not misroute a
+    # valid base model (regression: substring matching broke these).
+    assert detect_family("/models/edited/z-image-turbo-Q4_K_M.gguf").name == "z-image"
+    assert detect_family("unsloth/Z-Image-Edition-GGUF").name == "z-image"
+    assert detect_family("/models/kontextual/z-image-turbo-Q4_K_M.gguf").name == "z-image"
+    # Supported edit families still resolve (edit / kontext are whole tokens there).
+    assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF").name == "qwen-image-edit"
+    assert detect_family("unsloth/FLUX.1-Kontext-dev-GGUF").name == "flux.1-kontext"
+    # Unsupported variants sharing only a base arch keyword are still rejected.
+    assert detect_family("unsloth/Qwen-Image-Layered-GGUF") is None
+    assert detect_family("unsloth/Qwen-Image-2512-Inpaint") is None
+
+
+def test_detect_family_edit_keyword_scoped_to_basename():
+    from core.inference.diffusion_families import detect_family_for_pick
+
+    # A parent directory named `edit`/`inpaint` must NOT poison a valid pick: only
+    # the model id / filename basename is scanned for reject keywords. A direct
+    # local pick arrives as (parent_dir, filename).
+    assert detect_family("/models/edit") is None  # the dir alone is ambiguous
+    assert detect_family_for_pick("/models/edit", "Z-Image-Turbo-Q4.gguf").name == "z-image"
+    assert detect_family_for_pick("/models/inpaint", "qwen-image-2512-Q4.gguf").name == "qwen-image"
+    # A genuinely unsupported variant keyword in the FILENAME still rejects.
+    assert detect_family_for_pick("/models/misc", "Qwen-Image-Layered-Q4.gguf") is None
+
+
 def test_detect_family_override():
     assert detect_family("local/path", override = "z-image").name == "z-image"
     assert detect_family("local/path", override = "zimage").name == "z-image"
@@ -1489,10 +1525,30 @@ def test_validate_load_request(tmp_path):
         backend.validate_load_request("some-org/Z-Image", gguf_filename = "model.safetensors")
     with pytest.raises(ValueError, match = "family"):
         backend.validate_load_request("meta/Llama-3", gguf_filename = "q.gguf")
+    # A family-looking repo paired with a non-GGUF single-file name is rejected here,
+    # BEFORE the route evicts chat and hands over the GPU (the background load would
+    # otherwise be the first to notice README.md is not a checkpoint).
+    with pytest.raises(ValueError, match = r"\.gguf"):
+        backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "README.md")
     assert (
         backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf").name
         == "z-image"
     )
+    # A kind/extension mismatch fails fast here, before the route evicts chat + grabs the
+    # GPU only to fail in the background from_single_file path.
+    with pytest.raises(ValueError, match = ".gguf"):
+        backend.validate_load_request(
+            "unsloth/Z-Image-Turbo-GGUF", gguf_filename = "model.safetensors", model_kind = "gguf"
+        )
+    with pytest.raises(ValueError, match = "gguf"):
+        backend.validate_load_request(
+            "unsloth/Qwen-Image-2512-FP8", gguf_filename = "q.gguf", model_kind = "single_file"
+        )
+    # A remote "*-GGUF" repo loaded as a full pipeline (no single-file name) is a single-file
+    # GGUF repo, so from_pretrained would find no pipeline manifest and fail after chat is
+    # already evicted; reject it here before the GPU handoff.
+    with pytest.raises(ValueError, match = "GGUF"):
+        backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", model_kind = "pipeline")
     # A local path with a missing child fails here (before any GPU/network work).
     with pytest.raises(FileNotFoundError):
         backend.validate_load_request(

studio/backend/tests/test_diffusion_controlnet.py

@@ -71,6 +71,7 @@ def test_resolve_controlnet_local(tmp_path, monkeypatch):
     cn = d / "my-cn"
     cn.mkdir()
     (cn / "config.json").write_text("{}")
+    (cn / "diffusion_pytorch_model.safetensors").write_bytes(b"x")  # a loadable weight
     monkeypatch.setattr(dc, "controlnets_dir", lambda: d)
     entries = {e.id for e in dc.list_controlnets()}
     assert "my-cn" in entries
@@ -78,6 +79,21 @@ def test_resolve_controlnet_local(tmp_path, monkeypatch):
     assert r.is_local and r.path == str(cn)
 
 
+def test_scan_local_skips_config_only_folder(tmp_path, monkeypatch):
+    # A folder with config.json but no weight/index (interrupted copy) must NOT be
+    # advertised: it would otherwise fail deep in from_pretrained as a generic 500.
+    d = tmp_path / "controlnets"
+    d.mkdir()
+    incomplete = d / "incomplete-cn"
+    incomplete.mkdir()
+    (incomplete / "config.json").write_text("{}")
+    monkeypatch.setattr(dc, "controlnets_dir", lambda: d)
+    assert "incomplete-cn" not in {e.id for e in dc.list_controlnets()}
+    # A sharded weight index counts as a loadable weight.
+    (incomplete / "diffusion_pytorch_model.safetensors.index.json").write_text("{}")
+    assert "incomplete-cn" in {e.id for e in dc.list_controlnets()}
+
+
 def test_preprocess_control_passthrough_and_canny():
     from PIL import Image
 

studio/backend/tests/test_diffusion_lora.py

@@ -44,12 +44,14 @@ def test_inject_prompt_tags_validated_weight_overrides_user_typed():
     assert dl.inject_prompt_tags("a cat <lora:style:1>", [r]) == "a cat <lora:style:0.8>"
 
 
-def test_inject_prompt_tags_keeps_unselected_user_tags():
+def test_inject_prompt_tags_strips_unselected_user_tags():
     r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8)
-    # A user tag for an alias that is NOT one of the selected adapters is left untouched.
+    # A user tag for an alias that is NOT selected is stripped: only selected adapters are
+    # materialized in the managed --lora-model-dir, so sd-cli would drop the dead tag anyway;
+    # removing it keeps the prompt clean and unambiguous.
     out = dl.inject_prompt_tags("a cat <lora:other:0.5>", [r])
-    assert "<lora:other:0.5>" in out
-    assert "<lora:style:0.8>" in out
+    assert "<lora:other:0.5>" not in out
+    assert out == "a cat <lora:style:0.8>"
 
 
 def test_inject_prompt_tags_empty_returns_prompt():
@@ -83,6 +85,45 @@ def test_supports_lora_matrix():
     assert not dl.supports_lora(
         engine = "diffusers", family = "flux.1", model_kind = "gguf", transformer_quant = None
     )
+    # A torch.compile'd diffusers transformer (Speed=default/max) can't take a non-hotswap
+    # adapter: diffusers needs the adapter loaded before compilation.
+    assert not dl.supports_lora(
+        engine = "diffusers",
+        family = "flux.1",
+        model_kind = "pipeline",
+        transformer_quant = None,
+        compiled = True,
+    )
+    # compiled is diffusers-only; the native path ignores it.
+    assert dl.supports_lora(
+        engine = "sd_cpp",
+        family = "flux.1",
+        model_kind = "gguf",
+        transformer_quant = None,
+        compiled = True,
+    )
+
+
+def test_resolve_specs_maps_cancelled_to_diffusion_sentinel(tmp_path, monkeypatch):
+    # A Hub download cancelled mid-flight raises RuntimeError("Cancelled"); resolve_specs
+    # must convert it to the diffusion cancellation sentinel so the route maps it to 409,
+    # not a generic 500 server-error toast.
+    def _boom(spec_id, weight, **kw):
+        raise RuntimeError("Cancelled")
+
+    monkeypatch.setattr(dl, "resolve_one", _boom)
+    with pytest.raises(RuntimeError) as ei:
+        dl.resolve_specs([("a", 1.0)])
+    assert str(ei.value) == dl.DIFFUSION_CANCELLED_MSG
+
+    # A non-cancellation RuntimeError is left untouched.
+    def _other(spec_id, weight, **kw):
+        raise RuntimeError("disk full")
+
+    monkeypatch.setattr(dl, "resolve_one", _other)
+    with pytest.raises(RuntimeError) as ei2:
+        dl.resolve_specs([("a", 1.0)])
+    assert str(ei2.value) == "disk full"
 
 
 def test_materialize_native_dir_symlinks_and_breaks_collisions(tmp_path):
@@ -192,6 +233,12 @@ def test_lora_spec_and_request_validation():
         LoraSpec(id = "a", weight = -0.1)
     # default weight
     assert LoraSpec(id = "a").weight == 1.0
+    # duplicate ids are rejected: repeating an id would load the same adapter as several
+    # distinct suffixed adapters and stack its effect past the per-adapter weight bound.
+    with pytest.raises(Exception):
+        DiffusionGenerateRequest(
+            prompt = "x", loras = [{"id": "a", "weight": 0.5}, {"id": "a", "weight": 1.0}]
+        )
 
 
 # ── Diffusers apply manager ─────────────────────────────────────────────────

studio/backend/tests/test_diffusion_training.py

@@ -189,6 +189,22 @@ def test_apply_event_transitions():
     assert svc.status()["status"] == "error" and svc.status()["message"] == "boom"
 
 
+def test_terminal_events_clear_model_load_flag():
+    # A stop or error during model load emits complete/error WITHOUT a preceding
+    # model_load_completed, so the terminal update must reset in_model_load or the
+    # client shows a stale loading indicator after the job ended.
+    svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+    svc._apply_event({"type": "model_load_started"})
+    assert svc.status()["in_model_load"] is True
+    svc._apply_event({"type": "complete", "stopped": True})
+    assert svc.status()["in_model_load"] is False and svc.status()["status"] == "stopped"
+
+    svc2 = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+    svc2._apply_event({"type": "model_load_started"})
+    svc2._apply_event({"type": "error", "message": "load failed"})
+    assert svc2.status()["in_model_load"] is False and svc2.status()["status"] == "error"
+
+
 # ── route wiring (mocked service) ─────────────────────────────────────────────
 class _FakeService:
     def __init__(self):

studio/backend/tests/test_sd_cpp_engine.py

@@ -135,6 +135,25 @@ def test_runtime_env_handles_missing_lib_path():
     assert env[var] == "/opt/sdcpp/bin"
 
 
+def test_terminate_reaps_killed_child():
+    # Cancellation/timeout paths call _terminate then immediately raise, so it must
+    # reap the killed child itself -- otherwise a burst of image cancellations leaves
+    # zombies until a later Popen cleanup. After _terminate the returncode is set
+    # (the child has been waited on), so nothing lingers.
+    import subprocess
+    proc = subprocess.Popen(
+        [sys.executable, "-c", "import time; time.sleep(30)"],
+        start_new_session = (os.name == "posix"),
+    )
+    try:
+        eng._terminate(proc)
+        assert proc.returncode is not None
+    finally:
+        if proc.poll() is None:
+            proc.kill()
+            proc.wait()
+
+
 # ── generate (fake subprocess) ──────────────────────────────────────────────
 
 

studio/frontend/src/app/routes/__root.tsx

@@ -83,6 +83,11 @@ const CHAT_ONLY_ALLOWED = new Set([
 function isChatOnlyAllowed(pathname: string): boolean {
   if (CHAT_ONLY_ALLOWED.has(pathname)) return true;
   if (pathname === "/data-recipes" || pathname.startsWith("/data-recipes/")) return true;
+  // Images runs on CPU/MPS via the native sd.cpp engine, which is exactly the
+  // no-GPU (chat-only) setup it was added for. The generic chat-only flag is about
+  // training/export needing a GPU, so it must not redirect /images away here or the
+  // native image path is unreachable on the hosts that need it.
+  if (pathname === "/images" || pathname.startsWith("/images/")) return true;
   return false;
 }
 
@@ -159,6 +164,13 @@ function RootLayout() {
     setImagesMounted(true);
   }
   const shouldMountImages = isImagesRoute || imagesMounted;
+  // Chat and Images both render their own full-height shell (a fixed top rail + an
+  // internally-scrolling body), so both want the chat-style layout: no outer pt-14
+  // inset and no outer scroll. Keying the layout off isChatRoute alone gave /images
+  // the non-chat pt-14 + outer overflow, pushing its picker down and clipping the
+  // bottom gallery. Treat them the same for the container padding/overflow only; the
+  // keep-alive mounts below stay keyed to each specific route.
+  const isChatLike = isChatRoute || isImagesRoute;
 
   useTrainingUnloadGuard();
   // Global export driver: streams worker logs and tracks status from any route
@@ -251,10 +263,10 @@ function RootLayout() {
           className="!min-h-0 h-[calc(100dvh-var(--studio-titlebar-height,0px))] overflow-hidden"
         >
           <AppSidebar />
-          <SidebarInset className={isChatRoute ? "overflow-hidden" : "overflow-y-auto"}>
+          <SidebarInset className={isChatLike ? "overflow-hidden" : "overflow-y-auto"}>
             <Navbar />
             <div
-              className={`relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"} ${isChatRoute ? "" : "pt-14 md:pt-[var(--studio-non-chat-content-top-inset,var(--studio-content-top-inset,0px))] md:[--studio-titlebar-height:var(--studio-non-chat-content-top-inset,var(--studio-content-top-inset,0px))]"}`}
+              className={`relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatLike ? "overflow-hidden" : "overflow-visible"} ${isChatLike ? "" : "pt-14 md:pt-[var(--studio-non-chat-content-top-inset,var(--studio-content-top-inset,0px))] md:[--studio-titlebar-height:var(--studio-non-chat-content-top-inset,var(--studio-content-top-inset,0px))]"}`}
             >
               {/* Stays mounted across navigation so an in-flight generation is
                   not cancelled when leaving /chat; hidden (not unmounted) off-route.
@@ -280,7 +292,7 @@ function RootLayout() {
                 <div
                   className={
                     isImagesRoute
-                      ? "flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-visible"
+                      ? "flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden"
                       : "hidden"
                   }
                   inert={!isImagesRoute || undefined}

studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx

@@ -1072,11 +1072,19 @@ const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint", "layered"] as const;
 // hidden even though their id contains an edit keyword. Mirrors the backend's
 // qwen-image-edit family in diffusion_families.py.
 const SUPPORTED_EDIT_KEYWORDS = ["qwen-image-edit", "kontext"] as const;
+// Match a keyword as a whole path/name segment (bounded by a separator or a string
+// edge), not a raw substring, so "edit" does not hide ".../edited/..." or an
+// "*-edition" repo and "kontext" does not hide ".../kontextual/...". These keywords
+// are literals of [a-z-], so no regex escaping is needed. Mirrors _token_in_needle in
+// diffusion_families.py.
+function idHasSegment(id: string, keyword: string): boolean {
+  return new RegExp(`(?:^|[-_./\\\\])${keyword}(?:$|[-_./\\\\])`).test(id);
+}
 function isImageEditModel(repoId: string | null | undefined): boolean {
   if (!repoId) return false;
   const id = repoId.toLowerCase();
-  if (SUPPORTED_EDIT_KEYWORDS.some((kw) => id.includes(kw))) return false;
-  return IMAGE_EDIT_KEYWORDS.some((kw) => id.includes(kw));
+  if (SUPPORTED_EDIT_KEYWORDS.some((kw) => idHasSegment(id, kw))) return false;
+  return IMAGE_EDIT_KEYWORDS.some((kw) => idHasSegment(id, kw));
 }
 
 // Gate an on-device model by the picker's task scope. With a filter (the Images

studio/frontend/src/features/images/api.ts

@@ -111,7 +111,9 @@ export interface ControlNetSpecInput {
   id: string;
   // Base64/data-URL control image (a source image or an already-made control map).
   image: string;
-  control_type: "passthrough" | "canny";
+  // "canny" preprocesses edges from a source image; any other type (passthrough, or a
+  // union type like depth/pose) is an already-made map the backend maps to a control mode.
+  control_type: string;
   strength: number;
   guidance_start?: number;
   guidance_end?: number;

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

@@ -241,6 +241,16 @@ const ASPECT_RATIOS: Record<string, [number, number]> = {
 };
 const ASPECT_OPTIONS = ["custom", ...Object.keys(ASPECT_RATIOS)];
 
+// Friendly labels for ControlNet control types. "canny" traces edges from a source image;
+// every other type is an already-made map (passthrough/depth/pose/...). Unknown types fall
+// back to a capitalized "(map)" label so a new backend type still renders.
+const CONTROL_TYPE_LABELS: Record<string, string> = {
+  passthrough: "Passthrough (already a map)",
+  canny: "Canny (trace edges)",
+  depth: "Depth (map)",
+  pose: "Pose (map)",
+};
+
 // Z-Image accepts 256–2048, in multiples of 16. Snap any value into range.
 const MIN_DIM = 256;
 const MAX_DIM = 2048;
@@ -775,6 +785,31 @@ async function buildOutpaint(
   mctx.fillStyle = "#000000"; // ...except the kept original (inset by the seam overlap).
   mctx.fillRect(l + ol, t + ot, w - ol - or, h - ot - ob);
 
+  // The grown canvas can exceed the backend's 4096px-per-side decode limit (e.g. a
+  // 2048px source at 100% on both sides -> 6144px), which would 400 the load. Scale the
+  // built pair down proportionally to fit, so Extend still returns an outpaint instead
+  // of failing. The backend also rounds to /16, so exact dims here are not required.
+  const MAX_SIDE = 4096;
+  const longest = Math.max(nw, nh);
+  if (longest > MAX_SIDE) {
+    const scale = MAX_SIDE / longest;
+    const sw = Math.max(1, Math.round(nw * scale));
+    const sh = Math.max(1, Math.round(nh * scale));
+    const scaleCanvas = (source: HTMLCanvasElement): HTMLCanvasElement => {
+      const dst = document.createElement("canvas");
+      dst.width = sw;
+      dst.height = sh;
+      const dctx = dst.getContext("2d");
+      if (!dctx) throw new Error("Could not scale the extended canvas");
+      dctx.drawImage(source, 0, 0, sw, sh);
+      return dst;
+    };
+    return {
+      image: scaleCanvas(ic).toDataURL("image/png"),
+      mask: scaleCanvas(mc).toDataURL("image/png"),
+    };
+  }
+
   return { image: ic.toDataURL("image/png"), mask: mc.toDataURL("image/png") };
 }
 
@@ -922,7 +957,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
   // per loaded family; applied at generate time only when a model + control image are set.
   const [controlnetId, setControlnetId] = useState<string>("");
   const [controlImage, setControlImage] = useState<string | null>(null);
-  const [controlType, setControlType] = useState<"passthrough" | "canny">("passthrough");
+  // Free-form: a union ControlNet advertises depth/pose/etc alongside the preprocessing
+  // "canny", and the backend maps the exact control_type to the union control_mode. The
+  // picker is built from the selected model's control_types, so it isn't limited to two.
+  const [controlType, setControlType] = useState<string>("passthrough");
   const [controlStrength, setControlStrength] = useState(0.7);
   const [availableControlNets, setAvailableControlNets] = useState<DiffusionControlNetInfo[]>([]);
   // Advanced options live in a right-docked panel (like Chat's settings panel). Closed by
@@ -980,6 +1018,12 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
   const loadToastId = useRef<string | number | null>(null);
   // Last load-progress signature shown, so a tick that moved nothing skips the toast.
   const lastLoadSig = useRef<string | null>(null);
+  // The quant to restore if the current optimistic swap fails. A same-repo quant
+  // change sets `quant` immediately for picker feedback; if the load then fails
+  // AFTER starting (an error/eviction during download), the old pipeline stays
+  // loaded, so the poll must roll the label back rather than advertise the failed
+  // quant. `{ prev }` distinguishes "revert to null" from "nothing pending".
+  const quantRevert = useRef<{ prev: string | null } | null>(null);
 
   const dismissLoadToast = useCallback(() => {
     if (loadToastId.current != null) toast.dismiss(loadToastId.current);
@@ -1063,6 +1107,25 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
     };
   }, [controlnetCapable, status?.family]);
 
+  // The control types offered for the selected ControlNet. A union model advertises
+  // several (canny/depth/pose/passthrough); a plain model advertises its own. Fall back
+  // to the preprocessing pair when nothing is selected.
+  const controlTypeOptions = useMemo(() => {
+    const cn = availableControlNets.find((c) => c.id === controlnetId);
+    const types = cn?.control_types?.length ? cn.control_types : ["passthrough", "canny"];
+    return types;
+  }, [availableControlNets, controlnetId]);
+
+  // Keep controlType valid for the selected model: if the current choice isn't among the
+  // model's advertised types, snap to the first (prefer passthrough when offered).
+  useEffect(() => {
+    if (!controlTypeOptions.includes(controlType)) {
+      setControlType(
+        controlTypeOptions.includes("passthrough") ? "passthrough" : controlTypeOptions[0],
+      );
+    }
+  }, [controlTypeOptions, controlType]);
+
   const selected = useMemo(
     () => images.find((i) => i.id === selectedId) ?? images[0] ?? null,
     [images, selectedId],
@@ -1169,6 +1232,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
     const m = matchAspect(image.width, image.height);
     setAspect(m.key);
     setPortrait(m.portrait);
+    // Restore selected LoRA adapters from the recipe ("id:weight" strings); split on the
+    // LAST colon so an id that itself contains ':' is preserved. Unparseable entries are
+    // skipped, and a recipe with no LoRAs clears the current selection so the restore
+    // reproduces the image faithfully rather than leaking a stale form selection.
+    const restoredLoras: LoraSpecInput[] = [];
+    for (const entry of image.loras ?? []) {
+      const idx = entry.lastIndexOf(":");
+      if (idx <= 0) continue;
+      const id = entry.slice(0, idx);
+      const weight = Number(entry.slice(idx + 1));
+      if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight });
+    }
+    setLoras(restoredLoras);
     toast.success("Settings restored to inputs");
   }, []);
 
@@ -1251,12 +1327,22 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         setStatus(await getDiffusionStatus());
         toast.success("Model loaded");
         setBusy(null);
+        // Load succeeded: the optimistic quant is now the real one, so drop the
+        // pending revert.
+        quantRevert.current = null;
         return;
       }
       if (p.phase === "error") {
         dismissLoadToast();
         toast.error(p.error || "Failed to load model");
         setBusy(null);
+        // A load that failed AFTER starting leaves the previous pipeline loaded, so
+        // roll the optimistic quant label back to what is actually loaded (status
+        // does not carry the quant, so refreshStatus alone can't correct it).
+        if (quantRevert.current) {
+          setQuant(quantRevert.current.prev);
+          quantRevert.current = null;
+        }
         // A failed load may have freed a previously-loaded model, so resync to
         // the real backend state (the synchronous failure path does the same).
         void refreshStatus();
@@ -1269,6 +1355,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         // busy stuck on "loading", deadening the picker and Generate button.
         dismissLoadToast();
         setBusy(null);
+        // Same optimistic-quant rollback as the error path: the swap did not take.
+        if (quantRevert.current) {
+          setQuant(quantRevert.current.prev);
+          quantRevert.current = null;
+        }
         void refreshStatus();
         return;
       }
@@ -1410,17 +1501,23 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         return;
       }
       // GGUF quant pick from the variant expander. Optimistic for instant picker
-      // feedback, but revert if the load fails to START (400/409/network): the
-      // selector must not advertise a quant that is not the loaded one. Poll-phase
-      // failures re-sync via refreshStatus.
+      // feedback, but revert if the load fails to START (400/409/network) or LATER
+      // during the poll (download/preflight error/eviction) -- in both cases the old
+      // pipeline stays loaded, so the selector must not advertise the failed quant.
+      // The poll owns the after-start revert via quantRevert; here we only handle
+      // the never-started case.
       if (meta.ggufVariant && meta.ggufFilename) {
         const prevQuant = quant;
+        quantRevert.current = { prev: prevQuant };
         setQuant(meta.ggufVariant);
         const dq = defaultsFor(id);
         setSteps(dq.steps);
         setGuidance(dq.guidance);
         void handleLoad(id, { kind: "gguf", filename: meta.ggufFilename }).then((started) => {
-          if (!started) setQuant(prevQuant);
+          if (!started) {
+            setQuant(prevQuant);
+            quantRevert.current = null;
+          }
         });
         return;
       }
@@ -1435,14 +1532,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         if (!filename.toLowerCase().endsWith(".gguf")) return;
         // A direct pick carries no curated variant label; surface the filename so
         // the selector stops advertising the previously loaded quant. Optimistic,
-        // reverted if the load fails to start (mirrors the curated branch above).
+        // reverted if the load fails to start OR fails later in the poll (mirrors the
+        // curated branch above; the poll owns the after-start revert via quantRevert).
         const prevQuant = quant;
+        quantRevert.current = { prev: prevQuant };
         setQuant(filename);
         const dq2 = defaultsFor(id);
         setSteps(dq2.steps);
         setGuidance(dq2.guidance);
         void handleLoad(dir, { kind: "gguf", filename }).then((started) => {
-          if (!started) setQuant(prevQuant);
+          if (!started) {
+            setQuant(prevQuant);
+            quantRevert.current = null;
+          }
         });
         return;
       }
@@ -2250,16 +2352,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
                     <ImageDropzone value={controlImage} onChange={setControlImage} />
                     <div className="flex items-center gap-2">
                       <span className="shrink-0 text-xs text-muted-foreground">Control type</span>
-                      <Select
-                        value={controlType}
-                        onValueChange={(v) => setControlType(v as "passthrough" | "canny")}
-                      >
+                      <Select value={controlType} onValueChange={setControlType}>
                         <SelectTrigger className="h-8 flex-1 text-xs">
                           <SelectValue />
                         </SelectTrigger>
                         <SelectContent>
-                          <SelectItem value="passthrough">Passthrough (already a map)</SelectItem>
-                          <SelectItem value="canny">Canny (trace edges)</SelectItem>
+                          {controlTypeOptions.map((t) => (
+                            <SelectItem key={t} value={t}>
+                              {CONTROL_TYPE_LABELS[t] ??
+                                `${t.charAt(0).toUpperCase()}${t.slice(1)} (map)`}
+                            </SelectItem>
+                          ))}
                         </SelectContent>
                       </Select>
                     </div>

studio/install_sd_cpp_prebuilt.py

@@ -262,6 +262,10 @@ def _maybe_fetch_windows_cudart(release: dict, chosen: str, target: Path) -> Non
     print(f"downloading CUDA runtime {cudart['name']} ...", flush = True)
     try:
         _download(cudart["browser_download_url"], dest)
+        # Verify integrity BEFORE extracting, like the main sd-cli archive: these DLLs are
+        # loaded into sd-cli.exe at runtime, so a corrupt/tampered runtime archive must be
+        # rejected rather than extracted next to the binary.
+        _verify_sha256(dest, cudart.get("digest"))
         with zipfile.ZipFile(dest) as zf:
             _safe_extractall(zf, target)
     finally: