(self, x, pos, cond)
| 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): |
nothing calls this directly
no test coverage detected