Fix picker dead-end for single-file repos, stale LoRA state after deferred compile, slash upload cap

Tag cached diffusion repos that ship no model_index.json with single_file in the
cached-models listing, and keep them out of the task-scoped On Device pickers
unless the curated catalog carries their artifact: the selection fall-through
loads uncataloged rows as a full pipeline and from_pretrained fails on a
single-file checkpoint repo after the GPU handoff.

Refresh diffusion status after a successful generation run on the Images page.
Speed Auto compiles the transformer on the third LoRA-free generation and flips
supports_lora to false; without the refresh the LoRA picker stayed enabled and
the next LoRA generation failed on the backend.

Match upload passthrough exact paths with trailing slashes normalized: the
trailing-slash variant of /api/train/diffusion/dataset reaches MaxBodyMiddleware
before the router's redirect_slashes 307, so it fell through to the default
/api/train body cap and 413ed large uploads. JSON sub-routes keep extra path
components after normalization and stay on the small cap.

studio/backend/main.py

@@ -767,9 +767,14 @@ _BODY_UPLOAD_PASSTHROUGH_EXACT_PATHS = (_DIFFUSION_DATASET_UPLOAD_PATH,)
 def _get_upload_passthrough_request_max_bytes(path: str) -> int:
     if path.startswith(_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX):
         return upload_request_limit_bytes(UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES)
+    # The trailing-slash variant (/api/train/diffusion/dataset/) reaches this middleware
+    # BEFORE the router's redirect_slashes 307, so it must resolve to the same upload cap
+    # as the canonical path or a large upload 413s on the default /api/train body cap.
+    # Stripping slashes cannot promote a JSON sub-route: those all keep extra path
+    # components after normalization and still miss the exact match.
     if (
         path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX)
-        or path == _DIFFUSION_DATASET_UPLOAD_PATH
+        or path.rstrip("/") == _DIFFUSION_DATASET_UPLOAD_PATH
     ):
         return upload_request_limit_bytes()
     return default_request_body_limit_bytes()
@@ -831,7 +836,10 @@ class MaxBodyMiddleware:
         self.upload_passthrough_exact_paths = upload_passthrough_exact_paths
 
     def _is_upload_passthrough(self, path: str) -> bool:
