| 43 | |
| 44 | |
| 45 | class CogDiTBlock(torch.nn.Module): |
| 46 | def __init__(self, dim, dim_cond, num_heads): |
| 47 | super().__init__() |
| 48 | self.norm1 = CogAdaLayerNorm(dim, dim_cond) |
| 49 | self.attn1 = Attention(q_dim=dim, num_heads=48, head_dim=dim//num_heads, bias_q=True, bias_kv=True, bias_out=True) |
| 50 | self.norm_q = torch.nn.LayerNorm((dim//num_heads,), eps=1e-06, elementwise_affine=True) |
| 51 | self.norm_k = torch.nn.LayerNorm((dim//num_heads,), eps=1e-06, elementwise_affine=True) |
| 52 | |
| 53 | self.norm2 = CogAdaLayerNorm(dim, dim_cond) |
| 54 | self.ff = torch.nn.Sequential( |
| 55 | torch.nn.Linear(dim, dim*4), |
| 56 | torch.nn.GELU(approximate="tanh"), |
| 57 | torch.nn.Linear(dim*4, dim) |
| 58 | ) |
| 59 | |
| 60 | |
| 61 | def apply_rotary_emb(self, x, freqs_cis): |
| 62 | cos, sin = freqs_cis # [S, D] |
| 63 | cos = cos[None, None] |
| 64 | sin = sin[None, None] |
| 65 | cos, sin = cos.to(x.device), sin.to(x.device) |
| 66 | x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] |
| 67 | x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) |
| 68 | out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) |
| 69 | return out |
| 70 | |
| 71 | |
| 72 | def process_qkv(self, q, k, v, image_rotary_emb, text_seq_length): |
| 73 | q = self.norm_q(q) |
| 74 | k = self.norm_k(k) |
| 75 | q[:, :, text_seq_length:] = self.apply_rotary_emb(q[:, :, text_seq_length:], image_rotary_emb) |
| 76 | k[:, :, text_seq_length:] = self.apply_rotary_emb(k[:, :, text_seq_length:], image_rotary_emb) |
| 77 | return q, k, v |
| 78 | |
| 79 | |
| 80 | def forward(self, hidden_states, prompt_emb, time_emb, image_rotary_emb): |
| 81 | # Attention |
| 82 | norm_hidden_states, norm_encoder_hidden_states, gate_a, gate_b = self.norm1( |
| 83 | hidden_states, prompt_emb, time_emb |
| 84 | ) |
| 85 | attention_io = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1) |
| 86 | attention_io = self.attn1( |
| 87 | attention_io, |
| 88 | qkv_preprocessor=lambda q, k, v: self.process_qkv(q, k, v, image_rotary_emb, prompt_emb.shape[1]) |
| 89 | ) |
| 90 | |
| 91 | hidden_states = hidden_states + gate_a * attention_io[:, prompt_emb.shape[1]:] |
| 92 | prompt_emb = prompt_emb + gate_b * attention_io[:, :prompt_emb.shape[1]] |
| 93 | |
| 94 | # Feed forward |
| 95 | norm_hidden_states, norm_encoder_hidden_states, gate_a, gate_b = self.norm2( |
| 96 | hidden_states, prompt_emb, time_emb |
| 97 | ) |
| 98 | ff_io = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1) |
| 99 | ff_io = self.ff(ff_io) |
| 100 | |
| 101 | hidden_states = hidden_states + gate_a * ff_io[:, prompt_emb.shape[1]:] |
| 102 | prompt_emb = prompt_emb + gate_b * ff_io[:, :prompt_emb.shape[1]] |