From e2fd90783e087947c90b463a174e2da4c0ffe486 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 04:38:28 +0000 Subject: [PATCH 1/3] Shrink the LTX-2.3 base pull to the components the assembly reads A 2.3 checkpoint (GGUF or single file) carries the DiT and, with its extras files, the connectors, both VAEs and the vocoder; only the 2.0 base repo's scheduler, text encoder and tokenizer are read. Detect 2.3 from the checkpoint header after the pull, re-estimate, and scope the base pre-download accordingly (about 6 GB less per fresh install). --- studio/backend/core/inference/video.py | 73 ++++++++++++++++++---- studio/backend/tests/test_video_backend.py | 20 ++++++ 2 files changed, 82 insertions(+), 11 deletions(-) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index f00ccdee6b..d91cc62737 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -310,15 +310,49 @@ 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 +367,13 @@ 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 +383,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 +398,12 @@ 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 +414,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 +429,14 @@ 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 +452,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_video_backend.py b/studio/backend/tests/test_video_backend.py index bc15bc0207..3eece7f807 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -509,3 +509,23 @@ 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 + ) From 69f48d2ffc72f8879d648eb6119f77a7704707c2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 04:39:05 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/video.py | 24 +++++++++++++--------- studio/backend/tests/test_video_backend.py | 3 +-- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index d91cc62737..4bd7751f7c 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -346,9 +346,7 @@ class VideoBackend: 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 - ) + 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. @@ -373,7 +371,12 @@ class VideoBackend: _LTX23_BASE_PREFIXES = ("scheduler/", "text_encoder/", "tokenizer/") @staticmethod - def _base_download_files(info: Any, kind: str, *, ltx23: bool = False) -> 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, @@ -398,11 +401,7 @@ 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) - ): + if ltx23 and "/" in name and not name.startswith(VideoBackend._LTX23_BASE_PREFIXES): continue files.append((name, int(size))) return files @@ -435,7 +434,12 @@ class VideoBackend: return None def _predownload_base( - self, base: str, hf_token: Optional[str], kind: str, *, ltx23: bool = False + 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. diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 3eece7f807..c34cee5bf9 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -526,6 +526,5 @@ def test_base_download_files_ltx23_keeps_only_shared_components(): 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 + n.startswith(("vae/", "connectors/", "latent_upsampler/", "transformer/")) for n in names ) From 62cae5fe71cacff56dd2ec72b881a6697b028cce Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 04:41:43 +0000 Subject: [PATCH 3/3] Load image pipelines from the prefetched snapshot instead of re-sweeping the hub 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. --- studio/backend/core/inference/diffusion.py | 41 +++++++++++++++---- .../backend/tests/test_diffusion_backend.py | 32 +++++++++++++++ 2 files changed, 64 insertions(+), 9 deletions(-) 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()