[bugfix] v2: correct SF-Wan + LTX2 2-stage SR sampling defaults (GPU frame-verified)
Two distilled few-step video models rendered incorrectly on GPU; root-caused via dense frame sampling (contact sheets) and fixed in the recipe cards. SF-Wan2.1 (self-forcing causal, wan_causal card): was oversaturated/overcooked. The distilled student is CFG-FREE (guidance 1.0, single forward/step) and denoises with the 4-step DMD schedule [1000,750,500,250] (warped by FlowShiftPolicy(5.0)), at a native causal block of 3 latent frames. Defaults were ClassicCFG@6.0 + 2 steps + block 2 -> overcooked AND under-denoised. Fixed: num_chunks=7, chunk_size=3, steps_per_chunk=4; SamplingDefaults num_steps=4, guidance_scale=1.0 (7x3=21 latent -> 81 frames). Renders a clean raccoon-in-sunflowers across all 81 frames. LTX2-Distilled 2-stage SR (ltx2 card + LTX2VAE): was temporally blocky. Root cause was an OOM-forced 57-frame reduction (only 8 latent temporal frames); the model is designed for 121 (16 latent frames). Enable VAE tiling in LTX2VAE so the 121-frame full-res decode fits the 80 GiB GPU; keep base cfg_scale=3.0 (drives brightness; cfg=1 washed out) with stg_scale=0.0 (v2's drop-text perturbation is not real skip-layer STG). Renders the on-prompt backyard shot, bright + temporally coherent. torch_backend.py: enable LTX2VAE tiling; clear pre-existing mypy no-untyped-call / yapf debt on the file (surfaced once per-file linting bypassed the duplicate-module flakiness) by annotating the helper/constructor/maker signatures. Tests: update the 3 affected CPU defaults/chunk-count tests. v2 mini: 240 passed.
This commit is contained in:
SolitaryThinker
@@ -20,6 +20,7 @@ on a CPU box never imports torch and the mini stays green. Wan2.1 + LTX-2 A/V GP
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -56,9 +57,8 @@ def _ensure_fastvideo_runtime() -> None:
|
||||
def _require_checkpoint(spec) -> str:
|
||||
ckpt = getattr(spec, "checkpoint", "") or ""
|
||||
if not ckpt:
|
||||
raise RuntimeError(
|
||||
f"component {spec.component_id!r}: set ComponentSpec.checkpoint to the weights path / HF id "
|
||||
f"for the GPU backend (load_id={spec.load_id!r}). See GPU_BRINGUP.md, risk A.")
|
||||
raise RuntimeError(f"component {spec.component_id!r}: set ComponentSpec.checkpoint to the weights path / HF id "
|
||||
f"for the GPU backend (load_id={spec.load_id!r}). See GPU_BRINGUP.md, risk A.")
|
||||
return ckpt
|
||||
|
||||
|
||||
@@ -68,16 +68,22 @@ def _model_root(spec) -> str:
|
||||
return os.path.dirname(os.path.normpath(_require_checkpoint(spec)))
|
||||
|
||||
|
||||
def _fastvideo_args(spec):
|
||||
def _fastvideo_args(spec: Any) -> Any:
|
||||
"""Build the FastVideoArgs the real loaders need (BRINGUP risk A). ``from_kwargs`` populates
|
||||
``pipeline_config`` (dit/vae/text-encoder configs + precisions) from the model root. Single-GPU,
|
||||
all offload/FSDP OFF — weights resident on one device, the simplest correct bring-up."""
|
||||
from v2.fastvideo_args import FastVideoArgs
|
||||
return FastVideoArgs.from_kwargs(
|
||||
model_path=_model_root(spec), num_gpus=1, tp_size=1, sp_size=1,
|
||||
dit_cpu_offload=False, text_encoder_cpu_offload=False, vae_cpu_offload=False,
|
||||
image_encoder_cpu_offload=False, dit_layerwise_offload=False,
|
||||
use_fsdp_inference=False, pin_cpu_memory=False)
|
||||
return FastVideoArgs.from_kwargs(model_path=_model_root(spec),
|
||||
num_gpus=1,
|
||||
tp_size=1,
|
||||
sp_size=1,
|
||||
dit_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
image_encoder_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
use_fsdp_inference=False,
|
||||
pin_cpu_memory=False)
|
||||
|
||||
|
||||
def load_component(loader_attr: str, path: str, args):
|
||||
@@ -93,7 +99,7 @@ def _device(platform) -> str:
|
||||
return "cuda" if platform.device == "cuda" else platform.device
|
||||
|
||||
|
||||
def _native_dtype(module):
|
||||
def _native_dtype(module: Any) -> Any:
|
||||
"""Keep each component at the precision its loader produced (Wan DiT bf16, VAE/text fp32) rather than
|
||||
the card's uniform fp32 — faithful to fastvideo and ~2x faster/smaller for the DiT."""
|
||||
try:
|
||||
@@ -102,12 +108,12 @@ def _native_dtype(module):
|
||||
return torch.float32
|
||||
|
||||
|
||||
def _to_torch(a, *, device, dtype):
|
||||
def _to_torch(a: Any, *, device: Any, dtype: Any) -> torch.Tensor | None:
|
||||
return None if a is None else torch.as_tensor(np.asarray(a), dtype=dtype, device=device)
|
||||
|
||||
|
||||
def _to_numpy(t):
|
||||
if hasattr(t, "sample"): # some heads wrap output in an object with .sample
|
||||
def _to_numpy(t: Any) -> np.ndarray:
|
||||
if hasattr(t, "sample"): # some heads wrap output in an object with .sample
|
||||
t = t.sample
|
||||
return t.detach().to("cpu", torch.float32).numpy()
|
||||
|
||||
@@ -119,21 +125,21 @@ class TorchComponent:
|
||||
"""Wraps a real ``fastvideo.models.*`` module to the mini's numpy duck-typed surface. Subclasses
|
||||
override only the forward semantics (``__call__`` / ``encode`` / ``decode`` / ``upsample`` / ...)."""
|
||||
|
||||
def __init__(self, module, *, device, dtype, eager: bool = True):
|
||||
def __init__(self, module: Any, *, device: Any, dtype: Any, eager: bool = True) -> None:
|
||||
self.device, self.dtype = device, dtype
|
||||
self.module = module.to(device=device, dtype=dtype).eval() if eager else module
|
||||
|
||||
# numpy<->torch marshalling at the loop boundary (ONE place — torch-native is a later swap) ------- #
|
||||
def _t(self, a, *, batch: bool = True):
|
||||
def _t(self, a: Any, *, batch: bool = True) -> Any:
|
||||
t = _to_torch(a, device=self.device, dtype=self.dtype)
|
||||
return t.unsqueeze(0) if (t is not None and batch) else t
|
||||
|
||||
@staticmethod
|
||||
def _n(t):
|
||||
def _n(t: Any) -> np.ndarray:
|
||||
return _to_numpy(t.squeeze(0))
|
||||
|
||||
@staticmethod
|
||||
def _ctx(current_timestep=0):
|
||||
def _ctx(current_timestep: Any = 0) -> Any:
|
||||
"""The FastVideo attention layer reads attn_metadata via get_forward_context(), so every forward
|
||||
runs inside set_forward_context(...). attn_metadata=None selects the dense SDPA path."""
|
||||
from v2.forward_context import set_forward_context
|
||||
@@ -143,16 +149,16 @@ class TorchComponent:
|
||||
def copy_from(self, other) -> None:
|
||||
self.module.load_state_dict(other.module.state_dict())
|
||||
|
||||
def blend_from(self, other, decay: float) -> None: # EMA / decayed-old-policy
|
||||
def blend_from(self, other, decay: float) -> None: # EMA / decayed-old-policy
|
||||
with torch.no_grad():
|
||||
for p, q in zip(self.module.parameters(), other.module.parameters()):
|
||||
for p, q in zip(self.module.parameters(), other.module.parameters(), strict=False):
|
||||
p.mul_(decay).add_(q, alpha=1.0 - decay)
|
||||
|
||||
def clone(self):
|
||||
def clone(self) -> TorchComponent:
|
||||
import copy
|
||||
c = self.__class__.__new__(self.__class__)
|
||||
c.__dict__.update(self.__dict__) # tokenizer / shapes / sibling refs shared
|
||||
c.module = copy.deepcopy(self.module) # independent weights
|
||||
c.__dict__.update(self.__dict__) # tokenizer / shapes / sibling refs shared
|
||||
c.module = copy.deepcopy(self.module) # independent weights
|
||||
return c
|
||||
|
||||
def mse_grad_step(self, *a, **k):
|
||||
@@ -168,7 +174,13 @@ class WanDiT(TorchComponent):
|
||||
(wanvideo.py): forward(hidden_states[B,C,T,H,W], encoder_hidden_states, timestep,
|
||||
encoder_hidden_states_image=None) -> velocity (bare tensor)."""
|
||||
|
||||
def __init__(self, module, *, device, dtype, offload_group=None, component_id="transformer"):
|
||||
def __init__(self,
|
||||
module: Any,
|
||||
*,
|
||||
device: Any,
|
||||
dtype: Any,
|
||||
offload_group: Any = None,
|
||||
component_id: str = "transformer") -> None:
|
||||
# Wan2.2 MoE (A14B): two 14B experts don't both fit one 80GB GPU. With ``offload_group`` set, keep
|
||||
# this expert on CPU and bring only the *active* one onto the GPU on demand (single swap at the
|
||||
# boundary, not per-step thrash). Single-expert Wan stays resident (offload_group=None).
|
||||
@@ -200,22 +212,25 @@ class WanDiT(TorchComponent):
|
||||
def __call__(self, latent, text_embed, sigma, context=None, *, cond=None):
|
||||
self._ensure_resident()
|
||||
hs = self._t(latent)
|
||||
if cond is not None: # i2v: concat [noise (16ch) ; mask+cond_latent (20ch)] -> 36ch DiT input
|
||||
if cond is not None: # i2v: concat [noise (16ch) ; mask+cond_latent (20ch)] -> 36ch DiT input
|
||||
hs = torch.cat([hs, self._t(cond)], dim=1)
|
||||
ehs = self._t(text_embed)
|
||||
# timestep = sigma*1000 (BRINGUP risk B). Causal model needs per-latent-frame [B, num_frames].
|
||||
ts = float(sigma) * NUM_TRAIN_TIMESTEPS
|
||||
timestep = (torch.full((1, hs.shape[2]), ts, device=self.device) if self.causal
|
||||
else torch.tensor([ts], device=self.device))
|
||||
img = None if self.causal else self._t(context) # i2v image embedding for standard Wan
|
||||
timestep = (torch.full(
|
||||
(1, hs.shape[2]), ts, device=self.device) if self.causal else torch.tensor([ts], device=self.device))
|
||||
img = None if self.causal else self._t(context) # i2v image embedding for standard Wan
|
||||
with self._ctx():
|
||||
velocity = self.module(hidden_states=hs, encoder_hidden_states=ehs, timestep=timestep,
|
||||
velocity = self.module(hidden_states=hs,
|
||||
encoder_hidden_states=ehs,
|
||||
timestep=timestep,
|
||||
encoder_hidden_states_image=img)
|
||||
return self._n(velocity) # rectified-flow velocity (BRINGUP risk C)
|
||||
return self._n(velocity) # rectified-flow velocity (BRINGUP risk C)
|
||||
|
||||
def clone(self):
|
||||
def clone(self) -> WanDiT:
|
||||
c = super().clone()
|
||||
c.offload_group, c._on_gpu = None, True # standalone resident copy
|
||||
assert isinstance(c, WanDiT)
|
||||
c.offload_group, c._on_gpu = None, True # standalone resident copy
|
||||
return c
|
||||
|
||||
|
||||
@@ -225,7 +240,7 @@ class LTX2DiT(TorchComponent):
|
||||
returns the flow-match velocity ``(x_t - x0)/sigma`` the loop integrates. Joint A/V (audio_latent
|
||||
given) -> (video_velocity, audio_velocity) in one forward."""
|
||||
|
||||
def __init__(self, module, *, device, dtype):
|
||||
def __init__(self, module: Any, *, device: Any, dtype: Any) -> None:
|
||||
super().__init__(module, device=device, dtype=dtype)
|
||||
from v2.models.dits.ltx2 import VideoLatentShape
|
||||
self._VideoLatentShape = VideoLatentShape
|
||||
@@ -240,19 +255,23 @@ class LTX2DiT(TorchComponent):
|
||||
video_sigma = torch.tensor([s], device=self.device, dtype=torch.float32)
|
||||
av_kwargs: dict = {}
|
||||
au = None
|
||||
if audio_latent is not None: # joint A/V: audio latent [1,8,T,16] + audio text
|
||||
if audio_latent is not None: # joint A/V: audio latent [1,8,T,16] + audio text
|
||||
from v2.models.audio.ltx2_audio_vae import AudioLatentShape
|
||||
au = self._t(audio_latent)
|
||||
aeh = self._t(audio_text)
|
||||
atok = self.module.audio_patchifier.get_token_count(AudioLatentShape.from_torch_shape(tuple(au.shape)))
|
||||
av_kwargs = dict(
|
||||
audio_hidden_states=au, audio_encoder_hidden_states=aeh,
|
||||
audio_timestep=torch.full((1, atok, 1), s, device=self.device, dtype=torch.float32),
|
||||
audio_sigma=torch.tensor([s], device=self.device, dtype=torch.float32))
|
||||
av_kwargs = dict(audio_hidden_states=au,
|
||||
audio_encoder_hidden_states=aeh,
|
||||
audio_timestep=torch.full((1, atok, 1), s, device=self.device, dtype=torch.float32),
|
||||
audio_sigma=torch.tensor([s], device=self.device, dtype=torch.float32))
|
||||
with self._ctx(current_timestep=s):
|
||||
out = self.module(hidden_states=hs, encoder_hidden_states=ehs, timestep=timestep,
|
||||
video_sigma=video_sigma, encoder_attention_mask=None, **av_kwargs)
|
||||
if au is not None: # LTX-2 DiT predicts x0; integrate velocity
|
||||
out = self.module(hidden_states=hs,
|
||||
encoder_hidden_states=ehs,
|
||||
timestep=timestep,
|
||||
video_sigma=video_sigma,
|
||||
encoder_attention_mask=None,
|
||||
**av_kwargs)
|
||||
if au is not None: # LTX-2 DiT predicts x0; integrate velocity
|
||||
denoised_v, denoised_a = out
|
||||
vel_v = (hs.float() - denoised_v.float()) / max(s, 1e-6)
|
||||
vel_a = (au.float() - denoised_a.float()) / max(s, 1e-6)
|
||||
@@ -268,11 +287,10 @@ class WanVAE(TorchComponent):
|
||||
"""The DiT operates in NORMALIZED latent space, so encode applies ``(z - mean)/std`` and decode
|
||||
inverts it; ``AutoencoderKLWan.encode/decode`` operate in RAW latent space (BRINGUP risk D)."""
|
||||
|
||||
def _mean_invstd(self, like):
|
||||
mean = torch.tensor(self.module.latents_mean, device=like.device,
|
||||
dtype=torch.float32).view(1, -1, 1, 1, 1)
|
||||
inv_std = (1.0 / torch.tensor(self.module.latents_std, device=like.device,
|
||||
dtype=torch.float32)).view(1, -1, 1, 1, 1)
|
||||
def _mean_invstd(self, like: Any) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
mean = torch.tensor(self.module.latents_mean, device=like.device, dtype=torch.float32).view(1, -1, 1, 1, 1)
|
||||
inv_std = (1.0 / torch.tensor(self.module.latents_std, device=like.device, dtype=torch.float32)).view(
|
||||
1, -1, 1, 1, 1)
|
||||
return mean, inv_std
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -281,20 +299,28 @@ class WanVAE(TorchComponent):
|
||||
dist = self.module.encode(x)
|
||||
z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist)
|
||||
mean, inv_std = self._mean_invstd(z)
|
||||
return self._n((z.float() - mean) * inv_std) # -> the normalized latent the DiT expects
|
||||
return self._n((z.float() - mean) * inv_std) # -> the normalized latent the DiT expects
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, latent):
|
||||
z = self._t(latent).float()
|
||||
mean, inv_std = self._mean_invstd(z)
|
||||
z = z / inv_std + mean # invert encode normalization -> raw latent
|
||||
video = self.module.decode(z.to(self.dtype)) # -> video [B,3,T,H,W] in [-1,1]
|
||||
z = z / inv_std + mean # invert encode normalization -> raw latent
|
||||
video = self.module.decode(z.to(self.dtype)) # -> video [B,3,T,H,W] in [-1,1]
|
||||
return self._n(video)
|
||||
|
||||
|
||||
class LTX2VAE(TorchComponent):
|
||||
"""LTX-2 VideoDecoder un-normalizes internally (per-channel stats); the adapter just marshals."""
|
||||
|
||||
def __init__(self, module: Any, *, device: Any, dtype: Any) -> None:
|
||||
super().__init__(module, device=device, dtype=dtype)
|
||||
# Tiled decode (spatial tiles concatenated): the LTX-2 CausalVideoAutoencoder decode of a full-res,
|
||||
# full-length (121f / 1536x1024) latent allocates >80GB of conv activations otherwise. fastvideo
|
||||
# defaults vae_tiling=True for LTX-2; mirror it so the full clip fits on one GPU.
|
||||
if hasattr(self.module, "enable_tiling"):
|
||||
self.module.enable_tiling()
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, latent):
|
||||
video = self.module.decode(self._t(latent))
|
||||
@@ -315,7 +341,7 @@ class LTX2VAE(TorchComponent):
|
||||
class T5Encoder(TorchComponent):
|
||||
"""(U)MT5 text_encoder.encode(text) -> embedding[text_len, dim] (numpy out)."""
|
||||
|
||||
def __init__(self, module, tokenizer, *, device, dtype, max_length: int = 512):
|
||||
def __init__(self, module: Any, tokenizer: Any, *, device: Any, dtype: Any, max_length: int = 512) -> None:
|
||||
super().__init__(module, device=device, dtype=dtype)
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
@@ -340,14 +366,17 @@ class Gemma(TorchComponent):
|
||||
"""LTX-2 Gemma text encoder. ``encode`` -> video projection; ``encode_av`` -> (video, audio): the 2.3
|
||||
connector emits SEPARATE projections (audio in ``hidden_states[0]`` only when output_hidden_states)."""
|
||||
|
||||
def __init__(self, module, tokenizer, *, device, dtype, max_length: int = 256):
|
||||
def __init__(self, module: Any, tokenizer: Any, *, device: Any, dtype: Any, max_length: int = 256) -> None:
|
||||
super().__init__(module, device=device, dtype=dtype)
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
|
||||
def _tok(self, text):
|
||||
toks = self.tokenizer(text or "", return_tensors="pt", max_length=self.max_length,
|
||||
truncation=True, padding="max_length")
|
||||
def _tok(self, text: Any) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
toks = self.tokenizer(text or "",
|
||||
return_tensors="pt",
|
||||
max_length=self.max_length,
|
||||
truncation=True,
|
||||
padding="max_length")
|
||||
return toks.input_ids.to(self.device), toks.attention_mask.to(self.device)
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -365,7 +394,7 @@ class Gemma(TorchComponent):
|
||||
out = self.module(input_ids=ids, attention_mask=mask, output_hidden_states=True)
|
||||
video = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
|
||||
hs = getattr(out, "hidden_states", None)
|
||||
audio = hs[0] if hs else video # separate audio connector projection
|
||||
audio = hs[0] if hs else video # separate audio connector projection
|
||||
return _to_numpy(video.squeeze(0)), _to_numpy(audio.squeeze(0))
|
||||
|
||||
|
||||
@@ -375,7 +404,7 @@ class CLIPImageEncoder(TorchComponent):
|
||||
preprocesses, the encoder returns ``last_hidden_state``. BRINGUP: written-not-run; GPU-verify the
|
||||
processor/encoder subfolders + dtype against a real i2v checkpoint (see GPU_BRINGUP.md)."""
|
||||
|
||||
def __init__(self, module, processor, *, device, dtype):
|
||||
def __init__(self, module: Any, processor: Any, *, device: Any, dtype: Any) -> None:
|
||||
super().__init__(module, device=device, dtype=dtype)
|
||||
self.processor = processor
|
||||
|
||||
@@ -395,9 +424,9 @@ class LTX2Upsampler(TorchComponent):
|
||||
"""upsample(latent[C,T,H,W]) -> [C,T,2H,2W]. Applies the repo's ``upsample_video``: un_normalize (via
|
||||
the video VAE's per_channel_statistics) -> learned 2x upsample -> normalize."""
|
||||
|
||||
def __init__(self, module, stats_owner, *, device, dtype):
|
||||
def __init__(self, module: Any, stats_owner: Any, *, device: Any, dtype: Any) -> None:
|
||||
super().__init__(module, device=device, dtype=dtype)
|
||||
self.stats_owner = stats_owner # exposes per_channel_statistics (the VAE decoder/encoder)
|
||||
self.stats_owner = stats_owner # exposes per_channel_statistics (the VAE decoder/encoder)
|
||||
|
||||
@torch.no_grad()
|
||||
def upsample(self, latent):
|
||||
@@ -414,9 +443,9 @@ class LTX2AudioVAE(TorchComponent):
|
||||
"""audio_vae.decode(audio_latent[8,T,16]) -> waveform. Runs AudioDecoder (latent->mel) then the
|
||||
Vocoder (mel->waveform). ``sample_rate`` = the vocoder's output rate (so a saver writes the right wav)."""
|
||||
|
||||
def __init__(self, module, vocoder, *, device, dtype):
|
||||
def __init__(self, module: Any, vocoder: Any, *, device: Any, dtype: Any) -> None:
|
||||
super().__init__(module, device=device, dtype=dtype)
|
||||
self.vocoder = vocoder # LTX2Vocoder | None
|
||||
self.vocoder = vocoder # LTX2Vocoder | None
|
||||
self.sample_rate = int(getattr(getattr(vocoder, "module", None), "output_sample_rate", 24000))
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -436,7 +465,7 @@ class LTX2AudioVAE(TorchComponent):
|
||||
# --------------------------------------------------------------------------- #
|
||||
# build_component — one dispatch (replaces the six build_torch_* + trampolines) #
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _explicit_adapter(spec, module, device, dtype, *extra):
|
||||
def _explicit_adapter(spec: Any, module: Any, device: Any, dtype: Any, *extra: Any) -> Any:
|
||||
"""Build the card's explicitly-declared adapter (``ComponentSpec.adapter='module:Class'``) if set —
|
||||
the seam that lets a NEW architecture's recipe carry its own TorchComponent subclass without editing
|
||||
the dispatch here (so a port is a self-contained recipe package). Constructed as
|
||||
@@ -449,7 +478,7 @@ def _explicit_adapter(spec, module, device, dtype, *extra):
|
||||
return getattr(importlib.import_module(mod), cls)(module, *extra, device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _make_dit(spec, instance, platform, args):
|
||||
def _make_dit(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
||||
module = load_component("TransformerLoader", spec.checkpoint, args)
|
||||
device, dtype = _device(platform), _native_dtype(module)
|
||||
explicit = _explicit_adapter(spec, module, device, dtype)
|
||||
@@ -467,7 +496,7 @@ def _make_dit(spec, instance, platform, args):
|
||||
return WanDiT(module, device=device, dtype=dtype, offload_group=grp, component_id=spec.component_id)
|
||||
|
||||
|
||||
def _make_vae(spec, instance, platform, args):
|
||||
def _make_vae(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
||||
module = load_component("VAELoader", spec.checkpoint, args)
|
||||
device, dtype = _device(platform), _native_dtype(module)
|
||||
explicit = _explicit_adapter(spec, module, device, dtype)
|
||||
@@ -477,26 +506,26 @@ def _make_vae(spec, instance, platform, args):
|
||||
return cls(module, device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _make_text_encoder(spec, instance, platform, args):
|
||||
def _make_text_encoder(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
||||
module = load_component("TextEncoderLoader", spec.checkpoint, args)
|
||||
tokenizer = load_component("TokenizerLoader", os.path.join(_model_root(spec), "tokenizer"), args)
|
||||
device, dtype = _device(platform), _native_dtype(module)
|
||||
explicit = _explicit_adapter(spec, module, device, dtype, tokenizer) # adapter cls(module, tokenizer, ...)
|
||||
explicit = _explicit_adapter(spec, module, device, dtype, tokenizer) # adapter cls(module, tokenizer, ...)
|
||||
if explicit is not None:
|
||||
return explicit
|
||||
cls = Gemma if "Gemma" in type(module).__name__ else T5Encoder
|
||||
return cls(module, tokenizer, device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _make_image_encoder(spec, instance, platform, args):
|
||||
module = load_component("ImageEncoderLoader", spec.checkpoint, args) # CLIP vision
|
||||
def _make_image_encoder(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
||||
module = load_component("ImageEncoderLoader", spec.checkpoint, args) # CLIP vision
|
||||
# the HF image processor is a sibling subfolder (BRINGUP: confirm path on a real i2v checkpoint)
|
||||
processor = load_component("ImageProcessorLoader", os.path.join(_model_root(spec), "image_processor"), args)
|
||||
return CLIPImageEncoder(module, processor, device=_device(platform), dtype=_native_dtype(module))
|
||||
|
||||
|
||||
def _make_upsampler(spec, instance, platform, args):
|
||||
module = load_component("UpsamplerLoader", spec.checkpoint, args) # LTX2LatentUpsampler
|
||||
def _make_upsampler(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
||||
module = load_component("UpsamplerLoader", spec.checkpoint, args) # LTX2LatentUpsampler
|
||||
vae = instance.component("vae")
|
||||
dtype = getattr(vae, "dtype", _native_dtype(module))
|
||||
vae_module = getattr(vae, "module", None)
|
||||
@@ -505,33 +534,37 @@ def _make_upsampler(spec, instance, platform, args):
|
||||
return LTX2Upsampler(module, stats_owner, device=_device(platform), dtype=dtype)
|
||||
|
||||
|
||||
def _make_audio_vae(spec, instance, platform, args):
|
||||
module = load_component("AudioDecoderLoader", spec.checkpoint, args) # LTX2AudioDecoder
|
||||
voc = instance.component("vocoder") # chained: decoder -> vocoder
|
||||
def _make_audio_vae(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
||||
module = load_component("AudioDecoderLoader", spec.checkpoint, args) # LTX2AudioDecoder
|
||||
voc = instance.component("vocoder") # chained: decoder -> vocoder
|
||||
return LTX2AudioVAE(module, voc, device=_device(platform), dtype=_native_dtype(module))
|
||||
|
||||
|
||||
def _make_vocoder(spec, instance, platform, args):
|
||||
module = load_component("VocoderLoader", spec.checkpoint, args) # LTX2Vocoder
|
||||
def _make_vocoder(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
||||
module = load_component("VocoderLoader", spec.checkpoint, args) # LTX2Vocoder
|
||||
return LTX2Vocoder(module, device=_device(platform), dtype=_native_dtype(module))
|
||||
|
||||
|
||||
_MAKERS = {
|
||||
"dit": _make_dit, "vae": _make_vae, "text_encoder": _make_text_encoder,
|
||||
"image_encoder": _make_image_encoder, "upsampler": _make_upsampler,
|
||||
"audio_vae": _make_audio_vae, "vocoder": _make_vocoder,
|
||||
"dit": _make_dit,
|
||||
"vae": _make_vae,
|
||||
"text_encoder": _make_text_encoder,
|
||||
"image_encoder": _make_image_encoder,
|
||||
"upsampler": _make_upsampler,
|
||||
"audio_vae": _make_audio_vae,
|
||||
"vocoder": _make_vocoder,
|
||||
}
|
||||
|
||||
|
||||
def build_component(spec, instance, platform):
|
||||
def build_component(spec: Any, instance: Any, platform: Any) -> TorchComponent:
|
||||
"""The single cuda component builder (registered for every kind in ``torch_cuda.py``). Shared prefix
|
||||
— checkpoint check, dist-init, FastVideoArgs — then dispatch by ``spec.kind`` to its maker."""
|
||||
_require_checkpoint(spec) # fail fast on a mis-stamped card, before any dist init
|
||||
_require_checkpoint(spec) # fail fast on a mis-stamped card, before any dist init
|
||||
_ensure_fastvideo_runtime()
|
||||
args = _fastvideo_args(spec)
|
||||
try:
|
||||
maker = _MAKERS[spec.kind]
|
||||
except KeyError:
|
||||
raise RuntimeError(f"torch backend: no builder for component kind {spec.kind!r} "
|
||||
f"(have {sorted(_MAKERS)})")
|
||||
f"(have {sorted(_MAKERS)})") from None
|
||||
return maker(spec, instance, platform, args)
|
||||
|
||||
+207
-101
@@ -39,64 +39,110 @@ def build_ltx2_card(model_id: str = "ltx2-2stage-distilled") -> ModelCard:
|
||||
cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1.5e-4, per_unit_seconds=1.2e-7)
|
||||
|
||||
def base_factory():
|
||||
return LTX2DenoiseLoop(loop_id="ltx2_base", stage="base", sigmas=BASE_SIGMAS,
|
||||
cfg_scale=3.0, stg_scale=0.1, cost=cost, input_slot=None, seed_offset=0)
|
||||
# The frame-jump/blockiness was the OOM-forced 57-frame reduction (only 8 latent temporal frames);
|
||||
# at the model's intended 121 (16 latent frames, VAE tiling enabled) the shot is temporally coherent.
|
||||
# CFG drives brightness/contrast here (cfg=3 -> mean~86, cfg=1 -> washed out ~66), so keep cfg=3.
|
||||
# STG=0.0: the v2 "drop text" perturbation is NOT real skip-layer STG, so leave it off.
|
||||
return LTX2DenoiseLoop(loop_id="ltx2_base",
|
||||
stage="base",
|
||||
sigmas=BASE_SIGMAS,
|
||||
cfg_scale=3.0,
|
||||
stg_scale=0.0,
|
||||
cost=cost,
|
||||
input_slot=None,
|
||||
seed_offset=0)
|
||||
|
||||
def refine_factory():
|
||||
return LTX2DenoiseLoop(loop_id="ltx2_refine", stage="refine", sigmas=REFINE_SIGMAS,
|
||||
cfg_scale=1.0, stg_scale=0.0, cost=cost,
|
||||
input_slot="ltx_upsampled", seed_offset=1000,
|
||||
audio_input_slot="ltx_audio") # read threaded audio in T2VS (else unused)
|
||||
return LTX2DenoiseLoop(loop_id="ltx2_refine",
|
||||
stage="refine",
|
||||
sigmas=REFINE_SIGMAS,
|
||||
cfg_scale=1.0,
|
||||
stg_scale=0.0,
|
||||
cost=cost,
|
||||
input_slot="ltx_upsampled",
|
||||
seed_offset=1000,
|
||||
audio_input_slot="ltx_audio") # read threaded audio in T2VS (else unused)
|
||||
|
||||
components = {
|
||||
"text_encoder": ComponentSpec(component_id="text_encoder", kind="text_encoder",
|
||||
load_id="transformers:AutoModel", # Gemma
|
||||
factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}),
|
||||
"vae": ComponentSpec(component_id="vae", kind="vae",
|
||||
load_id="fastvideo.models.vaes.ltx2vae:VideoDecoder",
|
||||
factory=lambda inst: ToyVAE(), required_for={"t2v"}),
|
||||
"text_encoder":
|
||||
ComponentSpec(
|
||||
component_id="text_encoder",
|
||||
kind="text_encoder",
|
||||
load_id="transformers:AutoModel", # Gemma
|
||||
factory=lambda inst: ToyTextEncoder(),
|
||||
required_for={"t2v"}),
|
||||
"vae":
|
||||
ComponentSpec(component_id="vae",
|
||||
kind="vae",
|
||||
load_id="fastvideo.models.vaes.ltx2vae:VideoDecoder",
|
||||
factory=lambda inst: ToyVAE(),
|
||||
required_for={"t2v"}),
|
||||
# learned 2x spatial latent upsampler between the base and refine stages (real:
|
||||
# LTX2LatentUpsampler; CPU toy: nearest-neighbor). Applied in program.py:_upsample.
|
||||
"spatial_upsampler": ComponentSpec(
|
||||
component_id="spatial_upsampler", kind="upsampler",
|
||||
load_id="fastvideo.models.upsamplers.ltx2_upsampler:LTX2LatentUpsampler",
|
||||
factory=lambda inst: ToyUpsampler(), required_for={"t2v"}),
|
||||
"transformer": ComponentSpec(component_id="transformer", kind="dit",
|
||||
load_id="fastvideo.models.dits.ltx2:LTX2Transformer3DModel",
|
||||
factory=lambda inst: ToyDiT(seed=seed),
|
||||
resident_for=["ltx2_base", "ltx2_refine"], required_for={"t2v"}),
|
||||
"spatial_upsampler":
|
||||
ComponentSpec(component_id="spatial_upsampler",
|
||||
kind="upsampler",
|
||||
load_id="fastvideo.models.upsamplers.ltx2_upsampler:LTX2LatentUpsampler",
|
||||
factory=lambda inst: ToyUpsampler(),
|
||||
required_for={"t2v"}),
|
||||
"transformer":
|
||||
ComponentSpec(component_id="transformer",
|
||||
kind="dit",
|
||||
load_id="fastvideo.models.dits.ltx2:LTX2Transformer3DModel",
|
||||
factory=lambda inst: ToyDiT(seed=seed),
|
||||
resident_for=["ltx2_base", "ltx2_refine"],
|
||||
required_for={"t2v"}),
|
||||
# audio branch: lazy for T2V (not loaded), required for T2VS (joint audio+video, §9.11)
|
||||
"audio_vae": ComponentSpec(component_id="audio_vae", kind="audio_vae",
|
||||
load_id="fastvideo.models.audio.ltx2_audio_vae:AudioVAE",
|
||||
factory=lambda inst: ToyAudioVAE(),
|
||||
optional_for={"t2v"}, required_for={"t2vs"}),
|
||||
"audio_vae":
|
||||
ComponentSpec(component_id="audio_vae",
|
||||
kind="audio_vae",
|
||||
load_id="fastvideo.models.audio.ltx2_audio_vae:AudioVAE",
|
||||
factory=lambda inst: ToyAudioVAE(),
|
||||
optional_for={"t2v"},
|
||||
required_for={"t2vs"}),
|
||||
}
|
||||
loops = {
|
||||
"ltx2_base": LoopSpec(loop_id="ltx2_base", kind=LoopKind.DIFFUSION_DENOISE,
|
||||
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost,
|
||||
shared_weight_components=["transformer"], cache_policy=["feature"],
|
||||
loop_factory=base_factory),
|
||||
"ltx2_refine": LoopSpec(loop_id="ltx2_refine", kind=LoopKind.DIFFUSION_DENOISE,
|
||||
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost,
|
||||
shared_weight_components=["transformer"], cache_policy=["feature"],
|
||||
loop_factory=refine_factory),
|
||||
"ltx2_base":
|
||||
LoopSpec(loop_id="ltx2_base",
|
||||
kind=LoopKind.DIFFUSION_DENOISE,
|
||||
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
|
||||
step_cost_model=cost,
|
||||
shared_weight_components=["transformer"],
|
||||
cache_policy=["feature"],
|
||||
loop_factory=base_factory),
|
||||
"ltx2_refine":
|
||||
LoopSpec(loop_id="ltx2_refine",
|
||||
kind=LoopKind.DIFFUSION_DENOISE,
|
||||
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
|
||||
step_cost_model=cost,
|
||||
shared_weight_components=["transformer"],
|
||||
cache_policy=["feature"],
|
||||
loop_factory=refine_factory),
|
||||
}
|
||||
card = ModelCard(
|
||||
model_id=model_id, family="ltx2",
|
||||
components=components, loops=loops,
|
||||
model_id=model_id,
|
||||
family="ltx2",
|
||||
components=components,
|
||||
loops=loops,
|
||||
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.TEXT_TO_VIDEO_SOUND,
|
||||
Capability.VAE_DECODE),
|
||||
recipe=RecipeSpec(method="distilled", parents=["ltx2.3-base"], assumes_loop="ltx2_base",
|
||||
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
|
||||
recipe=RecipeSpec(method="distilled",
|
||||
parents=["ltx2.3-base"],
|
||||
assumes_loop="ltx2_base",
|
||||
assumes_precision="float32",
|
||||
consistency_required=ConsistencyLevel.C1),
|
||||
sampling_defaults=SamplingDefaults( # two-stage distilled: 8-step base (+ 2-step refine in loop)
|
||||
num_steps=8, guidance_scale=1.0, height=1024, width=1536, num_frames=121, fps=24,
|
||||
num_steps=8,
|
||||
guidance_scale=1.0,
|
||||
height=1024,
|
||||
width=1536,
|
||||
num_frames=121,
|
||||
fps=24,
|
||||
negative_prompt=""),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24,
|
||||
reuse_across_requests=True)},
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
|
||||
default_plan=ParallelPlan.single()),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
|
||||
@@ -111,43 +157,71 @@ def build_ltx2_base_card(model_id: str = "ltx2-single-stage") -> ModelCard:
|
||||
cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1.5e-4, per_unit_seconds=1.2e-7)
|
||||
|
||||
def single_factory():
|
||||
return LTX2DenoiseLoop(loop_id="ltx2_single", stage="single", sigmas=[1.0, 0.0],
|
||||
cfg_scale=3.0, stg_scale=0.0, cost=cost, input_slot=None,
|
||||
seed_offset=0, full_res=True, request_steps=True, shift=1.0)
|
||||
return LTX2DenoiseLoop(loop_id="ltx2_single",
|
||||
stage="single",
|
||||
sigmas=[1.0, 0.0],
|
||||
cfg_scale=3.0,
|
||||
stg_scale=0.0,
|
||||
cost=cost,
|
||||
input_slot=None,
|
||||
seed_offset=0,
|
||||
full_res=True,
|
||||
request_steps=True,
|
||||
shift=1.0)
|
||||
|
||||
components = {
|
||||
"text_encoder": ComponentSpec(component_id="text_encoder", kind="text_encoder",
|
||||
load_id="transformers:AutoModel", # Gemma
|
||||
factory=lambda inst: ToyTextEncoder(), required_for={"t2v"}),
|
||||
"vae": ComponentSpec(component_id="vae", kind="vae",
|
||||
load_id="fastvideo.models.vaes.ltx2vae:VideoDecoder",
|
||||
factory=lambda inst: ToyVAE(), required_for={"t2v"}),
|
||||
"transformer": ComponentSpec(component_id="transformer", kind="dit",
|
||||
load_id="fastvideo.models.dits.ltx2:LTX2Transformer3DModel",
|
||||
factory=lambda inst: ToyDiT(seed=seed),
|
||||
resident_for=["ltx2_single"], required_for={"t2v"}),
|
||||
"text_encoder":
|
||||
ComponentSpec(
|
||||
component_id="text_encoder",
|
||||
kind="text_encoder",
|
||||
load_id="transformers:AutoModel", # Gemma
|
||||
factory=lambda inst: ToyTextEncoder(),
|
||||
required_for={"t2v"}),
|
||||
"vae":
|
||||
ComponentSpec(component_id="vae",
|
||||
kind="vae",
|
||||
load_id="fastvideo.models.vaes.ltx2vae:VideoDecoder",
|
||||
factory=lambda inst: ToyVAE(),
|
||||
required_for={"t2v"}),
|
||||
"transformer":
|
||||
ComponentSpec(component_id="transformer",
|
||||
kind="dit",
|
||||
load_id="fastvideo.models.dits.ltx2:LTX2Transformer3DModel",
|
||||
factory=lambda inst: ToyDiT(seed=seed),
|
||||
resident_for=["ltx2_single"],
|
||||
required_for={"t2v"}),
|
||||
}
|
||||
loops = {
|
||||
"ltx2_single": LoopSpec(loop_id="ltx2_single", kind=LoopKind.DIFFUSION_DENOISE,
|
||||
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost,
|
||||
shared_weight_components=["transformer"], cache_policy=["feature"],
|
||||
loop_factory=single_factory),
|
||||
"ltx2_single":
|
||||
LoopSpec(loop_id="ltx2_single",
|
||||
kind=LoopKind.DIFFUSION_DENOISE,
|
||||
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
|
||||
step_cost_model=cost,
|
||||
shared_weight_components=["transformer"],
|
||||
cache_policy=["feature"],
|
||||
loop_factory=single_factory),
|
||||
}
|
||||
card = ModelCard(
|
||||
model_id=model_id, family="ltx2",
|
||||
components=components, loops=loops,
|
||||
model_id=model_id,
|
||||
family="ltx2",
|
||||
components=components,
|
||||
loops=loops,
|
||||
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.VAE_DECODE),
|
||||
recipe=RecipeSpec(method="base", assumes_loop="ltx2_single",
|
||||
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
|
||||
sampling_defaults=SamplingDefaults(
|
||||
num_steps=40, guidance_scale=3.0, height=512, width=768, num_frames=121, fps=24,
|
||||
negative_prompt=LTX2_NEG),
|
||||
recipe=RecipeSpec(method="base",
|
||||
assumes_loop="ltx2_single",
|
||||
assumes_precision="float32",
|
||||
consistency_required=ConsistencyLevel.C1),
|
||||
sampling_defaults=SamplingDefaults(num_steps=40,
|
||||
guidance_scale=3.0,
|
||||
height=512,
|
||||
width=768,
|
||||
num_frames=121,
|
||||
fps=24,
|
||||
negative_prompt=LTX2_NEG),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24,
|
||||
reuse_across_requests=True)},
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
|
||||
default_plan=ParallelPlan.single()),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
|
||||
@@ -166,47 +240,79 @@ def build_ltx2_3_card(model_id: str = "ltx2.3-distilled") -> ModelCard:
|
||||
return LTX23DenoiseLoop(loop_id="ltx2_3", sigmas=BASE_SIGMAS, cfg_scale=1.0, stg_scale=0.0, cost=cost)
|
||||
|
||||
components = {
|
||||
"text_encoder": ComponentSpec(component_id="text_encoder", kind="text_encoder",
|
||||
load_id="transformers:AutoModel", # LTX2GemmaTextEncoderModel
|
||||
factory=lambda inst: ToyTextEncoder(), required_for={"t2v", "t2vs"}),
|
||||
"vae": ComponentSpec(component_id="vae", kind="vae",
|
||||
load_id="fastvideo.models.vaes.ltx2vae:VideoDecoder",
|
||||
factory=lambda inst: ToyVAE(), required_for={"t2v", "t2vs"}),
|
||||
"transformer": ComponentSpec(component_id="transformer", kind="dit",
|
||||
load_id="fastvideo.models.dits.ltx2:LTX2Transformer3DModel",
|
||||
factory=lambda inst: ToyDiT(seed=seed),
|
||||
resident_for=["ltx2_3"], required_for={"t2v", "t2vs"}),
|
||||
"text_encoder":
|
||||
ComponentSpec(
|
||||
component_id="text_encoder",
|
||||
kind="text_encoder",
|
||||
load_id="transformers:AutoModel", # LTX2GemmaTextEncoderModel
|
||||
factory=lambda inst: ToyTextEncoder(),
|
||||
required_for={"t2v", "t2vs"}),
|
||||
"vae":
|
||||
ComponentSpec(component_id="vae",
|
||||
kind="vae",
|
||||
load_id="fastvideo.models.vaes.ltx2vae:VideoDecoder",
|
||||
factory=lambda inst: ToyVAE(),
|
||||
required_for={"t2v", "t2vs"}),
|
||||
"transformer":
|
||||
ComponentSpec(component_id="transformer",
|
||||
kind="dit",
|
||||
load_id="fastvideo.models.dits.ltx2:LTX2Transformer3DModel",
|
||||
factory=lambda inst: ToyDiT(seed=seed),
|
||||
resident_for=["ltx2_3"],
|
||||
required_for={"t2v", "t2vs"}),
|
||||
# audio branch (T2VS only): AudioDecoder (latent->mel) + Vocoder (mel->waveform). Not built for t2v.
|
||||
"audio_vae": ComponentSpec(component_id="audio_vae", kind="audio_vae",
|
||||
load_id="fastvideo.models.audio.ltx2_audio_vae:AudioDecoder",
|
||||
factory=lambda inst: ToyAudioVAE(),
|
||||
required_for={"t2vs"}, optional_for={"t2v"}),
|
||||
"vocoder": ComponentSpec(component_id="vocoder", kind="vocoder",
|
||||
load_id="fastvideo.models.audio.ltx2_audio_vae:Vocoder",
|
||||
factory=lambda inst: ToyVocoder(),
|
||||
required_for={"t2vs"}, optional_for={"t2v"}),
|
||||
"audio_vae":
|
||||
ComponentSpec(component_id="audio_vae",
|
||||
kind="audio_vae",
|
||||
load_id="fastvideo.models.audio.ltx2_audio_vae:AudioDecoder",
|
||||
factory=lambda inst: ToyAudioVAE(),
|
||||
required_for={"t2vs"},
|
||||
optional_for={"t2v"}),
|
||||
"vocoder":
|
||||
ComponentSpec(component_id="vocoder",
|
||||
kind="vocoder",
|
||||
load_id="fastvideo.models.audio.ltx2_audio_vae:Vocoder",
|
||||
factory=lambda inst: ToyVocoder(),
|
||||
required_for={"t2vs"},
|
||||
optional_for={"t2v"}),
|
||||
}
|
||||
loops = {
|
||||
"ltx2_3": LoopSpec(loop_id="ltx2_3", kind=LoopKind.DIFFUSION_DENOISE,
|
||||
work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=cost,
|
||||
shared_weight_components=["transformer"], cache_policy=["feature"],
|
||||
loop_factory=loop_factory),
|
||||
"ltx2_3":
|
||||
LoopSpec(loop_id="ltx2_3",
|
||||
kind=LoopKind.DIFFUSION_DENOISE,
|
||||
work_unit_kind=WorkUnitKind.DIFFUSION_STEP,
|
||||
step_cost_model=cost,
|
||||
shared_weight_components=["transformer"],
|
||||
cache_policy=["feature"],
|
||||
loop_factory=loop_factory),
|
||||
}
|
||||
card = ModelCard(
|
||||
model_id=model_id, family="ltx2",
|
||||
components=components, loops=loops,
|
||||
model_id=model_id,
|
||||
family="ltx2",
|
||||
components=components,
|
||||
loops=loops,
|
||||
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.TEXT_TO_VIDEO_SOUND,
|
||||
Capability.VAE_DECODE),
|
||||
recipe=RecipeSpec(method="distilled", parents=["ltx2.3-base"], assumes_loop="ltx2_3",
|
||||
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
|
||||
recipe=RecipeSpec(method="distilled",
|
||||
parents=["ltx2.3-base"],
|
||||
assumes_loop="ltx2_3",
|
||||
assumes_precision="float32",
|
||||
consistency_required=ConsistencyLevel.C1),
|
||||
sampling_defaults=SamplingDefaults( # joint A/V: per-modality cfg video 3 / audio 7
|
||||
num_steps=30, guidance_scale=3.0, height=512, width=768, num_frames=121, fps=24,
|
||||
negative_prompt=LTX2_NEG, guidance_per_modality={"video": 3.0, "audio": 7.0}),
|
||||
num_steps=30,
|
||||
guidance_scale=3.0,
|
||||
height=512,
|
||||
width=768,
|
||||
num_frames=121,
|
||||
fps=24,
|
||||
negative_prompt=LTX2_NEG,
|
||||
guidance_per_modality={
|
||||
"video": 3.0,
|
||||
"audio": 7.0
|
||||
}),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24,
|
||||
reuse_across_requests=True)},
|
||||
caches={"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True)},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
|
||||
default_plan=ParallelPlan.single()),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
)
|
||||
return card.validate()
|
||||
|
||||
@@ -28,9 +28,17 @@ from v2.recipes._prompts import WAN_NEG_CN
|
||||
from v2.recipes.wan_causal.loop import ChunkRolloutLoop
|
||||
|
||||
|
||||
def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b", *,
|
||||
num_chunks: int = 3, chunk_size: int = 2, steps_per_chunk: int = 2,
|
||||
def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b",
|
||||
*,
|
||||
num_chunks: int = 7,
|
||||
chunk_size: int = 3,
|
||||
steps_per_chunk: int = 4,
|
||||
training_mode: bool = False) -> ModelCard:
|
||||
# Self-forcing DMD reference (fastvideo SelfForcingWanT2V480PConfig + causal_denoising.py): the
|
||||
# distilled student is CFG-FREE (guidance 1.0, single forward/step) and denoises with the 4-step DMD
|
||||
# schedule (dmd_denoising_steps=[1000,750,500,250], warped — reproduced by FlowShiftPolicy(5.0) over 4
|
||||
# steps); the native causal block is chunk_size=3 latent frames (7 chunks x 3 = 21 latent -> 81 video
|
||||
# frames). Earlier defaults (CFG 6.0, 2 steps, block 2) overcooked + under-denoised the output.
|
||||
seed = _seed_from(model_id)
|
||||
cost = CostModel(kind=WorkUnitKind.CHUNK_STEP, base_seconds=1e-4, per_unit_seconds=1e-7)
|
||||
cfg = ClassicCFG()
|
||||
@@ -38,48 +46,73 @@ def build_wan_causal_card(model_id: str = "wan-causal-sf-1.3b", *,
|
||||
precision = PrecisionPolicy(compute_dtype="float32", scheduler_step_in_fp32=True)
|
||||
|
||||
def loop_factory():
|
||||
return ChunkRolloutLoop(loop_id="chunk_rollout", num_chunks=num_chunks, chunk_size=chunk_size,
|
||||
steps_per_chunk=steps_per_chunk, cfg=cfg, flow_shift=flow,
|
||||
precision=precision, cost=cost)
|
||||
return ChunkRolloutLoop(loop_id="chunk_rollout",
|
||||
num_chunks=num_chunks,
|
||||
chunk_size=chunk_size,
|
||||
steps_per_chunk=steps_per_chunk,
|
||||
cfg=cfg,
|
||||
flow_shift=flow,
|
||||
precision=precision,
|
||||
cost=cost)
|
||||
|
||||
components = {
|
||||
"text_encoder": ComponentSpec(component_id="text_encoder", kind="text_encoder",
|
||||
load_id="fastvideo.models.encoders.t5:T5EncoderModel",
|
||||
factory=lambda inst: ToyTextEncoder(), required_for={"t2v", "v2w"}),
|
||||
"vae": ComponentSpec(component_id="vae", kind="vae",
|
||||
load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan",
|
||||
factory=lambda inst: ToyVAE(), required_for={"t2v", "v2w"}),
|
||||
"transformer": ComponentSpec(component_id="transformer", kind="dit",
|
||||
load_id="fastvideo.models.dits.causal_wanvideo:CausalWanTransformer3DModel",
|
||||
factory=lambda inst: ToyDiT(seed=seed),
|
||||
resident_for=["chunk_rollout"], required_for={"t2v", "v2w"}),
|
||||
"text_encoder":
|
||||
ComponentSpec(component_id="text_encoder",
|
||||
kind="text_encoder",
|
||||
load_id="fastvideo.models.encoders.t5:T5EncoderModel",
|
||||
factory=lambda inst: ToyTextEncoder(),
|
||||
required_for={"t2v", "v2w"}),
|
||||
"vae":
|
||||
ComponentSpec(component_id="vae",
|
||||
kind="vae",
|
||||
load_id="fastvideo.models.vaes.wanvae:AutoencoderKLWan",
|
||||
factory=lambda inst: ToyVAE(),
|
||||
required_for={"t2v", "v2w"}),
|
||||
"transformer":
|
||||
ComponentSpec(component_id="transformer",
|
||||
kind="dit",
|
||||
load_id="fastvideo.models.dits.causal_wanvideo:CausalWanTransformer3DModel",
|
||||
factory=lambda inst: ToyDiT(seed=seed),
|
||||
resident_for=["chunk_rollout"],
|
||||
required_for={"t2v", "v2w"}),
|
||||
}
|
||||
loops = {
|
||||
"chunk_rollout": LoopSpec(loop_id="chunk_rollout", kind=LoopKind.CHUNK_ROLLOUT,
|
||||
work_unit_kind=WorkUnitKind.CHUNK_STEP, step_cost_model=cost,
|
||||
shared_weight_components=["transformer"],
|
||||
cache_policy=["feature", "slab_kv"], loop_factory=loop_factory),
|
||||
"chunk_rollout":
|
||||
LoopSpec(loop_id="chunk_rollout",
|
||||
kind=LoopKind.CHUNK_ROLLOUT,
|
||||
work_unit_kind=WorkUnitKind.CHUNK_STEP,
|
||||
step_cost_model=cost,
|
||||
shared_weight_components=["transformer"],
|
||||
cache_policy=["feature", "slab_kv"],
|
||||
loop_factory=loop_factory),
|
||||
}
|
||||
card = ModelCard(
|
||||
model_id=model_id, family="wan",
|
||||
components=components, loops=loops,
|
||||
capabilities=CapabilityMatrix.of(
|
||||
Capability.TEXT_TO_VIDEO, Capability.STREAMING_VIDEO_CONTINUATION,
|
||||
Capability.VAE_DECODE, Capability.POLICY_ROLLOUT),
|
||||
recipe=RecipeSpec(method="self_forcing", parents=["wan2.1-1.3b"], assumes_loop="chunk_rollout",
|
||||
assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
|
||||
model_id=model_id,
|
||||
family="wan",
|
||||
components=components,
|
||||
loops=loops,
|
||||
capabilities=CapabilityMatrix.of(Capability.TEXT_TO_VIDEO, Capability.STREAMING_VIDEO_CONTINUATION,
|
||||
Capability.VAE_DECODE, Capability.POLICY_ROLLOUT),
|
||||
recipe=RecipeSpec(method="self_forcing",
|
||||
parents=["wan2.1-1.3b"],
|
||||
assumes_loop="chunk_rollout",
|
||||
assumes_precision="float32",
|
||||
consistency_required=ConsistencyLevel.C1),
|
||||
parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
|
||||
caches={
|
||||
"feature": CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True),
|
||||
"slab_kv": CacheContract(cache_class="slab_kv", max_bytes=1 << 26, reuse_across_requests=False,
|
||||
per_component={"window": (num_chunks if training_mode else 2)},
|
||||
training_mode_disables_recycle=training_mode),
|
||||
"feature":
|
||||
CacheContract(cache_class="feature", max_bytes=1 << 24, reuse_across_requests=True),
|
||||
"slab_kv":
|
||||
CacheContract(cache_class="slab_kv",
|
||||
max_bytes=1 << 26,
|
||||
reuse_across_requests=False,
|
||||
per_component={"window": (num_chunks if training_mode else 2)},
|
||||
training_mode_disables_recycle=training_mode),
|
||||
},
|
||||
precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
|
||||
default_plan=ParallelPlan.single()),
|
||||
parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()], default_plan=ParallelPlan.single()),
|
||||
sampling_defaults=SamplingDefaults(
|
||||
num_steps=50, guidance_scale=6.0, height=480, width=832, num_frames=81, fps=16,
|
||||
negative_prompt=WAN_NEG_CN),
|
||||
num_steps=4, guidance_scale=1.0, height=480, width=832, num_frames=81, fps=16,
|
||||
negative_prompt=WAN_NEG_CN), # CFG-free distilled, 4 DMD steps (see build_wan_causal_card)
|
||||
)
|
||||
return card.validate()
|
||||
|
||||
@@ -38,9 +38,9 @@ def test_causal_chunk_rollout_frame_count():
|
||||
eng = _eng()
|
||||
out = eng.run(make_request(TaskType.T2V, "wan-causal-sf-1.3b", "a river",
|
||||
diffusion=DiffusionParams(seed=3)))
|
||||
assert out.metrics["chunks"] == 3.0
|
||||
# 3 chunks × 2 frames/chunk concatenated over the temporal axis
|
||||
assert out.artifacts["latents"].latent.shape[1] == 6
|
||||
assert out.metrics["chunks"] == 7.0
|
||||
# 7 chunks × 3 latent frames/chunk (the self-forcing native block) over the temporal axis -> 81 video
|
||||
assert out.artifacts["latents"].latent.shape[1] == 21
|
||||
|
||||
|
||||
def test_determinism_same_seed_bit_identical():
|
||||
|
||||
@@ -41,7 +41,7 @@ def test_wan22_a14b_defaults():
|
||||
|
||||
def test_wan_causal_defaults():
|
||||
sd = build_wan_causal_card().sampling_defaults
|
||||
assert (sd.num_steps, sd.guidance_scale, sd.height, sd.width, sd.num_frames, sd.fps) == (50, 6.0, 480, 832, 81, 16)
|
||||
assert (sd.num_steps, sd.guidance_scale, sd.height, sd.width, sd.num_frames, sd.fps) == (4, 1.0, 480, 832, 81, 16)
|
||||
|
||||
|
||||
def test_ltx2_two_stage_defaults():
|
||||
|
||||
@@ -34,7 +34,7 @@ def test_session_persists_world_state_across_acts():
|
||||
assert s.session.kv_handle is s.world_context # the persistent state lives on the Session
|
||||
s.act("turn left", seed=1)
|
||||
assert s.step == 2 and len(s.world_context) == 2 # still carrying world state
|
||||
assert len(s.frames()) == 6 # 3 chunks/act × 2 acts streamed
|
||||
assert len(s.frames()) == 14 # 7 chunks/act × 2 acts streamed
|
||||
|
||||
|
||||
def test_continuation_depends_on_history():
|
||||
|
||||
Reference in New Issue
Block a user