diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 23b761dd32..28dd3f599e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -176,7 +176,8 @@ class DiffusionBackend: fam = detect_family(repo_id, family_override) if fam is None: raise ValueError( - f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)." + f"'{repo_id}' isn't a supported image-generation model. " + f"Supported: Z-Image, Qwen-Image, FLUX.1, FLUX.2-klein." ) with self._lock: @@ -221,10 +222,12 @@ class DiffusionBackend: expected, base_files = self._estimate_download_bytes( kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token") ) - loading = self._loading - if loading is not None: - loading.base_repo = base - loading.expected_bytes = expected + with self._lock: + # Stamp progress only if this load is still current; a superseding + # load (or unload) has its own token and its own _LoadingState. + if self._load_token == token and self._loading is not None: + self._loading.base_repo = base + self._loading.expected_bytes = expected # Download outside the lock so unload()/an eviction can preempt the # multi-GB pull; load_pipeline below then assembles from the cache. self._prefetch_files( @@ -246,9 +249,13 @@ class DiffusionBackend: if self._load_token != token: return logger.error("diffusion.load_failed: %s", exc) + # Redact native paths: this error is surfaced verbatim via the + # load-progress poll, and Studio can run as a shared server. + from utils.native_path_leases import redact_native_paths + with self._lock: if self._load_token == token and self._loading is not None: - self._loading.error = str(exc) + self._loading.error = redact_native_paths(str(exc)) def load_progress(self) -> dict[str, Any]: """Phase + downloaded/total bytes for the in-flight load (cache-scan based).""" @@ -330,7 +337,8 @@ class DiffusionBackend: fam = detect_family(repo_id, family_override) if fam is None: raise ValueError( - f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)." + f"'{repo_id}' isn't a supported image-generation model. " + f"Supported: Z-Image, Qwen-Image, FLUX.1, FLUX.2-klein." ) base = _resolve_base_repo(repo_id, base_repo, fam, hf_token) device, dtype = self._pick_device_and_dtype() diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 4c486e9a96..e38595bb01 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -308,7 +308,8 @@ def test_load_without_gguf_raises(): def test_load_unknown_family_raises(): backend = DiffusionBackend() - with pytest.raises(ValueError): + # User-facing message: names the supported models, no internal-API jargon. + with pytest.raises(ValueError, match = "isn't a supported image-generation model"): backend.load_pipeline("some/unrecognised-repo", gguf_filename = "x.gguf") @@ -497,3 +498,28 @@ def test_prefetch_downloads_gguf_and_base(monkeypatch, tmp_path): backend._prefetch_files(str(tmp_path), "model.gguf", "base/repo", ["vae/x.safetensors"], None) assert all(repo != str(tmp_path) for repo, _ in calls) assert ("base/repo", "vae/x.safetensors") in calls + + +def test_run_load_does_not_stamp_superseded_progress(fake_runtime, monkeypatch): + # A worker whose load is superseded mid-resolve must not stamp its progress + # (base_repo / expected_bytes) onto the new load's _LoadingState. + backend = DiffusionBackend() + backend._loading = _LoadingState(repo_id = "unsloth/Z-Image-Turbo-GGUF", base_repo = "seed") + backend._load_token = 5 + monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None) + + def supersede_then_estimate(*a, **k): + backend._load_token = 6 # a newer begin_load bumped the token mid-resolve + return (99999, []) + + monkeypatch.setattr( + DiffusionBackend, "_estimate_download_bytes", staticmethod(supersede_then_estimate) + ) + monkeypatch.setattr(DiffusionBackend, "_prefetch_files", lambda self, *a, **k: None) + monkeypatch.setattr(DiffusionBackend, "load_pipeline", lambda self, **k: None) + + backend._run_load( + repo_id = "unsloth/Z-Image-Turbo-GGUF", gguf_filename = "m.gguf", base_repo = None, _load_token = 5 + ) + assert backend._loading.expected_bytes == 0 + assert backend._loading.base_repo == "seed" diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index e60fb3237a..5de10ede19 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -225,7 +225,7 @@ def test_generate_without_load_returns_409(client): def test_load_unknown_family_returns_400(client, monkeypatch): def _raise(*a, **k): - raise ValueError("Could not infer a diffusion family for 'x/y'.") + raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.") backend = _FakeBackend() backend.begin_load = _raise @@ -234,7 +234,7 @@ def test_load_unknown_family_returns_400(client, monkeypatch): "/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"} ) assert resp.status_code == 400 - assert "family" in resp.json()["detail"] + assert "isn't a supported image-generation model" in resp.json()["detail"] def test_load_progress_route(client): 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 a49acda408..5dfc22be14 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1580,9 +1580,10 @@ export function HubModelPicker({ // the chat classifier marks image tasks "unsupported". const isChatSupported = useCallback( (r: HfModelResult) => { - // Image tab: keep task-matching results, but drop editing checkpoints the - // backend rejects (so they don't appear in Hub search only to 400 on load). - if (task && taskMatchesFilter(r.pipelineTag, task)) return !isImageEditModel(r.id); + // Image tab (task set): only task-matching, non-editing results. Anything + // else (e.g. a chat GGUF surfaced by a typed query) is dropped rather than + // falling through to the chat classifier and appearing as loadable. + if (task) return taskMatchesFilter(r.pipelineTag, task) && !isImageEditModel(r.id); return ( classifyUnslothSupport({ modelId: r.id, diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index c832e38e4f..f7766bf9a6 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -115,6 +115,9 @@ const ASPECT_OPTIONS = ["custom", ...Object.keys(ASPECT_RATIOS)]; // Z-Image accepts 256–2048, in multiples of 16. Snap any value into range. const MIN_DIM = 256; const MAX_DIM = 2048; +// Sequential runs are a frontend loop the backend never sees, so cap them here +// (every other generate field is bounded by the backend's request validators). +const MAX_RUNS = 128; function snapDim(value: number): number { if (!Number.isFinite(value)) return 1024; return Math.min(MAX_DIM, Math.max(MIN_DIM, Math.round(value / 16) * 16)); @@ -604,6 +607,9 @@ export function ImagesPage() { dismissLoadToast(); toast.error(p.error || "Failed to load model"); setBusy(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(); return; } // Include bytes_total: the estimate lands as a 0→real jump while phase and @@ -617,7 +623,7 @@ export function ImagesPage() { // Transient poll failure: keep trying. } pollTimer.current = setTimeout(() => void pollLoadProgress(), 1000); - }, [dismissLoadToast]); + }, [dismissLoadToast, refreshStatus]); const handleLoad = useCallback( async (repoId: string, ggufFilename: string) => { @@ -700,6 +706,12 @@ export function ImagesPage() { const w = snapDim(width); const h = snapDim(height); + // Bound the run loop: the number input accepts typed values past the slider + // max (e.g. 10000) or non-numeric ones (NaN), and runs are a frontend loop + // with no backend limit. + const runs = Math.max(1, Math.min(Number.isFinite(count) ? count : 1, MAX_RUNS)); + if (runs !== count) setCount(runs); + setBusy("generating"); setGenDone(0); setGenStep(null); @@ -719,7 +731,7 @@ export function ImagesPage() { } }, 300); try { - for (let i = 0; i < count; i++) { + for (let i = 0; i < runs; i++) { const res = await generateDiffusionImage({ prompt: prompt.trim(), // Only send a negative prompt when guidance uses it, so the recipe @@ -856,7 +868,7 @@ export function ImagesPage() { hint="How many times to repeat the generation, one after another. Each run uses the next seed, so the images differ and can be reproduced." value={count} min={1} - max={128} + max={MAX_RUNS} step={1} onChange={setCount} />