| 241 | |
| 242 | |
| 243 | class JointTransformerBlock(torch.nn.Module): |
| 244 | def __init__(self, dim, num_attention_heads, use_rms_norm=False, dual=False): |
| 245 | super().__init__() |
| 246 | self.norm1_a = AdaLayerNorm(dim, dual=dual) |
| 247 | self.norm1_b = AdaLayerNorm(dim) |
| 248 | |
| 249 | self.attn = JointAttention(dim, dim, num_attention_heads, dim // num_attention_heads, use_rms_norm=use_rms_norm) |
| 250 | if dual: |
| 251 | self.attn2 = SingleAttention(dim, num_attention_heads, dim // num_attention_heads, use_rms_norm=use_rms_norm) |
| 252 | |
| 253 | self.norm2_a = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 254 | self.ff_a = torch.nn.Sequential( |
| 255 | torch.nn.Linear(dim, dim*4), |
| 256 | torch.nn.GELU(approximate="tanh"), |
| 257 | torch.nn.Linear(dim*4, dim) |
| 258 | ) |
| 259 | |
| 260 | self.norm2_b = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 261 | self.ff_b = torch.nn.Sequential( |
| 262 | torch.nn.Linear(dim, dim*4), |
| 263 | torch.nn.GELU(approximate="tanh"), |
| 264 | torch.nn.Linear(dim*4, dim) |
| 265 | ) |
| 266 | |
| 267 | |
| 268 | def forward(self, hidden_states_a, hidden_states_b, temb): |
| 269 | if self.norm1_a.dual: |
| 270 | norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a, norm_hidden_states_a_2, gate_msa_a_2 = self.norm1_a(hidden_states_a, emb=temb) |
| 271 | else: |
| 272 | norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a = self.norm1_a(hidden_states_a, emb=temb) |
| 273 | norm_hidden_states_b, gate_msa_b, shift_mlp_b, scale_mlp_b, gate_mlp_b = self.norm1_b(hidden_states_b, emb=temb) |
| 274 | |
| 275 | # Attention |
| 276 | attn_output_a, attn_output_b = self.attn(norm_hidden_states_a, norm_hidden_states_b) |
| 277 | |
| 278 | # Part A |
| 279 | hidden_states_a = hidden_states_a + gate_msa_a * attn_output_a |
| 280 | if self.norm1_a.dual: |
| 281 | hidden_states_a = hidden_states_a + gate_msa_a_2 * self.attn2(norm_hidden_states_a_2) |
| 282 | norm_hidden_states_a = self.norm2_a(hidden_states_a) * (1 + scale_mlp_a) + shift_mlp_a |
| 283 | hidden_states_a = hidden_states_a + gate_mlp_a * self.ff_a(norm_hidden_states_a) |
| 284 | |
| 285 | # Part B |
| 286 | hidden_states_b = hidden_states_b + gate_msa_b * attn_output_b |
| 287 | norm_hidden_states_b = self.norm2_b(hidden_states_b) * (1 + scale_mlp_b) + shift_mlp_b |
| 288 | hidden_states_b = hidden_states_b + gate_mlp_b * self.ff_b(norm_hidden_states_b) |
| 289 | |
| 290 | return hidden_states_a, hidden_states_b |
| 291 | |
| 292 | |
| 293 | |