ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Load image pipelines from the prefetched snapshot instead of re-sweeping the hub

Browse source 
The prefetch already scopes the file list (no packaged root singles, no
dtype-variant twins, no ONNX/Flax exports), but from_pretrained was then
called with the hub id, and its own snapshot sweep re-downloaded the
skipped files anyway: 24 GB per FLUX.1 repo and 65 GB on FLUX.2-dev, as
found in the blob cache. Return the snapshot dir from the prefetch (keyed
on the pipeline manifest) and hand it to every pipeline-assembly
from_pretrained site; any prefetch failure keeps the hub id and the old
behavior.
This commit is contained in:
Daniel Han 2026-07-05 04:41:43 +00:00
commit 62cae5fe71
2 changed files with 64 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

41
studio/backend/core/inference/diffusion.py
View file

@ -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

32
studio/backend/tests/test_diffusion_backend.py
View file

@ -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()