| 401 | |
| 402 | |
| 403 | class MMDoubleStreamBlockComponent(torch.nn.Module): |
| 404 | def __init__(self, hidden_size=3072, heads_num=24, mlp_width_ratio=4): |
| 405 | super().__init__() |
| 406 | self.heads_num = heads_num |
| 407 | |
| 408 | self.mod = ModulateDiT(hidden_size) |
| 409 | self.norm1 = torch.nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 410 | |
| 411 | self.to_qkv = torch.nn.Linear(hidden_size, hidden_size * 3) |
| 412 | self.norm_q = RMSNorm(dim=hidden_size // heads_num, eps=1e-6) |
| 413 | self.norm_k = RMSNorm(dim=hidden_size // heads_num, eps=1e-6) |
| 414 | self.to_out = torch.nn.Linear(hidden_size, hidden_size) |
| 415 | |
| 416 | self.norm2 = torch.nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 417 | self.ff = torch.nn.Sequential( |
| 418 | torch.nn.Linear(hidden_size, hidden_size * mlp_width_ratio), |
| 419 | torch.nn.GELU(approximate="tanh"), |
| 420 | torch.nn.Linear(hidden_size * mlp_width_ratio, hidden_size) |
| 421 | ) |
| 422 | |
| 423 | def forward(self, hidden_states, conditioning, freqs_cis=None, token_replace_vec=None, tr_token=None): |
| 424 | mod1_shift, mod1_scale, mod1_gate, mod2_shift, mod2_scale, mod2_gate = self.mod(conditioning).chunk(6, dim=-1) |
| 425 | if token_replace_vec is not None: |
| 426 | assert tr_token is not None |
| 427 | tr_mod1_shift, tr_mod1_scale, tr_mod1_gate, tr_mod2_shift, tr_mod2_scale, tr_mod2_gate = self.mod(token_replace_vec).chunk(6, dim=-1) |
| 428 | else: |
| 429 | tr_mod1_shift, tr_mod1_scale, tr_mod1_gate, tr_mod2_shift, tr_mod2_scale, tr_mod2_gate = None, None, None, None, None, None |
| 430 | |
| 431 | norm_hidden_states = self.norm1(hidden_states) |
| 432 | norm_hidden_states = modulate(norm_hidden_states, shift=mod1_shift, scale=mod1_scale, |
| 433 | tr_shift=tr_mod1_shift, tr_scale=tr_mod1_scale, tr_token=tr_token) |
| 434 | qkv = self.to_qkv(norm_hidden_states) |
| 435 | q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) |
| 436 | |
| 437 | q = self.norm_q(q) |
| 438 | k = self.norm_k(k) |
| 439 | |
| 440 | if freqs_cis is not None: |
| 441 | q, k = apply_rotary_emb(q, k, freqs_cis, head_first=False) |
| 442 | return (q, k, v), (mod1_gate, mod2_shift, mod2_scale, mod2_gate), (tr_mod1_gate, tr_mod2_shift, tr_mod2_scale, tr_mod2_gate) |
| 443 | |
| 444 | def process_ff(self, hidden_states, attn_output, mod, mod_tr=None, tr_token=None): |
| 445 | mod1_gate, mod2_shift, mod2_scale, mod2_gate = mod |
| 446 | if mod_tr is not None: |
| 447 | tr_mod1_gate, tr_mod2_shift, tr_mod2_scale, tr_mod2_gate = mod_tr |
| 448 | else: |
| 449 | tr_mod1_gate, tr_mod2_shift, tr_mod2_scale, tr_mod2_gate = None, None, None, None |
| 450 | hidden_states = hidden_states + apply_gate(self.to_out(attn_output), mod1_gate, tr_mod1_gate, tr_token) |
| 451 | x = self.ff(modulate(self.norm2(hidden_states), shift=mod2_shift, scale=mod2_scale, tr_shift=tr_mod2_shift, tr_scale=tr_mod2_scale, tr_token=tr_token)) |
| 452 | hidden_states = hidden_states + apply_gate(x, mod2_gate, tr_mod2_gate, tr_token) |
| 453 | return hidden_states |
| 454 | |
| 455 | |
| 456 | class MMDoubleStreamBlock(torch.nn.Module): |