From 9e8d51584204ffd0def4a67f8268a3aeb79af88d Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Fri, 19 Jun 2026 11:57:48 +0000 Subject: [PATCH] [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. --- v2/platform/backends/torch_backend.py | 191 +++++++++------- v2/recipes/ltx2/card.py | 308 +++++++++++++++++--------- v2/recipes/wan_causal/card.py | 101 ++++++--- v2/tests/test_inference.py | 6 +- v2/tests/test_sampling_defaults.py | 2 +- v2/tests/test_world_session.py | 2 +- 6 files changed, 391 insertions(+), 219 deletions(-) diff --git a/v2/platform/backends/torch_backend.py b/v2/platform/backends/torch_backend.py index 13cc367c5..22c4bf9d3 100644 --- a/v2/platform/backends/torch_backend.py +++ b/v2/platform/backends/torch_backend.py @@ -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) diff --git a/v2/recipes/ltx2/card.py b/v2/recipes/ltx2/card.py index bdcb6341a..86b775451 100644 --- a/v2/recipes/ltx2/card.py +++ b/v2/recipes/ltx2/card.py @@ -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() diff --git a/v2/recipes/wan_causal/card.py b/v2/recipes/wan_causal/card.py index fff0b94c0..0cef45f46 100644 --- a/v2/recipes/wan_causal/card.py +++ b/v2/recipes/wan_causal/card.py @@ -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() diff --git a/v2/tests/test_inference.py b/v2/tests/test_inference.py index bc45df963..3dfe2b0ea 100644 --- a/v2/tests/test_inference.py +++ b/v2/tests/test_inference.py @@ -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(): diff --git a/v2/tests/test_sampling_defaults.py b/v2/tests/test_sampling_defaults.py index 406e0425e..c4f047c7d 100644 --- a/v2/tests/test_sampling_defaults.py +++ b/v2/tests/test_sampling_defaults.py @@ -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(): diff --git a/v2/tests/test_world_session.py b/v2/tests/test_world_session.py index c1c6c3e25..96d1b0c89 100644 --- a/v2/tests/test_world_session.py +++ b/v2/tests/test_world_session.py @@ -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():