Guard superseded loads, redact load errors, drop chat models from the image picker, clamp runs

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()

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"

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):

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,

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}
           />