ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/tests/test_video_backend.py
Daniel Han 6485e68147 Harden video load path: early GGUF-repo rejection, family fallback parity, rollback and teardown fixes 
Review follow-ups on the video inference backend:

- validate_load_request now rejects a -GGUF repo picked as a diffusers
  pipeline (no gguf_filename) up front, instead of failing minutes later
  in from_pretrained after the GPU owner was already evicted.
- New _detect_load_family helper shared by validate_load_request and
  _run_load: when the repo id alone does not carry the family, fall back
  to detecting it from the picked GGUF filename, so both paths agree.
- routes/video.py now threads base_repo into validate_load_request so an
  untrusted companion repo is refused before the arbiter handoff.
- unload() now drains _generate_lock before _teardown_state so a
  cancelled clip actually exits the denoise loop before the VRAM is
  reported free.
- load_pipeline re-checks the load token after the generate-lock barrier
  and raises if the load was superseded while waiting.
- Pre-commit global mutations (backend flags, gguf compile installs) are
  registered per load token and rolled back in _run_load's error path
  via _rollback_precommit_globals, so a failed load no longer leaks
  process-wide state.
- fp32 memory estimates now apply a 2x dtype scale on non-CPU devices
  for pipeline, single-file and companion sizes (bf16 tables assume
  2 bytes/param); GGUF quant estimates stay unscaled.

Tests: GGUF-repo-as-pipeline rejection, _detect_load_family fallback and
override semantics; fake route backend accepts base_repo. 66 passed
across test_video_backend, test_video_routes, test_video_families,
test_video_gallery.
2026-07-05 07:55:26 +00:00

