From c249d0c50a70ecebbed2652d281da0ca95e3abce Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 10 Jul 2026 08:06:20 +0000 Subject: [PATCH] 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 | 12 ++++- studio/backend/routes/models.py | 33 ++++++++++--- .../backend/tests/test_cached_gguf_routes.py | 42 ++++++++++++++++ studio/backend/tests/test_middleware.py | 48 +++++++++++++++++++ .../assistant-ui/model-selector/pickers.tsx | 6 ++- .../src/features/chat/api/chat-api.ts | 4 ++ .../src/features/images/images-page.tsx | 7 ++- 7 files changed, 141 insertions(+), 11 deletions(-) diff --git a/studio/backend/main.py b/studio/backend/main.py index 5a65b62e57..92d9753753 100644 --- a/studio/backend/main.py +++ b/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 ) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1cee0fd548..a3cb144e5a 100644 --- a/studio/backend/routes/models.py +++ b/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)) diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 674aabe231..dd72d68c37 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/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"] diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index b2fd1c3613..9372a93b5f 100644 --- a/studio/backend/tests/test_middleware.py +++ b/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 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 8ae2111d1e..3949c42498 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/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, diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index acf3044077..c12ee92c44 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/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( diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index c22766392d..6dd755869d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/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