From 8772f81e7521a3a1baa4127197267d56cc36b199 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:46:50 +0000 Subject: [PATCH] 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 | 109 +++++++++++++----- .../core/inference/diffusion_families.py | 25 +++- .../backend/core/inference/sd_cpp_backend.py | 22 +++- .../backend/tests/test_diffusion_backend.py | 33 ++++++ .../assistant-ui/model-selector/pickers.tsx | 12 +- .../src/features/images/images-page.tsx | 25 ++++ 6 files changed, 183 insertions(+), 43 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 16f11a392b..8f7291cffa 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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 diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 1323ef1b1d..b8eaf312f0 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/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 diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 34ae1e98e3..3059dc440d 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/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"], } diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 214193d065..d1cbf513d0 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/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( diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index c5c562490c..b8d75b2eb6 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/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 diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index b6453ffc64..d23b2c85f0 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/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") }; }