601 lines
23 KiB
Python
Raw Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""VideoBackend lifecycle on a faked torch/diffusers runtime (CPU-only, offline).
Mirrors test_diffusion_backend's fake_runtime pattern: explicit fake signatures so
the signature-gated kwargs actually exercise, sys.modules stubs so no real ML
stack loads."""
import contextlib
import sys
import types
import pytest
from core.inference.video import (
VideoBackend,
_detect_load_family,
get_video_backend,
resolve_video_model_kind,
)
from core.inference.video_families import VIDEO_NOT_LOADED_MSG
class _FakeDtype:
def __init__(self, name: str) -> None:
self._name = name
def __repr__(self) -> str:
return f"torch.{self._name}"
__str__ = __repr__
class _FakeGenerator:
def __init__(self, device = None) -> None:
self.device = device
self.manual = None
def seed(self) -> int:
return 4242
def manual_seed(self, value: int):
self.manual = value
return self
class _FakeVae:
def __init__(self) -> None:
self.tiled = False
def enable_tiling(self) -> None:
self.tiled = True
class _FakePipe:
def __init__(self) -> None:
self.moved_to = None
self.vae = _FakeVae()
self.last_kwargs = None
self._interrupt = False
def to(self, device):
self.moved_to = device
return self
def enable_vae_tiling(self) -> None:
self.vae.tiled = True
# Explicit signature so generate()'s signature-gated kwargs (negative_prompt,
# frame_rate, callback) actually engage; **kwargs would defeat the gates.
def __call__(
self,
*,
prompt = None,
negative_prompt = None,
num_inference_steps = None,
guidance_scale = None,
width = None,
height = None,
num_frames = None,
frame_rate = None,
generator = None,
callback_on_step_end = None,
**kwargs,
):
self.last_kwargs = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"num_inference_steps": num_inference_steps,
"guidance_scale": guidance_scale,
"width": width,
"height": height,
"num_frames": num_frames,
"frame_rate": frame_rate,
**kwargs,
}
if callback_on_step_end is not None:
for step in range(int(num_inference_steps or 1)):
callback_on_step_end(self, step, 0, {})
if self._interrupt:
break
frames = [[object() for _ in range(int(num_frames or 1))]]
return types.SimpleNamespace(frames = frames, audio = None)
class _FakePipeline:
last: dict = {}
@classmethod
def from_pretrained(cls, base, **kwargs):
_FakePipeline.last = {"base": base, **kwargs}
return _FakePipe()
class _FakeTransformer:
last: dict = {}
@classmethod
def from_single_file(cls, path, **kwargs):
_FakeTransformer.last = {"path": path, **kwargs}
return object()
@pytest.fixture
def fake_runtime(monkeypatch):
torch = types.ModuleType("torch")
torch.bfloat16 = _FakeDtype("bfloat16")
torch.float16 = _FakeDtype("float16")
torch.float32 = _FakeDtype("float32")
torch.Generator = _FakeGenerator
torch.cuda = types.SimpleNamespace(is_available = lambda: False)
torch.backends = types.SimpleNamespace(mps = None)
torch.inference_mode = lambda: contextlib.nullcontext()
diffusers = types.ModuleType("diffusers")
diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype)
diffusers.LTX2Pipeline = _FakePipeline
diffusers.LTX2VideoTransformer3DModel = _FakeTransformer
monkeypatch.setitem(sys.modules, "torch", torch)
monkeypatch.setitem(sys.modules, "diffusers", diffusers)
monkeypatch.setattr("core.inference.video.clear_gpu_cache", lambda: None)
# MP4 encode needs real frames + PyAV; the backend contract under test is the
# byte handoff, so stub the encoder.
monkeypatch.setattr(
VideoBackend, "_encode_mp4", staticmethod(lambda frames, fps, audio, pipe: b"MP4")
)
_FakePipeline.last = {}
_FakeTransformer.last = {}
yield
def _load_gguf(backend, tmp_path):
(tmp_path / "model.gguf").write_bytes(b"weights")
return backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.gguf",
base_repo = "Lightricks/LTX-2",
family_override = "ltx-2",
)
def test_resolve_kind():
assert resolve_video_model_kind("x.gguf", None) == "gguf"
assert resolve_video_model_kind("x.safetensors", None) == "single_file"
assert resolve_video_model_kind(None, None) == "pipeline"
with pytest.raises(ValueError):
resolve_video_model_kind(None, "bogus")
def test_validate_rejects_unknown_and_untrusted():
backend = VideoBackend()
with pytest.raises(ValueError, match = "not a supported"):
backend.validate_load_request("someorg/some-image-model")
# A known family but an untrusted repo id must not open from_pretrained.
with pytest.raises(ValueError, match = "limited to"):
backend.validate_load_request("evil/ltx-2-repack")
# GGUF loads stay open to any repo (single-file read, no pickle).
fam = backend.validate_load_request(
"anyorg/ltx-2-GGUF", gguf_filename = "x.gguf", model_kind = "gguf"
)
assert fam.name == "ltx-2"
with pytest.raises(ValueError, match = "filename"):
backend.validate_load_request("unsloth/LTX-2.3-GGUF", model_kind = "gguf")
def test_validate_gates_base_repo_and_local_paths(tmp_path):
backend = VideoBackend()
# An arbitrary remote base_repo must not reach from_pretrained via a GGUF pick.
with pytest.raises(ValueError, match = "base_repo"):
backend.validate_load_request(
"unsloth/LTX-2.3-GGUF",
gguf_filename = "x.gguf",
model_kind = "gguf",
base_repo = "evil/companions",
)
# The family base and local dirs stay allowed.
fam = backend.validate_load_request(
"unsloth/LTX-2.3-GGUF",
gguf_filename = "x.gguf",
model_kind = "gguf",
base_repo = "Lightricks/LTX-2",
)
assert fam.name == "ltx-2"
# A local dir without the picked checkpoint fails BEFORE the GPU handoff.
with pytest.raises(ValueError):
backend.validate_load_request(
str(tmp_path), gguf_filename = "missing.gguf", family_override = "ltx-2"
)
# A path-shaped repo id that does not exist fails validation too.
with pytest.raises(ValueError, match = "does not exist"):
backend.validate_load_request(
str(tmp_path / "nope" / "model.gguf"),
gguf_filename = "model.gguf",
family_override = "ltx-2",
)
def test_validate_rejects_gguf_repo_as_pipeline():
backend = VideoBackend()
# A -GGUF repo with no quant filename resolves to the pipeline kind and would
# only fail minutes later in from_pretrained, AFTER evicting the GPU owner.
with pytest.raises(ValueError, match = "pick one of its .gguf files"):
backend.validate_load_request("unsloth/LTX-2.3-GGUF")
with pytest.raises(ValueError, match = "pick one of its .gguf files"):
backend.validate_load_request("unsloth/Wan2.2-TI2V-5B-GGUF/")
def test_detect_load_family_filename_fallback():
# Repo id alone carries the family.
fam = _detect_load_family("Lightricks/LTX-2", None, None)
assert fam is not None and fam.name == "ltx-2"
# Repo id is opaque but the picked filename carries it: fall back to the
# combined path so validate and _run_load agree on the family.
fam = _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", None)
assert fam is not None and fam.name == "ltx-2"
# No filename and no recognisable repo id: no family.
assert _detect_load_family("someorg/quants", None, None) is None
# An explicit override resolves by name/alias and skips the filename fallback:
# a bogus override stays None even when the filename would have matched.
fam = _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", "ltxv")
assert fam is not None and fam.name == "ltx-2"
assert _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", "bogus") is None
def test_load_generate_unload_gguf(fake_runtime, tmp_path):
backend = VideoBackend()
status = _load_gguf(backend, tmp_path)
assert status["loaded"] is True and status["family"] == "ltx-2"
assert status["model_kind"] == "gguf"
assert status["has_audio"] is True
# The GGUF transformer is dequant-configured and assembled onto the base repo.
assert _FakeTransformer.last["path"].endswith("model.gguf")
assert _FakeTransformer.last["quantization_config"][0] == "quant"
assert _FakePipeline.last["base"] == "Lightricks/LTX-2"
assert "transformer" in _FakePipeline.last
# Video decode is the memory peak: tiling is always on.
assert status["vae_tiling"] is True
assert status["defaults"]["frame_step"] == 8
result = backend.generate(
prompt = "a sloth surfing", width = 1000, height = 700, num_frames = 120, fps = 24
)
call = backend._state.pipe.last_kwargs
# Shape snapping happened BEFORE the pipe call: /32 sizes, 8k+1 frames.
assert (call["width"], call["height"]) == (992, 672)
assert call["num_frames"] == 113
assert call["frame_rate"] == 24.0
assert result["mp4_bytes"] == b"MP4"
assert result["num_frames"] == 113 and result["fps"] == 24
assert result["has_audio"] is False # fake pipe returned no audio track
assert 0 <= result["seed"] < 2**53
status = backend.unload()
assert status["loaded"] is False
def test_generate_defaults_from_variant(fake_runtime, tmp_path):
# A distilled GGUF pick defaults to the few-step no-CFG schedule.
(tmp_path / "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf").write_bytes(b"w")
backend = VideoBackend()
backend.load_pipeline(
str(tmp_path),
gguf_filename = "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf",
base_repo = "Lightricks/LTX-2",
family_override = "ltx-2",
)
backend.generate(prompt = "a sloth")
call = backend._state.pipe.last_kwargs
assert call["num_inference_steps"] == 8
assert call["guidance_scale"] == 1.0
def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
# FBCache residuals live on the long-lived DiT(s) and survive a generation, so
# the next clip at a new resolution would crash on stale state. generate must
# reset them when a cache is engaged (diffusers 0.39 exposes
# _reset_stateful_cache on the transformer; reset_stateful_hooks only exists on
# the HookRegistry) and must not touch an uncached load. transformer_2 (the Wan
# dual expert) resets too when present.
import dataclasses
(tmp_path / "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf").write_bytes(b"w")
backend = VideoBackend()
backend.load_pipeline(
str(tmp_path),
gguf_filename = "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf",
base_repo = "Lightricks/LTX-2",
family_override = "ltx-2",
)
resets = []
backend._state.pipe.transformer = types.SimpleNamespace(
_reset_stateful_cache = lambda: resets.append("transformer")
)
backend._state.pipe.transformer_2 = types.SimpleNamespace(
_reset_stateful_cache = lambda: resets.append("transformer_2")
)
# No cache engaged -> no reset.
backend.generate(prompt = "a sloth")
assert resets == []
# Cache engaged -> both resident DiTs reset before the pipe call.
backend._state = dataclasses.replace(backend._state, transformer_cache = "fbcache")
backend.generate(prompt = "a sloth")
assert resets == ["transformer", "transformer_2"]
def test_is_ltx23_checkpoint_gguf(monkeypatch, tmp_path):
# diffusers maps every LTX-2 single file to the 2.0 config; a 2.3 checkpoint
# (9-row modulation tables in the header) must be detected so the loader
# routes to the full 2.3 assembly. A 2.0 header must not, and an unreadable
# header must fall back to the stock path (False), never raise.
from core.inference.video_ltx2 import is_ltx23_checkpoint
def _reader_for(shapes):
tensors = [types.SimpleNamespace(name = n, shape = s) for n, s in shapes.items()]
return lambda path: types.SimpleNamespace(tensors = tensors)
gguf = types.ModuleType("gguf")
# GGUF headers store dims in GGML (reversed) order.
gguf.GGUFReader = _reader_for(
{
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 9),
}
)
monkeypatch.setitem(sys.modules, "gguf", gguf)
path = tmp_path / "ltx23.gguf"
path.write_bytes(b"x")
assert is_ltx23_checkpoint(path) is True
gguf.GGUFReader = _reader_for(
{
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 6),
}
)
assert is_ltx23_checkpoint(path) is False
def _boom(path):
raise RuntimeError("bad magic")
gguf.GGUFReader = _boom
assert is_ltx23_checkpoint(path) is False
def test_is_ltx23_checkpoint_safetensors(monkeypatch, tmp_path):
from core.inference.video_ltx2 import is_ltx23_checkpoint
class _FakeSlice:
def __init__(self, shape):
self._shape = shape
def get_shape(self):
return self._shape
class _FakeSafe:
def __init__(self, shapes):
self._shapes = shapes
def __enter__(self):
return self
def __exit__(self, *args):
return False
def keys(self):
return list(self._shapes)
def get_slice(self, name):
return _FakeSlice(self._shapes[name])
shapes = {
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (9, 4096),
}
safetensors = types.ModuleType("safetensors")
safetensors.safe_open = lambda path, framework = None: _FakeSafe(shapes)
monkeypatch.setitem(sys.modules, "safetensors", safetensors)
path = tmp_path / "ltx23.safetensors"
path.write_bytes(b"x")
assert is_ltx23_checkpoint(path) is True
def test_ltx23_split_and_variant(tmp_path):
# Pure functions: combined-checkpoint partitioning and companion-set choice.
from core.inference.video_ltx2 import _split_checkpoint, checkpoint_variant
state = {
"model.diffusion_model.transformer_blocks.0.attn1.to_q.weight": 1,
"model.diffusion_model.video_embeddings_connector.learnable_registers": 2,
"model.diffusion_model.prompt_adaln_single.linear.weight": 3,
"text_embedding_projection.video_aggregate_embed.weight": 4,
"vae.decoder.conv_in.weight": 5,
"audio_vae.encoder.conv_in.weight": 6,
"vocoder.bwe_generator.conv_pre.weight": 7,
}
groups = _split_checkpoint(state)
assert set(groups["dit"]) == {
"transformer_blocks.0.attn1.to_q.weight",
"prompt_adaln_single.linear.weight",
}
assert set(groups["connectors"]) == {
"video_embeddings_connector.learnable_registers",
"text_embedding_projection.video_aggregate_embed.weight",
}
assert groups["vae"] == {"decoder.conv_in.weight": 5}
assert groups["audio_vae"] == {"encoder.conv_in.weight": 6}
assert groups["vocoder"] == {"bwe_generator.conv_pre.weight": 7}
assert checkpoint_variant("x/ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf") == "distilled"
assert checkpoint_variant("x/ltx-2.3-22b-dev-Q8_0.gguf") == "dev"
def test_ltx23_scaled_fp8_refused(monkeypatch, tmp_path):
# The Lightricks fp8 files carry .weight_scale/.input_scale companions; a
# plain dtype cast would silently corrupt them, so the loader must refuse
# with a pointer to the supported GGUF path.
from core.inference import video_ltx2
# Stub the module tree so this also runs under the CI sim, which blocks the
# real diffusers import.
diffusers = types.ModuleType("diffusers")
diffusers.LTX2Pipeline = object
loaders = types.ModuleType("diffusers.loaders")
sfu = types.ModuleType("diffusers.loaders.single_file_utils")
sfu.load_single_file_checkpoint = lambda path: {
"model.diffusion_model.transformer_blocks.0.attn1.to_q.weight": object(),
"model.diffusion_model.transformer_blocks.0.attn1.to_q.weight_scale": object(),
}
diffusers.loaders = loaders
loaders.single_file_utils = sfu
monkeypatch.setitem(sys.modules, "diffusers", diffusers)
monkeypatch.setitem(sys.modules, "diffusers.loaders", loaders)
monkeypatch.setitem(sys.modules, "diffusers.loaders.single_file_utils", sfu)
monkeypatch.setitem(sys.modules, "transformers", types.ModuleType("transformers"))
path = tmp_path / "ltx-2.3-22b-distilled-fp8.safetensors"
path.write_bytes(b"x")
with pytest.raises(ValueError, match = "scaled fp8"):
video_ltx2.load_ltx23_pipeline(
path, base_repo = "Lightricks/LTX-2", torch_dtype = None, is_gguf = False
)
def test_generate_without_load_raises(fake_runtime):
backend = VideoBackend()
with pytest.raises(RuntimeError, match = VIDEO_NOT_LOADED_MSG):
backend.generate(prompt = "x")
def test_generate_progress_and_cancel_idle(fake_runtime):
backend = VideoBackend()
assert backend.generate_progress() == {"active": False}
assert backend.cancel_generate() is False
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("tokenizer/chat_template.jinja", 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
# The standalone chat template must survive the whitelist: apply_chat_template
# reads it at generation time and it is not embedded in tokenizer_config.json.
assert "tokenizer/chat_template.jinja" in files
assert sum(files.values()) == 10 + 1 + 20 + 18 + 25 + 25 + 3 + 1 + 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()
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
)
def test_predownload_base_honors_cancel_between_files(monkeypatch):
# A warm-cache sweep returns each file instantly without consulting the event,
# so the loop must check it explicitly or an unload mid-predownload is ignored.
backend = VideoBackend()
backend._cancel_event.set()
calls: list = []
monkeypatch.setattr(
"utils.hf_xet_fallback.hf_hub_download_with_xet_fallback",
lambda repo, fn, tok, **kw: (calls.append(fn), f"/cache/{fn}")[1],
)
class _Api:
def __init__(self, token = None):
pass
def model_info(
self,
repo,
files_metadata = True,
):
return types.SimpleNamespace(
siblings = [
_sibling("model_index.json", 1),
_sibling("vae/diffusion_pytorch_model.safetensors", 2),
]
)
import huggingface_hub
monkeypatch.setattr(huggingface_hub, "HfApi", _Api)
with pytest.raises(RuntimeError, match = "cancelled"):
backend._predownload_base("base/repo", None, "pipeline")
assert calls == []
Reference in a new issue View git blame Copy permalink