| 8 | |
| 9 | |
| 10 | class VaceWanAttentionBlock(WanAttentionBlock): |
| 11 | |
| 12 | def __init__(self, |
| 13 | cross_attn_type, |
| 14 | dim, |
| 15 | ffn_dim, |
| 16 | num_heads, |
| 17 | window_size=(-1, -1), |
| 18 | qk_norm=True, |
| 19 | cross_attn_norm=False, |
| 20 | eps=1e-6, |
| 21 | block_id=0): |
| 22 | super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size, |
| 23 | qk_norm, cross_attn_norm, eps) |
| 24 | self.block_id = block_id |
| 25 | if block_id == 0: |
| 26 | self.before_proj = nn.Linear(self.dim, self.dim) |
| 27 | nn.init.zeros_(self.before_proj.weight) |
| 28 | nn.init.zeros_(self.before_proj.bias) |
| 29 | self.after_proj = nn.Linear(self.dim, self.dim) |
| 30 | nn.init.zeros_(self.after_proj.weight) |
| 31 | nn.init.zeros_(self.after_proj.bias) |
| 32 | |
| 33 | def forward(self, c, x, **kwargs): |
| 34 | if self.block_id == 0: |
| 35 | c = self.before_proj(c) + x |
| 36 | |
| 37 | c = super().forward(c, **kwargs) |
| 38 | c_skip = self.after_proj(c) |
| 39 | return c, c_skip |
| 40 | |
| 41 | |
| 42 | class BaseWanAttentionBlock(WanAttentionBlock): |