"""UMT5 text encoder used by Wan2.2 (custom diffsynth weight layout). Architecture is a stripped UMT5: tied vocab, GELU-gated FFN, T5-style relative position bias, RMS-style layer norm. Matches the keys in `models_t5_umt5-xxl-enc-bf16.pth`. """ import html import math import re import ftfy import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoTokenizer def _fp16_clamp(x): if x.dtype == torch.float16 and torch.isinf(x).any(): c = torch.finfo(x.dtype).max - 1000 x = torch.clamp(x, min=-c, max=c) return 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 _T5LayerNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): x = x * torch.rsqrt(x.float().pow(2).mean(dim=-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 _T5Attention(nn.Module): def __init__(self, dim, dim_attn, num_heads): super().__init__() self.num_heads = num_heads self.head_dim = 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.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) out = torch.einsum('bnij,bjnc->binc', attn, v).reshape(b, -1, n * c) return self.o(out) class _T5FeedForward(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 _T5SelfAttention(nn.Module): def __init__(self, dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos): super().__init__() self.norm1 = _T5LayerNorm(dim) self.attn = _T5Attention(dim, dim_attn, num_heads) self.norm2 = _T5LayerNorm(dim) self.ffn = _T5FeedForward(dim, dim_ffn) self.pos_embedding = None if shared_pos else _T5RelativeEmbedding( num_buckets, num_heads, bidirectional=True) def forward(self, x, mask=None, pos_bias=None): e = pos_bias if self.pos_embedding is None else self.pos_embedding(x.size(1), x.size(1)) x = _fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) x = _fp16_clamp(x + self.ffn(self.norm2(x))) return x class _T5RelativeEmbedding(nn.Module): def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): super().__init__() self.num_buckets = num_buckets self.bidirectional = bidirectional self.max_dist = max_dist self.embedding = nn.Embedding(num_buckets, num_heads) def forward(self, lq, lk): device = self.embedding.weight.device rel = torch.arange(lk, device=device).unsqueeze(0) - torch.arange(lq, device=device).unsqueeze(1) rel = self._bucket(rel) return self.embedding(rel).permute(2, 0, 1).unsqueeze(0).contiguous() def _bucket(self, rel): if self.bidirectional: n = self.num_buckets // 2 buckets = (rel > 0).long() * n rel = rel.abs() else: n = self.num_buckets buckets = 0 rel = -torch.min(rel, torch.zeros_like(rel)) max_exact = n // 2 large = max_exact + (torch.log(rel.float() / max_exact) / math.log(self.max_dist / max_exact) * (n - max_exact)).long() large = torch.min(large, torch.full_like(large, n - 1)) buckets += torch.where(rel < max_exact, rel, large) return buckets class WanTextEncoder(nn.Module): """UMT5 encoder used by Wan2.2-TI2V-5B; loads `models_t5_umt5-xxl-enc-bf16.pth`.""" def __init__(self, vocab=256384, dim=4096, dim_attn=4096, dim_ffn=10240, num_heads=64, num_layers=24, num_buckets=32, shared_pos=False): super().__init__() self.token_embedding = nn.Embedding(vocab, dim) self.pos_embedding = _T5RelativeEmbedding( num_buckets, num_heads, bidirectional=True) if shared_pos else None self.blocks = nn.ModuleList([ _T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos) for _ in range(num_layers) ]) self.norm = _T5LayerNorm(dim) self.shared_pos = shared_pos def forward(self, ids, mask=None): x = self.token_embedding(ids) e = self.pos_embedding(x.size(1), x.size(1)) if self.shared_pos else None for blk in self.blocks: x = blk(x, mask=mask, pos_bias=e) return self.norm(x) def _whitespace_clean(text: str) -> str: """Bit-identical port of `diffsynth.models.wan_video_text_encoder. whitespace_clean(basic_clean(text))`. Run on every prompt before tokenizing — both the training-time pipeline and the cache precompute do this, so skipping it makes the negative prompt's T5 embedding diverge (e.g. the Chinese fullwidth comma `，` U+FF0C → ASCII `,` U+002C swap maps to a completely different UMT5 token id). Don't drop this.""" text = ftfy.fix_text(text) text = html.unescape(html.unescape(text)) text = text.strip() text = re.sub(r"\s+", " ", text) return text.strip() class WanTokenizer: """Wraps HF AutoTokenizer with the (return_mask, max_length) interface used by Wan.""" def __init__(self, path: str, seq_len: int = 512): self.tokenizer = AutoTokenizer.from_pretrained(path) self.seq_len = seq_len def __call__(self, text): if isinstance(text, str): text = [text] text = [_whitespace_clean(t) for t in text] out = self.tokenizer(text, return_tensors='pt', padding='max_length', truncation=True, max_length=self.seq_len, add_special_tokens=True) return out.input_ids, out.attention_mask