"""One-off script to precompute Wan UMT5-XXL text embeddings for Any4D + Wan training. Usage: python3 -m custom.wan.build_default_text_wan \\ --prompt "a high quality video of driving, realistic, cinematic" \\ --umt5_tokenizer_path /path/to/Wan2.1-Fun-1.3B-InP/google/umt5-xxl \\ --umt5_weights_path /path/to/Wan2.1-Fun-1.3B-InP/models_t5_umt5-xxl-enc-bf16.pth \\ --out custom/wan/default_text_wan.pkl Produces a pickle with BOTH positive and negative UMT5-XXL embeddings: t5_text_embeddings: (1, 512, 4096) bf16 — positive prompt t5_text_mask: (1, 512) bf16 — attention mask t5_text_embeddings_neg: (1, 512, 4096) bf16 — negative prompt ('' by default) t5_text_mask_neg: (1, 512) bf16 Critical correctness details (match diffsynth's WanPrompter.encode_prompt): - Use Wan's own UMT5 implementation (Wan's checkpoint keys don't match HF T5). - After encoding, zero positions beyond the real sequence length — without this, padding-token embeddings leak into cross-attention and degrade generation. - The negative embedding is UMT5('') + zero-beyond-seq-len. NOT a raw zero tensor (a zero tensor still passes through dit.text_embedding to produce spurious activations). """ import argparse import math import pickle import torch import torch.nn as nn import torch.nn.functional as F MAX_LENGTH = 512 # ---- Wan UMT5-XXL encoder (inlined from DiffSynth-Studio) ----------------- class _T5LayerNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.dim, self.eps = dim, eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): x = x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) if self.weight.dtype in (torch.float16, torch.bfloat16): x = x.type_as(self.weight) return self.weight * x class _GELU(nn.Module): def forward(self, x): return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x.pow(3)))) class _T5Attn(nn.Module): def __init__(self, dim, dim_attn, num_heads): super().__init__() self.dim, self.dim_attn = dim, dim_attn self.num_heads, self.head_dim = num_heads, dim_attn // num_heads self.q = nn.Linear(dim, dim_attn, bias=False) self.k = nn.Linear(dim, dim_attn, bias=False) self.v = nn.Linear(dim, dim_attn, bias=False) self.o = nn.Linear(dim_attn, dim, bias=False) def forward(self, x, mask=None, pos_bias=None): b, n, c = x.size(0), self.num_heads, self.head_dim q = self.q(x).view(b, -1, n, c) k = self.k(x).view(b, -1, n, c) v = self.v(x).view(b, -1, n, c) attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) if pos_bias is not None: attn_bias = attn_bias + pos_bias if mask is not None: mask = mask.view(b, 1, 1, -1) if mask.ndim == 2 else mask.unsqueeze(1) attn_bias = attn_bias.masked_fill(mask == 0, torch.finfo(x.dtype).min) attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias attn = F.softmax(attn.float(), dim=-1).type_as(attn) x = torch.einsum('bnij,bjnc->binc', attn, v).reshape(b, -1, n * c) return self.o(x) class _T5FFN(nn.Module): def __init__(self, dim, dim_ffn): super().__init__() self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), _GELU()) self.fc1 = nn.Linear(dim, dim_ffn, bias=False) self.fc2 = nn.Linear(dim_ffn, dim, bias=False) def forward(self, x): return self.fc2(self.fc1(x) * self.gate(x)) class _T5RelEmbed(nn.Module): def __init__(self, num_buckets=32, num_heads=64, bidirectional=True, max_dist=128): super().__init__() self.num_buckets, self.num_heads = num_buckets, num_heads self.bidirectional, self.max_dist = bidirectional, max_dist self.embedding = nn.Embedding(num_buckets, num_heads) def forward(self, lq, lk): dev = self.embedding.weight.device rel = torch.arange(lk, device=dev).unsqueeze(0) - torch.arange(lq, device=dev).unsqueeze(1) nb = self.num_buckets // 2 if self.bidirectional else self.num_buckets rel_buckets = (rel > 0).long() * nb if self.bidirectional else torch.zeros_like(rel) rel = rel.abs() if self.bidirectional else (-torch.min(rel, torch.zeros_like(rel))) max_exact = nb // 2 rel_large = max_exact + (torch.log(rel.float() / max_exact) / math.log(self.max_dist / max_exact) * (nb - max_exact)).long() rel_large = torch.min(rel_large, torch.full_like(rel_large, nb - 1)) rel_buckets = rel_buckets + torch.where(rel < max_exact, rel, rel_large) return self.embedding(rel_buckets).permute(2, 0, 1).unsqueeze(0).contiguous() class _T5Block(nn.Module): def __init__(self, dim=4096, dim_ffn=10240, num_heads=64, num_buckets=32): super().__init__() self.norm1 = _T5LayerNorm(dim) self.attn = _T5Attn(dim, dim, num_heads) self.norm2 = _T5LayerNorm(dim) self.ffn = _T5FFN(dim, dim_ffn) self.pos_embedding = _T5RelEmbed(num_buckets, num_heads) def forward(self, x, mask=None): e = self.pos_embedding(x.size(1), x.size(1)) x = x + self.attn(self.norm1(x), mask=mask, pos_bias=e) x = x + self.ffn(self.norm2(x)) return x class WanUMT5Encoder(nn.Module): def __init__(self, vocab=256384, dim=4096, dim_ffn=10240, num_heads=64, num_layers=24, num_buckets=32): super().__init__() self.token_embedding = nn.Embedding(vocab, dim) self.blocks = nn.ModuleList([ _T5Block(dim, dim_ffn, num_heads, num_buckets) for _ in range(num_layers) ]) self.norm = _T5LayerNorm(dim) def forward(self, ids, mask=None): x = self.token_embedding(ids) for block in self.blocks: x = block(x, mask=mask) return self.norm(x) # ---- Text encoding with Wan-correct seq-length zeroing -------------------- def _encode_with_zeroing(model, tok, prompt: str, device, dtype): import ftfy, html, re text = ftfy.fix_text(prompt) text = html.unescape(html.unescape(text)) text = re.sub(r'\s+', ' ', text).strip() enc = tok([text], padding='max_length', truncation=True, max_length=MAX_LENGTH, return_tensors='pt', add_special_tokens=True) ids = enc['input_ids'].to(device) mask = enc['attention_mask'].to(device) seq_lens = mask.gt(0).sum(dim=1).long() with torch.no_grad(): emb = model(ids, mask=mask) # (1, 512, 4096) for i, v in enumerate(seq_lens): emb[i, v:] = 0 # zero beyond real prompt (matches WanPrompter) return emb.to(dtype=dtype).cpu(), mask.to(dtype=dtype).cpu() def build_default_text(prompt: str, negative_prompt: str, tokenizer_path: str, weights_path: str, out_path: str, device: str = 'cuda', dtype: torch.dtype = torch.bfloat16): from transformers import AutoTokenizer print(f'Loading UMT5 tokenizer from: {tokenizer_path}') tok = AutoTokenizer.from_pretrained(tokenizer_path) print(f'Loading UMT5 weights from: {weights_path}') model = WanUMT5Encoder().to(dtype=dtype).eval() sd = torch.load(weights_path, map_location='cpu', weights_only=True) miss, surp = model.load_state_dict(sd, strict=False) print(f'UMT5 load: missing={len(miss)} surplus={len(surp)}') if miss: print(f' first missing: {miss[:3]}') model = model.to(device) print(f'Encoding positive: {prompt!r}') emb_pos, mask_pos = _encode_with_zeroing(model, tok, prompt, device, dtype) print(f' pos nonzero-tokens: {mask_pos.sum().item()}') print(f'Encoding negative: {negative_prompt!r}') emb_neg, mask_neg = _encode_with_zeroing(model, tok, negative_prompt, device, dtype) print(f' neg nonzero-tokens: {mask_neg.sum().item()}') payload = { 't5_text_embeddings': emb_pos, 't5_text_mask': mask_pos, 't5_text_embeddings_neg': emb_neg, 't5_text_mask_neg': mask_neg, 'prompt': prompt, 'negative_prompt': negative_prompt, } with open(out_path, 'wb') as f: pickle.dump(payload, f) print(f'Wrote {out_path}') def main(): p = argparse.ArgumentParser() p.add_argument('--prompt', type=str, default='a high quality video of driving, realistic, cinematic') p.add_argument('--negative_prompt', type=str, default='') p.add_argument('--umt5_tokenizer_path', type=str, required=True, help='Dir with Wan UMT5 tokenizer files (spiece.model, tokenizer.json, ...).') p.add_argument('--umt5_weights_path', type=str, required=True, help='Path to models_t5_umt5-xxl-enc-bf16.pth') p.add_argument('--out', type=str, default='custom/wan/default_text_wan.pkl') p.add_argument('--device', type=str, default='cuda') p.add_argument('--dtype', type=str, default='bfloat16', choices=['bfloat16', 'float16', 'float32']) args = p.parse_args() dtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16, 'float32': torch.float32}[args.dtype] build_default_text(args.prompt, args.negative_prompt, args.umt5_tokenizer_path, args.umt5_weights_path, args.out, device=args.device, dtype=dtype) if __name__ == '__main__': main()