| 459 | return x + skip |
| 460 | |
| 461 | class ShiftedWindowSelfAttentionBlock(nn.Module): |
| 462 | def __init__(self, d_model, d_head, cond_features, window_size, window_shift, dropout=0.0): |
| 463 | super().__init__() |
| 464 | self.d_head = d_head |
| 465 | self.n_heads = d_model // d_head |
| 466 | self.window_size = window_size |
| 467 | self.window_shift = window_shift |
| 468 | self.norm = AdaRMSNorm(d_model, cond_features) |
| 469 | self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False)) |
| 470 | self.scale = nn.Parameter(torch.full([self.n_heads], 10.0)) |
| 471 | self.pos_emb = AxialRoPE(d_head // 2, self.n_heads) |
| 472 | self.dropout = nn.Dropout(dropout) |
| 473 | self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False))) |
| 474 | |
| 475 | def extra_repr(self): |
| 476 | return f"d_head={self.d_head}, window_size={self.window_size}, window_shift={self.window_shift}" |
| 477 | |
| 478 | def forward(self, x, pos, cond): |
| 479 | skip = x |
| 480 | x = self.norm(x, cond) |
| 481 | qkv = self.qkv_proj(x) |
| 482 | q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh h w e", t=3, e=self.d_head) |
| 483 | q, k = scale_for_cosine_sim(q, k, self.scale[:, None, None, None], 1e-6) |
| 484 | theta = self.pos_emb(pos).movedim(-2, -4) |
| 485 | q = apply_rotary_emb_(q, theta) |
| 486 | k = apply_rotary_emb_(k, theta) |
| 487 | x = apply_window_attention(self.window_size, self.window_shift, q, k, v, scale=1.0) |
| 488 | x = rearrange(x, "n nh h w e -> n h w (nh e)") |
| 489 | x = self.dropout(x) |
| 490 | x = self.out_proj(x) |
| 491 | return x + skip |
| 492 | |
| 493 | |
| 494 | class FeedForwardBlock(nn.Module): |