diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 5ca18831e2..c42bc4ccac 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -94,6 +94,10 @@ class DiffusionBackend: # superseded (a new load) or cancelled (unload, incl. an arbiter eviction) # neither commits its pipeline nor stamps progress onto the current load. self._load_token = 0 + # Set by unload() to abort an in-flight download (which runs without the + # lock, like the chat backend), so an eviction/unload can preempt a slow + # load instead of blocking on the lock for the whole download. + self._cancel_event = threading.Event() # The callback mutates this and generate_progress() reads it, both without # the lock (generate holds it for the whole call), so polling stays live. self._gen: Optional[_GenState] = None @@ -106,7 +110,12 @@ class DiffusionBackend: import torch if torch.cuda.is_available(): - return "cuda", torch.bfloat16 + # BF16 needs Ampere+ (compute capability >= 8); pre-Ampere cards + # (Turing/Volta/Pascal) only emulate it, so use FP16 there. (Checked by + # capability, not torch.cuda.is_bf16_supported(), which returns True via + # emulation on those cards and would still pick BF16.) + dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16 + return "cuda", dtype mps = getattr(torch.backends, "mps", None) if mps is not None and mps.is_available(): return "mps", torch.float16 @@ -120,6 +129,24 @@ class DiffusionBackend: return hf_hub_download(repo_id, gguf_filename, token = hf_token) + def _prefetch_files( + self, repo_id: str, gguf_filename: Optional[str], base: str, base_files: list[str], hf_token: Optional[str] + ) -> None: + """Pre-download the GGUF + the given ``base_files`` into the HF cache, + WITHOUT the lock and honoring ``_cancel_event``, so load_pipeline's + from_single_file / from_pretrained hit the cache and the heavy download can + be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``.""" + from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback + + # GGUF transformer (hub repos only; a local path is already on disk). + if gguf_filename and not Path(repo_id).expanduser().exists(): + hf_hub_download_with_xet_fallback(repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event) + # Base repo (VAE / text-encoder / scheduler); list comes from the estimate. + for rfilename in base_files: + if self._cancel_event.is_set(): + raise RuntimeError("Cancelled") + hf_hub_download_with_xet_fallback(base, rfilename, hf_token, cancel_event = self._cancel_event) + # ── Background load + progress ───────────────────────────────────────── def begin_load( @@ -149,6 +176,9 @@ class DiffusionBackend: raise RuntimeError("A diffusion load is already in progress.") self._load_token += 1 token = self._load_token + # Best-effort download preemption only; the token (not this event) is + # the real guard that a superseded worker can't commit its pipeline. + self._cancel_event.clear() # Seed with the family fallback; the worker resolves the real base # (a network lookup) and updates this, so begin_load never blocks. self._loading = _LoadingState(repo_id = repo_id, base_repo = fam.base_repo) @@ -179,15 +209,18 @@ class DiffusionBackend: kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token") ) kwargs["base_repo"] = base + 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 = self._estimate_download_bytes( - kwargs["repo_id"], - kwargs.get("gguf_filename"), - base, - kwargs.get("hf_token"), - ) + 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( + kwargs["repo_id"], kwargs.get("gguf_filename"), base, base_files, kwargs.get("hf_token") + ) self.load_pipeline(**kwargs) with self._lock: # Only clear the marker if this load is still the current one; a @@ -225,22 +258,26 @@ class DiffusionBackend: @staticmethod def _estimate_download_bytes( repo_id: str, gguf_filename: Optional[str], base_repo: str, hf_token: Optional[str] - ) -> int: + ) -> tuple[int, list[str]]: + """Total download size for the progress bar, plus the base-repo files to + fetch (the prefetch reuses this list, so the base is listed only once).""" from huggingface_hub import HfApi api = HfApi() total = 0 + base_files: list[str] = [] try: if gguf_filename: 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) - total += sum( - s.size or 0 for s in base_info.siblings if _base_file_downloaded(s.rfilename) - ) + for s in base_info.siblings: + if _base_file_downloaded(s.rfilename): + base_files.append(s.rfilename) + total += s.size or 0 except Exception as exc: # noqa: BLE001 — estimate is best-effort logger.warning("diffusion.size_estimate_failed: %s", exc) - return total + return total, base_files @staticmethod def _cache_bytes(repo_id: str) -> int: @@ -306,6 +343,9 @@ class DiffusionBackend: torch_dtype = dtype, config = base, subfolder = "transformer", + # Forward the token: the config is fetched from the (possibly gated) + # base repo before from_pretrained gets a chance to authenticate. + token = hf_token, ) pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} @@ -427,6 +467,9 @@ class DiffusionBackend: } def unload(self) -> dict[str, Any]: + # Abort an in-flight download (it runs without the lock and checks this), + # so unload/an eviction returns promptly instead of waiting it out. + self._cancel_event.set() with self._lock: self._unload_locked() # Cancel any in-flight load (its worker checks this token before diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 621ea6a899..65bc0602ac 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -216,6 +216,7 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): gguf_filename = "model.gguf", base_repo = "base/repo", family_override = "z-image", + hf_token = "hf_secret", ) assert status["loaded"] is True assert status["family"] == "z-image" @@ -226,6 +227,8 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): # Transformer built from the local GGUF, pipeline assembled from the base repo. assert _FakeTransformer.last["path"] == str((tmp_path / "model.gguf").resolve()) assert _FakeTransformer.last["subfolder"] == "transformer" + # The token reaches the (possibly gated) base config fetch and the pipeline. + assert _FakeTransformer.last["token"] == "hf_secret" assert _FakePipeline.last["base"] == "base/repo" assert "transformer" in _FakePipeline.last @@ -398,10 +401,12 @@ def test_generate_qwen_uses_true_cfg_scale(fake_runtime, tmp_path): def test_begin_load_rejects_concurrent(monkeypatch): backend = DiffusionBackend() - # The worker resolves the base via a network lookup; stub it so the test is offline. + # The worker resolves the base + downloads, both over the network; stub them + # so the test is offline. monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None) + monkeypatch.setattr(DiffusionBackend, "_prefetch_files", lambda self, *a, **k: None) monkeypatch.setattr( - DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0) + DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: (0, [])) ) # Block the spawned worker so the load stays "in progress". monkeypatch.setattr( @@ -429,3 +434,60 @@ def test_unload_cancels_in_flight_load(fake_runtime): base_repo = fam.base_repo, _load_token = token, ) + + +def test_pick_dtype_bf16_only_on_ampere(fake_runtime, monkeypatch): + # BF16 only on Ampere+ (cc >= 8); pre-Ampere cards must fall back to FP16. + torch = sys.modules["torch"] + backend = DiffusionBackend() + monkeypatch.setattr(torch.cuda, "is_available", lambda: True, raising = False) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 0), raising = False) + assert backend._pick_device_and_dtype() == ("cuda", torch.bfloat16) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5), raising = False) + assert backend._pick_device_and_dtype() == ("cuda", torch.float16) + + +def test_unload_sets_cancel_event(fake_runtime): + # unload signals an in-flight download (which runs without the lock) to abort. + backend = DiffusionBackend() + assert not backend._cancel_event.is_set() + backend.unload() + assert backend._cancel_event.is_set() + + +def test_prefetch_aborts_when_cancelled(tmp_path): + # A prefetch interrupted by unload (cancel event set) raises rather than + # downloading the whole base, so the load can be preempted mid-download. + backend = DiffusionBackend() + backend._cancel_event.set() + # Local gguf path so the transformer download is skipped; the base loop hits + # the cancel check on its first file (no network). + (tmp_path / "model.gguf").write_bytes(b"x") + with pytest.raises(RuntimeError, match = "Cancelled"): + backend._prefetch_files( + str(tmp_path), "model.gguf", "Tongyi-MAI/Z-Image-Turbo", + ["vae/diffusion_pytorch_model.safetensors"], None, + ) + + +def test_prefetch_downloads_gguf_and_base(monkeypatch, tmp_path): + backend = DiffusionBackend() + calls: list = [] + monkeypatch.setattr( + "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback", + lambda repo, fn, tok, **k: (calls.append((repo, fn)), f"/cache/{fn}")[1], + ) + # Hub repo: the GGUF transformer and each base file are fetched. + backend._prefetch_files( + "unsloth/Z-Image-Turbo-GGUF", "model.gguf", "base/repo", + ["vae/x.safetensors", "text_encoder/y.safetensors"], "hf_tok", + ) + assert ("unsloth/Z-Image-Turbo-GGUF", "model.gguf") in calls + assert ("base/repo", "vae/x.safetensors") in calls + assert ("base/repo", "text_encoder/y.safetensors") in calls + # Local GGUF path: the transformer download is skipped, base still fetched. + calls.clear() + (tmp_path / "model.gguf").write_bytes(b"x") + 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 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 b0fc7177c4..a49acda408 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -985,11 +985,28 @@ export const IMAGE_GEN_TASKS = [ "image-text-to-image", ] as const; +// Editing/inpaint checkpoints are tagged image-to-image but need an input image, +// which the text-to-image backend rejects (mirrors its _EDIT_KEYWORDS). Hidden by +// id so they don't show in the Images picker only to 400 on load. Keeping the +// image-to-image task itself is required: some supported models (FLUX.2-klein) +// carry that tag too. +const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint"] as const; +function isImageEditModel(repoId: string | null | undefined): boolean { + if (!repoId) return false; + const id = repoId.toLowerCase(); + return IMAGE_EDIT_KEYWORDS.some((kw) => id.includes(kw)); +} + // 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); +// page) keep only matching, non-editing 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, + repoId: string | null | undefined, + filter: HfTaskFilter, +): boolean { + if (filter) return taskMatchesFilter(repoTask, filter) && !isImageEditModel(repoId); return !(repoTask != null && (IMAGE_GEN_TASKS as readonly string[]).includes(repoTask)); } @@ -1563,7 +1580,9 @@ export function HubModelPicker({ // the chat classifier marks image tasks "unsupported". const isChatSupported = useCallback( (r: HfModelResult) => { - if (task && taskMatchesFilter(r.pipelineTag, task)) return true; + // 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); return ( classifyUnslothSupport({ modelId: r.id, @@ -1744,7 +1763,7 @@ export function HubModelPicker({ const sortedCachedGguf = useMemo( () => sortCachedRepos( - cachedGguf.filter((c) => passesTaskGate(c.task, task)), + cachedGguf.filter((c) => passesTaskGate(c.task, c.repo_id, task)), downloadedSort, loadTimes, ), @@ -1758,7 +1777,7 @@ export function HubModelPicker({ task ? [] : sortCachedRepos( - cachedModels.filter((c) => passesTaskGate(c.task, task)), + cachedModels.filter((c) => passesTaskGate(c.task, c.repo_id, task)), downloadedSort, loadTimes, ), @@ -1779,7 +1798,7 @@ export function HubModelPicker({ sortLocalModels( lmStudioModels.filter( (m) => - passesTaskGate(m.task, task) && + passesTaskGate(m.task, m.model_id ?? m.id, task) && localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), ), @@ -1797,7 +1816,7 @@ export function HubModelPicker({ sortLocalModels( localDirModels.filter( (m) => - passesTaskGate(m.task, task) && + passesTaskGate(m.task, m.model_id ?? m.id, task) && (!chatOnly || localModelIsGguf(m) || (isMac && localModelIsMlx(m))) && @@ -1824,7 +1843,7 @@ export function HubModelPicker({ sortLocalModels( customFolderModels.filter( (m) => - passesTaskGate(m.task, task) && + passesTaskGate(m.task, m.model_id ?? m.id, task) && localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), ), diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index dae4656677..c832e38e4f 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -37,6 +37,7 @@ import type { ModelSelectorChangeMeta, } from "@/components/assistant-ui/model-selector/types"; import { ModelLoadDescription } from "@/features/chat/components/model-load-status"; +import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store"; import { formatBytes, formatEta } from "@/features/hub/lib/format"; import { cn } from "@/lib/utils"; import { toast } from "@/lib/toast"; @@ -630,7 +631,12 @@ export function ImagesPage() { try { // Returns immediately — the load runs in the background; we poll for it. // The backend infers the family + base diffusers repo from the repo id. - await loadDiffusionModel({ model_path: repoId, gguf_filename: ggufFilename }); + // Forward the saved HF token so gated bases (FLUX dev/klein) can download. + await loadDiffusionModel({ + model_path: repoId, + gguf_filename: ggufFilename, + hf_token: hfApiToken(getHfToken()), + }); } catch (err) { dismissLoadToast(); toast.error(err instanceof Error ? err.message : "Failed to start load");