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; + } }); } },