From 7937374bd8608ed32fbc3271db28444b5abaa6ff Mon Sep 17 00:00:00 2001 From: Tai An Date: Fri, 11 Sep 2026 18:11:23 -0700 Subject: [PATCH] fix(wan): un-shadow XLMRoberta's SelfAttention/AttentionBlock `diffsynth/models/wan_video_image_encoder.py` merges two upstream modules (the XLM-RoBERTa text tower and the CLIP vision tower) into one file, and both define module-level `SelfAttention` and `AttentionBlock`. The CLIP versions come later, so they win the module-level binding and the XLM-RoBERTa classes at the top of the file are unreachable. `XLMRoberta.__init__` resolves `AttentionBlock` at call time and therefore builds CLIP blocks with XLM-RoBERTa's argument order: AttentionBlock(dim, num_heads, post_norm, dropout, eps) # -> CLIP signature (dim, mlp_ratio, num_heads, post_norm, causal, ...) which silently yields `mlp_ratio=16`, `num_heads=True` (so `head_dim=1024` instead of 64) and 37.8M parameters per block instead of 12.6M. The assert `dim % num_heads == 0` passes because `1024 % True == 0`. The mismatch only surfaces on the forward pass, as `TypeError: AttentionBlock.forward() takes 2 positional arguments but 3 were given`, since the CLIP block's forward takes no mask. Rename the two shadowed definitions to `XLMRobertaSelfAttention` and `XLMRobertaAttentionBlock`. The CLIP/vision path used by `WanImageEncoder` is untouched: the module-level bindings of `SelfAttention` and `AttentionBlock` are byte-identical before and after this change. Signed-off-by: Tai An Co-Authored-By: Claude Opus 5 (1M context) --- diffsynth/models/wan_video_image_encoder.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/diffsynth/models/wan_video_image_encoder.py b/diffsynth/models/wan_video_image_encoder.py index 37d17d6a1..28ab6e5d9 100644 --- a/diffsynth/models/wan_video_image_encoder.py +++ b/diffsynth/models/wan_video_image_encoder.py @@ -11,7 +11,7 @@ from .wan_video_dit import flash_attention -class SelfAttention(nn.Module): +class XLMRobertaSelfAttention(nn.Module): def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5): assert dim % num_heads == 0 @@ -50,7 +50,7 @@ def forward(self, x, mask): return x -class AttentionBlock(nn.Module): +class XLMRobertaAttentionBlock(nn.Module): def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5): super().__init__() @@ -60,7 +60,7 @@ def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5): self.eps = eps # layers - self.attn = SelfAttention(dim, num_heads, dropout, eps) + self.attn = XLMRobertaSelfAttention(dim, num_heads, dropout, eps) self.norm1 = nn.LayerNorm(dim, eps=eps) self.ffn = nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), @@ -112,7 +112,7 @@ def __init__(self, # blocks self.blocks = nn.ModuleList([ - AttentionBlock(dim, num_heads, post_norm, dropout, eps) + XLMRobertaAttentionBlock(dim, num_heads, post_norm, dropout, eps) for _ in range(num_layers) ])