import torch
from typing import List, Union, Optional, Dict, Any, Callable
from diffusers.models.attention_processor import Attention, F
from .lora_controller import enable_lora
def attn_forward(
attn: Attention,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor = None,
condition_latents: torch.FloatTensor = None,
attention_mask: Optional[torch.FloatTensor] = None,
image_rotary_emb: Optional[torch.Tensor] = None,
cond_rotary_emb: Optional[torch.Tensor] = None,
model_config: Optional[Dict[str, Any]] = {},
) -> torch.FloatTensor:
batch_size, _, _ = (
hidden_states.shape
if encoder_hidden_states is None
else encoder_hidden_states.shape
)
with enable_lora(
(attn.to_q, attn.to_k, attn.to_v), model_config.get("latent_lora", False)
):
# `sample` projections.
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
if attn.norm_q is not None:
query = attn.norm_q(query)
if attn.norm_k is not None:
key = attn.norm_k(key)
# the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
if encoder_hidden_states is not None:
# `context` projections.
encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
batch_size, -1, attn.heads, head_dim
).transpose(1, 2)
encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
batch_size, -1, attn.heads, head_dim
).transpose(1, 2)
encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
batch_size, -1, attn.heads, head_dim
).transpose(1, 2)
if attn.norm_added_q is not None:
encoder_hidden_states_query_proj = attn.norm_added_q(
encoder_hidden_states_query_proj
)
if attn.norm_added_k is not None:
encoder_hidden_states_key_proj = attn.norm_added_k(
encoder_hidden_states_key_proj
)
# attention
query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
if image_rotary_emb is not None:
from diffusers.models.embeddings import apply_rotary_emb
query = apply_rotary_emb(query, image_rotary_emb)
key = apply_rotary_emb(key, image_rotary_emb)
if condition_latents is not None:
cond_query = attn.to_q(condition_latents)
cond_key = attn.to_k(condition_latents)
cond_value = attn.to_v(condition_latents)
cond_query = cond_query.view(batch_size, -1, attn.heads, head_dim).transpose(
1, 2
)
cond_key = cond_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
cond_value = cond_value.view(batch_size, -1, attn.heads, head_dim).transpose(
1, 2
)
if attn.norm_q is not None:
cond_query = attn.norm_q(cond_query)
if attn.norm_k is not None:
cond_key = attn.norm_k(cond_key)
if cond_rotary_emb is not None:
cond_query = apply_rotary_emb(cond_query, cond_rotary_emb)
cond_key = apply_rotary_emb(cond_key, cond_rotary_emb)
if condition_latents is not None:
query = torch.cat([query, cond_query], dim=2)
key = torch.cat([key, cond_key], dim=2)
value = torch.cat([value, cond_value], dim=2)
if not model_config.get("union_cond_attn", True):
# If we don't want to use the union condition attention, we need to mask the attention
# between the hidden states and the condition latents
attention_mask = torch.ones(
query.shape[2], key.shape[2], device=query.device, dtype=torch.bool
)
condition_n = cond_query.shape[2]
attention_mask[-condition_n:, :-condition_n] = False
attention_mask[:-condition_n, -condition_n:] = False
if hasattr(attn, "c_factor"):
attention_mask = torch.zeros(
query.shape[2], key.shape[2], device=query.device, dtype=query.dtype
)
condition_n = cond_query.shape[2]
bias = torch.log(attn.c_factor[0])
attention_mask[-condition_n:, :-condition_n] = bias
attention_mask[:-condition_n, -condition_n:] = bias
hidden_states = F.scaled_dot_product_attention(
query, key, value, dropout_p=0.0, is_causal=False, attn_mask=attention_mask
)
hidden_states = hidden_states.transpose(1, 2).reshape(
batch_size, -1, attn.heads * head_dim
)
hidden_states = hidden_states.to(query.dtype)
if encoder_hidden_states is not None:
if condition_latents is not None:
encoder_hidden_states, hidden_states, condition_latents = (
hidden_states[:, : encoder_hidden_states.shape[1]],
hidden_states[
:, encoder_hidden_states.shape[1] : -condition_latents.shape[1]
],
hidden_states[:, -condition_latents.shape[1] :],
)
else:
encoder_hidden_states, hidden_states = (
hidden_states[:, : encoder_hidden_states.shape[1]],
hidden_states[:, encoder_hidden_states.shape[1] :],
)
with enable_lora((attn.to_out[0],), model_config.get("latent_lora", False)):
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
if condition_latents is not None:
condition_latents = attn.to_out[0](condition_latents)
condition_latents = attn.to_out[1](condition_latents)
return (
(hidden_states, encoder_hidden_states, condition_latents)
if condition_latents is not None
else (hidden_states, encoder_hidden_states)
)
elif condition_latents is not None:
# if there are condition_latents, we need to separate the hidden_states and the condition_latents
hidden_states, condition_latents = (
hidden_states[:, : -condition_latents.shape[1]],
hidden_states[:, -condition_latents.shape[1] :],
)
return hidden_states, condition_latents
else:
return hidden_states
def block_forward(
self,
hidden_states: torch.FloatTensor,
encoder_hidden_states: torch.FloatTensor,
condition_latents: torch.FloatTensor,
temb: torch.FloatTensor,
cond_temb: torch.FloatTensor,
cond_rotary_emb=None,
image_rotary_emb=None,
model_config: Optional[Dict[str, Any]] = {},
):
use_cond = condition_latents is not None
with enable_lora((self.norm1.linear,), model_config.get("latent_lora", False)):
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
hidden_states, emb=temb
)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = (
self.norm1_context(encoder_hidden_states, emb=temb)
)
if use_cond:
(
norm_condition_latents,
cond_gate_msa,
cond_shift_mlp,
cond_scale_mlp,
cond_gate_mlp,
) = self.norm1(condition_latents, emb=cond_temb)
# Attention.
result = attn_forward(
self.attn,
model_config=model_config,
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
condition_latents=norm_condition_latents if use_cond else None,
image_rotary_emb=image_rotary_emb,
cond_rotary_emb=cond_rotary_emb if use_cond else None,
)
attn_output, context_attn_output = result[:2]
cond_attn_output = result[2] if use_cond else None
# Process attention outputs for the `hidden_states`.
# 1. hidden_states
attn_output = gate_msa.unsqueeze(1) * attn_output
hidden_states = hidden_states + attn_output
# 2. encoder_hidden_states
context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
encoder_hidden_states = encoder_hidden_states + context_attn_output
# 3. condition_latents
if use_cond:
cond_attn_output = cond_gate_msa.unsqueeze(1) * cond_attn_output
condition_latents = condition_latents + cond_attn_output
if model_config.get("add_cond_attn", False):
hidden_states += cond_attn_output
# LayerNorm + MLP.
# 1. hidden_states
norm_hidden_states = self.norm2(hidden_states)
norm_hidden_states = (
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
)
# 2. encoder_hidden_states
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
norm_encoder_hidden_states = (
norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
)
# 3. condition_latents
if use_cond:
norm_condition_latents = self.norm2(condition_latents)
norm_condition_latents = (
norm_condition_latents * (1 + cond_scale_mlp[:, None])
+ cond_shift_mlp[:, None]
)
# Feed-forward.
with enable_lora((self.ff.net[2],), model_config.get("latent_lora", False)):
# 1. hidden_states
ff_output = self.ff(norm_hidden_states)
ff_output = gate_mlp.unsqueeze(1) * ff_output
# 2. encoder_hidden_states
context_ff_output = self.ff_context(norm_encoder_hidden_states)
context_ff_output = c_gate_mlp.unsqueeze(1) * context_ff_output
# 3. condition_latents
if use_cond:
cond_ff_output = self.ff(norm_condition_latents)
cond_ff_output = cond_gate_mlp.unsqueeze(1) * cond_ff_output
# Process feed-forward outputs.
hidden_states = hidden_states + ff_output
encoder_hidden_states = encoder_hidden_states + context_ff_output
if use_cond:
condition_latents = condition_latents + cond_ff_output
# Clip to avoid overflow.
if encoder_hidden_states.dtype == torch.float16:
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
return encoder_hidden_states, hidden_states, condition_latents if use_cond else None
def single_block_forward(
self,
hidden_states: torch.FloatTensor,
temb: torch.FloatTensor,
image_rotary_emb=None,
condition_latents: torch.FloatTensor = None,
cond_temb: torch.FloatTensor = None,
cond_rotary_emb=None,
model_config: Optional[Dict[str, Any]] = {},
):
using_cond = condition_latents is not None
residual = hidden_states
with enable_lora(
(
self.norm.linear,
self.proj_mlp,
),
model_config.get("latent_lora", False),
):
norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
if using_cond:
residual_cond = condition_latents
norm_condition_latents, cond_gate = self.norm(condition_latents, emb=cond_temb)
mlp_cond_hidden_states = self.act_mlp(self.proj_mlp(norm_condition_latents))
attn_output = attn_forward(
self.attn,
model_config=model_config,
hidden_states=norm_hidden_states,
image_rotary_emb=image_rotary_emb,
**(
{
"condition_latents": norm_condition_latents,
"cond_rotary_emb": cond_rotary_emb if using_cond else None,
}
if using_cond
else {}
),
)
if using_cond:
attn_output, cond_attn_output = attn_output
with enable_lora((self.proj_out,), model_config.get("latent_lora", False)):
hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
gate = gate.unsqueeze(1)
hidden_states = gate * self.proj_out(hidden_states)
hidden_states = residual + hidden_states
if using_cond:
condition_latents = torch.cat([cond_attn_output, mlp_cond_hidden_states], dim=2)
cond_gate = cond_gate.unsqueeze(1)
condition_latents = cond_gate * self.proj_out(condition_latents)
condition_latents = residual_cond + condition_latents
if hidden_states.dtype == torch.float16:
hidden_states = hidden_states.clip(-65504, 65504)
return hidden_states if not using_cond else (hidden_states, condition_latents)