diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 09b2fa1b42..2d69c9bb56 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -439,11 +439,17 @@ class DiffusionBackend: base: str, base_files: list[str], hf_token: Optional[str], - ) -> None: + ) -> Optional[str]: """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")``.""" + be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``. + + Returns the base repo's local snapshot dir when the prefetched set includes + the pipeline manifest, so from_pretrained can load from disk instead of + re-sweeping the hub (its own sweep also pulls files the scoped list skips, + e.g. the 24 GB packaged root singles in each FLUX.1 repo); None otherwise + (estimate failure, config-only base, local repo) -> hub id as before.""" from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback # GGUF transformer (hub repos only; a local path is already on disk). @@ -452,12 +458,16 @@ class DiffusionBackend: repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event ) # Base repo (VAE / text-encoder / scheduler); list comes from the estimate. + snapshot_root: Optional[str] = None for rfilename in base_files: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") - hf_hub_download_with_xet_fallback( + local = hf_hub_download_with_xet_fallback( base, rfilename, hf_token, cancel_event = self._cancel_event ) + if rfilename == "model_index.json": + snapshot_root = str(Path(local).parent) + return snapshot_root def validate_load_request( self, @@ -674,7 +684,7 @@ class DiffusionBackend: 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( + kwargs["_base_local_dir"] = self._prefetch_files( kwargs["repo_id"], kwargs.get("gguf_filename"), base, @@ -887,6 +897,7 @@ class DiffusionBackend: transformer_cache_threshold: Optional[float] = None, model_kind: Optional[str] = None, _load_token: Optional[int] = None, + _base_local_dir: Optional[str] = None, ) -> dict[str, Any]: # A blank / whitespace-only token must degrade to anonymous access, not be passed # as an explicit credential (from_single_file / from_pretrained / the Hub client @@ -1003,6 +1014,7 @@ class DiffusionBackend: transformer_quant, transformer_quant_fast_accum, fam = fam, + base_local_dir = _base_local_dir, prequant_path = transformer_prequant_path, ) except Exception as exc: # noqa: BLE001 — fall back to the GGUF build @@ -1033,7 +1045,12 @@ class DiffusionBackend: pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} if hf_token: pipe_kwargs["token"] = hf_token - pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs) + # The prefetched snapshot dir keeps from_pretrained off the hub: + # its own snapshot sweep re-downloads files the scoped prefetch + # skipped (root packaged singles, e.g. 24 GB per FLUX.1 repo). + pipe = pipeline_cls.from_pretrained( + _base_local_dir or repo_id, **pipe_kwargs + ) elif kind == "single_file" and fam.single_file_is_pipeline: # A single-file SDXL-style checkpoint is the WHOLE pipeline # (U-Net + VAE + both text encoders), not a transformer-only file, @@ -1069,7 +1086,9 @@ class DiffusionBackend: pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer} if hf_token: pipe_kwargs["token"] = hf_token - pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + pipe = pipeline_cls.from_pretrained( + _base_local_dir or base, **pipe_kwargs + ) # Resolve the effective speed mode: GGUF models default to the # near-lossless `default` profile (compile is ~2.2x and sits below @@ -1301,6 +1320,7 @@ class DiffusionBackend: *, fam: Optional[DiffusionFamily] = None, prequant_path: Optional[str] = None, + base_local_dir: Optional[str] = None, ) -> tuple[Any, str]: """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``. @@ -1348,7 +1368,7 @@ class DiffusionBackend: ) if transformer is not None: pipe = self._assemble_pipe( - pipeline_cls, base, transformer, dtype, hf_token, device + pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir ) return pipe, scheme @@ -1356,7 +1376,9 @@ class DiffusionBackend: transformer = transformer_cls.from_pretrained( base, subfolder = "transformer", torch_dtype = dtype, token = hf_token ) - pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device) + pipe = self._assemble_pipe( + pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir + ) scheme = quantize_transformer( pipe, target, @@ -1377,13 +1399,14 @@ class DiffusionBackend: dtype: Any, hf_token: Optional[str], device: str, + base_local_dir: Optional[str] = None, ) -> Any: """Assemble the diffusers pipeline around ``transformer`` and place it on ``device`` (a no-op for an already-placed pre-quantized transformer; it moves the companions).""" pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} if hf_token: pipe_kwargs["token"] = hf_token - pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs) pipe.to(device) return pipe diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 6baa0a8a0a..1dccc2d0ac 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2121,3 +2121,35 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path): backend.generate(prompt = "a sloth") backend.generate(prompt = "another sloth") assert resets == [True, True] + + +def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch): + # The prefetched pipeline manifest's directory is the local snapshot root; a + # config-only base list (no manifest) returns None so the hub id stays in use. + backend = DiffusionBackend() + monkeypatch.setattr( + "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback", + lambda repo, fn, tok, **k: f"/cache/snap/{fn}", + ) + root = backend._prefetch_files( + "base/repo", None, "base/repo", ["model_index.json", "vae/x.safetensors"], None + ) + assert root == "/cache/snap" + assert ( + backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) + is None + ) + + +def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path): + # With a prefetched snapshot, from_pretrained must receive the local dir -- + # its own hub sweep would re-download the root packaged singles the scoped + # prefetch skips (24 GB per FLUX.1 repo). + backend = DiffusionBackend() + backend.load_pipeline( + "unsloth/Qwen-Image-2512-bnb-4bit", + model_kind = "pipeline", + _base_local_dir = str(tmp_path), + ) + assert _FakePipeline.last["base"] == str(tmp_path) + backend.unload()