unsloth/studio/backend/core/inference/video_families.py
Daniel Han eee191d30e Tighten diffusion comments (second pass)
Collapse the remaining multi-line comment blocks in the video page, training routes and service, sd.cpp server and installer, memory and speed planners, and the shared request models. Comments only, no code or behaviour changes.
2026-07-27 12:01:44 +00:00

275 lines
15 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
"""Pure helpers for text-to-video model identification.
The video registry mirrors ``diffusion_families`` (no torch/diffusers imports, so
everything unit-tests without the heavy runtime) but is a SEPARATE registry with a
separate backend: video pipelines take frame/fps arguments, return frame stacks
(and, for LTX-2, synchronized audio) instead of PIL images, and their artifacts are
MP4s. Keeping the registries apart means neither picker can mis-route a checkpoint
to the wrong engine.
A video checkpoint published as a single-file GGUF only carries the DiT weights;
the VAE / text encoder / connectors / vocoder come from the companion diffusers
base repo, exactly like the image GGUF path.
"""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from typing import Optional
# Runtime->route contract: routes match these EXACTLY for a 409 instead of a 500.
VIDEO_NOT_LOADED_MSG = "No video model is loaded."
VIDEO_CANCELLED_MSG = "Video generation was cancelled."
VIDEO_GENERATION_BUSY_MSG = "A video generation is already in progress."
@dataclass(frozen = True)
class VideoFamily:
name: str
pipeline_class: str
transformer_class: str
base_repo: str
# Pipeline kwarg carrying the guidance value.
cfg_kwarg: str = "guidance_scale"
# The pipe attribute holding the denoiser (all current video families are DiTs).
denoiser_attr: str = "transformer"
# Extra lowercased substrings (besides ``name``) that map a repo id here.
aliases: tuple[str, ...] = field(default_factory = tuple)
# True when the pipeline returns synchronized audio (LTX-2): export muxes the track and size estimates count the audio VAE + vocoder.
has_audio: bool = False
# Wan2.2-A14B dual-expert MoE: a second DiT (transformer_2) handles the low-noise steps with its own guidance kwarg. None/False for single-DiT.
transformer2_class: Optional[str] = None
is_moe: bool = False
cfg2_kwarg: Optional[str] = None
# HunyuanVideo-1.5 guidance: __call__ takes NO guidance kwarg; CFG lives on a ``guider`` component whose guidance_scale is set per request. When True, generate() writes pipe.guider.
guidance_via_guider: bool = False
# Generation defaults + shape. ``frame_step`` is the temporal compression: a valid frame count is k*frame_step + 1, so requests are snapped BEFORE latents are allocated.
default_steps: int = 40
default_guidance: float = 4.0
default_num_frames: int = 121
default_fps: int = 24
frame_step: int = 8
# Width/height must be divisible by this (LTX-2's pipeline rejects non-/32).
resolution_multiple: int = 32
# (width, height) UI presets, landscape first; the first is the default.
resolution_presets: tuple[tuple[int, int], ...] = ((768, 512),)
# Component bf16-RESIDENT sizes in decimal GB (denoiser(s), text encoder, VAE + audio companions): what sits on device after the dtype cast, not the download size.
bf16_components_gb: Optional[tuple[float, float, float]] = None
# True when the DiT compiles cleanly with regional torch.compile (declares _repeated_blocks).
supports_torch_compile: bool = True
# Video DiTs are bf16-native, so fp16 promotes to float32; defaults True.
fp16_incompatible: bool = True
# Wan VAE decodes in float32 (loading it bf16 causes banding / black frames), so when True the loader pins it back to fp32. Its bf16_components_gb term is already the fp32 size.
vae_force_fp32: bool = False
# Curated GGUF repo for the picker (the DiT as single-file GGUF quants).
gguf_repo: Optional[str] = None
# Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples; same semantics as DiffusionFamily.te_prequant_repos.
te_prequant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
_FAMILIES: tuple[VideoFamily, ...] = (
# LTX-2 (diffusers >= 0.39): ~19B single-stream video DiT generating synchronized audio + video in one pass. The Gemma3-12B text encoder is stored fp32 on the hub (~49 GB download, ~24 GB resident as bf16).
# Base repo carries the dev config (40 steps, CFG 4); distilled runs few-step.
VideoFamily(
name = "ltx-2",
pipeline_class = "LTX2Pipeline",
transformer_class = "LTX2VideoTransformer3DModel",
base_repo = "Lightricks/LTX-2",
aliases = ("ltx-2.3", "ltx2", "ltx-video", "ltxv", "ltx"),
has_audio = True,
default_steps = 40,
default_guidance = 4.0,
default_num_frames = 121,
default_fps = 24,
frame_step = 8,
resolution_multiple = 32,
# 768x512 native default; 1216x704 the card's quality target; 704x1216 vertical.
resolution_presets = ((768, 512), (1216, 704), (704, 1216), (512, 768)),
# transformer 37.8 bf16; Gemma3-12B TE ~24.4 bf16 RESIDENT (the hub stores it fp32, ~49 GB download, but the pipeline loads torch_dtype=bf16); VAE 2.4 + connectors 2.9 + audio 0.2. The old 50.4 figure double-counted the fp32 store and pushed auto toward offload.
bf16_components_gb = (37.8, 24.4, 5.5),
gguf_repo = "unsloth/LTX-2.3-GGUF",
# Pre-cast Gemma3-12B TE (fp32 ~49 GB on the hub, pre-cast ~13.2 GB): the biggest download win.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/LTX-2-FP8"),),
),
# Wan2.2-TI2V-5B (diffusers >= 0.35, verified on 0.39): ~5B single-stream video DiT (UMT5 text encoder). No audio, no second expert, so single-DiT. Wan VAE temporal compression 4 gives valid frame counts 4k+1. Pipeline defaults 50 steps / CFG 5; UI presets target 720p at 24 fps.
VideoFamily(
name = "wan2.2-ti2v-5b",
pipeline_class = "WanPipeline",
transformer_class = "WanTransformer3DModel",
base_repo = "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
# "wan2.2-5b"/"wan-ti2v" are the picker/GGUF short ids; "wan2.2-ti2v" catches the repo stem.
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 ~5s; on the 4k+1 lattice (121 = 4*30 + 1) it needs no snapping.
default_num_frames = 121,
default_fps = 24,
# Wan VAE temporal factor 4, so valid counts are 4k+1.
frame_step = 4,
# TI2V-5B VAE is 16x spatial + patch 2, so WanPipeline floors H/W to 32; snap to 32 so the recorded size matches the rendered clip.
resolution_multiple = 32,
# 720p-class presets (all /32); first is the default the loader plans against.
resolution_presets = ((1280, 704), (704, 1280), (960, 960), (832, 480)),
# bf16-RESIDENT. transformer + VAE ship FP32 on disk (index 20.0 GB = 5B x 4), so bf16 transformer ~10.0; UMT5 TE ships bf16 (11.4); VAE runs fp32 (2.8).
bf16_components_gb = (10.0, 11.4, 2.8),
vae_force_fp32 = True,
gguf_repo = "QuantStack/Wan2.2-TI2V-5B-GGUF",
),
# Wan2.2-T2V-A14B (diffusers >= 0.35, verified on 0.39): the dual-expert MoE. Both transformers are WanTransformer3DModel with boundary_ratio 0.875; the pipeline routes high-noise steps through transformer (guidance_scale) and low-noise through transformer_2 (guidance_scale_2, accepted only when boundary_ratio is set), so cfg2_kwarg is threaded ONLY here.
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,
# is_moe drives the dual-DiT optimisation layers (speed/attention/cache/quant on BOTH); cfg2_kwarg names the pipeline kwarg for 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 ~5s (81 = 4*20 + 1), the A14B card's default clip.
default_num_frames = 81,
default_fps = 16, # A14B runs at 16 fps (vs TI2V-5B's 24)
frame_step = 4,
resolution_multiple = 16,
# 480p + 720p presets (landscape + vertical). A14B's VAE is 8x so multiple 16 renders 720 exactly (unlike TI2V-5B's 16x VAE, which floors 720 to 704).
resolution_presets = ((1280, 720), (832, 480), (480, 832), (720, 1280)),
# bf16-RESIDENT. Each expert ships FP32 (index 57.15 GB = 14.3B x 4), so ~28.6 bf16 each and ~57.2 for BOTH (the headline before offload), NOT the 114.3 fp32 sum. UMT5 TE bf16 (11.4); VAE fp32 (0.5).
bf16_components_gb = (57.2, 11.4, 0.5),
vae_force_fp32 = True,
# No gguf_repo: community GGUFs split the experts, and a single-file load covers only one.
),
# HunyuanVideo-1.5 (diffusers >= 0.39): 8.3B DiT, Qwen2.5-VL text encoder + ByT5 glyph encoder. Three quirks: (1) __call__ has NO guidance kwarg, CFG lives on the ``guider``; (2) NO callback_on_step_end (generate() wraps scheduler.step); (3) tencent's repo has no model_index.json, so only the community Diffusers repacks load.
VideoFamily(
name = "hunyuanvideo-1.5",
pipeline_class = "HunyuanVideo15Pipeline",
transformer_class = "HunyuanVideo15Transformer3DModel",
base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
# No bare "hunyuanvideo" alias: it would also claim the incompatible 1.0 repos.
aliases = ("hunyuanvideo-1-5", "hunyuanvideo1.5", "hunyuanvideo1-5", "hv15"),
has_audio = False,
guidance_via_guider = True,
default_steps = 50,
default_guidance = 6.0,
# 121 frames at 24 fps ~5s, the pipeline's own default.
default_num_frames = 121,
default_fps = 24,
# HV15 VAE compresses 16x spatial / 4x temporal, patch-1, so sizes snap /16, frames 4k+1.
frame_step = 4,
resolution_multiple = 16,
# 480p-class presets (the base is the 480p variant): landscape, vertical, square.
resolution_presets = ((832, 480), (480, 832), (624, 624)),
# DiT fp32 on disk (32.0 to 16.6 bf16); VAE (4.7 to 2.4); Qwen2.5-VL TE bf16 14.0 + ByT5 0.8.
bf16_components_gb = (16.6, 14.8, 2.4),
),
# The 720p t2v repack: same architecture/quirks/footprint as the 480p entry, only the trained resolution differs. Own family so a 720p load defaults to 720p sizes. Its full-path alias out-lengths (and outranks) the generic "hunyuanvideo-1.5" token for this repo only.
VideoFamily(
name = "hunyuanvideo-1.5-720p",
pipeline_class = "HunyuanVideo15Pipeline",
transformer_class = "HunyuanVideo15Transformer3DModel",
base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v",
aliases = ("hunyuanvideo-1.5-diffusers-720p_t2v", "hv15-720p"),
has_audio = False,
guidance_via_guider = True,
default_steps = 50,
default_guidance = 6.0,
default_num_frames = 121,
default_fps = 24,
frame_step = 4,
resolution_multiple = 16,
# 720p-class presets: landscape, vertical, square (all /16).
resolution_presets = ((1280, 720), (720, 1280), (960, 960)),
bf16_components_gb = (16.6, 14.8, 2.4),
),
)
def _token_in_needle(token: str, needle: str) -> bool:
"""Whole path/name segment match, as in diffusion_families (a short alias like
'ltx' must not match inside an unrelated word)."""
return re.search(r"(?:^|[-_./\\])" + re.escape(token) + r"(?:$|[-_./\\])", needle) is not None
def detect_video_family(repo_id: str, override: Optional[str] = None) -> Optional[VideoFamily]:
"""Resolve a ``VideoFamily`` from a repo id, or an explicit override.
Same contract as ``diffusion_families.detect_family``: an override matches a
name/alias exactly; otherwise the longest name/alias appearing as a whole
segment of the repo id wins.
"""
if override:
key = override.strip().lower()
for fam in _FAMILIES:
if key == fam.name or key in fam.aliases:
return fam
return None
needle = repo_id.lower()
best: Optional[tuple[VideoFamily, int]] = None
for fam in _FAMILIES:
for token in (fam.name, *fam.aliases):
if _token_in_needle(token, needle) and (best is None or len(token) > best[1]):
best = (fam, len(token))
return best[0] if best else None
def supported_video_family_names() -> tuple[str, ...]:
return tuple(fam.name for fam in _FAMILIES)
def resolve_video_base_repo(fam: VideoFamily, base_repo: Optional[str]) -> str:
"""The companion diffusers repo: caller-supplied if given, else the family fallback."""
base = (base_repo or "").strip()
return base or fam.base_repo
def snap_num_frames(fam: VideoFamily, num_frames: int) -> int:
"""The nearest valid frame count at or below the request (k * frame_step + 1).
Video latents are allocated as (num_frames - 1) / temporal_compression + 1, so
an off-lattice count wastes a partial latent frame at best and trips shape
checks at worst; snapping mirrors the image path's silent /16 size snap.
"""
step = max(1, fam.frame_step)
return max(1, ((max(1, num_frames) - 1) // step) * step + 1)
def snap_video_size(fam: VideoFamily, width: int, height: int) -> tuple[int, int]:
"""Width/height floored to the family's required multiple (minimum one unit)."""
multiple = max(1, fam.resolution_multiple)
snap = lambda v: max(multiple, (max(1, v) // multiple) * multiple) # noqa: E731
return snap(width), snap(height)
# Default (steps, guidance) per checkpoint variant, matched by substring (picked id then base repo), most specific first: distilled LTX-2.3 runs few-step CFG-off, the dev base wants 40/4.
_VIDEO_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
("distilled", 8, 1.0),
("ltx", 40, 4.0),
# Wan2.2 pipelines default to 50 steps / CFG 5.0; both TI2V-5B and A14B share these.
("wan", 50, 5.0),
# HunyuanVideo-1.5: 50 steps with the guider's shipped CFG 6.0.
("hunyuanvideo", 50, 6.0),
)
def default_video_generation_params(
*identifiers: Optional[str], fallback: tuple[int, float] = (40, 4.0)
) -> tuple[int, float]:
"""Default ``(steps, guidance)`` for a loaded video model; the first identifier
naming a known variant wins, so a GGUF filename ('...distilled...Q4_K_M.gguf')
beats the family base repo. ``fallback`` is used when no identifier names a variant --
callers pass the resolved family's own default so a Wan model loaded from an opaque local
path under an explicit family_override still gets 50/5.0, not the hardcoded LTX 40/4.0."""
for identifier in identifiers:
needle = (identifier or "").lower()
for key, steps, guidance in _VIDEO_GENERATION_DEFAULTS:
# Match the key as a name segment: reject a preceding ASCII letter so "swan-video" or "taiwan-clips" does not false-match "wan". Trailing chars stay free.
if re.search(r"(?<![a-z])" + re.escape(key), needle):
return steps, guidance
return fallback