From adebdfc9f4f3fe538db98d53f92f463aa78e7ae9 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 24 Jun 2026 17:17:28 -0300 Subject: [PATCH] Fix image load progress, hide diffusion models from chat picker, and gallery polish --- studio/backend/core/inference/diffusion.py | 29 +++- studio/backend/routes/models.py | 15 ++ .../backend/tests/test_cached_gguf_routes.py | 39 ++++- .../backend/tests/test_diffusion_backend.py | 31 +++- .../assistant-ui/model-selector/pickers.tsx | 47 ++++-- .../src/features/chat/api/chat-api.ts | 3 + .../src/features/images/images-page.tsx | 137 ++++++++++++------ 7 files changed, 237 insertions(+), 64 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index cf5f153e85..0f92e3bba3 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -177,10 +177,12 @@ class DiffusionBackend: downloaded = self._cache_bytes(loading.repo_id) + self._cache_bytes(loading.base_repo) expected = loading.expected_bytes - # Downloads done but pipeline still dequantising / moving to GPU. + # Downloads done but pipeline still dequantising / moving to GPU. The cache + # scan can slightly exceed the estimate (extra cached quants, blob padding), + # so clamp the reported bytes/fraction so the bar never overshoots 100%. if expected > 0 and downloaded >= expected * 0.999: - return _progress("finalizing", downloaded, expected, 1.0) - fraction = downloaded / expected if expected > 0 else 0.0 + return _progress("finalizing", min(downloaded, expected), expected, 1.0) + fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0 return _progress("downloading", downloaded, expected, fraction) @staticmethod @@ -196,12 +198,8 @@ class DiffusionBackend: info = api.model_info(repo_id, files_metadata = True, token = hf_token) total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename) base_info = api.model_info(base_repo, files_metadata = True, token = hf_token) - # The transformer subfolder is supplied by the GGUF, so from_pretrained - # never downloads it — exclude it from the expected total. total += sum( - s.size or 0 - for s in base_info.siblings - if not s.rfilename.startswith("transformer/") + s.size or 0 for s in base_info.siblings if _base_file_downloaded(s.rfilename) ) except Exception as exc: # noqa: BLE001 — estimate is best-effort logger.warning("diffusion.size_estimate_failed: %s", exc) @@ -373,6 +371,21 @@ class DiffusionBackend: } +def _base_file_downloaded(rfilename: str) -> bool: + """True for base-repo files ``from_pretrained`` actually fetches. + + The transformer is supplied by the GGUF, and repo docs (``assets/``, the + top-level README/PDF/images) are never downloaded — counting them would peg + the progress estimate above what lands on disk, so the bar would sit short of + 100% for the whole pipeline-load phase instead of advancing to "finalizing". + """ + if rfilename.startswith("transformer/"): + return False + if "/" not in rfilename: # top-level: only the pipeline manifest is fetched + return rfilename == "model_index.json" + return not rfilename.startswith("assets/") + + def _progress( phase: Optional[str], bytes_downloaded: int = 0, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 92b8824016..7aa1a3a580 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3215,6 +3215,20 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): return {"cached": []} +def _repo_is_diffusers(repo_info) -> bool: + """A diffusers pipeline repo (e.g. a Z-Image / FLUX base) carries a top-level + model_index.json. These render images rather than chat, so the chat picker + hides them — mirroring how cached diffusion GGUFs are classified by arch.""" + 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 + + @router.get("/cached-models") async def list_cached_models(current_subject: str = Depends(get_current_subject)): """List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" @@ -3261,6 +3275,7 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) row = { "repo_id": repo_id, "size_bytes": total_size, + "task": "text-to-image" if _repo_is_diffusers(repo_info) else None, } # Keep the newest timestamp across duplicate caches; # attach only when known so absent rows sort as oldest. diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index b2ead305ba..046eb2543d 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -83,6 +83,7 @@ def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monk "size_bytes": 5_000, "cache_path": str(repo.repo_path), "has_vision": False, + "task": None, } ] @@ -105,6 +106,7 @@ def test_list_cached_gguf_matches_extension_case_insensitively(monkeypatch, tmp_ "size_bytes": 7_000, "cache_path": str(repo.repo_path), "has_vision": False, + "task": None, } ] @@ -200,6 +202,7 @@ def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans(monkeypatch, "size_bytes": 6_000, "cache_path": str(larger.repo_path), "has_vision": False, + "task": None, } ] @@ -230,6 +233,7 @@ def test_list_cached_gguf_dedupes_shared_blobs_across_revisions(monkeypatch, tmp "size_bytes": 5_000, "cache_path": str(repo.repo_path), "has_vision": False, + "task": None, } ] @@ -280,6 +284,7 @@ def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypa "size_bytes": 5_000, "cache_path": str(mixed.repo_path), "has_vision": False, + "task": None, } ] @@ -307,6 +312,7 @@ def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path): "size_bytes": 5_000, "cache_path": str(partial.repo_path), "has_vision": False, + "task": None, } ] @@ -343,6 +349,7 @@ def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(monkeypat "size_bytes": 5_000, "cache_path": str(healthy.repo_path), "has_vision": False, + "task": None, } ] @@ -390,7 +397,35 @@ def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) - assert result["cached"] == [{"repo_id": "Org/MmprojAux", "size_bytes": 15_000}] + assert result["cached"] == [{"repo_id": "Org/MmprojAux", "size_bytes": 15_000, "task": None}] + + +def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch, tmp_path): + """A cached diffusers pipeline repo (model_index.json present) is tagged + text-to-image so the chat picker hides it, while a plain checkpoint isn't.""" + diffusion = _repo( + "Tongyi-MAI/Z-Image-Turbo", + [_file("model_index.json", 1_000), _file("text_encoder/model.safetensors", 9_000)], + tmp_path / "models--Tongyi-MAI--Z-Image-Turbo", + ) + checkpoint = _repo( + "unsloth/Llama-3.2-1B-Instruct", + [_file("config.json", 1_000), _file("model.safetensors", 9_000)], + tmp_path / "models--unsloth--Llama-3.2-1B-Instruct", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [diffusion, checkpoint])], + ) + + result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) + by_repo = {c["repo_id"]: c["task"] for c in result["cached"]} + assert by_repo == { + "Tongyi-MAI/Z-Image-Turbo": "text-to-image", + "unsloth/Llama-3.2-1B-Instruct": None, + } def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeypatch, tmp_path): @@ -419,6 +454,7 @@ def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeyp "size_bytes": 5_000, "cache_path": str(vision_repo.repo_path), "has_vision": True, + "task": None, } ] @@ -473,6 +509,7 @@ def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_pat "size_bytes": 5_000, "cache_path": str(tmp_path / "active"), "has_vision": False, + "task": None, } ] diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 5b1fb5a342..289fa367b7 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -15,7 +15,11 @@ import types import pytest -from core.inference.diffusion import DiffusionBackend, encode_png_base64 +from core.inference.diffusion import ( + DiffusionBackend, + _base_file_downloaded, + encode_png_base64, +) from core.inference.diffusion_families import ( detect_family, resolve_base_repo, @@ -270,6 +274,31 @@ def test_load_progress_downloading_then_finalizing(monkeypatch): assert backend.load_progress()["phase"] == "finalizing" # 1000/1000 +def test_base_file_downloaded_excludes_undownloaded(): + # Counted: the pipeline manifest + component subfolders from_pretrained fetches. + assert _base_file_downloaded("model_index.json") + assert _base_file_downloaded("text_encoder/model-00001-of-00003.safetensors") + assert _base_file_downloaded("vae/diffusion_pytorch_model.safetensors") + # Excluded: the GGUF supplies the transformer; docs/assets and top-level files + # are never downloaded, so counting them would peg the bar short of 100%. + assert not _base_file_downloaded("transformer/diffusion_pytorch_model-00001-of-00003.safetensors") + assert not _base_file_downloaded("assets/Z-Image-Gallery.pdf") + assert not _base_file_downloaded("README.md") + assert not _base_file_downloaded(".gitattributes") + + +def test_load_progress_fraction_clamped(monkeypatch): + # The cache scan can exceed the estimate (e.g. a second cached quant); the + # reported fraction must still clamp to 1.0 rather than overshoot. + backend = DiffusionBackend() + backend._loading = _LoadingState(repo_id = "r", base_repo = "b", expected_bytes = 1000) + monkeypatch.setattr(DiffusionBackend, "_cache_bytes", staticmethod(lambda repo: 900)) + p = backend.load_progress() # summed 1800 > expected 1000 + assert p["phase"] == "finalizing" + assert p["fraction"] == 1.0 + assert p["bytes_downloaded"] == 1000 # clamped to the estimate + + def test_begin_load_rejects_concurrent(monkeypatch): backend = DiffusionBackend() monkeypatch.setattr(DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0)) 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 d995887ca9..1e3b7f257f 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -975,6 +975,22 @@ function taskMatchesFilter(repoTask: string | null | undefined, filter: HfTaskFi return repoTask != null && (wanted as readonly string[]).includes(repoTask); } +// Image-generation pipeline tasks: handled by the Images page, never loadable as +// chat models. The backend reports "text-to-image" for diffusion-arch GGUFs. +const IMAGE_GEN_TASKS: readonly string[] = [ + "text-to-image", + "image-to-image", + "image-text-to-image", +]; + +// Gate an on-device model by the picker's task scope. With a filter (the Images +// page) keep only matching tasks; with no filter (chat) drop image-generation +// models so a downloaded diffusion GGUF doesn't show up as a loadable chat model. +function passesTaskGate(repoTask: string | null | undefined, filter: HfTaskFilter): boolean { + if (filter) return taskMatchesFilter(repoTask, filter); + return !(repoTask != null && IMAGE_GEN_TASKS.includes(repoTask)); +} + // Module-level caches so re-mounting the popover shows results instantly let _cachedGgufCache: CachedGgufRepo[] = []; let _cachedModelsCache: CachedModelRepo[] = []; @@ -1720,22 +1736,30 @@ export function HubModelPicker({ return map; }, [results, recommendedSearch.results]); - // Ordered by the On Device dropdown (recent/download date/size/name). When a - // task filter is set (e.g. Images page), drop cached repos whose architecture - // is not that task so On Device matches the Recommended tab. + // Ordered by the On Device dropdown (recent/download date/size/name). The task + // gate keeps the Images picker to diffusion GGUFs and, conversely, hides those + // diffusion GGUFs from the chat picker (where they aren't loadable models). const sortedCachedGguf = useMemo( () => sortCachedRepos( - task ? cachedGguf.filter((c) => taskMatchesFilter(c.task, task)) : cachedGguf, + cachedGguf.filter((c) => passesTaskGate(c.task, task)), downloadedSort, loadTimes, ), [cachedGguf, downloadedSort, loadTimes, task], ); // Non-GGUF (safetensors) cached repos aren't single-file diffusion GGUFs, so - // hide them entirely when a task filter is active (the Images picker). + // hide them entirely when a task filter is active (the Images picker). In chat, + // drop cached diffusers pipeline repos (a Z-Image / FLUX base) the same way. const sortedCachedModels = useMemo( - () => (task ? [] : sortCachedRepos(cachedModels, downloadedSort, loadTimes)), + () => + task + ? [] + : sortCachedRepos( + cachedModels.filter((c) => passesTaskGate(c.task, task)), + downloadedSort, + loadTimes, + ), [cachedModels, downloadedSort, loadTimes, task], ); // Each local section's search is scoped to its own models (matched by name). @@ -1745,14 +1769,15 @@ export function HubModelPicker({ normalizeForSearch( `${m.model_id ?? ""} ${m.display_name} ${m.id}`, ).includes(localQuery); - // A task filter (the Images page) wants diffusion GGUFs only, so local models - // are filtered to that task (by the GGUF architecture the backend reports). + // The Images page wants diffusion GGUFs only; chat wants everything but those. + // passesTaskGate handles both directions off the GGUF architecture the backend + // reports as each model's task. const sortedLmStudio = useMemo( () => sortLocalModels( lmStudioModels.filter( (m) => - (!task || taskMatchesFilter(m.task, task)) && + passesTaskGate(m.task, task) && localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), ), @@ -1770,7 +1795,7 @@ export function HubModelPicker({ sortLocalModels( localDirModels.filter( (m) => - (!task || taskMatchesFilter(m.task, task)) && + passesTaskGate(m.task, task) && (!chatOnly || localModelIsGguf(m) || (isMac && localModelIsMlx(m))) && @@ -1797,7 +1822,7 @@ export function HubModelPicker({ sortLocalModels( customFolderModels.filter( (m) => - (!task || taskMatchesFilter(m.task, task)) && + passesTaskGate(m.task, task) && localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), ), diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index ee7ca1950c..a5e82d62c5 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -319,6 +319,9 @@ export interface CachedModelRepo { /** Epoch seconds of the newest downloaded weight file; sorts Downloaded * newest-first. Optional for older-backend compatibility. */ last_modified?: number; + /** HF pipeline task: "text-to-image" for a cached diffusers pipeline repo + * (model_index.json present), so the chat picker can hide it. Absent = chat. */ + task?: string | null; } export async function listCachedModels(): Promise { diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 24cffd87e2..595f997fce 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -89,28 +89,64 @@ function downloadImage(src: string, seed: number) { link.click(); } -// Title + percent/label for the chat progress component, from a progress poll. -function progressView(p: DiffusionLoadProgress): { - title: string; - progressPercent: number | null; - progressLabel: string | null; -} { - const title = p.phase === "finalizing" ? "Loading to GPU…" : "Downloading model…"; - if (p.bytes_total > 0) { - return { - title, - progressPercent: p.fraction * 100, - progressLabel: `${formatBytes(p.bytes_downloaded)} of ${formatBytes(p.bytes_total)}`, - }; - } - // Unknown total: show bytes counting up, no bar. +// The chat tab's model-load toast styling, reused verbatim so the diffusion +// load toast is visually identical (persistent, progress bar, same chrome). +const LOAD_TOAST_CLASSNAMES = { + toast: "chat-model-load-toast items-center gap-2.5", + content: "gap-0.5 flex-1 min-w-0", + title: "leading-5", + description: "mt-0 w-full", +} as const; + +// Render the chat ModelLoadDescription for a progress poll. The base text- +// encoder/VAE repo downloads alongside the GGUF, so the total can exceed the +// picked quant's size — that's the real one-time download. +function loadToastDescription(p: DiffusionLoadProgress) { + // "Downloading" only when bytes actually remain to fetch — a cached model (or + // the pre-estimate window, total still 0) shouldn't claim a download. + const downloading = p.bytes_total > 0 && p.bytes_downloaded < p.bytes_total * 0.999; + const title = downloading + ? "Downloading model…" + : p.phase === "finalizing" + ? "Loading to GPU…" + : "Starting model…"; + const hasTotal = p.bytes_total > 0; + return ( + 0 + ? `${formatBytes(p.bytes_downloaded)} downloaded` + : null + } + /> + ); +} + +// Toast args mirroring chat's: persistent, closeable, content in `description`. +// Pass `id` to update the existing toast in place instead of stacking a new one. +function loadToastArgs(p: DiffusionLoadProgress, id?: string | number) { return { - title, - progressPercent: null, - progressLabel: p.bytes_downloaded > 0 ? `${formatBytes(p.bytes_downloaded)} downloaded` : null, + ...(id != null ? { id } : {}), + description: loadToastDescription(p), + duration: Infinity, + closeButton: true, + classNames: LOAD_TOAST_CLASSNAMES, }; } +const IDLE_PROGRESS: DiffusionLoadProgress = { + phase: null, + bytes_downloaded: 0, + bytes_total: 0, + fraction: 0, + error: null, +}; + // Mirrors the Train page's SliderRow (studio/sections/params-section.tsx): // label + standard Slider + number input, same classes. function SliderField({ @@ -179,13 +215,19 @@ export function ImagesPage() { const [busy, setBusy] = useState(null); const [status, setStatus] = useState(null); - const [loadProgress, setLoadProgress] = useState(null); const [results, setResults] = useState(() => imageSession.results); const [selectedId, setSelectedId] = useState(() => imageSession.selectedId); // Stable, ever-increasing ids: results are prepended, so an array index would // re-key every existing image on each new generation. const nextResultId = useRef(imageSession.nextId); const pollTimer = useRef | null>(null); + // The persistent load toast's id, so each poll updates it in place (chat-style). + const loadToastId = useRef(null); + + const dismissLoadToast = useCallback(() => { + if (loadToastId.current != null) toast.dismiss(loadToastId.current); + loadToastId.current = null; + }, []); // Persist the gallery to module scope so a tab switch doesn't drop it. useEffect(() => { @@ -217,36 +259,39 @@ export function ImagesPage() { }; }, [refreshStatus]); - // Poll load-progress until the background load reaches "ready" or "error". + // Poll load-progress until the background load reaches "ready" or "error", + // updating the persistent toast in place each tick. const pollLoadProgress = useCallback(async () => { try { const p = await getDiffusionLoadProgress(); - setLoadProgress(p); if (p.phase === "ready") { + dismissLoadToast(); setStatus(await getDiffusionStatus()); toast.success("Model loaded"); setBusy(null); - setLoadProgress(null); return; } if (p.phase === "error") { + dismissLoadToast(); toast.error(p.error || "Failed to load model"); setBusy(null); - setLoadProgress(null); return; } + if (loadToastId.current != null) toast(null, loadToastArgs(p, loadToastId.current)); } catch { // Transient poll failure: keep trying. } pollTimer.current = setTimeout(() => void pollLoadProgress(), 1000); - }, []); + }, [dismissLoadToast]); const handleLoad = useCallback( async (ggufFilename: string) => { // Cancel any prior poll loop so two can't run at once. if (pollTimer.current) clearTimeout(pollTimer.current); setBusy("loading"); - setLoadProgress(null); + // Show the chat-style toast immediately; the poll updates it by id. + dismissLoadToast(); + loadToastId.current = toast(null, loadToastArgs(IDLE_PROGRESS)); try { // Returns immediately — the load runs in the background; we poll for it. await loadDiffusionModel({ @@ -255,6 +300,7 @@ export function ImagesPage() { family_override: MODEL.family, }); } catch (err) { + dismissLoadToast(); toast.error(err instanceof Error ? err.message : "Failed to start load"); setBusy(null); void refreshStatus(); @@ -262,7 +308,7 @@ export function ImagesPage() { } void pollLoadProgress(); }, - [pollLoadProgress, refreshStatus], + [pollLoadProgress, refreshStatus, dismissLoadToast], ); // The chat picker emits (modelId, picked quant); load its matching GGUF. @@ -346,7 +392,8 @@ export function ImagesPage() { return (
- {/* ── Top: the chat-style model selector (header padding mirrors chat) ── */} + {/* ── Top: the chat-style model selector (header padding mirrors chat). The + load progress shows in a chat-style toast, not here. ── */}
+ ) : busy === "generating" ? ( + // First image (no past result to keep showing): spin in place. +
+ +

Generating…

+
) : (
@@ -452,33 +505,31 @@ export function ImagesPage() {

)} - {busy === "generating" && ( -
- -
- )} - {/* Download progress, lower-right (chat-style). */} - {busy === "loading" && loadProgress && ( -
- -
- )}
- {results.length > 0 && ( + {(results.length > 0 || busy === "generating") && (
+ {/* In-progress generation: a placeholder tile at the front so past + images stay visible and browsable while the new one renders. */} + {busy === "generating" && ( +
+ +
+ )} {results.map((r) => ( ))}