(self,
dim,
num_heads,
window_size=(-1, -1),
qk_norm=True,
eps=1e-6)
| 258 | class WanSelfAttention(nn.Module): |
| 259 | |
| 260 | def __init__(self, |
| 261 | dim, |
| 262 | num_heads, |
| 263 | window_size=(-1, -1), |
| 264 | qk_norm=True, |
| 265 | eps=1e-6): |
| 266 | assert dim % num_heads == 0 |
| 267 | super().__init__() |
| 268 | self.dim = dim |
| 269 | self.num_heads = num_heads |
| 270 | self.head_dim = dim // num_heads |
| 271 | self.window_size = window_size |
| 272 | self.qk_norm = qk_norm |
| 273 | self.eps = eps |
| 274 | |
| 275 | # layers |
| 276 | self.q = nn.Linear(dim, dim) |
| 277 | self.k = nn.Linear(dim, dim) |
| 278 | self.v = nn.Linear(dim, dim) |
| 279 | self.o = nn.Linear(dim, dim) |
| 280 | self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() |
| 281 | self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() |
| 282 | |
| 283 | def forward(self, x, seq_lens, freqs): |
| 284 | b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim |
nothing calls this directly
no test coverage detected