[feat]: align the DiT with official — vendored WanModel (pinned 9737cba) + official weights + official autocast regime
This commit is contained in:
SolitaryThinker
@@ -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
|
||||
|
||||
+3
-2
@@ -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)
|
||||
|
||||
+11
-1
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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",
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user