-        return path in self.upload_passthrough_exact_paths or any(
+        # Exact paths also match their trailing-slash variant: the middleware runs
+        # before the router's redirect_slashes 307, and a JSON sub-route can never
+        # normalize down to the exact path (it keeps extra components).
+        return path.rstrip("/") in self.upload_passthrough_exact_paths or any(
             path.startswith(p) for p in self.upload_passthrough_prefixes
         )
 

studio/backend/routes/models.py

@@ -36,6 +36,11 @@ class CachedModelRepo(BaseModel):
     # weights). The picker must not treat a partial base repo as a usable download, or an
     # On Device click routes to a fresh multi-GB re-download instead of the complete GGUF.
     partial: Optional[bool] = None
+    # True for a diffusion-tagged repo with NO top-level model_index.json: a single-file
+    # checkpoint that needs from_single_file + a filename. The task-scoped pickers must not
+    # offer it as a pipeline load (from_pretrained on it fails after the GPU handoff)
+    # unless the curated catalog carries its artifact.
+    single_file: Optional[bool] = None
 
 
 class CachedModelsResponse(BaseModel):
@@ -3380,6 +3385,20 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
         return {"cached": []}
 
 
+def _repo_has_pipeline_index(repo_info) -> bool:
+    """Whether the cached snapshot carries a model_index.json, i.e. is loadable as a
+    full diffusers pipeline (from_pretrained). Single-file / ComfyUI checkpoints ship
+    none and need a checkpoint filename + from_single_file instead."""
+    try:
+        for rev in repo_info.revisions:
+            for f in rev.files:
+                if f.file_name == "model_index.json" or f.file_name.endswith("/model_index.json"):
+                    return True
+    except Exception:
+        pass
+    return False
+
+
 def _repo_is_diffusers(repo_info) -> bool:
     """True for an image-diffusion repo, so the chat picker hides it (it renders
     images, not chat) and the Images picker claims it — mirroring how cached
@@ -3390,13 +3409,8 @@ def _repo_is_diffusers(repo_info) -> bool:
     Qwen-Image or a z-image .safetensors) ship none. For those, fall back to the
     repo id resolving to a known diffusion family — the same resolver the Images
     backend loads from — so they don't surface as loadable chat models."""
-    try:
-        for rev in repo_info.revisions:
-            for f in rev.files:
-                if f.file_name == "model_index.json" or f.file_name.endswith("/model_index.json"):
-                    return True
-    except Exception:
-        pass
+    if _repo_has_pipeline_index(repo_info):
+        return True
     try:
         from core.inference.diffusion_families import detect_family
         if detect_family(getattr(repo_info, "repo_id", "") or "") is not None:
@@ -3502,6 +3516,11 @@ async def list_cached_models(
                         }
                         if is_partial:
                             row["partial"] = True
+                        # Flag diffusion repos with no pipeline index: loadable only via
+                        # from_single_file with a checkpoint filename, so the pickers must
+                        # not offer them as pipeline loads unless the catalog carries them.
+                        if row["task"] is not None and not _repo_has_pipeline_index(repo_info):
+                            row["single_file"] = True
                         # Keep the newest timestamp across duplicate caches;
                         # attach only when known so absent rows sort as oldest.
                         lm = max(last_modified, (existing or {}).get("last_modified", 0.0))

studio/backend/tests/test_cached_gguf_routes.py

@@ -1076,3 +1076,45 @@ def test_cached_repo_partial_scopes_probe_to_snapshot_dir(monkeypatch):
 
     monkeypatch.setattr(scan, "is_snapshot_partial", _boom)
     assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is False
+
+
+def test_list_cached_models_flags_single_file_diffusion_repos(monkeypatch, tmp_path):
+    # A diffusion-tagged repo with NO top-level model_index.json is a single-file
+    # checkpoint: the task pickers must not offer it as a pipeline load (from_pretrained
+    # fails on it), so the row carries single_file=True. A full pipeline repo (has
+    # model_index.json) and a chat repo (task None) carry no flag.
+    single = _repo(
+        "unsloth/Qwen-Image-fp8-single",
+        [_file("qwen-image-fp8.safetensors", 10_000)],
+        tmp_path / "models--unsloth--Qwen-Image-fp8-single",
+    )
+    pipeline = _repo(
+        "unsloth/Qwen-Image-pipeline",
+        [_file("model_index.json", 10), _file("transformer/model.safetensors", 10_000)],
+        tmp_path / "models--unsloth--Qwen-Image-pipeline",
+    )
+    chat = _repo(
+        "Org/ChatRepo",
+        [_file("model.safetensors", 10_000)],
+        tmp_path / "models--Org--ChatRepo",
+    )
+
+    monkeypatch.setattr(
+        models_route,
+        "_cached_repo_task",
+        lambda repo_info: (
+            "text-to-image" if "Qwen-Image" in repo_info.repo_id else None
+        ),
+    )
+    monkeypatch.setattr(
+        models_route,
+        "_all_hf_cache_scans",
+        lambda: [SimpleNamespace(repos = [single, pipeline, chat])],
+    )
+
+    result = asyncio.run(models_route.list_cached_models(current_subject = "test-user"))
+
+    rows = {r["repo_id"]: r for r in result["cached"]}
+    assert rows["unsloth/Qwen-Image-fp8-single"].get("single_file") is True
+    assert "single_file" not in rows["unsloth/Qwen-Image-pipeline"]
+    assert "single_file" not in rows["Org/ChatRepo"]

studio/backend/tests/test_middleware.py

@@ -204,6 +204,54 @@ class TestMaxBodyMiddleware:
                 default_request_body_limit_bytes()
             ), path
 
+    def test_diffusion_dataset_trailing_slash_gets_upload_cap(self, main_module):
+        # The trailing-slash variant reaches the middleware BEFORE the router's
+        # redirect_slashes 307, so it must resolve to the same passthrough + upload cap
+        # as the canonical path or a large upload 413s on the default /api/train cap.
+        # JSON sub-routes keep extra components after normalization, so they stay capped.
+        from utils.upload_limits import (
+            default_request_body_limit_bytes,
+            upload_request_limit_bytes,
+        )
+
+        slashed = "/api/train/diffusion/dataset/"
+        assert main_module._get_upload_passthrough_request_max_bytes(slashed) == (
+            upload_request_limit_bytes()
+        )
+        # End-to-end through the middleware: a body over the default cap but under the
+        # upload cap passes through on both the canonical and the slashed path.
+        app = _make_protected_app(
+            128,
+            main_module,
+            upload_passthrough_max_bytes_getter = lambda _p: 1024,
+            upload_passthrough_exact_paths = ("/api/train/diffusion/dataset",),
+        )
+
+        @app.post("/api/train/diffusion/dataset")
+        async def upload(request: Request):
+            body = await request.body()
+            return {"total": len(body)}
+
+        c = TestClient(app)
+        for path in ("/api/train/diffusion/dataset", "/api/train/diffusion/dataset/"):
+            r = c.post(
+                path,
+                content = b"x" * 512,
+                headers = {"content-type": "application/octet-stream"},
+            )
+            assert r.status_code == 200, path
+            assert r.json()["total"] == 512, path
+        # A slashed JSON sub-route is still NOT passthrough: over-cap body is rejected.
+        r = c.post(
+            "/api/train/diffusion/dataset/import-example/",
+            content = b"x" * 512,
+            headers = {"content-type": "application/octet-stream"},
+        )
+        assert r.status_code == 413
+        assert main_module._get_upload_passthrough_request_max_bytes(
+            "/api/train/diffusion/dataset/import-example/"
+        ) == default_request_body_limit_bytes()
+
     def test_v1_surface_is_body_protected(self, main_module):
         # /images/generations is mounted at both /api/inference and /v1; the /v1 alias (and every
         # other /v1 POST route) must be body-capped via the /v1 blanket prefix, or an unbounded

studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx

@@ -2273,8 +2273,12 @@ export function HubModelPicker({
             // Gate on a curated ARTIFACT (artifactForRepoId, what loadSpecFor resolves), not a
             // group-key match: a base / uncurated-quant sibling (Qwen/Qwen-Image-2512) matches
             // the group by key but has no loadable artifact and dead-ends at the trust gate.
+            // An unsloth repo must also be a full pipeline (not single_file): the selection
+            // fall-through loads uncataloged rows as kind "pipeline", and from_pretrained on
+            // a single-file checkpoint repo (no model_index.json) fails after the handoff.
+            // Curated single-file artifacts stay: loadSpecFor carries their filename.
             (!task ||
-              isUnslothRepoId(c.repo_id) ||
+              (isUnslothRepoId(c.repo_id) && !c.single_file) ||
               (catalog ? artifactForRepoId(c.repo_id, catalog) !== null : false)),
         ),
         downloadedSort,

studio/frontend/src/features/chat/api/chat-api.ts

@@ -326,6 +326,10 @@ export interface CachedModelRepo {
   /** True when the snapshot is incomplete (a cancelled/partial download). Such a
    * repo must not count as downloaded, or a click re-downloads the full weights. */
   partial?: boolean;
+  /** True for a diffusion repo with no model_index.json: a single-file checkpoint that
+   * loads only via from_single_file + a checkpoint filename. Task pickers must not offer
+   * it as a pipeline load unless the curated catalog carries its artifact. */
+  single_file?: boolean;
 }
 
 export async function listCachedModels(

studio/frontend/src/features/images/images-page.tsx

@@ -1903,6 +1903,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         res.images.forEach((image) => void ensureSrc(image));
         setGenDone(i + 1);
       }
+      // A generation can change server-side status: Speed=Auto compiles the
+      // transformer on the 3rd LoRA-free run (supports_lora flips to false), so
+      // without a refresh the LoRA picker stays enabled and the next LoRA run
+      // fails on the backend. Cheap status GET; also picks up any other drift.
+      if (isMounted.current) void refreshStatus();
     } catch (err) {
       toast.error(err instanceof Error ? err.message : "Image generation failed");
     } finally {
@@ -1912,7 +1917,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
       setGenDone(null);
       setGenStep(null);
     }
-  }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, controlnetCapable, controlnetId, controlImage, controlType, controlStrength, ensureSrc]);
+  }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, controlnetCapable, controlnetId, controlImage, controlType, controlStrength, ensureSrc, refreshStatus]);
 
   // Keep the active workflow valid for the loaded model: an edit-only model (Qwen-Image-
   // Edit) has no Create/Transform tabs, a base model has no Edit tab. Snap to the first