From 2000e8a0fffcf71b2f4ba8bcb7cc14033803edc0 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 22 Jul 2026 15:20:45 -0700 Subject: [PATCH] =?UTF-8?q?[feat]:=20align=20the=20DiT=20with=20official?= =?UTF-8?q?=20=E2=80=94=20vendored=20WanModel=20(pinned=209737cba)=20+=20o?= =?UTF-8?q?fficial=20weights=20+=20official=20autocast=20regime?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- AGENTS.md | 20 +- fastvideo2/card.py | 5 +- fastvideo2/engine.py | 12 +- fastvideo2/loading.py | 5 +- fastvideo2/wan21/anchor.py | 7 +- fastvideo2/wan21/attention.py | 200 +++++++++++ fastvideo2/wan21/card.py | 7 +- fastvideo2/wan21/loop.py | 29 +- fastvideo2/wan21/model.py | 640 ++++++++++++++++++++++++++++++++++ fastvideo2/wan21/reference.py | 45 +-- 10 files changed, 923 insertions(+), 47 deletions(-) create mode 100644 fastvideo2/wan21/attention.py create mode 100644 fastvideo2/wan21/model.py diff --git a/AGENTS.md b/AGENTS.md index 05bab9890..090809ea2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -24,8 +24,9 @@ surfaces. Read `README.md` first. paths. If it can't round-trip through JSON, it doesn't belong on a card. 2. **Import direction is one-way:** `card` → `loop`/`pipeline` → `engine` → `verify`. Family packages depend on core, never the reverse. - `wan21/reference.py` imports **nothing** from fastvideo2 and nothing from - fastvideo2 outside `verify.py` may import it. + `wan21/reference.py` imports only the vendored official model file + (`wan21/model.py`, itself standalone) — never core/runtime modules — and + nothing outside `verify.py` may import the reference. 3. **Loop modules import torch-free** (torch inside methods) so contracts validate anywhere; `import fastvideo2` must work without torch installed. 4. **Model-specific inputs are typed** (`WanForwardInputs`); never add an @@ -35,12 +36,15 @@ surfaces. Read `README.md` first. `reference.py` to make a failing gate pass — say so instead. 6. **One catalog.** New servable ⇒ card constant + registry entry. No parallel model lists. -7. **Official implementations are the numerics authority.** diffusers (or any - port) is a backend whose fidelity is certified against captured official - goldens (`verify --anchor`), never assumed. The official repo enters this - codebase only as a commit hash in a goldens manifest — never as a - dependency, submodule, or vendored code. When conventions conflict, - official wins; goldens are re-captured deliberately, like re-blessing. +7. **Official implementations are the numerics authority.** Where fidelity is + the requirement, run the authors' modeling code: the Wan DiT is vendored + VERBATIM from the pinned official commit (`wan21/model.py`, provenance in + its header; re-vendor deliberately, like re-capturing goldens). Ports + (diffusers, etc.) may serve components only with anchor certification + (`verify --anchor`), never on trust. The official repo never becomes a + package dependency or submodule, and the anchor always compares against + captured goldens + hashes, not live upstream code. When conventions + conflict, official wins. 8. **One environment for goldens and gates.** The supported env is python 3.12 + torch 2.12 (the fastvideo cluster venv). Goldens are captured with that same env — official code rides `PYTHONPATH`, its extra deps go to a pip diff --git a/fastvideo2/card.py b/fastvideo2/card.py index 01d3e4c0c..2f5ae446f 100644 --- a/fastvideo2/card.py +++ b/fastvideo2/card.py @@ -48,9 +48,10 @@ class ComponentSpec: """One weight-bearing (or processing) component of the artifact.""" component_id: str kind: str # dit | vae | text_encoder | tokenizer - module: str # loader reference, e.g. "diffusers:WanTransformer3DModel" - subfolder: str # subfolder in the diffusers checkpoint layout + module: str # loader reference, e.g. "fastvideo2.wan21.model:WanModel" + subfolder: str # subfolder in the checkpoint layout ("" = repo root) dtype: str = "bf16" # bf16 | fp32 | "" (dtype-less, e.g. tokenizer) + source: str = "" # weights repo override; "" = the card-level `weights` @dataclass(frozen=True) diff --git a/fastvideo2/engine.py b/fastvideo2/engine.py index 21ea85c9f..75f957b1b 100644 --- a/fastvideo2/engine.py +++ b/fastvideo2/engine.py @@ -74,6 +74,7 @@ class Instance: self.device = device self.root = resolve_weights(card, root) self._components: dict[str, Any] = {} + self._source_roots: dict[str, str] = {} self._loops: dict[str, Any] = {} def component(self, component_id: str) -> Any: @@ -81,9 +82,18 @@ class Instance: spec = self.card.components.get(component_id) if spec is None: raise KeyError(f"component {component_id!r} not declared on card {self.card.model_id!r}") - self._components[component_id] = load_component(spec, self.root, self.device) + root = self._source_root(spec.source) if spec.source else self.root + self._components[component_id] = load_component(spec, root, self.device) return self._components[component_id] + def _source_root(self, source: str) -> str: + """Resolve a per-component weights source (e.g. the official-layout + transformer repo) through the same snapshot cache as card weights.""" + if source not in self._source_roots: + from huggingface_hub import snapshot_download + self._source_roots[source] = snapshot_download(source) + return self._source_roots[source] + def loop(self, loop_id: str) -> Any: if loop_id not in self._loops: spec = self.card.loops.get(loop_id) diff --git a/fastvideo2/loading.py b/fastvideo2/loading.py index 49680424a..e5566e6ae 100644 --- a/fastvideo2/loading.py +++ b/fastvideo2/loading.py @@ -43,12 +43,13 @@ def load_component(spec: ComponentSpec, root: str, device: str = "cpu") -> Any: raise CardError(f"component {spec.component_id!r}: bad module ref {spec.module!r}") cls = getattr(importlib.import_module(lib), clsname) + sub = {"subfolder": spec.subfolder} if spec.subfolder else {} if spec.kind == "tokenizer": - return cls.from_pretrained(root, subfolder=spec.subfolder) + return cls.from_pretrained(root, **sub) import torch dtype = getattr(torch, _DTYPES[spec.dtype]) - module = cls.from_pretrained(root, subfolder=spec.subfolder, torch_dtype=dtype) + module = cls.from_pretrained(root, torch_dtype=dtype, **sub) # NOTE: no blanket .to(dtype) here. Loaders keep selected submodules fp32 # inside a bf16 model on purpose (Wan: time_embedder, norms), so a loaded # module is legitimately mixed-dtype. Consumers must not infer a compute diff --git a/fastvideo2/wan21/anchor.py b/fastvideo2/wan21/anchor.py index 3deb424cd..738b702f1 100644 --- a/fastvideo2/wan21/anchor.py +++ b/fastvideo2/wan21/anchor.py @@ -68,10 +68,11 @@ def fastvideo2_adapter(root: str, device: str) -> dict[str, Callable]: def dit_forward(x: np.ndarray, t: float, context: np.ndarray) -> np.ndarray: from fastvideo2.wan21.loop import WanForwardInputs dit = inst.component("transformer") - xt = torch.from_numpy(x).to(device, torch.bfloat16) - ct = torch.from_numpy(context).to(device, torch.bfloat16)[None] + xt = torch.from_numpy(x).to(device, torch.float32) + ct = torch.from_numpy(context).to(device, torch.float32)[None] tt = torch.tensor([t], device=device, dtype=torch.float32) - with torch.no_grad(): + # official precision regime, same as the production loop + with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16): out = WanForwardInputs(xt, tt, ct).forward(dit) return out.to(torch.float32).cpu().numpy() diff --git a/fastvideo2/wan21/attention.py b/fastvideo2/wan21/attention.py new file mode 100644 index 000000000..6e9740e42 --- /dev/null +++ b/fastvideo2/wan21/attention.py @@ -0,0 +1,200 @@ +# Vendored from the official Wan2.1 implementation (wan/modules/attention.py, +# same commit/license as model.py). Local modification (only this): when flash +# attention is unavailable, inputs are fp32 (the exact-math anchor path), or we +# are not on CUDA, fall back to torch SDPA *honoring q_lens/k_lens masking and +# the input dtype* — the upstream SDPA fallback drops masking and force-casts. +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import torch + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +import warnings + + +def _sdpa_attention(q, k, v, q_lens=None, k_lens=None, softmax_scale=None, causal=False): + """Mask- and dtype-honoring SDPA fallback (local modification).""" + qh, kh, vh = (u.transpose(1, 2) for u in (q, k, v)) # [B, H, L, C] + mask = None + if k_lens is not None: + idx = torch.arange(kh.shape[2], device=kh.device)[None] + mask = (idx < k_lens.to(kh.device)[:, None])[:, None, None] + out = torch.nn.functional.scaled_dot_product_attention( + qh, kh, vh, attn_mask=mask, is_causal=causal, scale=softmax_scale) + return out.transpose(1, 2) + +__all__ = [ + 'flash_attention', + 'attention', +] + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + if (not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE) + or q.dtype == torch.float32 or q.device.type != 'cuda'): + return _sdpa_attention(q, k, v, q_lens=q_lens, k_lens=k_lens, + softmax_scale=softmax_scale, causal=causal) + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == 'cuda' and q.size(-1) <= 256 + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor( + [lq] * b, dtype=torch.int32).to( + device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor( + [lk] * b, dtype=torch.int32).to( + device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + 'Flash attention 3 is not available, use flash attention 2 instead.' + ) + + # apply attention + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + seqused_q=None, + seqused_k=None, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic)[0].unflatten(0, (b, lq)) + else: + assert FLASH_ATTN_2_AVAILABLE + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic).unflatten(0, (b, lq)) + + # output + return x.type(out_dtype) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE: + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=fa_version, + ) + else: + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' + ) + attn_mask = None + + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + + out = out.transpose(1, 2).contiguous() + return out diff --git a/fastvideo2/wan21/card.py b/fastvideo2/wan21/card.py index b1197ae7f..f4bdc4fe1 100644 --- a/fastvideo2/wan21/card.py +++ b/fastvideo2/wan21/card.py @@ -25,9 +25,12 @@ WAN21_T2V_1_3B = ModelCard( "text_encoder": ComponentSpec("text_encoder", kind="text_encoder", module="transformers:UMT5EncoderModel", subfolder="text_encoder", dtype="bf16"), + # The DiT is the OFFICIAL modeling code (vendored verbatim) running the + # official-layout weights — alignment by construction, not by porting. "transformer": ComponentSpec("transformer", kind="dit", - module="diffusers:WanTransformer3DModel", - subfolder="transformer", dtype="bf16"), + module="fastvideo2.wan21.model:WanModel", + subfolder="", dtype="bf16", + source="Wan-AI/Wan2.1-T2V-1.3B"), # Wan VAE runs fp32 (bf16 decode visibly degrades output). "vae": ComponentSpec("vae", kind="vae", module="diffusers:AutoencoderKLWan", diff --git a/fastvideo2/wan21/loop.py b/fastvideo2/wan21/loop.py index 9840c3b07..65273e1f2 100644 --- a/fastvideo2/wan21/loop.py +++ b/fastvideo2/wan21/loop.py @@ -32,14 +32,21 @@ class WanForwardInputs: """Everything one Wan DiT forward consumes — typed, so a new conditioning channel is a new field the adapter must handle, never a silently-dropped kwarg.""" - latent: Any # [1, C, T, h, w], transformer dtype + latent: Any # [1, C, T, h, w], fp32 (official convention: cast under autocast) timestep: Any # [1] float tensor, 0..1000 scale (sigma * 1000) - text: Any # [1, 512, 4096] UMT5 embeddings + text: Any # [1, 512, 4096] UMT5 embeddings (the zero-pad is load-bearing) def forward(self, transformer: Any) -> Any: - """The one place the diffusers calling convention lives.""" - return transformer(hidden_states=self.latent, timestep=self.timestep, - encoder_hidden_states=self.text, return_dict=False)[0] + """The one place the official calling convention lives: list-based, + unbatched, with the max sequence length derived from the patch grid. + Callers provide official's precision regime (autocast bf16 on CUDA; + plain fp32 for the exact-math path).""" + pt, ph, pw = transformer.patch_size + _, _, f, h, w = self.latent.shape + seq_len = (f // pt) * (h // ph) * (w // pw) + out = transformer([self.latent[0]], t=self.timestep, + context=[self.text[0]], seq_len=seq_len)[0] + return out[None] class WanDenoiseLoop: @@ -92,13 +99,17 @@ class WanDenoiseLoop: import torch dit = instance.component("transformer") t = torch.tensor([sigma * 1000.0], device=x.device, dtype=torch.float32) - x_in = x.to(dit_dtype) - with torch.no_grad(): - v_cond = WanForwardInputs(x_in, t, text).forward(dit).to(torch.float32) + # Official precision regime: fp32 latents under autocast in the + # card-declared compute dtype (the model's own fp32 islands — + # rope, time embedding, modulation — are autocast-disabled inside). + autocast = torch.amp.autocast("cuda", dtype=dit_dtype, + enabled=x.is_cuda and dit_dtype != torch.float32) + with torch.no_grad(), autocast: + v_cond = WanForwardInputs(x, t, text).forward(dit).to(torch.float32) if guidance == 1.0: v = v_cond else: - v_neg = WanForwardInputs(x_in, t, neg).forward(dit).to(torch.float32) + v_neg = WanForwardInputs(x, t, neg).forward(dit).to(torch.float32) v = v_neg + guidance * (v_cond - v_neg) # Euler step in fp32: x' = x + (sigma_next - sigma) * v return {"latents": x + (sigma_next - sigma) * v} diff --git a/fastvideo2/wan21/model.py b/fastvideo2/wan21/model.py new file mode 100644 index 000000000..cba123c5f --- /dev/null +++ b/fastvideo2/wan21/model.py @@ -0,0 +1,640 @@ +# Vendored VERBATIM from the official Wan2.1 implementation: +# https://github.com/Wan-Video/Wan2.1 @ 9737cba9c1c3c4d04b33fcad41c111989865d315 +# (wan/modules/model.py, Apache-2.0, Copyright The Alibaba Wan Team) +# This IS fastvideo2's production Wan DiT: per the authority rule, alignment +# with official is achieved by running the authors' modeling code, not a port. +# Local modifications (only these): +# * absolute import of the vendored attention module; +# Re-vendor deliberately (like re-capturing goldens) when bumping the pinned +# commit; the anchor gate detects drift either way. +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import math + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin + +from fastvideo2.wan21.attention import flash_attention + +__all__ = ['WanModel'] + +T5_CONTEXT_TOKEN_NUMBER = 512 +FIRST_LAST_FRAME_CONTEXT_TOKEN_NUMBER = 257 * 2 + + +def sinusoidal_embedding_1d(dim, position): + # preprocess + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer( + position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +@amp.autocast(enabled=False) +def rope_params(max_seq_len, dim, theta=10000): + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim))) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +@amp.autocast(enabled=False) +def rope_apply(x, grid_sizes, freqs): + n, c = x.size(2), x.size(3) // 2 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape( + seq_len, n, -1, 2)) + freqs_i = torch.cat([ + freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], + dim=-1).reshape(seq_len, 1, -1) + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).float() + + +class WanRMSNorm(nn.Module): + + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return super().forward(x.float()).type_as(x) + + +class WanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, seq_lens, grid_sizes, freqs): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + x = flash_attention( + q=rope_apply(q, grid_sizes, freqs), + k=rope_apply(k, grid_sizes, freqs), + v=v, + k_lens=seq_lens, + window_size=self.window_size) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanT2VCrossAttention(WanSelfAttention): + + def forward(self, x, context, context_lens): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.q(x)).view(b, -1, n, d) + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + # compute attention + x = flash_attention(q, k, v, k_lens=context_lens) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanI2VCrossAttention(WanSelfAttention): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + super().__init__(dim, num_heads, window_size, qk_norm, eps) + + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + # self.alpha = nn.Parameter(torch.zeros((1, ))) + self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, context, context_lens): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + image_context_length = context.shape[1] - T5_CONTEXT_TOKEN_NUMBER + context_img = context[:, :image_context_length] + context = context[:, image_context_length:] + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.q(x)).view(b, -1, n, d) + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d) + v_img = self.v_img(context_img).view(b, -1, n, d) + img_x = flash_attention(q, k_img, v_img, k_lens=None) + # compute attention + x = flash_attention(q, k, v, k_lens=context_lens) + + # output + x = x.flatten(2) + img_x = img_x.flatten(2) + x = x + img_x + x = self.o(x) + return x + + +WAN_CROSSATTENTION_CLASSES = { + 't2v_cross_attn': WanT2VCrossAttention, + 'i2v_cross_attn': WanI2VCrossAttention, +} + + +class WanAttentionBlock(nn.Module): + + def __init__(self, + cross_attn_type, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, + num_heads, + (-1, -1), + qk_norm, + eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + assert e.dtype == torch.float32 + with amp.autocast(dtype=torch.float32): + e = (self.modulation + e).chunk(6, dim=1) + assert e[0].dtype == torch.float32 + + # self-attention + y = self.self_attn( + self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, grid_sizes, + freqs) + with amp.autocast(dtype=torch.float32): + x = x + y * e[2] + + # cross-attention & ffn function + def cross_attn_ffn(x, context, context_lens, e): + x = x + self.cross_attn(self.norm3(x), context, context_lens) + y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3]) + with amp.autocast(dtype=torch.float32): + x = x + y * e[5] + return x + + x = cross_attn_ffn(x, context, context_lens, e) + return x + + +class Head(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, C] + """ + assert e.dtype == torch.float32 + with amp.autocast(dtype=torch.float32): + e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + e[1]) + e[0])) + return x + + +class MLPProj(torch.nn.Module): + + def __init__(self, in_dim, out_dim, flf_pos_emb=False): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim), + torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim), + torch.nn.LayerNorm(out_dim)) + if flf_pos_emb: # NOTE: we only use this for `flf2v` + self.emb_pos = nn.Parameter( + torch.zeros(1, FIRST_LAST_FRAME_CONTEXT_TOKEN_NUMBER, 1280)) + + def forward(self, image_embeds): + if hasattr(self, 'emb_pos'): + bs, n, d = image_embeds.shape + image_embeds = image_embeds.view(-1, 2 * n, d) + image_embeds = image_embeds + self.emb_pos + clip_extra_context_tokens = self.proj(image_embeds) + return clip_extra_context_tokens + + +class WanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size' + ] + _no_split_modules = ['WanAttentionBlock'] + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + r""" + Initialize the diffusion model backbone. + + Args: + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) or 'flf2v' (first-last-frame-to-video) or 'vace' + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + window_size (`tuple`, *optional*, defaults to (-1, -1)): + Window size for local attention (-1 indicates global attention) + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v', 'flf2v', 'vace'] + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn' + self.blocks = nn.ModuleList([ + WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, + window_size, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ]) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + self.freqs = torch.cat([ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)) + ], + dim=1) + + if model_type == 'i2v' or model_type == 'flf2v': + self.img_emb = MLPProj(1280, dim, flf_pos_emb=model_type == 'flf2v') + + # initialize weights + self.init_weights() + + def forward( + self, + x, + t, + context, + seq_len, + clip_fea=None, + y=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode or first-last-frame-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v' or self.model_type == 'flf2v': + assert clip_fea is not None and y is not None + # params + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([ + torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], + dim=1) for u in x + ]) + + # time embeddings + with amp.autocast(dtype=torch.float32): + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t).float()) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)) + assert e.dtype == torch.float32 and e0.dtype == torch.float32 + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat( + [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + if clip_fea is not None: + context_clip = self.img_emb(clip_fea) # bs x 257 (x2) x dim + context = torch.concat([context_clip, context], dim=1) + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens) + + for block in self.blocks: + x = block(x, **kwargs) + + # head + x = self.head(x, e) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + return [u.float() for u in x] + + def unpatchify(self, x, grid_sizes): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + List[Tensor]: + Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] + """ + + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum('fhwpqrc->cfphqwr', u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/fastvideo2/wan21/reference.py b/fastvideo2/wan21/reference.py index 1697bc2bc..d56d7ca53 100644 --- a/fastvideo2/wan21/reference.py +++ b/fastvideo2/wan21/reference.py @@ -1,10 +1,12 @@ """Wan2.1 T2V reference — the oracle, in one readable file. -This is the complete generation path in plain eager PyTorch on stock -diffusers/transformers components: encode the prompt with UMT5, run N -classifier-free-guided flow-match Euler steps on the DiT, decode with the Wan -VAE. No fastvideo2 imports, no runtime, no policies — copy this file out of -the repo and it still runs. +This is the complete generation path in plain eager PyTorch: encode the +prompt with UMT5 (HF transformers — anchored bit-identical to official), run N +classifier-free-guided flow-match Euler steps on the OFFICIAL WanModel (the +vendored authors' modeling code in ``fastvideo2/wan21/model.py`` — the only +fastvideo2 import here, itself a standalone file), decode with the Wan VAE +(diffusers — anchored to 3e-4). No runtime, no policies — copy this file plus +``model.py``/``attention.py`` out of the repo and they still run. It serves three roles: * the textbook: what Wan2.1 T2V *is*, in execution order; @@ -34,7 +36,8 @@ import argparse from dataclasses import dataclass, field from typing import Any -WEIGHTS = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +WEIGHTS = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # text encoder / tokenizer / VAE +DIT_WEIGHTS = "Wan-AI/Wan2.1-T2V-1.3B" # the official transformer artifact NEGATIVE = ("色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，" "低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，" "毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走") @@ -47,22 +50,24 @@ class ReferenceResult: trajectory: list = field(default_factory=list) # per-step fp32 latents (cpu), when captured -def load_models(root: str | None = None, device: str = "cuda"): - """Load the four components exactly as declared on the card: DiT and UMT5 - in bf16, VAE in fp32.""" +def load_models(root: str | None = None, device: str = "cuda", dit_root: str | None = None): + """Load the four components exactly as declared on the card: the official + WanModel (bf16, official-layout weights), UMT5 bf16, VAE fp32.""" import torch - from diffusers import AutoencoderKLWan, WanTransformer3DModel + from diffusers import AutoencoderKLWan from transformers import AutoTokenizer, UMT5EncoderModel + + from fastvideo2.wan21.model import WanModel # the vendored official DiT if root is None: from huggingface_hub import snapshot_download root = snapshot_download(WEIGHTS) + if dit_root is None: + from huggingface_hub import snapshot_download + dit_root = snapshot_download(DIT_WEIGHTS) tokenizer = AutoTokenizer.from_pretrained(root, subfolder="tokenizer") - # torch_dtype (not a post-hoc .to()) so the loader's kept-in-fp32 islands - # (Wan's time_embedder and norms inside the bf16 DiT) survive as intended. text_encoder = UMT5EncoderModel.from_pretrained( root, subfolder="text_encoder", torch_dtype=torch.bfloat16).to(device).eval() - dit = WanTransformer3DModel.from_pretrained( - root, subfolder="transformer", torch_dtype=torch.bfloat16).to(device).eval() + dit = WanModel.from_pretrained(dit_root, torch_dtype=torch.bfloat16).to(device).eval() vae = AutoencoderKLWan.from_pretrained( root, subfolder="vae", torch_dtype=torch.float32).to(device).eval() return tokenizer, text_encoder, dit, vae @@ -133,20 +138,20 @@ def generate(prompt: str, x = torch.randn(shape, generator=gen, device=device, dtype=torch.float32) # 3) CFG flow-match Euler: v = v_neg + g*(v_cond - v_neg); x += dsigma * v. + # The DiT runs the official convention: fp32 latents under autocast + # bf16, list-based call, seq_len from the (1,2,2) patch grid. sigmas = flow_sigmas(num_steps, shift) + seq_len = x.shape[2] * (x.shape[3] // 2) * (x.shape[4] // 2) trajectory: list = [] for i in range(num_steps): sigma, sigma_next = sigmas[i], sigmas[i + 1] t = torch.tensor([sigma * 1000.0], device=device, dtype=torch.float32) - x_in = x.to(torch.bfloat16) - with torch.no_grad(): - v_cond = dit(hidden_states=x_in, timestep=t, encoder_hidden_states=text, - return_dict=False)[0].to(torch.float32) + with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16): + v_cond = dit([x[0]], t=t, context=[text[0]], seq_len=seq_len)[0][None].to(torch.float32) if guidance_scale == 1.0: v = v_cond else: - v_neg = dit(hidden_states=x_in, timestep=t, encoder_hidden_states=neg, - return_dict=False)[0].to(torch.float32) + v_neg = dit([x[0]], t=t, context=[neg[0]], seq_len=seq_len)[0][None].to(torch.float32) v = v_neg + guidance_scale * (v_cond - v_neg) x = x + (sigma_next - sigma) * v if capture_trajectory: