| 107 | |
| 108 | |
| 109 | class FluxJointTransformerBlock(torch.nn.Module): |
| 110 | def __init__(self, dim, num_attention_heads): |
| 111 | super().__init__() |
| 112 | self.norm1_a = AdaLayerNorm(dim) |
| 113 | self.norm1_b = AdaLayerNorm(dim) |
| 114 | |
| 115 | self.attn = FluxJointAttention(dim, dim, num_attention_heads, dim // num_attention_heads) |
| 116 | |
| 117 | self.norm2_a = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 118 | self.ff_a = torch.nn.Sequential( |
| 119 | torch.nn.Linear(dim, dim*4), |
| 120 | torch.nn.GELU(approximate="tanh"), |
| 121 | torch.nn.Linear(dim*4, dim) |
| 122 | ) |
| 123 | |
| 124 | self.norm2_b = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 125 | self.ff_b = torch.nn.Sequential( |
| 126 | torch.nn.Linear(dim, dim*4), |
| 127 | torch.nn.GELU(approximate="tanh"), |
| 128 | torch.nn.Linear(dim*4, dim) |
| 129 | ) |
| 130 | |
| 131 | |
| 132 | def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None): |
| 133 | norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a = self.norm1_a(hidden_states_a, emb=temb) |
| 134 | norm_hidden_states_b, gate_msa_b, shift_mlp_b, scale_mlp_b, gate_mlp_b = self.norm1_b(hidden_states_b, emb=temb) |
| 135 | |
| 136 | # Attention |
| 137 | attn_output_a, attn_output_b = self.attn(norm_hidden_states_a, norm_hidden_states_b, image_rotary_emb, attn_mask, ipadapter_kwargs_list) |
| 138 | |
| 139 | # Part A |
| 140 | hidden_states_a = hidden_states_a + gate_msa_a * attn_output_a |
| 141 | norm_hidden_states_a = self.norm2_a(hidden_states_a) * (1 + scale_mlp_a) + shift_mlp_a |
| 142 | hidden_states_a = hidden_states_a + gate_mlp_a * self.ff_a(norm_hidden_states_a) |
| 143 | |
| 144 | # Part B |
| 145 | hidden_states_b = hidden_states_b + gate_msa_b * attn_output_b |
| 146 | norm_hidden_states_b = self.norm2_b(hidden_states_b) * (1 + scale_mlp_b) + shift_mlp_b |
| 147 | hidden_states_b = hidden_states_b + gate_mlp_b * self.ff_b(norm_hidden_states_b) |
| 148 | |
| 149 | return hidden_states_a, hidden_states_b |
| 150 | |
| 151 | |
| 152 | |