776 lines
28 KiB
Python
776 lines
28 KiB
Python
# 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, 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()
|
|
|
|
|
|
# ── Wan2.2 fakes: a per-DiT trackable transformer so the dual-DiT optimisation
|
|
# tests can assert speed / cache / attention engaged on BOTH experts, plus two
|
|
# pipeline fakes -- single-DiT (TI2V-5B) and dual-DiT MoE (A14B). The MoE __call__
|
|
# carries guidance_scale_2 so the cfg2 signature-gate actually exercises; the
|
|
# single-DiT __call__ omits it so the gate proves it is NOT threaded there.
|
|
|
|
|
|
class _FakeWanDiT:
|
|
"""One Wan denoiser. Records which optimisation helpers touched it (the loader
|
|
applies each once per expert on an MoE load), so a test can prove BOTH experts
|
|
were covered. compile_repeated_blocks / enable_cache / set_attention_backend are
|
|
exactly the attribute names the imported helpers look for."""
|
|
|
|
def __init__(self) -> None:
|
|
self.compiled = False
|
|
self.cache_config = None
|
|
self.attention = None
|
|
|
|
def compile_repeated_blocks(self, **kwargs) -> None:
|
|
self.compiled = True
|
|
|
|
def enable_cache(self, config) -> None:
|
|
self.cache_config = config
|
|
|
|
def set_attention_backend(self, backend) -> None:
|
|
self.attention = backend
|
|
|
|
|
|
class _FakeWanVae:
|
|
def __init__(self) -> None:
|
|
self.tiled = False
|
|
|
|
def enable_tiling(self) -> None:
|
|
self.tiled = True
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
|
|
class _FakeWanPipeBase:
|
|
"""Shared Wan pipeline state. Subclasses provide the __call__ with the right
|
|
explicit signature (with/without guidance_scale_2) so the generate() cfg2 and
|
|
frame_rate signature-gates actually exercise -- ``**kwargs`` alone would hide the
|
|
parameter names inspect.signature reads."""
|
|
|
|
moe: bool = False
|
|
|
|
def __init__(self) -> None:
|
|
self.vae = _FakeWanVae()
|
|
self.transformer = _FakeWanDiT()
|
|
self.transformer_2 = _FakeWanDiT() if self.moe else None
|
|
self.components = {"transformer": self.transformer, "vae": self.vae}
|
|
if self.moe:
|
|
self.components["transformer_2"] = self.transformer_2
|
|
self.moved_to = None
|
|
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
|
|
|
|
def _finish(self, num_inference_steps, num_frames, callback_on_step_end):
|
|
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 _FakeWanPipeSingle(_FakeWanPipeBase):
|
|
"""Single-DiT Wan pipeline (TI2V-5B): NO guidance_scale_2 in the signature, so the
|
|
cfg2 gate must not thread it."""
|
|
|
|
moe = False
|
|
|
|
def __call__(
|
|
self,
|
|
*,
|
|
prompt = None,
|
|
negative_prompt = None,
|
|
num_inference_steps = None,
|
|
guidance_scale = None,
|
|
width = None,
|
|
height = None,
|
|
num_frames = 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,
|
|
**kwargs,
|
|
}
|
|
return self._finish(num_inference_steps, num_frames, callback_on_step_end)
|
|
|
|
|
|
class _FakeWanPipeMoE(_FakeWanPipeBase):
|
|
"""Dual-DiT MoE Wan pipeline (A14B): guidance_scale_2 IS in the signature, matching
|
|
WanPipeline.__call__ in diffusers 0.39, so the cfg2 gate threads it."""
|
|
|
|
moe = True
|
|
|
|
def __call__(
|
|
self,
|
|
*,
|
|
prompt = None,
|
|
negative_prompt = None,
|
|
num_inference_steps = None,
|
|
guidance_scale = None,
|
|
guidance_scale_2 = None,
|
|
width = None,
|
|
height = None,
|
|
num_frames = 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,
|
|
"guidance_scale_2": guidance_scale_2,
|
|
"width": width,
|
|
"height": height,
|
|
"num_frames": num_frames,
|
|
**kwargs,
|
|
}
|
|
return self._finish(num_inference_steps, num_frames, callback_on_step_end)
|
|
|
|
|
|
class _FakeWanPipelineSingle:
|
|
"""WanPipeline fake (from_pretrained). One class serves both families and picks the
|
|
single-DiT / dual-DiT pipe by the repo id, exactly as diffusers dispatches on the
|
|
repo's model_index.json (A14B lists transformer_2, TI2V-5B does not)."""
|
|
|
|
last: dict = {}
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, repo, **kwargs):
|
|
_FakeWanPipelineSingle.last = {"repo": repo, **kwargs}
|
|
moe = "a14b" in str(repo).lower()
|
|
return _FakeWanPipeMoE() if moe else _FakeWanPipeSingle()
|
|
|
|
|
|
@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
|
|
# Wan2.2: one pipeline class serves both families (it dispatches on the repo id).
|
|
diffusers.WanPipeline = _FakeWanPipelineSingle
|
|
diffusers.WanTransformer3DModel = _FakeTransformer
|
|
diffusers.FirstBlockCacheConfig = lambda threshold = None: ("fbcache", threshold)
|
|
|
|
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_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()
|
|
|
|
|
|
# ── Wan2.2 ─────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_load_wan_ti2v_5b_pipeline(fake_runtime):
|
|
# A full-pipeline load of the single-DiT TI2V-5B repo: WanPipeline.from_pretrained,
|
|
# no audio, tiling forced on, and the 4k+1 frame lattice surfaced.
|
|
backend = VideoBackend()
|
|
status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
|
|
assert status["loaded"] is True
|
|
assert status["family"] == "wan2.2-ti2v-5b"
|
|
assert status["model_kind"] == "pipeline"
|
|
assert status["has_audio"] is False
|
|
assert status["vae_tiling"] is True
|
|
assert status["defaults"]["frame_step"] == 4
|
|
assert status["transformer_quant"] is None
|
|
assert _FakeWanPipelineSingle.last["repo"] == "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
|
|
|
|
|
def test_wan_frame_snapping_4k_plus_1(fake_runtime):
|
|
# Wan snaps num_frames to 4k+1 (temporal factor 4), unlike LTX-2's 8k+1.
|
|
backend = VideoBackend()
|
|
backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
|
|
backend.generate(prompt = "a sloth", width = 1000, height = 700, num_frames = 120)
|
|
call = backend._state.pipe.last_kwargs
|
|
assert call["num_frames"] == 117 # 4*29 + 1
|
|
# /16 spatial snap for Wan (spatial 8 * patch 2).
|
|
assert (call["width"], call["height"]) == (992, 688)
|
|
|
|
|
|
def test_wan_ti2v_defaults_applied(fake_runtime):
|
|
# No steps/guidance passed -> the Wan pipeline defaults (50 / 5.0).
|
|
backend = VideoBackend()
|
|
backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
|
|
backend.generate(prompt = "a sloth")
|
|
call = backend._state.pipe.last_kwargs
|
|
assert call["num_inference_steps"] == 50
|
|
assert call["guidance_scale"] == 5.0
|
|
|
|
|
|
def test_wan_ti2v_does_not_thread_cfg2(fake_runtime):
|
|
# The single-DiT TI2V pipeline has no guidance_scale_2 in its signature, so a
|
|
# request value must NOT be threaded (WanPipeline raises on it when boundary_ratio
|
|
# is None), even if the caller passes guidance_2.
|
|
backend = VideoBackend()
|
|
backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
|
|
backend.generate(prompt = "a sloth", guidance_2 = 3.5)
|
|
call = backend._state.pipe.last_kwargs
|
|
assert "guidance_scale_2" not in call
|
|
|
|
|
|
def test_wan_a14b_dual_dit_pipeline_loads(fake_runtime):
|
|
# The A14B repo builds a dual-DiT MoE pipeline (transformer + transformer_2).
|
|
backend = VideoBackend()
|
|
status = backend.load_pipeline("Wan-AI/Wan2.2-T2V-A14B-Diffusers", model_kind = "pipeline")
|
|
assert status["loaded"] is True and status["family"] == "wan2.2-t2v-a14b"
|
|
pipe = backend._state.pipe
|
|
assert pipe.transformer is not None and pipe.transformer_2 is not None
|
|
|
|
|
|
def test_wan_a14b_cfg2_threaded_when_signature_has_it(fake_runtime):
|
|
# The MoE pipeline's __call__ carries guidance_scale_2, so an explicit guidance_2
|
|
# is threaded through as that kwarg (the cfg2_kwarg the family declares).
|
|
backend = VideoBackend()
|
|
backend.load_pipeline("Wan-AI/Wan2.2-T2V-A14B-Diffusers", model_kind = "pipeline")
|
|
backend.generate(prompt = "a sloth", guidance = 5.0, guidance_2 = 3.0)
|
|
call = backend._state.pipe.last_kwargs
|
|
assert call["guidance_scale"] == 5.0
|
|
assert call["guidance_scale_2"] == 3.0
|
|
|
|
# A None guidance_2 must NOT be threaded, so the pipeline defaults it itself.
|
|
backend.generate(prompt = "a sloth", guidance = 5.0)
|
|
call2 = backend._state.pipe.last_kwargs
|
|
assert call2["guidance_scale_2"] is None
|
|
|
|
|
|
def test_wan_a14b_step_cache_applies_to_both_dits(fake_runtime):
|
|
# A dual-DiT MoE load must engage the step cache on BOTH experts, not just the
|
|
# first: transformer_2 handles the low-noise steps and would otherwise run uncached.
|
|
backend = VideoBackend()
|
|
status = backend.load_pipeline(
|
|
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
|
model_kind = "pipeline",
|
|
transformer_cache = "fbcache",
|
|
)
|
|
pipe = backend._state.pipe
|
|
assert pipe.transformer.cache_config is not None
|
|
assert pipe.transformer_2.cache_config is not None
|
|
assert status["transformer_cache"] == "fbcache"
|
|
|
|
|
|
def test_wan_a14b_attention_applies_to_both_dits(fake_runtime):
|
|
# An explicit attention backend must be set on both experts.
|
|
backend = VideoBackend()
|
|
backend.load_pipeline(
|
|
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
|
model_kind = "pipeline",
|
|
attention_backend = "cudnn",
|
|
)
|
|
pipe = backend._state.pipe
|
|
assert pipe.transformer.attention is not None
|
|
assert pipe.transformer_2.attention is not None
|
|
# Both experts got the SAME kernel.
|
|
assert pipe.transformer.attention == pipe.transformer_2.attention
|
|
|
|
|
|
def test_wan_ti2v_single_dit_only_touches_one(fake_runtime):
|
|
# A single-DiT load must not fabricate a second expert or try to optimise one.
|
|
backend = VideoBackend()
|
|
backend.load_pipeline(
|
|
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
|
model_kind = "pipeline",
|
|
transformer_cache = "fbcache",
|
|
)
|
|
pipe = backend._state.pipe
|
|
assert pipe.transformer_2 is None
|
|
assert pipe.transformer.cache_config is not None
|
|
|
|
|
|
def test_wan_a14b_dense_quant_applies_to_both_dits(fake_runtime, monkeypatch):
|
|
# transformer_quant on a pipeline load quantises the dense DiT(s). On CPU the real
|
|
# dense path is unsupported, so stub the two quant seams to record which pipe view
|
|
# each helper saw: BOTH experts must be quantised (via the _SecondDiTView proxy),
|
|
# and status must report the engaged scheme.
|
|
import core.inference.video as video_mod
|
|
|
|
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
|
quantised = []
|
|
|
|
def _fake_quant(
|
|
view,
|
|
target,
|
|
*,
|
|
mode,
|
|
family,
|
|
logger = None,
|
|
):
|
|
# The helper reads view.transformer; record the object it would quantise so the
|
|
# test proves the second expert was reached through the proxy.
|
|
quantised.append(view.transformer)
|
|
return "int8"
|
|
|
|
monkeypatch.setattr(video_mod, "quantize_transformer", _fake_quant)
|
|
|
|
backend = VideoBackend()
|
|
status = backend.load_pipeline(
|
|
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
|
model_kind = "pipeline",
|
|
transformer_quant = "int8",
|
|
)
|
|
pipe = backend._state.pipe
|
|
# Both experts were passed to quantize_transformer, in that order.
|
|
assert quantised == [pipe.transformer, pipe.transformer_2]
|
|
assert status["transformer_quant"] == "int8"
|
|
|
|
|
|
def test_wan_ti2v_dense_quant_applies_to_single_dit(fake_runtime, monkeypatch):
|
|
# A single-DiT pipeline load quantises exactly one transformer.
|
|
import core.inference.video as video_mod
|
|
|
|
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
|
quantised = []
|
|
|
|
def _fake_quant(
|
|
view,
|
|
target,
|
|
*,
|
|
mode,
|
|
family,
|
|
logger = None,
|
|
):
|
|
quantised.append(view.transformer)
|
|
return "fp8"
|
|
|
|
monkeypatch.setattr(video_mod, "quantize_transformer", _fake_quant)
|
|
|
|
backend = VideoBackend()
|
|
status = backend.load_pipeline(
|
|
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
|
model_kind = "pipeline",
|
|
transformer_quant = "fp8",
|
|
)
|
|
assert quantised == [backend._state.pipe.transformer]
|
|
assert status["transformer_quant"] == "fp8"
|
|
|
|
|
|
def test_wan_validate_trusted_repos(fake_runtime):
|
|
# The two Wan base repos are trusted for non-GGUF (pipeline) loads; an unrelated
|
|
# repo carrying the family name is not.
|
|
backend = VideoBackend()
|
|
fam = backend.validate_load_request("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
|
|
assert fam.name == "wan2.2-ti2v-5b"
|
|
fam2 = backend.validate_load_request("Wan-AI/Wan2.2-T2V-A14B-Diffusers", model_kind = "pipeline")
|
|
assert fam2.name == "wan2.2-t2v-a14b"
|
|
with pytest.raises(ValueError, match = "limited to"):
|
|
backend.validate_load_request("evil/wan2.2-ti2v-5b-repack", model_kind = "pipeline")
|
|
# A bad transformer_quant scheme is rejected cheaply at validate time.
|
|
with pytest.raises(ValueError, match = "transformer_quant"):
|
|
backend.validate_load_request(
|
|
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
|
model_kind = "pipeline",
|
|
transformer_quant = "bogus",
|
|
)
|