Add files via upload
This commit is contained in:
HM-RunningHub
+333
@@ -0,0 +1,333 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user