| 367 | |
| 368 | |
| 369 | class SelfAttentionBlock(nn.Module): |
| 370 | def __init__(self, d_model, d_head, cond_features, dropout=0.0): |
| 371 | super().__init__() |
| 372 | self.d_head = d_head |
| 373 | self.n_heads = d_model // d_head |
| 374 | self.norm = AdaRMSNorm(d_model, cond_features) |
| 375 | self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False)) |
| 376 | self.scale = nn.Parameter(torch.full([self.n_heads], 10.0)) |
| 377 | self.pos_emb = AxialRoPE(d_head // 2, self.n_heads) |
| 378 | self.dropout = nn.Dropout(dropout) |
| 379 | self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False))) |
| 380 | |
| 381 | def extra_repr(self): |
| 382 | return f"d_head={self.d_head}," |
| 383 | |
| 384 | def forward(self, x, pos, cond): |
| 385 | skip = x |
| 386 | x = self.norm(x, cond) |
| 387 | qkv = self.qkv_proj(x) |
| 388 | pos = rearrange(pos, "... h w e -> ... (h w) e").to(qkv.dtype) |
| 389 | theta = self.pos_emb(pos) |
| 390 | if use_flash_2(qkv): |
| 391 | qkv = rearrange(qkv, "n h w (t nh e) -> n (h w) t nh e", t=3, e=self.d_head) |
| 392 | qkv = scale_for_cosine_sim_qkv(qkv, self.scale, 1e-6) |
| 393 | theta = torch.stack((theta, theta, torch.zeros_like(theta)), dim=-3) |
| 394 | qkv = apply_rotary_emb_(qkv, theta) |
| 395 | flops_shape = qkv.shape[-5], qkv.shape[-2], qkv.shape[-4], qkv.shape[-1] |
| 396 | flops.op(flops.op_attention, flops_shape, flops_shape, flops_shape) |
| 397 | x = flash_attn.flash_attn_qkvpacked_func(qkv, softmax_scale=1.0) |
| 398 | x = rearrange(x, "n (h w) nh e -> n h w (nh e)", h=skip.shape[-3], w=skip.shape[-2]) |
| 399 | else: |
| 400 | q, k, v = rearrange(qkv, "n h w (t nh e) -> t n nh (h w) e", t=3, e=self.d_head) |
| 401 | q, k = scale_for_cosine_sim(q, k, self.scale[:, None, None], 1e-6) |
| 402 | theta = theta.movedim(-2, -3) |
| 403 | q = apply_rotary_emb_(q, theta) |
| 404 | k = apply_rotary_emb_(k, theta) |
| 405 | flops.op(flops.op_attention, q.shape, k.shape, v.shape) |
| 406 | x = F.scaled_dot_product_attention(q, k, v, scale=1.0) |
| 407 | x = rearrange(x, "n nh (h w) e -> n h w (nh e)", h=skip.shape[-3], w=skip.shape[-2]) |
| 408 | x = self.dropout(x) |
| 409 | x = self.out_proj(x) |
| 410 | return x + skip |
| 411 | |
| 412 | class NeighborhoodSelfAttentionBlock(nn.Module): |
| 413 | def __init__(self, d_model, d_head, cond_features, kernel_size, dropout=0.0): |