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

Merge branch 'video-inference' into video-tab

Browse source
This commit is contained in:
Daniel Han 2026-07-05 03:18:04 +00:00
commit d281de0923
2 changed files with 150 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

98
studio/backend/core/inference/video.py
View file

@ -308,8 +308,8 @@ class VideoBackend:
self._loading.expected_bytes = expected
# The GGUF/single-file checkpoint downloads outside the lock so an
# unload/eviction can preempt the multi-GB pull; the pipeline
# companions download inside from_pretrained (which resumes from the
# cache, so a cancelled pull costs nothing).
# companions pre-download the same way (scoped file list, cancellable,
# resumes from the cache so a cancelled pull costs nothing).
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(
@ -318,6 +318,9 @@ class VideoBackend:
kwargs.get("hf_token"),
cancel_event = self._cancel_event,
)
kwargs["_base_local_dir"] = self._predownload_base(
base, kwargs.get("hf_token"), kind
)
self.load_pipeline(**kwargs)
with self._lock:
if self._load_token == token:
@ -332,6 +335,32 @@ class VideoBackend:
if self._load_token == token and self._loading is not None:
self._loading.error = redact_native_paths(str(exc))
@staticmethod
def _base_download_files(info: Any, kind: str) -> 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,
so the two can never disagree. Excluded on purpose:
- root-level packaged checkpoints (ComfyUI-style singles; 170 GB of the LTX-2
repo) -- the diffusers pipeline only reads per-component subfolders;
- 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."""
files: list[tuple[str, int]] = []
for sibling in info.siblings or []:
name, size = sibling.rfilename, sibling.size or 0
if not name.endswith((".safetensors", ".json", ".model", ".txt")):
continue
if "/" not in name and name.endswith(".safetensors"):
continue
if kind != "pipeline" and name.startswith("transformer/"):
continue
if name.startswith("text_encoder/diffusion_pytorch_model"):
continue
files.append((name, int(size)))
return files
def _estimate_download_bytes(
self,
repo_id: str,
@ -353,22 +382,48 @@ class VideoBackend:
total += int(sibling.size)
if base and not Path(base).expanduser().exists():
info = api.model_info(base, files_metadata = True)
for sibling in info.siblings or []:
name, size = sibling.rfilename, sibling.size or 0
if not name.endswith((".safetensors", ".json", ".model", ".txt")):
continue
# A GGUF/single-file load replaces the base repo's DiT, and the
# LTX-2 base repo ships its text encoder TWICE (two shard
# namings); count only what from_pretrained will pull.
if kind != "pipeline" and name.startswith("transformer/"):
continue
if name.startswith("text_encoder/diffusion_pytorch_model"):
continue
total += int(size)
total += sum(size for _, size in self._base_download_files(info, kind))
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]:
"""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
whole 314 GB repo (root packaged checkpoints plus a second 50 GB text-encoder
shard set) when ~93 GB is used. Downloading the scoped file list ourselves is
also cancellable per file, and handing the local dir to from_pretrained skips
diffusers' own expected-files sweep. None -> caller keeps the hub id (local
path, non-diffusers layout, or any metadata failure: from_pretrained then
resolves the repo exactly as before)."""
try:
if not base or Path(base).expanduser().exists():
return None
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)
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
snapshot_root: Optional[Path] = None
for name, _ in files:
local = Path(
hf_hub_download_with_xet_fallback(
base, name, hf_token, cancel_event = self._cancel_event
)
)
if name == "model_index.json":
snapshot_root = local.parent
return str(snapshot_root) if snapshot_root is not None else None
except Exception as exc: # noqa: BLE001 -- fall back to from_pretrained's own pull
if self._cancel_event.is_set():
raise
logger.warning("video.predownload_fallback: %s", exc)
return None
def _cache_bytes(self, repo_id: Optional[str]) -> int:
"""Bytes of ``repo_id`` currently in the HF blob cache (progress polling)."""
if not repo_id:
@ -397,6 +452,11 @@ class VideoBackend:
phase = "downloading"
if expected and downloaded >= expected:
phase = "finalizing"
# The cache scan counts every blob of the repo(s), including files a
# previous (or broader) pull left behind that this load never reads, so
# the raw counter can exceed the scoped estimate. Clamp: everything the
# load needs is present, which is what the bar reports.
downloaded = expected
return _progress(
phase,
downloaded_bytes = int(downloaded),
@ -420,6 +480,7 @@ class VideoBackend:
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]:
import diffusers
import torch
@ -506,7 +567,10 @@ class VideoBackend:
if hf_token:
pipe_kwargs["token"] = hf_token
if kind == "pipeline":
pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
# The pre-downloaded snapshot dir keeps from_pretrained off the hub (its
# own snapshot sweep would also pull the repo's packaged root checkpoints
# and duplicate text-encoder shards); hub id when pre-download was skipped.
pipe = pipeline_cls.from_pretrained(_base_local_dir or repo_id, **pipe_kwargs)
else:
transformer_cls = getattr(diffusers, fam.transformer_class)
# checkpoint_path was already resolved (and downloaded) by the memory
@ -536,7 +600,9 @@ class VideoBackend:
)
else:
transformer = transformer_cls.from_single_file(str(checkpoint_path), **sf_kwargs)
pipe = pipeline_cls.from_pretrained(base, transformer = transformer, **pipe_kwargs)
pipe = pipeline_cls.from_pretrained(
_base_local_dir or base, transformer = transformer, **pipe_kwargs
)
if _load_token is not None and _load_token != self._load_token:
del pipe

68
studio/backend/tests/test_video_backend.py
View file

@ -441,3 +441,71 @@ def test_generate_progress_and_cancel_idle(fake_runtime):
def test_singleton():
assert get_video_backend() is get_video_backend()
def _sibling(name, size):
return types.SimpleNamespace(rfilename = name, size = size)
_LTX2_SIBLINGS = [
_sibling("model_index.json", 10),
_sibling("ltx-2-19b-packaged-fp8.safetensors", 170),
_sibling("transformer/config.json", 1),
_sibling("transformer/diffusion_pytorch_model-00001-of-00002.safetensors", 20),
_sibling("transformer/diffusion_pytorch_model-00002-of-00002.safetensors", 18),
_sibling("text_encoder/model-00001-of-00002.safetensors", 25),
_sibling("text_encoder/model-00002-of-00002.safetensors", 25),
_sibling("text_encoder/diffusion_pytorch_model-00001-of-00002.safetensors", 25),
_sibling("text_encoder/diffusion_pytorch_model-00002-of-00002.safetensors", 25),
_sibling("vae/diffusion_pytorch_model.safetensors", 3),
_sibling("tokenizer/tokenizer.model", 1),
_sibling("assets/example.mp4", 500),
]
def test_base_download_files_scopes_pipeline_pull():
# A pipeline load skips the packaged root checkpoint, the duplicate
# text-encoder shard naming, and non-weight assets -- and keeps everything else.
info = types.SimpleNamespace(siblings = _LTX2_SIBLINGS)
files = dict(VideoBackend._base_download_files(info, "pipeline"))
assert "ltx-2-19b-packaged-fp8.safetensors" not in files
assert "text_encoder/diffusion_pytorch_model-00001-of-00002.safetensors" not in files
assert "assets/example.mp4" not in files
assert files["text_encoder/model-00001-of-00002.safetensors"] == 25
assert files["transformer/diffusion_pytorch_model-00001-of-00002.safetensors"] == 20
assert sum(files.values()) == 10 + 1 + 20 + 18 + 25 + 25 + 3 + 1
def test_base_download_files_gguf_drops_transformer():
# A GGUF/single-file checkpoint replaces the DiT: the base transformer never pulls.
info = types.SimpleNamespace(siblings = _LTX2_SIBLINGS)
names = [n for n, _ in VideoBackend._base_download_files(info, "gguf")]
assert not any(n.startswith("transformer/") for n in names)
assert "text_encoder/model-00001-of-00002.safetensors" in names
def test_load_progress_clamps_overshoot(fake_runtime, monkeypatch):
# The cache scan counts blobs a broader previous pull left behind; the reported
# counter must never exceed the scoped estimate (no "282 GB of 263 GB").
backend = VideoBackend()
backend._loading = types.SimpleNamespace(
repo_id = "Lightricks/LTX-2", base_repo = None, expected_bytes = 100, error = None
)
monkeypatch.setattr(VideoBackend, "_cache_bytes", lambda self, repo: 150)
progress = backend.load_progress()
assert progress["phase"] == "finalizing"
assert progress["downloaded_bytes"] == 100
assert progress["expected_bytes"] == 100
def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path):
# When the scoped pre-download produced a local snapshot, from_pretrained must
# receive that dir (keeping diffusers' own broader snapshot sweep off the hub).
backend = VideoBackend()
backend.load_pipeline(
"Lightricks/LTX-2",
model_kind = "pipeline",
_base_local_dir = str(tmp_path),
)
assert _FakePipeline.last["base"] == str(tmp_path)
backend.unload()