diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 117c20248c..5f4dfb9859 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -487,11 +487,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). @@ -500,12 +506,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, @@ -728,7 +738,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, @@ -941,6 +951,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 @@ -1122,6 +1133,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 @@ -1168,7 +1180,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 (packaged root singles, 24 GB per FLUX.1). + 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, @@ -1209,7 +1226,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 @@ -1534,6 +1553,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)``. @@ -1581,7 +1601,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 @@ -1589,7 +1609,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, @@ -1610,13 +1632,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/core/inference/video.py b/studio/backend/core/inference/video.py index f00ccdee6b..4bd7751f7c 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -310,15 +310,47 @@ class VideoBackend: # unload/eviction can preempt the multi-GB pull; the pipeline # companions pre-download the same way (scoped file list, cancellable, # resumes from the cache so a cancelled pull costs nothing). + checkpoint_local: Optional[Path] = None if kwargs.get("gguf_filename") and not Path(kwargs["repo_id"]).expanduser().exists(): from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback - hf_hub_download_with_xet_fallback( - kwargs["repo_id"], - kwargs["gguf_filename"], - kwargs.get("hf_token"), - cancel_event = self._cancel_event, + checkpoint_local = Path( + hf_hub_download_with_xet_fallback( + kwargs["repo_id"], + kwargs["gguf_filename"], + kwargs.get("hf_token"), + cancel_event = self._cancel_event, + ) ) - kwargs["_base_local_dir"] = self._predownload_base(base, kwargs.get("hf_token"), kind) + # An LTX-2.3 checkpoint replaces the base VAEs/vocoder/connectors too, so + # its base pull shrinks to scheduler + text encoder + tokenizer; the + # estimate is recomputed to match (detectable only once the checkpoint + # header is on disk, hence after the pull above). + ltx23 = False + if fam is not None and fam.name == "ltx-2" and kind != "pipeline": + from .video_ltx2 import is_ltx23_checkpoint + + probe = checkpoint_local + if probe is None: + root = Path(kwargs["repo_id"]).expanduser() + probe = root if root.is_file() else None + ltx23 = probe is not None and is_ltx23_checkpoint(probe) + if ltx23: + expected = self._estimate_download_bytes( + kwargs["repo_id"], + kwargs.get("gguf_filename"), + base, + kwargs.get("hf_token"), + kind, + ltx23 = True, + ) + with self._lock: + if self._load_token == token and self._loading is not None: + self._loading.expected_bytes = expected + base_local = self._predownload_base(base, kwargs.get("hf_token"), kind, ltx23 = ltx23) + # The 2.3 assembly pulls per component from the hub id (its snapshot here + # deliberately lacks the base VAEs), so it only gets the warmed cache; the + # generic from_pretrained paths get the complete local snapshot. + kwargs["_base_local_dir"] = None if ltx23 else base_local self.load_pipeline(**kwargs) with self._lock: if self._load_token == token: @@ -333,8 +365,18 @@ class VideoBackend: if self._load_token == token and self._loading is not None: self._loading.error = redact_native_paths(str(exc)) + # Base-repo subfolders an LTX-2.3 assembly reads: the checkpoint (plus the GGUF + # repo's extras files) supplies the DiT, connectors, both VAEs and the vocoder, + # so only the 2.0 base's scheduler / text encoder / tokenizer are pulled. + _LTX23_BASE_PREFIXES = ("scheduler/", "text_encoder/", "tokenizer/") + @staticmethod - def _base_download_files(info: Any, kind: str) -> list[tuple[str, int]]: + def _base_download_files( + info: Any, + kind: str, + *, + ltx23: bool = False, + ) -> list[tuple[str, int]]: """The (rfilename, size) list a load actually needs from the base repo. Single source of truth for the progress estimate AND the scoped pre-download, @@ -344,7 +386,10 @@ class VideoBackend: - the duplicate ``text_encoder/diffusion_pytorch_model*`` shard set (the LTX-2 base repo ships its text encoder twice; transformers loads the ``model-*`` naming via the shard index); - - ``transformer/`` when a GGUF/single-file checkpoint replaces the DiT.""" + - ``transformer/`` when a GGUF/single-file checkpoint replaces the DiT; + - everything but scheduler / text encoder / tokenizer for an LTX-2.3 + checkpoint (``ltx23``), whose VAEs/vocoder/connectors come from the + checkpoint and its extras, not the 2.0 base.""" files: list[tuple[str, int]] = [] for sibling in info.siblings or []: name, size = sibling.rfilename, sibling.size or 0 @@ -356,6 +401,8 @@ class VideoBackend: continue if name.startswith("text_encoder/diffusion_pytorch_model"): continue + if ltx23 and "/" in name and not name.startswith(VideoBackend._LTX23_BASE_PREFIXES): + continue files.append((name, int(size))) return files @@ -366,6 +413,7 @@ class VideoBackend: base: str, hf_token: Optional[str], kind: str, + ltx23: bool = False, ) -> Optional[int]: """Total bytes this load will pull (checkpoint + companions), or None.""" try: @@ -380,12 +428,19 @@ class VideoBackend: total += int(sibling.size) if base and not Path(base).expanduser().exists(): info = api.model_info(base, files_metadata = True) - total += sum(size for _, size in self._base_download_files(info, kind)) + total += sum(size for _, size in self._base_download_files(info, kind, ltx23 = ltx23)) return total or None except Exception: # noqa: BLE001 -- progress totals are best-effort only return None - def _predownload_base(self, base: str, hf_token: Optional[str], kind: str) -> Optional[str]: + def _predownload_base( + self, + base: str, + hf_token: Optional[str], + kind: str, + *, + ltx23: bool = False, + ) -> Optional[str]: """Pull exactly the base-repo files the load needs; return the local snapshot dir. A bare ``from_pretrained(repo_id)`` snapshot of Lightricks/LTX-2 downloads the @@ -401,7 +456,7 @@ class VideoBackend: from huggingface_hub import HfApi info = HfApi(token = hf_token or None).model_info(base, files_metadata = True) - files = self._base_download_files(info, kind) + files = self._base_download_files(info, kind, ltx23 = ltx23) if not any(name == "model_index.json" for name, _ in files): return None from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index b796cb0004..623c210b7a 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2194,3 +2194,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() diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index bc15bc0207..c34cee5bf9 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -509,3 +509,22 @@ def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path): ) assert _FakePipeline.last["base"] == str(tmp_path) backend.unload() + + +def test_base_download_files_ltx23_keeps_only_shared_components(): + # A 2.3 checkpoint supplies the DiT, connectors, both VAEs and the vocoder, so + # the base pull shrinks to scheduler + text encoder + tokenizer (+ root manifest). + siblings = _LTX2_SIBLINGS + [ + _sibling("scheduler/scheduler_config.json", 1), + _sibling("connectors/diffusion_pytorch_model.safetensors", 3), + _sibling("latent_upsampler/diffusion_pytorch_model.safetensors", 1), + ] + info = types.SimpleNamespace(siblings = siblings) + names = [n for n, _ in VideoBackend._base_download_files(info, "gguf", ltx23 = True)] + assert "model_index.json" in names + assert "scheduler/scheduler_config.json" in names + assert "text_encoder/model-00001-of-00002.safetensors" in names + assert "tokenizer/tokenizer.model" in names + assert not any( + n.startswith(("vae/", "connectors/", "latent_upsampler/", "transformer/")) for n in names + )