Add Wan2.2 text-to-video families to the video backend
Register two new video families and wire them through the backend, routes, and frontend picker: wan2.2-ti2v-5b (single DiT) and wan2.2-t2v-a14b (the dual-expert MoE). Both share diffusers' WanPipeline + WanTransformer3DModel + AutoencoderKLWan, which the VideoFamily dataclass already reserved fields for (transformer2_class, is_moe, cfg2_kwarg). Verified against the installed diffusers 0.39.0 before writing code: - WanPipeline, WanTransformer3DModel, and AutoencoderKLWan are all exported from top-level diffusers 0.39.0. - WanPipeline.__call__ (pipeline_wan.py:383) defaults to num_frames=81, num_inference_steps=50, guidance_scale=5.0. guidance_scale_2 DOES exist in 0.39 (line 392) and its check_inputs raises if it is passed when the pipeline's boundary_ratio is None (line 322), so the second guidance is threaded ONLY for the MoE family and only when inspect.signature accepts it (the same gate frame_rate already uses). - The Wan VAE temporal factor is 4 (autoencoder_kl_wan.py scale_factor_temporal), and the pipeline snaps num_frames to 4k+1 (line 493), so frame_step is 4, unlike LTX-2's 8k+1. Sizes patchify at spatial 8 * patch 2 = 16, so resolution_multiple is 16. - boundary_ratio and transformer_2 come from model_index.json: TI2V-5B ships boundary_ratio=null and transformer_2=[null,null] (single DiT), while A14B ships boundary_ratio=0.875 and transformer_2=WanTransformer3DModel (dual DiT). boundary_ratio lives in the pipeline config, so it needs no per-call plumbing. - WanTransformer3DModel declares _repeated_blocks=["WanTransformerBlock"] and inherits CacheMixin (transformer_wan.py:508/551), so regional compile and First-Block-Cache both work. bf16-resident component sizes, measured from each diffusers repo's on-disk safetensors (all stored bf16), feed the auto memory table: TI2V-5B: transformer 20.0, UMT5 text encoder 11.4, VAE 2.8 GB. A14B: two experts 57.2 each (114.3 total), text encoder 11.4, VAE 0.5 GB. Backend changes make the optimisation layers dual-DiT aware: a small _SecondDiTView proxy presents transformer_2 as pipe.transformer so the existing single-DiT helpers (apply_speed_optims, apply_attention_backend, apply_step_cache, quantize_transformer) cover BOTH experts on an is_moe load without forking any helper; single-DiT loads are unchanged (views is just (pipe,)). The two Wan base repos are added to the trusted non-GGUF allowlist. A transformer_quant option is added to the load path, mirroring the image backend's dense torchao fast path: on a pipeline-kind load the dense DiT(s) are quantised in place onto the low-precision tensor cores and the engaged scheme is surfaced in status. generate() threads guidance_2 through the family's cfg2_kwarg when the loaded pipeline accepts it. Routes and Pydantic models gain the optional transformer_quant (load / status) and guidance_2 (generate) fields. The frontend picker gains the two Wan models with 50-step / CFG 5.0 defaults; fps is supplied per family by the backend. Tests extend the fake runtime with WanPipeline and per-DiT transformer fakes (single-DiT and dual-DiT), and cover family detection for both repos, 4k+1 frame snapping, default application, dual-DiT speed/cache/attention/quant coverage on both experts, cfg2 threading gated on the pipeline signature, trusted-repo validation, and the new route fields. Both the standard and the diffusers/torchao-blocked CI-sim runs are green.
This commit is contained in:
Daniel Han
parent
ffce9b8ba8
commit
9f83d387c2
8 changed files with 833 additions and 25 deletions
|
|
@ -55,6 +55,7 @@ from .diffusion_memory import (
|
|||
snapshot_device_memory,
|
||||
)
|
||||
from .diffusion_speed import (
|
||||
SPEED_DEFAULT,
|
||||
SPEED_OFF,
|
||||
apply_speed_optims,
|
||||
resolve_speed_mode,
|
||||
|
|
@ -62,6 +63,11 @@ from .diffusion_speed import (
|
|||
snapshot_backend_flags,
|
||||
)
|
||||
from .diffusion_auto_policy import build_resolved_record
|
||||
from .diffusion_transformer_quant import (
|
||||
dense_transformer_supported,
|
||||
normalize_transformer_quant,
|
||||
quantize_transformer,
|
||||
)
|
||||
from .video_families import (
|
||||
VIDEO_CANCELLED_MSG,
|
||||
VIDEO_NOT_LOADED_MSG,
|
||||
|
|
@ -90,6 +96,10 @@ _TRUSTED_NON_GGUF_VIDEO_REPOS = frozenset(
|
|||
"lightricks/ltx-2",
|
||||
"lightricks/ltx-2.3",
|
||||
"lightricks/ltx-2.3-fp8",
|
||||
# Wan2.2 official diffusers base repos (Wan-AI org): safetensors-only, no
|
||||
# remote code, so allowed as full (pipeline-kind) loads like the LTX-2 bases.
|
||||
"wan-ai/wan2.2-ti2v-5b-diffusers",
|
||||
"wan-ai/wan2.2-t2v-a14b-diffusers",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -151,6 +161,11 @@ class _VideoLoadState:
|
|||
backend_flags: Optional[dict] = None
|
||||
attention_backend: Optional[str] = None
|
||||
transformer_cache: Optional[str] = None
|
||||
# Dense transformer quant actually engaged ("int8" | "fp8" | "nvfp4" | "mxfp8") or
|
||||
# None. Mirrors the image backend's _LoadState.transformer_quant: on a pipeline-kind
|
||||
# load the dense DiT(s) can be torchao-quantised in place onto the low-precision
|
||||
# tensor cores; None means they run at their loaded (bf16) precision.
|
||||
transformer_quant: Optional[str] = None
|
||||
resolved: Optional[dict] = None
|
||||
|
||||
|
||||
|
|
@ -166,6 +181,63 @@ def _progress(phase: Optional[str], **extra: Any) -> dict[str, Any]:
|
|||
return {"phase": phase, **extra}
|
||||
|
||||
|
||||
# ── dual-DiT (Wan2.2-A14B MoE) helpers ────────────────────────────────────────
|
||||
#
|
||||
# The imported optimisation helpers (apply_speed_optims / apply_attention_backend /
|
||||
# apply_step_cache) and the dense quantiser all read ``pipe.transformer`` and act on
|
||||
# that ONE denoiser -- correct for every single-DiT family (LTX-2, Wan2.2-TI2V-5B).
|
||||
# Wan2.2-A14B is a dual-expert MoE: ``transformer`` handles the high-noise steps and
|
||||
# ``transformer_2`` the low-noise steps (pipeline_wan.py routes by boundary_ratio), so
|
||||
# an optimisation applied only to ``transformer`` would leave the second expert eager /
|
||||
# unquantised / on the wrong attention kernel for half the schedule. Rather than fork
|
||||
# each helper, present the second DiT to them AS ``pipe.transformer`` via a thin proxy
|
||||
# and call the helper a second time, so the helpers stay untouched and single-DiT loads
|
||||
# are bit-identical (the proxy is only built for is_moe families).
|
||||
|
||||
|
||||
def _transformer_names(pipe: Any, fam: VideoFamily) -> tuple[str, ...]:
|
||||
"""Attribute names of the denoiser(s) on ``pipe`` to optimise. Just
|
||||
("transformer",) for a single-DiT family; also "transformer_2" for an MoE family
|
||||
whose second expert is actually present (a checkpoint may ship only the first)."""
|
||||
names = ["transformer"]
|
||||
if fam.is_moe and getattr(pipe, "transformer_2", None) is not None:
|
||||
names.append("transformer_2")
|
||||
return tuple(names)
|
||||
|
||||
|
||||
class _SecondDiTView:
|
||||
"""A thin proxy that makes ``pipe.transformer_2`` look like ``pipe.transformer`` to a
|
||||
helper that hardcodes ``getattr(pipe, "transformer")``, while every other attribute
|
||||
(vae, components, __call__, ...) reads through to the real pipe unchanged.
|
||||
|
||||
This lets the existing single-DiT helpers optimise the second expert without a fork:
|
||||
``apply_speed_optims(_SecondDiTView(pipe), ...)`` compiles / caches / sets attention on
|
||||
``transformer_2``. Only ever wrapped around an MoE pipe (guarded by fam.is_moe)."""
|
||||
|
||||
def __init__(self, pipe: Any) -> None:
|
||||
# Store on the instance dict under a name __getattr__ never fires for.
|
||||
object.__setattr__(self, "_pipe", pipe)
|
||||
|
||||
@property
|
||||
def transformer(self) -> Any:
|
||||
return self._pipe.transformer_2
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
# Only reached for attributes not found on the instance/class (i.e. not
|
||||
# ``transformer`` / ``_pipe``), so everything else delegates to the real pipe.
|
||||
return getattr(object.__getattribute__(self, "_pipe"), name)
|
||||
|
||||
|
||||
def _views_for(pipe: Any, fam: VideoFamily) -> tuple[Any, ...]:
|
||||
"""The pipe view(s) to pass through the ``getattr(pipe, "transformer")`` helpers so
|
||||
they cover every denoiser: the real pipe (its ``transformer``), plus a
|
||||
``_SecondDiTView`` (its ``transformer_2``) for a dual-DiT MoE family. A single-DiT
|
||||
load returns just ``(pipe,)``, so its behaviour is unchanged."""
|
||||
if fam.is_moe and getattr(pipe, "transformer_2", None) is not None:
|
||||
return (pipe, _SecondDiTView(pipe))
|
||||
return (pipe,)
|
||||
|
||||
|
||||
class VideoBackend:
|
||||
"""One loaded video pipeline; loads swap it atomically (same model as images)."""
|
||||
|
||||
|
|
@ -189,6 +261,7 @@ class VideoBackend:
|
|||
gguf_filename: Optional[str] = None,
|
||||
family_override: Optional[str] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
) -> VideoFamily:
|
||||
"""Cheap, network-free validation shared by the route and the load path."""
|
||||
kind = resolve_video_model_kind(gguf_filename, model_kind)
|
||||
|
|
@ -210,6 +283,11 @@ class VideoBackend:
|
|||
)
|
||||
if kind in ("gguf", "single_file") and not gguf_filename:
|
||||
raise ValueError("A gguf/single_file load needs the checkpoint filename.")
|
||||
# Reject a malformed transformer_quant scheme cheaply, before the GPU handoff
|
||||
# (normalize_transformer_quant raises ValueError on an unknown scheme). It applies
|
||||
# only on pipeline-kind loads (the dense DiT from the base repo); an ignored value
|
||||
# on a gguf/single_file load is left to the loader, matching the image backend.
|
||||
normalize_transformer_quant(transformer_quant)
|
||||
_ensure_mp4_encoder_available()
|
||||
return fam
|
||||
|
||||
|
|
@ -228,6 +306,7 @@ class VideoBackend:
|
|||
attention_backend: Optional[str] = None,
|
||||
transformer_cache: Optional[str] = None,
|
||||
transformer_cache_threshold: Optional[float] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Validate, then run the (slow) load on a daemon thread. Returns at once."""
|
||||
|
|
@ -237,6 +316,7 @@ class VideoBackend:
|
|||
gguf_filename = gguf_filename,
|
||||
family_override = family_override,
|
||||
model_kind = model_kind,
|
||||
transformer_quant = transformer_quant,
|
||||
)
|
||||
with self._lock:
|
||||
if self._loading is not None and self._loading.error is None:
|
||||
|
|
@ -259,6 +339,7 @@ class VideoBackend:
|
|||
attention_backend = attention_backend,
|
||||
transformer_cache = transformer_cache,
|
||||
transformer_cache_threshold = transformer_cache_threshold,
|
||||
transformer_quant = transformer_quant,
|
||||
model_kind = model_kind,
|
||||
_load_token = token,
|
||||
),
|
||||
|
|
@ -398,6 +479,7 @@ class VideoBackend:
|
|||
attention_backend: Optional[str] = None,
|
||||
transformer_cache: Optional[str] = None,
|
||||
transformer_cache_threshold: Optional[float] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
_load_token: Optional[int] = None,
|
||||
) -> dict[str, Any]:
|
||||
|
|
@ -409,6 +491,7 @@ class VideoBackend:
|
|||
gguf_filename = gguf_filename,
|
||||
family_override = family_override,
|
||||
model_kind = model_kind,
|
||||
transformer_quant = transformer_quant,
|
||||
)
|
||||
kind = resolve_video_model_kind(gguf_filename, model_kind)
|
||||
base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo)
|
||||
|
|
@ -514,26 +597,87 @@ class VideoBackend:
|
|||
clear_gpu_cache()
|
||||
raise RuntimeError("Video load was cancelled or superseded.")
|
||||
|
||||
# For a dual-DiT MoE family (Wan2.2-A14B), every optimisation site below must
|
||||
# cover BOTH experts: ``views`` is (pipe, _SecondDiTView(pipe)) so a helper that
|
||||
# reads ``pipe.transformer`` runs once per denoiser. A single-DiT load resolves to
|
||||
# (pipe,), so it behaves exactly as before.
|
||||
views = _views_for(pipe, fam)
|
||||
|
||||
# ── dense transformer quant (opt-in, pipeline-kind only): load the dense bf16
|
||||
# DiT from the base repo and torchao-quantise it in place onto the low-precision
|
||||
# tensor cores, mirroring the image backend's transformer_quant fast path. Only
|
||||
# the pipeline kind materialises the dense weights (gguf/single_file already carry
|
||||
# their own precision), and only on CUDA + bf16. Best-effort: any failure leaves
|
||||
# the DiT dense. Quant must precede compile (dynamic quant is ~30x slower eager),
|
||||
# so it runs before apply_speed_optims below -- same order as diffusion.py.
|
||||
transformer_quant_engaged: Optional[str] = None
|
||||
if (
|
||||
kind == "pipeline"
|
||||
and normalize_transformer_quant(transformer_quant) is not None
|
||||
and dense_transformer_supported(target)
|
||||
):
|
||||
engaged = []
|
||||
for view in views:
|
||||
# quantize_transformer reads ``pipe.transformer`` and returns the scheme it
|
||||
# engaged (or None); pass each expert's view so both DiTs are quantised with
|
||||
# the same arch-chosen scheme. The family name drives the per-family deny
|
||||
# table (_FAMILY_SCHEME_DENY) exactly as on the image side.
|
||||
scheme = quantize_transformer(
|
||||
view,
|
||||
target,
|
||||
mode = transformer_quant,
|
||||
family = fam.name,
|
||||
logger = logger,
|
||||
)
|
||||
if scheme is not None:
|
||||
engaged.append(scheme)
|
||||
# Report the scheme only if it engaged on every DiT; a partial quant (one expert
|
||||
# dense, one quantised) would run the schedule at mismatched precision, so treat
|
||||
# anything short of full coverage as not engaged.
|
||||
if engaged and len(engaged) == len(views):
|
||||
transformer_quant_engaged = engaged[0]
|
||||
|
||||
# ── optimisation layers, in the image backend's order: speed profile
|
||||
# (compile must precede placement), attention (compile traces it),
|
||||
# placement/offload, then the step cache.
|
||||
effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf")
|
||||
# A torchao-quantised DiT must be compiled (eager dynamic quant is ~30x slower and
|
||||
# would lose to the bf16 it replaced), so force at least the regional-compile
|
||||
# profile when quant engaged and the effective speed was off, matching diffusion.py.
|
||||
if transformer_quant_engaged is not None and effective_speed == SPEED_OFF:
|
||||
logger.info(
|
||||
"video.transformer_quant: forcing speed_mode=default "
|
||||
"(quantized transformer must be compiled; eager is ~30x slower)"
|
||||
)
|
||||
effective_speed = SPEED_DEFAULT
|
||||
backend_flags = snapshot_backend_flags()
|
||||
attention_engaged = apply_attention_backend(
|
||||
pipe,
|
||||
select_attention_backend(
|
||||
target, attention_backend, speed_active = effective_speed != SPEED_OFF
|
||||
),
|
||||
logger = logger,
|
||||
)
|
||||
speed_optims = apply_speed_optims(
|
||||
pipe,
|
||||
target,
|
||||
is_gguf = kind == "gguf",
|
||||
family = fam,
|
||||
speed_mode = effective_speed,
|
||||
offload_active = plan.offload_policy != "none",
|
||||
)
|
||||
attention_engaged = None
|
||||
speed_optims: tuple = ()
|
||||
for view in views:
|
||||
# apply_attention_backend / apply_speed_optims both act on ``view.transformer``;
|
||||
# calling them once per view sets the kernel and compiles each expert. The
|
||||
# engaged values match across experts (same device/family/mode), so record the
|
||||
# first pass; a dense torchao transformer on the pipeline path is not a GGUF one,
|
||||
# so is_gguf keys off the load kind (gguf) AND no quant having engaged.
|
||||
gguf_transformer = kind == "gguf" and transformer_quant_engaged is None
|
||||
engaged = apply_attention_backend(
|
||||
view,
|
||||
select_attention_backend(
|
||||
target, attention_backend, speed_active = effective_speed != SPEED_OFF
|
||||
),
|
||||
logger = logger,
|
||||
)
|
||||
applied = apply_speed_optims(
|
||||
view,
|
||||
target,
|
||||
is_gguf = gguf_transformer,
|
||||
family = fam,
|
||||
speed_mode = effective_speed,
|
||||
offload_active = plan.offload_policy != "none",
|
||||
)
|
||||
if view is pipe:
|
||||
attention_engaged = engaged
|
||||
speed_optims = tuple(k for k, v in applied.items() if v)
|
||||
offload_policy, vae_tiling = apply_memory_plan(
|
||||
pipe, plan, device = device, logger = logger
|
||||
)
|
||||
|
|
@ -545,12 +689,18 @@ class VideoBackend:
|
|||
vae_tiling = True
|
||||
except Exception as exc: # noqa: BLE001 -- tiling is an optimisation only
|
||||
logger.warning("video.vae_tiling_failed: %s", exc)
|
||||
cache_engaged = apply_step_cache(
|
||||
pipe,
|
||||
mode = normalize_transformer_cache(transformer_cache),
|
||||
threshold = transformer_cache_threshold,
|
||||
logger = logger,
|
||||
)
|
||||
cache_engaged = None
|
||||
for view in views:
|
||||
# apply_step_cache engages First-Block-Cache on ``view.transformer``; run it per
|
||||
# expert so both denoisers cache. The engaged mode is identical across experts.
|
||||
engaged = apply_step_cache(
|
||||
view,
|
||||
mode = normalize_transformer_cache(transformer_cache),
|
||||
threshold = transformer_cache_threshold,
|
||||
logger = logger,
|
||||
)
|
||||
if view is pipe:
|
||||
cache_engaged = engaged
|
||||
|
||||
resolved = build_resolved_record(
|
||||
{
|
||||
|
|
@ -562,7 +712,9 @@ class VideoBackend:
|
|||
"speed_mode": (
|
||||
speed_mode,
|
||||
effective_speed,
|
||||
"GGUF video loads default to the near-lossless compile profile",
|
||||
"quantized transformer requires compile"
|
||||
if transformer_quant_engaged is not None
|
||||
else "GGUF video loads default to the near-lossless compile profile",
|
||||
),
|
||||
"attention_backend": (
|
||||
attention_backend,
|
||||
|
|
@ -574,6 +726,13 @@ class VideoBackend:
|
|||
cache_engaged or "off",
|
||||
"step cache engages on many-step schedules only",
|
||||
),
|
||||
"transformer_quant": (
|
||||
transformer_quant,
|
||||
transformer_quant_engaged or "off",
|
||||
"dense DiT(s) torchao-quantised onto the low-precision tensor cores"
|
||||
if transformer_quant_engaged is not None
|
||||
else "not engaged (dense bf16 DiT loaded)",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -599,11 +758,13 @@ class VideoBackend:
|
|||
backend_flags = backend_flags,
|
||||
attention_backend = attention_engaged,
|
||||
transformer_cache = cache_engaged,
|
||||
transformer_quant = transformer_quant_engaged,
|
||||
resolved = resolved,
|
||||
)
|
||||
logger.info(
|
||||
"video.loaded: %s (%s, %s, offload=%s, speed=%s)",
|
||||
"video.loaded: %s (%s, %s, offload=%s, speed=%s, quant=%s)",
|
||||
repo_id, fam.name, kind, offload_policy, effective_speed,
|
||||
transformer_quant_engaged or "off",
|
||||
)
|
||||
return self.status()
|
||||
|
||||
|
|
@ -638,6 +799,7 @@ class VideoBackend:
|
|||
fps: Optional[int] = None,
|
||||
steps: Optional[int] = None,
|
||||
guidance: Optional[float] = None,
|
||||
guidance_2: Optional[float] = None,
|
||||
seed: Optional[int] = None,
|
||||
) -> dict[str, Any]:
|
||||
import torch
|
||||
|
|
@ -685,6 +847,21 @@ class VideoBackend:
|
|||
# pipelines fix their own rate and fps only matters at export.
|
||||
if "frame_rate" in call_params:
|
||||
kwargs["frame_rate"] = float(out_fps)
|
||||
# Dual-DiT MoE (Wan2.2-A14B): the low-noise expert (transformer_2) has its
|
||||
# own guidance kwarg (cfg2_kwarg = "guidance_scale_2"). Thread it only when
|
||||
# the loaded family declares one AND the pipeline signature accepts it (the
|
||||
# same inspect.signature gate frame_rate uses), so a single-DiT pipeline is
|
||||
# never handed a kwarg its check_inputs would reject. WanPipeline raises if
|
||||
# guidance_scale_2 is passed to a pipeline with boundary_ratio=None
|
||||
# (pipeline_wan.py:322), so the gate is BOTH the family flag and the
|
||||
# signature: TI2V-5B has no cfg2_kwarg, so it never reaches here. A None
|
||||
# request lets the pipeline default it (to guidance_scale) itself.
|
||||
if (
|
||||
fam.cfg2_kwarg
|
||||
and fam.cfg2_kwarg in call_params
|
||||
and guidance_2 is not None
|
||||
):
|
||||
kwargs[fam.cfg2_kwarg] = float(guidance_2)
|
||||
|
||||
started = time.monotonic()
|
||||
self._gen = {
|
||||
|
|
@ -824,6 +1001,7 @@ class VideoBackend:
|
|||
"speed_optims": [],
|
||||
"attention_backend": None,
|
||||
"transformer_cache": None,
|
||||
"transformer_quant": None,
|
||||
"has_audio": False,
|
||||
"defaults": None,
|
||||
"resolved": None,
|
||||
|
|
@ -847,6 +1025,7 @@ class VideoBackend:
|
|||
"speed_optims": list(state.speed_optims),
|
||||
"attention_backend": state.attention_backend,
|
||||
"transformer_cache": state.transformer_cache,
|
||||
"transformer_quant": state.transformer_quant,
|
||||
"has_audio": fam.has_audio,
|
||||
"defaults": {
|
||||
"steps": default_steps,
|
||||
|
|
|
|||
|
|
@ -108,6 +108,88 @@ _FAMILIES: tuple[VideoFamily, ...] = (
|
|||
bf16_components_gb = (37.8, 50.4, 5.5),
|
||||
gguf_repo = "unsloth/LTX-2.3-GGUF",
|
||||
),
|
||||
# Wan2.2-TI2V-5B (diffusers >= 0.35, verified on 0.39): a ~5B single-stream
|
||||
# video DiT (WanPipeline + WanTransformer3DModel + AutoencoderKLWan + a UMT5
|
||||
# text encoder). No audio, no second expert -- its model_index.json ships
|
||||
# ``boundary_ratio: null`` and ``transformer_2: [null, null]``, so it is a
|
||||
# plain single-DiT family (is_moe left False). The Wan VAE has a temporal
|
||||
# compression of 4, so valid frame counts are 4k+1 (frame_step = 4), which
|
||||
# matches the pipeline's own ``num_frames % vae_scale_factor_temporal == 1``
|
||||
# check (pipeline_wan.py:493). The pipeline defaults to 50 steps / CFG 5, but
|
||||
# the 5B TI2V card ships the 720p-class few-step recipe, so the picker default
|
||||
# (see _VIDEO_GENERATION_DEFAULTS) uses the pipeline's 50/5 while the UI presets
|
||||
# target 720p at 24 fps (the model card's playback rate).
|
||||
VideoFamily(
|
||||
name = "wan2.2-ti2v-5b",
|
||||
pipeline_class = "WanPipeline",
|
||||
transformer_class = "WanTransformer3DModel",
|
||||
base_repo = "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
# "wan2.2-5b" and "wan-ti2v" are the short ids the picker / GGUF filenames
|
||||
# use; "wan2.2-ti2v" catches the diffusers repo stem without the "-5b".
|
||||
aliases = ("wan2.2-5b", "wan-ti2v", "wan2.2-ti2v", "wan-ti2v-5b"),
|
||||
has_audio = False,
|
||||
default_steps = 50,
|
||||
default_guidance = 5.0,
|
||||
# 121 frames at 24 fps is ~5s, the model card's headline clip length; on the
|
||||
# 4k+1 lattice (121 = 4*30 + 1) it needs no snapping.
|
||||
default_num_frames = 121,
|
||||
default_fps = 24,
|
||||
# Wan VAE temporal factor is 4 (autoencoder_kl_wan.py scale_factor_temporal),
|
||||
# so valid counts are 4k+1, unlike LTX-2's 8k+1.
|
||||
frame_step = 4,
|
||||
# The pipeline patchifies at vae_scale_factor_spatial (8) * patch_size (2) = 16,
|
||||
# and its check_inputs rejects non-/16 sizes (pipeline_wan.py:291); 16 keeps every
|
||||
# preset valid and the snap silent.
|
||||
resolution_multiple = 16,
|
||||
# 720p-class presets: 1280x704 landscape (the card's target), its vertical variant,
|
||||
# and a square. The first preset is the default the loader plans memory against.
|
||||
resolution_presets = ((1280, 704), (704, 1280), (960, 960), (832, 480)),
|
||||
# Measured from the diffusers repo (safetensors on disk, all stored bf16, so these
|
||||
# are the bf16-resident sizes): transformer 20.0, UMT5 text encoder 11.4, VAE 2.8.
|
||||
bf16_components_gb = (20.0, 11.4, 2.8),
|
||||
gguf_repo = "unsloth/Wan2.2-TI2V-5B-GGUF",
|
||||
),
|
||||
# Wan2.2-T2V-A14B (diffusers >= 0.35, verified on 0.39): the dual-expert MoE.
|
||||
# Its model_index.json lists BOTH ``transformer`` and ``transformer_2`` as
|
||||
# WanTransformer3DModel and sets ``boundary_ratio: 0.875``; the pipeline routes
|
||||
# the high-noise steps (timestep >= boundary) through ``transformer`` at
|
||||
# guidance_scale and the low-noise steps through ``transformer_2`` at
|
||||
# guidance_scale_2 (pipeline_wan.py:584-603). ``guidance_scale_2`` exists in
|
||||
# 0.39 (pipeline_wan.py:392) and is only accepted when boundary_ratio is set
|
||||
# (its check_inputs raises otherwise, pipeline_wan.py:322), so cfg2_kwarg is
|
||||
# threaded ONLY for this family. boundary_ratio itself lives in the pipeline
|
||||
# config (loaded from model_index.json), so no per-generation plumbing is needed.
|
||||
VideoFamily(
|
||||
name = "wan2.2-t2v-a14b",
|
||||
pipeline_class = "WanPipeline",
|
||||
transformer_class = "WanTransformer3DModel",
|
||||
base_repo = "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
aliases = ("wan2.2-14b", "wan-t2v", "wan2.2-t2v", "wan-t2v-a14b", "wan-a14b"),
|
||||
has_audio = False,
|
||||
# The second expert is the same class; is_moe drives the dual-DiT optimisation
|
||||
# layers (speed / attention / cache / quant apply to BOTH transformers), and
|
||||
# cfg2_kwarg names the pipeline kwarg carrying transformer_2's guidance.
|
||||
transformer2_class = "WanTransformer3DModel",
|
||||
is_moe = True,
|
||||
cfg2_kwarg = "guidance_scale_2",
|
||||
default_steps = 50,
|
||||
default_guidance = 5.0,
|
||||
# 81 frames at 16 fps is ~5s (81 = 4*20 + 1), the A14B card's default clip.
|
||||
default_num_frames = 81,
|
||||
# The A14B card runs at 16 fps (vs the 5B TI2V's 24), per its model_index /
|
||||
# model card; export uses this rate.
|
||||
default_fps = 16,
|
||||
frame_step = 4,
|
||||
resolution_multiple = 16,
|
||||
# 480p and 720p presets (landscape, vertical, square), the two resolutions the
|
||||
# A14B card documents. 832x480 is the native 480p; 1280x704 the 720p target.
|
||||
resolution_presets = ((1280, 704), (832, 480), (480, 832), (704, 1280)),
|
||||
# Measured from the diffusers repo (bf16 on disk): each DiT 57.2, so both experts
|
||||
# total ~114.3; UMT5 text encoder 11.4; VAE 0.5. The two-expert DiT total is the
|
||||
# memory headline (~114 GB bf16-resident before offload).
|
||||
bf16_components_gb = (114.3, 11.4, 0.5),
|
||||
gguf_repo = "unsloth/Wan2.2-T2V-A14B-GGUF",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -174,6 +256,12 @@ def snap_video_size(fam: VideoFamily, width: int, height: int) -> tuple[int, int
|
|||
_VIDEO_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
|
||||
("distilled", 8, 1.0),
|
||||
("ltx", 40, 4.0),
|
||||
# Wan2.2 pipelines default to 50 steps at CFG 5.0 (WanPipeline.__call__:
|
||||
# num_inference_steps = 50, guidance_scale = 5.0, verified in diffusers 0.39).
|
||||
# Both TI2V-5B and A14B share these; the substring "wan" catches the picked id
|
||||
# and the base repo. A future distilled Wan GGUF is caught by the "distilled"
|
||||
# row above (listed first), exactly as the LTX-2.3 distilled checkpoints are.
|
||||
("wan", 50, 5.0),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2328,6 +2328,19 @@ class VideoLoadRequest(BaseModel):
|
|||
description = "FBCache residual threshold (higher = skips more steps = faster, lower "
|
||||
"quality). null auto-picks the family default.",
|
||||
)
|
||||
transformer_quant: Optional[Literal["auto", "none", "off", "int8", "fp8", "nvfp4", "mxfp8"]] = (
|
||||
Field(
|
||||
None,
|
||||
description = "Quantise the dense DiT(s) on a full-pipeline load. On a diffusers "
|
||||
"pipeline load the dense bf16 transformer(s) are torchao-quantised in place onto "
|
||||
"the low-precision tensor cores (data-center fp8, consumer/Ampere int8), which is "
|
||||
"faster than running dense bf16. For a dual-expert MoE family (Wan2.2-A14B) BOTH "
|
||||
"experts are quantised with the same scheme. null/none/off keeps the DiT(s) at "
|
||||
"their loaded precision; an explicit scheme forces it. Needs CUDA + bf16; ignored "
|
||||
"on gguf/single_file loads (they carry their own precision). Mirrors the image "
|
||||
"backend's transformer_quant field.",
|
||||
)
|
||||
)
|
||||
|
||||
@field_validator("attention_backend", mode = "before")
|
||||
@classmethod
|
||||
|
|
@ -2368,6 +2381,15 @@ class VideoGenerateRequest(BaseModel):
|
|||
guidance: Optional[float] = Field(
|
||||
None, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale (default per model)"
|
||||
)
|
||||
guidance_2: Optional[float] = Field(
|
||||
None,
|
||||
ge = 0.0,
|
||||
le = 20.0,
|
||||
description = "Low-noise-stage guidance scale for a dual-expert MoE family (Wan2.2-A14B): "
|
||||
"the guidance the second transformer uses on the low-noise denoise steps. null lets the "
|
||||
"pipeline default it to the main guidance. Ignored by single-DiT families (their pipeline "
|
||||
"signature has no second guidance kwarg).",
|
||||
)
|
||||
# le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript
|
||||
# rounds integers above Number.MAX_SAFE_INTEGER -- a restored recipe would then
|
||||
# generate a different clip. Random seeds are already masked to this range.
|
||||
|
|
@ -2473,6 +2495,12 @@ class VideoStatusResponse(BaseModel):
|
|||
"_native_cudnn), or null for the default SDPA",
|
||||
)
|
||||
transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null")
|
||||
transformer_quant: Optional[str] = Field(
|
||||
None,
|
||||
description = "Dense transformer quant engaged on a pipeline load: int8 | fp8 | nvfp4 | "
|
||||
"mxfp8 | null (null = the DiT(s) run at their loaded bf16 precision). For a dual-expert "
|
||||
"MoE family both experts share the reported scheme.",
|
||||
)
|
||||
has_audio: bool = Field(
|
||||
False, description = "Whether the loaded family produces a synchronized audio track"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -94,6 +94,7 @@ async def load_video_model(
|
|||
gguf_filename = request.gguf_filename,
|
||||
family_override = request.family_override,
|
||||
model_kind = request.model_kind,
|
||||
transformer_quant = request.transformer_quant,
|
||||
)
|
||||
# Refuse while training is running: a multi-GB video pipeline would compete
|
||||
# with the training subprocess for VRAM. Mirrors the image-load guard.
|
||||
|
|
@ -121,6 +122,7 @@ async def load_video_model(
|
|||
attention_backend = request.attention_backend,
|
||||
transformer_cache = request.transformer_cache,
|
||||
transformer_cache_threshold = request.transformer_cache_threshold,
|
||||
transformer_quant = request.transformer_quant,
|
||||
model_kind = request.model_kind,
|
||||
)
|
||||
return VideoStatusResponse(**status_dict)
|
||||
|
|
@ -158,6 +160,7 @@ async def generate_video(
|
|||
fps = request.fps,
|
||||
steps = request.steps,
|
||||
guidance = request.guidance,
|
||||
guidance_2 = request.guidance_2,
|
||||
seed = request.seed,
|
||||
)
|
||||
except ValueError as exc:
|
||||
|
|
|
|||
|
|
@ -116,6 +116,163 @@ class _FakeTransformer:
|
|||
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")
|
||||
|
|
@ -131,6 +288,10 @@ def fake_runtime(monkeypatch):
|
|||
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)
|
||||
|
|
@ -341,3 +502,197 @@ def test_generate_progress_and_cancel_idle(fake_runtime):
|
|||
|
||||
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",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -39,11 +39,68 @@ def test_detect_override_and_unknown():
|
|||
assert detect_video_family("x", override = "ltx-2").name == "ltx-2"
|
||||
assert detect_video_family("x", override = "ltx2").name == "ltx-2"
|
||||
assert detect_video_family("x", override = "nope") is None
|
||||
assert detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") is None # V3 family
|
||||
# A short alias must not match inside an unrelated word.
|
||||
assert detect_video_family("someorg/deluxtreme-model") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"repo_id",
|
||||
[
|
||||
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"wan-ai/wan2.2-ti2v-5b-diffusers",
|
||||
"unsloth/Wan2.2-TI2V-5B-GGUF",
|
||||
"some/dir/wan2.2-ti2v-5b-Q4_K_M.gguf",
|
||||
],
|
||||
)
|
||||
def test_detect_wan_ti2v_5b(repo_id):
|
||||
# The TI2V-5B repo ids route to the single-DiT Wan family (no MoE, no audio).
|
||||
fam = detect_video_family(repo_id)
|
||||
assert fam is not None and fam.name == "wan2.2-ti2v-5b"
|
||||
assert fam.pipeline_class == "WanPipeline"
|
||||
assert fam.transformer_class == "WanTransformer3DModel"
|
||||
assert fam.is_moe is False
|
||||
assert fam.cfg2_kwarg is None
|
||||
assert fam.has_audio is False
|
||||
assert fam.frame_step == 4 # Wan VAE temporal factor is 4 (4k+1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"repo_id",
|
||||
[
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"wan-ai/wan2.2-t2v-a14b-diffusers",
|
||||
"unsloth/Wan2.2-T2V-A14B-GGUF",
|
||||
"some/dir/wan2.2-t2v-a14b-Q4_K_M.gguf",
|
||||
],
|
||||
)
|
||||
def test_detect_wan_t2v_a14b(repo_id):
|
||||
# The A14B repo ids route to the dual-expert MoE family: a second DiT + a second
|
||||
# guidance kwarg (guidance_scale_2, verified present in diffusers 0.39).
|
||||
fam = detect_video_family(repo_id)
|
||||
assert fam is not None and fam.name == "wan2.2-t2v-a14b"
|
||||
assert fam.pipeline_class == "WanPipeline"
|
||||
assert fam.transformer2_class == "WanTransformer3DModel"
|
||||
assert fam.is_moe is True
|
||||
assert fam.cfg2_kwarg == "guidance_scale_2"
|
||||
assert fam.has_audio is False
|
||||
assert fam.frame_step == 4
|
||||
|
||||
|
||||
def test_detect_wan_overrides():
|
||||
# Short aliases the picker / GGUF filenames use resolve to the right family.
|
||||
assert detect_video_family("x", override = "wan2.2-5b").name == "wan2.2-ti2v-5b"
|
||||
assert detect_video_family("x", override = "wan-ti2v").name == "wan2.2-ti2v-5b"
|
||||
assert detect_video_family("x", override = "wan2.2-14b").name == "wan2.2-t2v-a14b"
|
||||
assert detect_video_family("x", override = "wan-t2v").name == "wan2.2-t2v-a14b"
|
||||
|
||||
|
||||
def test_wan_and_ltx_do_not_cross_route():
|
||||
# LTX ids must never resolve to a Wan family and vice versa (separate engines).
|
||||
assert detect_video_family("Lightricks/LTX-2").name == "ltx-2"
|
||||
assert detect_video_family("unsloth/LTX-2.3-GGUF").name == "ltx-2"
|
||||
assert detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers").name == "wan2.2-ti2v-5b"
|
||||
|
||||
|
||||
def test_sentinels_are_video_specific():
|
||||
# The routes match these EXACTLY for 409s; they must not collide with the
|
||||
# image sentinels or a video 409 would be mis-attributed.
|
||||
|
|
@ -89,7 +146,50 @@ def test_generation_defaults_distilled_vs_dev():
|
|||
|
||||
|
||||
def test_supported_names():
|
||||
assert supported_video_family_names() == ("ltx-2",)
|
||||
assert supported_video_family_names() == (
|
||||
"ltx-2",
|
||||
"wan2.2-ti2v-5b",
|
||||
"wan2.2-t2v-a14b",
|
||||
)
|
||||
|
||||
|
||||
def test_wan_snap_num_frames_4k_plus_1():
|
||||
# Wan's temporal factor is 4, so valid counts are 4k+1 (not LTX-2's 8k+1).
|
||||
fam = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
|
||||
assert fam.frame_step == 4
|
||||
assert snap_num_frames(fam, 81) == 81 # 4*20 + 1, on-lattice
|
||||
assert snap_num_frames(fam, 121) == 121 # 4*30 + 1
|
||||
assert snap_num_frames(fam, 120) == 117 # floors to 4*29 + 1
|
||||
assert snap_num_frames(fam, 3) == 1 # below the first stride floors to 1
|
||||
assert snap_num_frames(fam, 5) == 5 # 4*1 + 1
|
||||
|
||||
|
||||
def test_wan_snap_video_size_16():
|
||||
# Wan patchifies at spatial factor 8 * patch 2 = 16; sizes floor to /16.
|
||||
fam = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers")
|
||||
assert fam.resolution_multiple == 16
|
||||
assert snap_video_size(fam, 1280, 704) == (1280, 704) # on-grid preset
|
||||
assert snap_video_size(fam, 1000, 700) == (992, 688)
|
||||
|
||||
|
||||
def test_wan_generation_defaults():
|
||||
# Both Wan families default to the pipeline's 50 steps / CFG 5.0.
|
||||
assert default_video_generation_params(None, "Wan-AI/Wan2.2-TI2V-5B-Diffusers") == (50, 5.0)
|
||||
assert default_video_generation_params(None, "Wan-AI/Wan2.2-T2V-A14B-Diffusers") == (50, 5.0)
|
||||
# A GGUF filename carrying the family name still lands on the Wan defaults.
|
||||
assert default_video_generation_params(
|
||||
"wan2.2-ti2v-5b-Q4_K_M.gguf", "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
||||
) == (50, 5.0)
|
||||
|
||||
|
||||
def test_wan_size_tables_present():
|
||||
ti2v = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
|
||||
a14b = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers")
|
||||
assert ti2v.bf16_components_gb is not None and a14b.bf16_components_gb is not None
|
||||
# The A14B DiT total (two experts) dwarfs the single TI2V-5B DiT.
|
||||
assert a14b.bf16_components_gb[0] > ti2v.bf16_components_gb[0] * 3
|
||||
# A portrait preset is offered for the 5B (a vertical option per the task).
|
||||
assert any(h > w for (w, h) in ti2v.resolution_presets)
|
||||
|
||||
|
||||
def test_family_size_table_present():
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@ def _unloaded_status():
|
|||
"speed_optims": [],
|
||||
"attention_backend": None,
|
||||
"transformer_cache": None,
|
||||
"transformer_quant": None,
|
||||
"has_audio": False,
|
||||
"defaults": None,
|
||||
"resolved": None,
|
||||
|
|
@ -70,6 +71,7 @@ class _FakeBackend:
|
|||
gguf_filename = None,
|
||||
family_override = None,
|
||||
model_kind = None,
|
||||
transformer_quant = None,
|
||||
):
|
||||
# Mirror the real backend's cheap validation so the route's
|
||||
# validate-before-evict ordering is exercised.
|
||||
|
|
@ -244,6 +246,41 @@ def test_load_threads_options_through_to_backend(client):
|
|||
assert kwargs.get("transformer_cache_threshold") == 0.1
|
||||
|
||||
|
||||
def test_load_threads_transformer_quant_and_guidance_2(client):
|
||||
# The new load-time transformer_quant field reaches the backend, and the new
|
||||
# per-generation guidance_2 field reaches generate() (dual-DiT MoE second guidance).
|
||||
resp = client.post(
|
||||
"/api/inference/video/load",
|
||||
json = {
|
||||
"model_path": "unsloth/LTX-2.3-GGUF",
|
||||
"gguf_filename": "q.gguf",
|
||||
"transformer_quant": "fp8",
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
kwargs = video_module.get_video_backend().last_load_kwargs
|
||||
assert kwargs.get("transformer_quant") == "fp8"
|
||||
|
||||
gen = client.post(
|
||||
"/api/inference/video/generate",
|
||||
json = {"prompt": "a sloth", "guidance": 5.0, "guidance_2": 3.0},
|
||||
)
|
||||
assert gen.status_code == 200
|
||||
|
||||
|
||||
def test_load_rejects_bad_transformer_quant_422(client):
|
||||
# transformer_quant is a Literal, so an unknown scheme is a 422 at request validation.
|
||||
resp = client.post(
|
||||
"/api/inference/video/load",
|
||||
json = {
|
||||
"model_path": "unsloth/LTX-2.3-GGUF",
|
||||
"gguf_filename": "q.gguf",
|
||||
"transformer_quant": "bogus",
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 422
|
||||
|
||||
|
||||
def test_load_progress_route(client):
|
||||
idle = client.get("/api/inference/video/load-progress")
|
||||
assert idle.status_code == 200 and idle.json()["phase"] is None
|
||||
|
|
|
|||
|
|
@ -67,6 +67,11 @@ import {
|
|||
type PipelineSpec = { kind: "pipeline"; filename?: string };
|
||||
const PIPELINE_MODELS: Record<string, PipelineSpec> = {
|
||||
"Lightricks/LTX-2": { kind: "pipeline" },
|
||||
// Wan2.2 diffusers base repos (no GGUF variant yet): loaded as full pipelines. TI2V-5B
|
||||
// is a single-DiT 720p-class model; T2V-A14B is the dual-expert MoE. The backend gates
|
||||
// these to the Wan-AI base repos (see _TRUSTED_NON_GGUF_VIDEO_REPOS).
|
||||
"Wan-AI/Wan2.2-TI2V-5B-Diffusers": { kind: "pipeline" },
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": { kind: "pipeline" },
|
||||
};
|
||||
|
||||
// A curated GGUF picker entry: isGguf true expands its .gguf files in the quant expander
|
||||
|
|
@ -91,6 +96,16 @@ const pipelineModel = (id: string, name: string, description: string): ModelOpti
|
|||
const VIDEO_MODELS: ModelOption[] = [
|
||||
ggufModel("unsloth/LTX-2.3-GGUF", "LTX 2.3 distilled"),
|
||||
pipelineModel("Lightricks/LTX-2", "LTX 2 (base, bf16)", "Text-to-video with audio · Safetensors"),
|
||||
pipelineModel(
|
||||
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"Wan 2.2 TI2V 5B",
|
||||
"Text-to-video 720p · Safetensors",
|
||||
),
|
||||
pipelineModel(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"Wan 2.2 T2V A14B (MoE)",
|
||||
"Text-to-video, dual-expert · Safetensors",
|
||||
),
|
||||
];
|
||||
|
||||
// Per-model generation defaults (steps + guidance), matched by repo-id substring, most
|
||||
|
|
@ -102,6 +117,9 @@ const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }>
|
|||
// "distilled" before the generic "ltx": the distilled model runs at 8 steps, guidance 1.
|
||||
{ match: "distilled", steps: 8, guidance: 1 },
|
||||
{ match: "ltx", steps: 40, guidance: 4 },
|
||||
// Wan2.2 pipelines default to 50 steps at CFG 5.0 (WanPipeline defaults, verified in
|
||||
// diffusers 0.39). The backend supplies the fps per family (24 for TI2V-5B, 16 for A14B).
|
||||
{ match: "wan", steps: 50, guidance: 5 },
|
||||
];
|
||||
|
||||
function defaultsFor(repoId: string): { steps: number; guidance: number } {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue