| 204 | |
| 205 | |
| 206 | class FluxSingleTransformerBlock(torch.nn.Module): |
| 207 | def __init__(self, dim, num_attention_heads): |
| 208 | super().__init__() |
| 209 | self.num_heads = num_attention_heads |
| 210 | self.head_dim = dim // num_attention_heads |
| 211 | self.dim = dim |
| 212 | |
| 213 | self.norm = AdaLayerNormSingle(dim) |
| 214 | self.to_qkv_mlp = torch.nn.Linear(dim, dim * (3 + 4)) |
| 215 | self.norm_q_a = RMSNorm(self.head_dim, eps=1e-6) |
| 216 | self.norm_k_a = RMSNorm(self.head_dim, eps=1e-6) |
| 217 | |
| 218 | self.proj_out = torch.nn.Linear(dim * 5, dim) |
| 219 | |
| 220 | |
| 221 | def apply_rope(self, xq, xk, freqs_cis): |
| 222 | xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) |
| 223 | xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) |
| 224 | xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] |
| 225 | xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] |
| 226 | return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) |
| 227 | |
| 228 | |
| 229 | def process_attention(self, hidden_states, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None): |
| 230 | batch_size = hidden_states.shape[0] |
| 231 | |
| 232 | qkv = hidden_states.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2) |
| 233 | q, k, v = qkv.chunk(3, dim=1) |
| 234 | q, k = self.norm_q_a(q), self.norm_k_a(k) |
| 235 | |
| 236 | q, k = self.apply_rope(q, k, image_rotary_emb) |
| 237 | |
| 238 | hidden_states = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) |
| 239 | hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim) |
| 240 | hidden_states = hidden_states.to(q.dtype) |
| 241 | if ipadapter_kwargs_list is not None: |
| 242 | hidden_states = interact_with_ipadapter(hidden_states, q, **ipadapter_kwargs_list) |
| 243 | return hidden_states |
| 244 | |
| 245 | |
| 246 | def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None): |
| 247 | residual = hidden_states_a |
| 248 | norm_hidden_states, gate = self.norm(hidden_states_a, emb=temb) |
| 249 | hidden_states_a = self.to_qkv_mlp(norm_hidden_states) |
| 250 | attn_output, mlp_hidden_states = hidden_states_a[:, :, :self.dim * 3], hidden_states_a[:, :, self.dim * 3:] |
| 251 | |
| 252 | attn_output = self.process_attention(attn_output, image_rotary_emb, attn_mask, ipadapter_kwargs_list) |
| 253 | mlp_hidden_states = torch.nn.functional.gelu(mlp_hidden_states, approximate="tanh") |
| 254 | |
| 255 | hidden_states_a = torch.cat([attn_output, mlp_hidden_states], dim=2) |
| 256 | hidden_states_a = gate.unsqueeze(1) * self.proj_out(hidden_states_a) |
| 257 | hidden_states_a = residual + hidden_states_a |
| 258 | |
| 259 | return hidden_states_a, hidden_states_b |
| 260 | |
| 261 | |
| 262 | |