From 253098ddf639203c2f4db616229e73aab0e63c24 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 04:37:03 +0000 Subject: [PATCH 01/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_sdxl.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index 2ebe7f2013..6bec0a4100 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -138,7 +138,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params(): class _GgufDenoiser: def parameters(self): - yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block + yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block yield types.SimpleNamespace(dtype = torch.bfloat16) # compute dtype vae = _FakeVae(dtype = torch.float32) @@ -151,9 +151,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params(): yield types.SimpleNamespace(dtype = torch.uint8) vae2 = _FakeVae(dtype = torch.float32) - DiffusionBackend._align_vae_dtype( - types.SimpleNamespace(transformer = _AllPacked(), vae = vae2) - ) + DiffusionBackend._align_vae_dtype(types.SimpleNamespace(transformer = _AllPacked(), vae = vae2)) assert vae2.moved_to is None From 4eb2c2225869ab528e2209b37b05ed6182f34641 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 04:37:50 +0000 Subject: [PATCH 02/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_sdxl.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index 2ebe7f2013..6bec0a4100 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -138,7 +138,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params(): class _GgufDenoiser: def parameters(self): - yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block + yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block yield types.SimpleNamespace(dtype = torch.bfloat16) # compute dtype vae = _FakeVae(dtype = torch.float32) @@ -151,9 +151,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params(): yield types.SimpleNamespace(dtype = torch.uint8) vae2 = _FakeVae(dtype = torch.float32) - DiffusionBackend._align_vae_dtype( - types.SimpleNamespace(transformer = _AllPacked(), vae = vae2) - ) + DiffusionBackend._align_vae_dtype(types.SimpleNamespace(transformer = _AllPacked(), vae = vae2)) assert vae2.moved_to is None From a3e23a8080fc1e46a68afc9b968a28f4ced6d923 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 04:38:21 +0000 Subject: [PATCH 03/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_sdxl.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index 2ebe7f2013..6bec0a4100 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -138,7 +138,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params(): class _GgufDenoiser: def parameters(self): - yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block + yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block yield types.SimpleNamespace(dtype = torch.bfloat16) # compute dtype vae = _FakeVae(dtype = torch.float32) @@ -151,9 +151,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params(): yield types.SimpleNamespace(dtype = torch.uint8) vae2 = _FakeVae(dtype = torch.float32) - DiffusionBackend._align_vae_dtype( - types.SimpleNamespace(transformer = _AllPacked(), vae = vae2) - ) + DiffusionBackend._align_vae_dtype(types.SimpleNamespace(transformer = _AllPacked(), vae = vae2)) assert vae2.moved_to is None From f6b12700711d4035faa982db7a7f5fa92e499afb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 04:38:52 +0000 Subject: [PATCH 04/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_sdxl.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py index 2ebe7f2013..6bec0a4100 100644 --- a/studio/backend/tests/test_diffusion_sdxl.py +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -138,7 +138,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params(): class _GgufDenoiser: def parameters(self): - yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block + yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block yield types.SimpleNamespace(dtype = torch.bfloat16) # compute dtype vae = _FakeVae(dtype = torch.float32) @@ -151,9 +151,7 @@ def test_align_vae_dtype_skips_gguf_packed_uint8_params(): yield types.SimpleNamespace(dtype = torch.uint8) vae2 = _FakeVae(dtype = torch.float32) - DiffusionBackend._align_vae_dtype( - types.SimpleNamespace(transformer = _AllPacked(), vae = vae2) - ) + DiffusionBackend._align_vae_dtype(types.SimpleNamespace(transformer = _AllPacked(), vae = vae2)) assert vae2.moved_to is None From 18f9510d11f5cc33ac9336de34d98aa699ac08fa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:41:23 +0000 Subject: [PATCH 05/15] Address a further round of Codex review findings on the image PR Backend: - validate_load_request rejects a non-.gguf single-file name before the GPU handoff, so a family-looking repo paired with README.md no longer evicts the chat model and only fails in the background load. - detect_family scopes the edit/kontext/inpaint keyword check to the model id or filename basename, not arbitrary parent directories, so a valid text-to-image file under a folder named edit is no longer rejected. - the images gallery listing skips records that fail schema validation, so one corrupt or hand-dropped PNG can no longer 500 the whole endpoint. - _terminate reaps the killed sd-cli child so cancellation and timeout paths do not leak zombie process-table entries. - the images load route gates the chat-eviction handoff on the resolved device being non-CPU, so a CPU-only diffusers fallback no longer evicts a resident chat model for a load that cannot use the GPU. Frontend: - treat Images as a chat-like full-height route (no outer padding or scroll) so its picker is not pushed down and the gallery is not clipped. - allow /images under the chat-only guard so the native CPU/MPS image path is reachable on the no-GPU hosts it was built for. - roll the optimistic quant label back when a same-repo swap fails after the load started, so the selector never advertises a quant that is not loaded. --- studio/backend/core/inference/diffusion.py | 8 ++++ .../core/inference/diffusion_families.py | 10 +++-- .../backend/core/inference/sd_cpp_engine.py | 8 ++++ studio/backend/routes/inference.py | 33 ++++++++++----- .../backend/tests/test_diffusion_backend.py | 19 +++++++++ studio/backend/tests/test_sd_cpp_engine.py | 20 +++++++++ studio/frontend/src/app/routes/__root.tsx | 18 ++++++-- .../src/features/images/images-page.tsx | 42 ++++++++++++++++--- 8 files changed, 136 insertions(+), 22 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 63914103bd..1680ce0288 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -290,6 +290,14 @@ class DiffusionBackend: raise ValueError( "gguf_filename is required: this backend loads single-file GGUF checkpoints only." ) + # Reject a non-GGUF single-file name (e.g. README.md, config.json) here, before + # the route hands the GPU over: without this a family-looking repo_id paired with + # a non-GGUF filename passes preflight, evicts the chat model, and only fails in + # the background from_single_file -- exactly the eviction this validation prevents. + if not gguf_filename.lower().endswith(".gguf"): + raise ValueError( + f"gguf_filename must name a .gguf single-file checkpoint; got '{gguf_filename}'." + ) fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: raise ValueError( diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 7257c69e9e..5beb8e36b6 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -170,9 +170,13 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff # Match edit keywords as whole id segments, not raw substrings, so a normal # text-to-image repo like ".../some-image-edition" isn't misread as an editing # checkpoint. Qwen-Image-Edit / FLUX.1-Kontext still match (edit/kontext are - # whole tokens there). Split on both path separators so a Windows local path - # is segmented too. - segments = set(re.split(r"[-_./\\]+", needle)) + # whole tokens there). Scope the 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 + # just because a parent folder happens to be named `edit`. The combined + # `repo_id/gguf_filename` fallback passes the filename as that last segment. + basename = re.split(r"[/\\]+", needle)[-1] + segments = set(re.split(r"[-_.]+", basename)) if any(kw in segments for kw in _EDIT_KEYWORDS): return None for fam in _FAMILIES: diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 77c3b457b0..182b25a1b7 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/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: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8c946484a2..67e541a8d8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11051,12 +11051,10 @@ async def load_diffusion_model( from core.inference.diffusion import get_diffusion_backend 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() @@ -11078,12 +11076,15 @@ async def load_diffusion_model( # installing the sd-cli binary if needed -- all BEFORE evicting chat, so a # native fallback never strands a half-loaded state. engine = await asyncio.to_thread(select_and_activate_engine, fam, hf_token = request.hf_token) - # 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 can't 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. @@ -11216,6 +11217,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)) @@ -11223,10 +11226,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") diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 1ea3d19e98..2801e96c3c 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -56,6 +56,20 @@ def test_detect_family_from_repo_id(): assert detect_family("meta-llama/Llama-3-8B") 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`/`kontext`/`inpaint` must NOT reject a valid + # text-to-image file: only the model id / filename basename is scanned for the + # edit keyword. A direct local pick arrives as (parent_dir, filename). + assert detect_family("/models/edit") is None # 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/kontext", "qwen-image-2512-Q4.gguf").name == "qwen-image" + # But a genuine editing checkpoint (keyword in the id/filename) is still rejected. + assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None + assert detect_family_for_pick("/models/misc", "Qwen-Image-Edit-2511-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" @@ -760,6 +774,11 @@ def test_validate_load_request(tmp_path): backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF") 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" diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 2c10699be1..4187eac7cf 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -135,6 +135,26 @@ 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) ────────────────────────────────────────────── diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 9b959c3939..220f66d277 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/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" > - +
{/* 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() {
(null); // Last load-progress signature shown, so a tick that moved nothing skips the toast. const lastLoadSig = useRef(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); @@ -644,12 +650,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(); @@ -662,6 +678,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; } @@ -751,15 +772,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (busy !== null) return; if (meta.ggufVariant && meta.ggufFilename) { // 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. + // 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. const prevQuant = quant; + quantRevert.current = { prev: prevQuant }; setQuant(meta.ggufVariant); const d = defaultsFor(id); setSteps(d.steps); setGuidance(d.guidance); void handleLoad(id, meta.ggufFilename).then((started) => { - if (!started) setQuant(prevQuant); + if (!started) { + setQuant(prevQuant); + quantRevert.current = null; + } }); return; } @@ -774,14 +801,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 d = defaultsFor(id); setSteps(d.steps); setGuidance(d.guidance); void handleLoad(dir, filename).then((started) => { - if (!started) setQuant(prevQuant); + if (!started) { + setQuant(prevQuant); + quantRevert.current = null; + } }); } }, From ca782d102e9f45300149d7c4a45436ba5f753ad9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 05:41:58 +0000 Subject: [PATCH 06/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_sd_cpp_engine.py | 1 - 1 file changed, 1 deletion(-) diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 4187eac7cf..5c9d3256af 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -141,7 +141,6 @@ def test_terminate_reaps_killed_child(): # 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"), From f58c3ddb07654743875b8ae09c67ca5a709d155a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:46:48 +0000 Subject: [PATCH 07/15] Count LR scheduler warmup/decay in optimizer steps, not micro-steps lr_sched.step() runs once per outer optimizer step (after the gradient accumulation inner loop), for train_steps total. The scheduler was configured with num_warmup_steps and num_training_steps multiplied by gradient_accumulation_steps, so with accumulation > 1 a warmup or non-constant schedule stretched past the run and never reached the intended decay. Count both in optimizer steps. --- studio/backend/core/training/diffusion_lora_trainer.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 70f7c8aedf..0e20b6ede5 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/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 From 8772f81e7521a3a1baa4127197267d56cc36b199 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:46:50 +0000 Subject: [PATCH 08/15] 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") }; } From c2b25feaeefe80236a3ff7ed1ffe3d4c52965531 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:47:50 +0000 Subject: [PATCH 09/15] Guard inference loads and worker lifetime against diffusion training Teach the chat and image load guards about an active diffusion (SDXL) LoRA job: a chat load is refused (its footprint cannot be fit-checked against the trainer) and an image load is refused outright, mirroring the existing LLM training guards, so a load can no longer allocate GPU memory alongside the trainer and undo the pre-start cleanup. Bind the diffusion trainer subprocess to the parent's lifetime and scrub the native path lease secret from it by running the child through run_without_native_path_secret, matching the inference/export/LLM workers, so a Studio crash or kill no longer leaves the trainer holding the GPU. Reset in_model_load on the complete and error terminal events: a stop or failure during model loading otherwise leaves the status reporting a stale loading indicator after the job has ended. --- .../training/diffusion_training_service.py | 29 +++++++++++++-- studio/backend/routes/inference.py | 36 ++++++++++++++++--- .../backend/tests/test_diffusion_training.py | 16 +++++++++ 3 files changed, 75 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 92621931f7..bc11b38a37 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -31,12 +31,26 @@ _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 +239,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 +253,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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5f1a5103a4..648f07ea2f 100644 --- a/studio/backend/routes/inference.py +++ b/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 = ( diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index c96dabb76f..35ff816046 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/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): From 722797897812b847b781f735e91561c6b4d2de53 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:48:06 +0000 Subject: [PATCH 10/15] Harden diffusion LoRA handling on the diffusers and native paths Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max): diffusers requires the adapter loaded before compilation, so applying one to the already-compiled module fails with adapter-key mismatches. The status gate now hides the picker and generate raises a clear message instead. Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the diffusion cancellation sentinel in resolve_specs, so an unload/superseding load during resolution maps to a 409 instead of a generic server error. Drop weight-0 LoRA rows before the native support gate so a request carrying only disabled adapters stays a no-op on families where native LoRA is unsupported, matching the diffusers path. Reject duplicate LoRA ids in the request model: both apply paths suffix colliding names, so a repeated id would stack the same adapter past its per-adapter weight bound. Strip all user-typed prompt tags on the native path (only the selected adapters are materialized in the managed lora-model-dir, so an unselected tag can never resolve), and restore saved LoRA selections from a gallery recipe so restore reproduces a LoRA image. --- studio/backend/core/inference/diffusion.py | 7 ++- .../backend/core/inference/diffusion_lora.py | 35 ++++++++++---- .../backend/core/inference/sd_cpp_backend.py | 9 +++- studio/backend/models/inference.py | 15 ++++++ studio/backend/tests/test_diffusion_lora.py | 48 +++++++++++++++++-- .../src/features/images/images-page.tsx | 13 +++++ 6 files changed, 110 insertions(+), 17 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 2bc54db955..8a445b852e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1372,11 +1372,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) @@ -1829,6 +1831,7 @@ class DiffusionBackend: family = state.family.name, model_kind = state.kind, transformer_quant = state.transformer_quant, + compiled = "compiled" in (getattr(state, "speed_optims", ()) or ()), ), } diff --git a/studio/backend/core/inference/diffusion_lora.py b/studio/backend/core/inference/diffusion_lora.py index 20b40aaf55..7db07dcd60 100644 --- a/studio/backend/core/inference/diffusion_lora.py +++ b/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 `` 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"" 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 diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 3a09acd0c5..3e92086df0 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -528,8 +528,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, @@ -541,7 +546,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: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 234dea7d7f..6392e19544 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1897,6 +1897,21 @@ class DiffusionGenerateRequest(BaseModel): "message when the loaded model or its quantisation can't apply LoRA.", ) + @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]]: diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py index e7df3da936..4b0fc32227 100644 --- a/studio/backend/tests/test_diffusion_lora.py +++ b/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 ", [r]) == "a cat " -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 ", [r]) - assert "" in out - assert "" in out + assert "" not in out + assert out == "a cat " def test_inject_prompt_tags_empty_returns_prompt(): @@ -83,6 +85,38 @@ 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 +226,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 ───────────────────────────────────────────────── diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index c206b8d0c0..863bdf5752 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1087,6 +1087,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"); }, []); From 3cedcfe781df93883a8953b020c4103c40449c0d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 05:48:20 +0000 Subject: [PATCH 11/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_training_service.py | 1 - 1 file changed, 1 deletion(-) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index bc11b38a37..d0b3d3220e 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -45,7 +45,6 @@ def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None: # 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 ) From 048d0422c7c6d240d02aa7c4cecc938444ef7cdc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:54:22 +0000 Subject: [PATCH 12/15] Harden ControlNet resolve, gallery metadata, and the control-type picker Check cancellation immediately after a ControlNet from_pretrained and before any device placement, so an unload/eviction that raced the download does not allocate several GB onto the GPU after the load was already cleared. Require a loadable weight or shard index (not just config.json) before a local ControlNet folder is advertised, so an interrupted copy is hidden instead of failing deep in from_pretrained as a generic 500. Do not record a strength-0 ControlNet in the gallery recipe: it is treated as disabled and skipped, so the image is unconditioned and the metadata must not claim a ControlNet was applied. Build the control-type picker from the selected ControlNet's advertised control_types instead of a hardcoded passthrough/canny pair, so a union model with a precomputed depth or pose map sends the correct control_mode. --- studio/backend/core/inference/diffusion.py | 7 +++ .../core/inference/diffusion_controlnet.py | 24 +++++++++- studio/backend/routes/inference.py | 5 +- .../tests/test_diffusion_controlnet.py | 16 +++++++ studio/frontend/src/features/images/api.ts | 4 +- .../src/features/images/images-page.tsx | 47 ++++++++++++++++--- 6 files changed, 93 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 03a1c39b78..0ededa6614 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1360,6 +1360,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 diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index 6b7cdec9fb..b318b040c8 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/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( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5f1a5103a4..daf895072b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11230,7 +11230,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, diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index e91abad053..5708c3a3ae 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/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 diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index e33da62113..ad38261705 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/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; diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 9d7107a927..9471882067 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -224,6 +224,16 @@ const ASPECT_RATIOS: Record = { }; 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 = { + 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; @@ -900,7 +910,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(""); const [controlImage, setControlImage] = useState(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("passthrough"); const [controlStrength, setControlStrength] = useState(0.7); const [availableControlNets, setAvailableControlNets] = useState([]); // Advanced options live in a right-docked panel (like Chat's settings panel). Closed by @@ -1027,6 +1040,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], @@ -2163,16 +2195,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
Control type - - Passthrough (already a map) - Canny (trace edges) + {controlTypeOptions.map((t) => ( + + {CONTROL_TYPE_LABELS[t] ?? + `${t.charAt(0).toUpperCase()}${t.slice(1)} (map)`} + + ))}
From 691bad30c4a73c58c9b9edec4a67b3599dbbe671 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 06:21:07 +0000 Subject: [PATCH 13/15] Address further Codex findings on the image-workflows PR - Persist the actual output image size in the gallery recipe instead of the request sliders: Transform/Inpaint/Edit derive the size from the uploaded image, Extend grows the canvas, and Upscale resizes it, so the sliders recorded (and later restored) the wrong dimensions for those workflows. - Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file name) in validate_load_request, so the unloadable pick fails before chat is evicted rather than deep in from_pretrained. - Only publish an image-conditioned from_pipe wrapper to the shared aux cache when the load is still current: from_pipe runs under the generate lock but not the state lock, so an unload racing its construction could otherwise cache a wrapper over torn-down modules that a later load would reuse. - Verify the Windows CUDA runtime archive checksum before extracting it, like the main sd-cli archive, so a corrupt or tampered runtime is rejected rather than extracted next to the binary. --- studio/backend/core/inference/diffusion.py | 19 ++++++++++++++++++- studio/backend/routes/inference.py | 9 +++++++-- .../backend/tests/test_diffusion_backend.py | 5 +++++ studio/install_sd_cpp_prebuilt.py | 4 ++++ 4 files changed, 34 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8f7291cffa..7bd7180b53 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -488,6 +488,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 ───────────────────────────────────────── @@ -1348,7 +1358,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 @staticmethod diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 084d892a07..30f35cb9de 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11199,8 +11199,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, diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index d1cbf513d0..ce1ced8352 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1461,6 +1461,11 @@ def test_validate_load_request(tmp_path): 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( diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 193c5e4670..639e3a894d 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/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: From ffed150070aeea0e6203d98e2988aefe30de1c6f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 06:37:32 +0000 Subject: [PATCH 14/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion.py | 1 + studio/backend/core/inference/diffusion_families.py | 4 +--- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 978432f8be..eb4bb33034 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1042,6 +1042,7 @@ class DiffusionBackend: install_arch_patches, uninstall_arch_patches, ) + try: if effective_speed != SPEED_OFF: install_compile_safe_patches() diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index b95278b037..08140f7957 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -258,9 +258,7 @@ def _token_in_needle(token: str, needle: str) -> bool: ('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 + return re.search(r"(?:^|[-_./\\])" + re.escape(token) + r"(?:$|[-_./\\])", needle) is not None def _best_family_match(needle: str) -> Optional[DiffusionFamily]: From 5a221223e6d14974f063787455435f33aac2952c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 06:39:19 +0000 Subject: [PATCH 15/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/models/inference.py | 4 +++- studio/backend/tests/test_diffusion_lora.py | 15 +++++++++++---- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 6392e19544..8926a6bb7e 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1908,7 +1908,9 @@ class DiffusionGenerateRequest(BaseModel): 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") + raise ValueError( + f"duplicate LoRA id '{spec.id}'; list each adapter at most once" + ) seen.add(spec.id) return value diff --git a/studio/backend/tests/test_diffusion_lora.py b/studio/backend/tests/test_diffusion_lora.py index 4b0fc32227..f99739eff9 100644 --- a/studio/backend/tests/test_diffusion_lora.py +++ b/studio/backend/tests/test_diffusion_lora.py @@ -88,13 +88,19 @@ def test_supports_lora_matrix(): # 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, + 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, + engine = "sd_cpp", + family = "flux.1", + model_kind = "gguf", + transformer_quant = None, + compiled = True, ) @@ -109,6 +115,7 @@ def test_resolve_specs_maps_cancelled_to_diffusion_sentinel(tmp_path, monkeypatc 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")