Merge branch 'video-inference' into video-tab

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
 

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

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

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
+    )