Address Codex review findings on the image-workflows PR

Keep diffusion.py importable without torch: the compile/arch patch modules
import torch at module level, so import them lazily at their load/unload
call sites instead of at module load. This restores the torchless contract
so get_diffusion_backend() works on a CPU/native sd.cpp install.

Match family reject keywords and aliases as whole path/name segments, not
raw substrings, so an unrelated word like edited, edition, or kontextual no
longer misroutes or hides a valid base image model, while supported edit
families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same
segment matching in the picker task filter.

Route FLUX.2-dev native guidance through --guidance like the other FLUX
families rather than --cfg-scale. Reject native upscale requests that have
no input image. Read image header dimensions and reject over-limit inputs
before decoding pixels, so a crafted small-payload image cannot spike
memory. Reject an upscale that would shrink the source below its input
size. Validate the model_kind against the filename extension before the
GPU handoff. Estimate a local diffusers pipeline's size from its on-disk
weights so auto memory planning does not skip offload and OOM. Report
workflows: [txt2img] from the native backend status so the Create tab
stays enabled for a loaded native model. Clamp the outpaint canvas to the
backend's 4096px decode limit.

Adds regression tests for segment matching and kind/extension validation.

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)
 
 
@@ -467,6 +464,14 @@ 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'.")
             if local_root.exists():
                 resolve_local_gguf_child(local_root, gguf_filename)
             elif path_shaped:
@@ -721,6 +726,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.
@@ -732,21 +761,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 ───────────────────────────────
@@ -997,6 +1012,17 @@ 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()
@@ -1246,7 +1272,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
@@ -1475,6 +1507,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
@@ -1671,6 +1712,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()
         # Drop the workflow pipes built around this load's modules so they don't pin the

studio/backend/core/inference/diffusion_families.py

@@ -251,14 +251,28 @@ _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
 
@@ -286,7 +300,8 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff
         # matched family does not itself declare, reject so the load fails fast + clearly.
         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, needle) and not any(kw in tok for tok in matched_tokens)
+            for kw in _EDIT_KEYWORDS
         ):
             return None
         return match

studio/backend/core/inference/sd_cpp_backend.py

@@ -179,7 +179,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
@@ -498,10 +498,18 @@ class SdCppDiffusionBackend:
 
         from PIL import Image
 
-        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."
             )
 
         cancel = threading.Event()
@@ -647,6 +655,7 @@ class SdCppDiffusionBackend:
                 "attention_backend": None,
                 "transformer_cache": None,
                 "engine": "sd_cpp",
+                "workflows": [],
             }
         return {
             "loaded": True,
@@ -670,6 +679,11 @@ class SdCppDiffusionBackend:
             "attention_backend": None,
             "transformer_cache": None,
             "engine": "sd_cpp",
+            # 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/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,21 @@ 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_override():
     assert detect_family("local/path", override = "z-image").name == "z-image"
     assert detect_family("local/path", override = "zimage").name == "z-image"
@@ -1428,6 +1451,16 @@ def test_validate_load_request(tmp_path):
         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 local path with a missing child fails here (before any GPU/network work).
     with pytest.raises(FileNotFoundError):
         backend.validate_load_request(

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/images-page.tsx

@@ -752,6 +752,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") };
 }