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)
|
||||
@@ -0,0 +1,128 @@
|
||||
import torch
|
||||
from typing import Optional, Union, List, Tuple
|
||||
from diffusers.pipelines import FluxPipeline
|
||||
from PIL import Image, ImageFilter
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
condition_dict = {
|
||||
"depth": 0,
|
||||
"canny": 1,
|
||||
"subject": 4,
|
||||
"coloring": 6,
|
||||
"deblurring": 7,
|
||||
"fill": 9,
|
||||
}
|
||||
|
||||
|
||||
class Condition(object):
|
||||
def __init__(
|
||||
self,
|
||||
condition_type: str,
|
||||
raw_img: Union[Image.Image, torch.Tensor] = None,
|
||||
condition: Union[Image.Image, torch.Tensor] = None,
|
||||
mask=None,
|
||||
) -> None:
|
||||
self.condition_type = condition_type
|
||||
assert raw_img is not None or condition is not None
|
||||
if raw_img is not None:
|
||||
self.condition = self.get_condition(condition_type, raw_img)
|
||||
else:
|
||||
self.condition = condition
|
||||
# TODO: Add mask support
|
||||
assert mask is None, "Mask not supported yet"
|
||||
|
||||
def get_condition(
|
||||
self, condition_type: str, raw_img: Union[Image.Image, torch.Tensor]
|
||||
) -> Union[Image.Image, torch.Tensor]:
|
||||
"""
|
||||
Returns the condition image.
|
||||
"""
|
||||
if condition_type == "depth":
|
||||
from transformers import pipeline
|
||||
import folder_paths
|
||||
import os
|
||||
|
||||
depth_pipe = pipeline(
|
||||
task="depth-estimation",
|
||||
# model="LiheYoung/depth-anything-small-hf",
|
||||
model=os.path.join(folder_paths.models_dir, 'flux', 'OminiControl', 'depth-anything-small-hf'),
|
||||
device="cuda",
|
||||
)
|
||||
source_image = raw_img.convert("RGB")
|
||||
condition_img = depth_pipe(source_image)["depth"].convert("RGB")
|
||||
return condition_img
|
||||
elif condition_type == "canny":
|
||||
img = np.array(raw_img)
|
||||
edges = cv2.Canny(img, 100, 200)
|
||||
edges = Image.fromarray(edges).convert("RGB")
|
||||
return edges
|
||||
elif condition_type == "subject":
|
||||
return raw_img
|
||||
elif condition_type == "coloring":
|
||||
return raw_img.convert("L").convert("RGB")
|
||||
elif condition_type == "deblurring":
|
||||
condition_image = (
|
||||
raw_img.convert("RGB")
|
||||
.filter(ImageFilter.GaussianBlur(10))
|
||||
.convert("RGB")
|
||||
)
|
||||
return condition_image
|
||||
elif condition_type == "fill":
|
||||
return raw_img.convert("RGB")
|
||||
return self.condition
|
||||
|
||||
@property
|
||||
def type_id(self) -> int:
|
||||
"""
|
||||
Returns the type id of the condition.
|
||||
"""
|
||||
return condition_dict[self.condition_type]
|
||||
|
||||
@classmethod
|
||||
def get_type_id(cls, condition_type: str) -> int:
|
||||
"""
|
||||
Returns the type id of the condition.
|
||||
"""
|
||||
return condition_dict[condition_type]
|
||||
|
||||
def _encode_image(self, pipe: FluxPipeline, cond_img: Image.Image) -> torch.Tensor:
|
||||
"""
|
||||
Encodes an image condition into tokens using the pipeline.
|
||||
"""
|
||||
cond_img = pipe.image_processor.preprocess(cond_img)
|
||||
# cond_img = cond_img.to(pipe.device).to(pipe.dtype)
|
||||
cond_img = cond_img.to('cuda').to(pipe.dtype)
|
||||
cond_img = pipe.vae.encode(cond_img).latent_dist.sample()
|
||||
cond_img = (
|
||||
cond_img - pipe.vae.config.shift_factor
|
||||
) * pipe.vae.config.scaling_factor
|
||||
cond_tokens = pipe._pack_latents(cond_img, *cond_img.shape)
|
||||
cond_ids = pipe._prepare_latent_image_ids(
|
||||
cond_img.shape[0],
|
||||
cond_img.shape[2],
|
||||
cond_img.shape[3],
|
||||
pipe.device,
|
||||
pipe.dtype,
|
||||
)
|
||||
return cond_tokens, cond_ids
|
||||
|
||||
def encode(self, pipe: FluxPipeline) -> Tuple[torch.Tensor, torch.Tensor, int]:
|
||||
"""
|
||||
Encodes the condition into tokens, ids and type_id.
|
||||
"""
|
||||
if self.condition_type in [
|
||||
"depth",
|
||||
"canny",
|
||||
"subject",
|
||||
"coloring",
|
||||
"deblurring",
|
||||
"fill",
|
||||
]:
|
||||
tokens, ids = self._encode_image(pipe, self.condition)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Condition type {self.condition_type} not implemented"
|
||||
)
|
||||
type_id = torch.ones_like(ids[:, :1]) * self.type_id
|
||||
return tokens, ids, type_id
|
||||
+300
@@ -0,0 +1,300 @@
|
||||
import torch
|
||||
import yaml, os
|
||||
from diffusers.pipelines import FluxPipeline
|
||||
from typing import List, Union, Optional, Dict, Any, Callable
|
||||
from .transformer import tranformer_forward
|
||||
from .condition import Condition
|
||||
|
||||
from diffusers.pipelines.flux.pipeline_flux import (
|
||||
FluxPipelineOutput,
|
||||
calculate_shift,
|
||||
retrieve_timesteps,
|
||||
np,
|
||||
)
|
||||
|
||||
|
||||
def prepare_params(
|
||||
prompt: Union[str, List[str]] = None,
|
||||
prompt_2: Optional[Union[str, List[str]]] = None,
|
||||
height: Optional[int] = 512,
|
||||
width: Optional[int] = 512,
|
||||
num_inference_steps: int = 28,
|
||||
timesteps: List[int] = None,
|
||||
guidance_scale: float = 3.5,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
output_type: Optional[str] = "pil",
|
||||
return_dict: bool = True,
|
||||
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
|
||||
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
||||
max_sequence_length: int = 512,
|
||||
**kwargs: dict,
|
||||
):
|
||||
return (
|
||||
prompt,
|
||||
prompt_2,
|
||||
height,
|
||||
width,
|
||||
num_inference_steps,
|
||||
timesteps,
|
||||
guidance_scale,
|
||||
num_images_per_prompt,
|
||||
generator,
|
||||
latents,
|
||||
prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
output_type,
|
||||
return_dict,
|
||||
joint_attention_kwargs,
|
||||
callback_on_step_end,
|
||||
callback_on_step_end_tensor_inputs,
|
||||
max_sequence_length,
|
||||
)
|
||||
|
||||
|
||||
def seed_everything(seed: int = 42):
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.manual_seed(seed)
|
||||
np.random.seed(seed)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(
|
||||
pipeline: FluxPipeline,
|
||||
conditions: List[Condition] = None,
|
||||
model_config: Optional[Dict[str, Any]] = {},
|
||||
condition_scale: float = 1.0,
|
||||
**params: dict,
|
||||
):
|
||||
|
||||
# print(params)
|
||||
# return
|
||||
# model_config = model_config or get_config(config_path).get("model", {})
|
||||
if condition_scale != 1:
|
||||
for name, module in pipeline.transformer.named_modules():
|
||||
if not name.endswith(".attn"):
|
||||
continue
|
||||
module.c_factor = torch.ones(1, 1) * condition_scale
|
||||
|
||||
self = pipeline
|
||||
(
|
||||
prompt,
|
||||
prompt_2,
|
||||
height,
|
||||
width,
|
||||
num_inference_steps,
|
||||
timesteps,
|
||||
guidance_scale,
|
||||
num_images_per_prompt,
|
||||
generator,
|
||||
latents,
|
||||
prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
output_type,
|
||||
return_dict,
|
||||
joint_attention_kwargs,
|
||||
callback_on_step_end,
|
||||
callback_on_step_end_tensor_inputs,
|
||||
max_sequence_length,
|
||||
) = prepare_params(**params)
|
||||
|
||||
height = height or self.default_sample_size * self.vae_scale_factor
|
||||
width = width or self.default_sample_size * self.vae_scale_factor
|
||||
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(
|
||||
prompt,
|
||||
prompt_2,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
|
||||
max_sequence_length=max_sequence_length,
|
||||
)
|
||||
|
||||
self._guidance_scale = guidance_scale
|
||||
self._joint_attention_kwargs = joint_attention_kwargs
|
||||
self._interrupt = False
|
||||
|
||||
# 2. Define call parameters
|
||||
if prompt is not None and isinstance(prompt, str):
|
||||
batch_size = 1
|
||||
elif prompt is not None and isinstance(prompt, list):
|
||||
batch_size = len(prompt)
|
||||
else:
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
device = self._execution_device
|
||||
# print('execution device --->', device)
|
||||
# device = 'cpu'
|
||||
|
||||
lora_scale = (
|
||||
self.joint_attention_kwargs.get("scale", None)
|
||||
if self.joint_attention_kwargs is not None
|
||||
else None
|
||||
)
|
||||
(
|
||||
prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
text_ids,
|
||||
) = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
prompt_2=prompt_2,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
|
||||
# print('-----> real prompt embeds:')
|
||||
# print(prompt_embeds[0])
|
||||
# print(pooled_prompt_embeds[0])
|
||||
# print(text_ids)
|
||||
|
||||
# 4. Prepare latent variables
|
||||
num_channels_latents = self.transformer.config.in_channels // 4
|
||||
latents, latent_image_ids = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
|
||||
# 4.1. Prepare conditions
|
||||
condition_latents, condition_ids, condition_type_ids = ([] for _ in range(3))
|
||||
use_condition = conditions is not None or []
|
||||
if use_condition:
|
||||
assert len(conditions) <= 1, "Only one condition is supported for now."
|
||||
pipeline.set_adapters(conditions[0].condition_type)
|
||||
for condition in conditions:
|
||||
# tokens, ids, type_id = condition.encode(self)
|
||||
tokens, ids, type_id = params['encoded_condition']
|
||||
tokens.to(torch.bfloat16)
|
||||
ids.to(torch.bfloat16)
|
||||
type_id.to(torch.bfloat16)
|
||||
condition_latents.append(tokens) # [batch_size, token_n, token_dim]
|
||||
condition_ids.append(ids) # [token_n, id_dim(3)]
|
||||
condition_type_ids.append(type_id) # [token_n, 1]
|
||||
condition_latents = torch.cat(condition_latents, dim=1)
|
||||
condition_ids = torch.cat(condition_ids, dim=0)
|
||||
if condition.condition_type == "subject":
|
||||
condition_ids[:, 2] += width // 16
|
||||
condition_type_ids = torch.cat(condition_type_ids, dim=0)
|
||||
|
||||
# 5. Prepare timesteps
|
||||
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
|
||||
image_seq_len = latents.shape[1]
|
||||
mu = calculate_shift(
|
||||
image_seq_len,
|
||||
self.scheduler.config.base_image_seq_len,
|
||||
self.scheduler.config.max_image_seq_len,
|
||||
self.scheduler.config.base_shift,
|
||||
self.scheduler.config.max_shift,
|
||||
)
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler,
|
||||
num_inference_steps,
|
||||
device,
|
||||
timesteps,
|
||||
sigmas,
|
||||
mu=mu,
|
||||
)
|
||||
num_warmup_steps = max(
|
||||
len(timesteps) - num_inference_steps * self.scheduler.order, 0
|
||||
)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
# 6. Denoising loop
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
if self.interrupt:
|
||||
continue
|
||||
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timestep = t.expand(latents.shape[0]).to(latents.dtype)
|
||||
|
||||
# handle guidance
|
||||
if self.transformer.config.guidance_embeds:
|
||||
guidance = torch.tensor([guidance_scale], device=device)
|
||||
guidance = guidance.expand(latents.shape[0])
|
||||
else:
|
||||
guidance = None
|
||||
noise_pred = tranformer_forward(
|
||||
self.transformer,
|
||||
model_config=model_config,
|
||||
# Inputs of the condition (new feature)
|
||||
condition_latents=condition_latents if use_condition else None,
|
||||
condition_ids=condition_ids if use_condition else None,
|
||||
condition_type_ids=condition_type_ids if use_condition else None,
|
||||
# Inputs to the original transformer
|
||||
hidden_states=latents,
|
||||
# YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
|
||||
timestep=timestep / 1000,
|
||||
guidance=guidance,
|
||||
pooled_projections=pooled_prompt_embeds,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
txt_ids=text_ids,
|
||||
img_ids=latent_image_ids,
|
||||
joint_attention_kwargs=self.joint_attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents_dtype = latents.dtype
|
||||
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
|
||||
|
||||
if latents.dtype != latents_dtype:
|
||||
if torch.backends.mps.is_available():
|
||||
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
|
||||
latents = latents.to(latents_dtype)
|
||||
|
||||
if callback_on_step_end is not None:
|
||||
callback_kwargs = {}
|
||||
for k in callback_on_step_end_tensor_inputs:
|
||||
callback_kwargs[k] = locals()[k]
|
||||
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
||||
|
||||
latents = callback_outputs.pop("latents", latents)
|
||||
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
||||
|
||||
# call the callback, if provided
|
||||
if i == len(timesteps) - 1 or (
|
||||
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
|
||||
):
|
||||
progress_bar.update()
|
||||
|
||||
if output_type == "latent":
|
||||
image = latents
|
||||
|
||||
else:
|
||||
latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
|
||||
latents = (
|
||||
latents / self.vae.config.scaling_factor
|
||||
) + self.vae.config.shift_factor
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
image = self.image_processor.postprocess(image, output_type=output_type)
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
if condition_scale != 1:
|
||||
for name, module in pipeline.transformer.named_modules():
|
||||
if not name.endswith(".attn"):
|
||||
continue
|
||||
del module.c_factor
|
||||
|
||||
if not return_dict:
|
||||
return (image,)
|
||||
|
||||
return FluxPipelineOutput(images=image)
|
||||
@@ -0,0 +1,75 @@
|
||||
from peft.tuners.tuners_utils import BaseTunerLayer
|
||||
from typing import List, Any, Optional, Type
|
||||
|
||||
|
||||
class enable_lora:
|
||||
def __init__(self, lora_modules: List[BaseTunerLayer], activated: bool) -> None:
|
||||
self.activated: bool = activated
|
||||
if activated:
|
||||
return
|
||||
self.lora_modules: List[BaseTunerLayer] = [
|
||||
each for each in lora_modules if isinstance(each, BaseTunerLayer)
|
||||
]
|
||||
self.scales = [
|
||||
{
|
||||
active_adapter: lora_module.scaling[active_adapter]
|
||||
for active_adapter in lora_module.active_adapters
|
||||
}
|
||||
for lora_module in self.lora_modules
|
||||
]
|
||||
|
||||
def __enter__(self) -> None:
|
||||
if self.activated:
|
||||
return
|
||||
|
||||
for lora_module in self.lora_modules:
|
||||
if not isinstance(lora_module, BaseTunerLayer):
|
||||
continue
|
||||
lora_module.scale_layer(0)
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[Any],
|
||||
) -> None:
|
||||
if self.activated:
|
||||
return
|
||||
for i, lora_module in enumerate(self.lora_modules):
|
||||
if not isinstance(lora_module, BaseTunerLayer):
|
||||
continue
|
||||
for active_adapter in lora_module.active_adapters:
|
||||
lora_module.scaling[active_adapter] = self.scales[i][active_adapter]
|
||||
|
||||
|
||||
class set_lora_scale:
|
||||
def __init__(self, lora_modules: List[BaseTunerLayer], scale: float) -> None:
|
||||
self.lora_modules: List[BaseTunerLayer] = [
|
||||
each for each in lora_modules if isinstance(each, BaseTunerLayer)
|
||||
]
|
||||
self.scales = [
|
||||
{
|
||||
active_adapter: lora_module.scaling[active_adapter]
|
||||
for active_adapter in lora_module.active_adapters
|
||||
}
|
||||
for lora_module in self.lora_modules
|
||||
]
|
||||
self.scale = scale
|
||||
|
||||
def __enter__(self) -> None:
|
||||
for lora_module in self.lora_modules:
|
||||
if not isinstance(lora_module, BaseTunerLayer):
|
||||
continue
|
||||
lora_module.scale_layer(self.scale)
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[Any],
|
||||
) -> None:
|
||||
for i, lora_module in enumerate(self.lora_modules):
|
||||
if not isinstance(lora_module, BaseTunerLayer):
|
||||
continue
|
||||
for active_adapter in lora_module.active_adapters:
|
||||
lora_module.scaling[active_adapter] = self.scales[i][active_adapter]
|
||||
@@ -0,0 +1,270 @@
|
||||
import torch
|
||||
from diffusers.pipelines import FluxPipeline
|
||||
from typing import List, Union, Optional, Dict, Any, Callable
|
||||
from .block import block_forward, single_block_forward
|
||||
from .lora_controller import enable_lora
|
||||
from diffusers.models.transformers.transformer_flux import (
|
||||
FluxTransformer2DModel,
|
||||
Transformer2DModelOutput,
|
||||
USE_PEFT_BACKEND,
|
||||
is_torch_version,
|
||||
scale_lora_layers,
|
||||
unscale_lora_layers,
|
||||
logger,
|
||||
)
|
||||
import numpy as np
|
||||
|
||||
|
||||
def prepare_params(
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor = None,
|
||||
pooled_projections: torch.Tensor = None,
|
||||
timestep: torch.LongTensor = None,
|
||||
img_ids: torch.Tensor = None,
|
||||
txt_ids: torch.Tensor = None,
|
||||
guidance: torch.Tensor = None,
|
||||
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
controlnet_block_samples=None,
|
||||
controlnet_single_block_samples=None,
|
||||
return_dict: bool = True,
|
||||
**kwargs: dict,
|
||||
):
|
||||
return (
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
pooled_projections,
|
||||
timestep,
|
||||
img_ids,
|
||||
txt_ids,
|
||||
guidance,
|
||||
joint_attention_kwargs,
|
||||
controlnet_block_samples,
|
||||
controlnet_single_block_samples,
|
||||
return_dict,
|
||||
)
|
||||
|
||||
|
||||
def tranformer_forward(
|
||||
transformer: FluxTransformer2DModel,
|
||||
condition_latents: torch.Tensor,
|
||||
condition_ids: torch.Tensor,
|
||||
condition_type_ids: torch.Tensor,
|
||||
model_config: Optional[Dict[str, Any]] = {},
|
||||
return_conditional_latents: bool = False,
|
||||
c_t=0,
|
||||
**params: dict,
|
||||
):
|
||||
self = transformer
|
||||
use_condition = condition_latents is not None
|
||||
use_condition_in_single_blocks = model_config.get(
|
||||
"use_condition_in_single_blocks", True
|
||||
)
|
||||
# if return_conditional_latents is True, use_condition and use_condition_in_single_blocks must be True
|
||||
assert not return_conditional_latents or (
|
||||
use_condition and use_condition_in_single_blocks
|
||||
), "`return_conditional_latents` is True, `use_condition` and `use_condition_in_single_blocks` must be True"
|
||||
|
||||
(
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
pooled_projections,
|
||||
timestep,
|
||||
img_ids,
|
||||
txt_ids,
|
||||
guidance,
|
||||
joint_attention_kwargs,
|
||||
controlnet_block_samples,
|
||||
controlnet_single_block_samples,
|
||||
return_dict,
|
||||
) = prepare_params(**params)
|
||||
|
||||
if joint_attention_kwargs is not None:
|
||||
joint_attention_kwargs = joint_attention_kwargs.copy()
|
||||
lora_scale = joint_attention_kwargs.pop("scale", 1.0)
|
||||
else:
|
||||
lora_scale = 1.0
|
||||
|
||||
if USE_PEFT_BACKEND:
|
||||
# weight the lora layers by setting `lora_scale` for each PEFT layer
|
||||
scale_lora_layers(self, lora_scale)
|
||||
else:
|
||||
if (
|
||||
joint_attention_kwargs is not None
|
||||
and joint_attention_kwargs.get("scale", None) is not None
|
||||
):
|
||||
logger.warning(
|
||||
"Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
|
||||
)
|
||||
with enable_lora((self.x_embedder,), model_config.get("latent_lora", False)):
|
||||
hidden_states = self.x_embedder(hidden_states)
|
||||
condition_latents = self.x_embedder(condition_latents) if use_condition else None
|
||||
|
||||
timestep = timestep.to(hidden_states.dtype) * 1000
|
||||
if guidance is not None:
|
||||
guidance = guidance.to(hidden_states.dtype) * 1000
|
||||
else:
|
||||
guidance = None
|
||||
temb = (
|
||||
self.time_text_embed(timestep, pooled_projections)
|
||||
if guidance is None
|
||||
else self.time_text_embed(timestep, guidance, pooled_projections)
|
||||
)
|
||||
cond_temb = (
|
||||
self.time_text_embed(torch.ones_like(timestep) * c_t * 1000, pooled_projections)
|
||||
if guidance is None
|
||||
else self.time_text_embed(
|
||||
torch.ones_like(timestep) * c_t * 1000, guidance, pooled_projections
|
||||
)
|
||||
)
|
||||
if hasattr(self, "cond_type_embed") and condition_type_ids is not None:
|
||||
cond_type_proj = self.time_text_embed.time_proj(condition_type_ids[0])
|
||||
cond_type_emb = self.cond_type_embed(cond_type_proj.to(dtype=cond_temb.dtype))
|
||||
cond_temb = cond_temb + cond_type_emb
|
||||
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
|
||||
|
||||
if txt_ids.ndim == 3:
|
||||
logger.warning(
|
||||
"Passing `txt_ids` 3d torch.Tensor is deprecated."
|
||||
"Please remove the batch dimension and pass it as a 2d torch Tensor"
|
||||
)
|
||||
txt_ids = txt_ids[0]
|
||||
if img_ids.ndim == 3:
|
||||
logger.warning(
|
||||
"Passing `img_ids` 3d torch.Tensor is deprecated."
|
||||
"Please remove the batch dimension and pass it as a 2d torch Tensor"
|
||||
)
|
||||
img_ids = img_ids[0]
|
||||
|
||||
ids = torch.cat((txt_ids, img_ids), dim=0)
|
||||
image_rotary_emb = self.pos_embed(ids)
|
||||
if use_condition:
|
||||
cond_ids = condition_ids
|
||||
cond_rotary_emb = self.pos_embed(cond_ids)
|
||||
|
||||
# hidden_states = torch.cat([hidden_states, condition_latents], dim=1)
|
||||
|
||||
for index_block, block in enumerate(self.transformer_blocks):
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = (
|
||||
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
)
|
||||
encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(block),
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb,
|
||||
image_rotary_emb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
encoder_hidden_states, hidden_states, condition_latents = block_forward(
|
||||
block,
|
||||
model_config=model_config,
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
condition_latents=condition_latents if use_condition else None,
|
||||
temb=temb,
|
||||
cond_temb=cond_temb if use_condition else None,
|
||||
cond_rotary_emb=cond_rotary_emb if use_condition else None,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
)
|
||||
|
||||
# controlnet residual
|
||||
if controlnet_block_samples is not None:
|
||||
interval_control = len(self.transformer_blocks) / len(
|
||||
controlnet_block_samples
|
||||
)
|
||||
interval_control = int(np.ceil(interval_control))
|
||||
hidden_states = (
|
||||
hidden_states
|
||||
+ controlnet_block_samples[index_block // interval_control]
|
||||
)
|
||||
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
|
||||
|
||||
for index_block, block in enumerate(self.single_transformer_blocks):
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = (
|
||||
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
)
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(block),
|
||||
hidden_states,
|
||||
temb,
|
||||
image_rotary_emb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
result = single_block_forward(
|
||||
block,
|
||||
model_config=model_config,
|
||||
hidden_states=hidden_states,
|
||||
temb=temb,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
**(
|
||||
{
|
||||
"condition_latents": condition_latents,
|
||||
"cond_temb": cond_temb,
|
||||
"cond_rotary_emb": cond_rotary_emb,
|
||||
}
|
||||
if use_condition_in_single_blocks and use_condition
|
||||
else {}
|
||||
),
|
||||
)
|
||||
if use_condition_in_single_blocks and use_condition:
|
||||
hidden_states, condition_latents = result
|
||||
else:
|
||||
hidden_states = result
|
||||
|
||||
# controlnet residual
|
||||
if controlnet_single_block_samples is not None:
|
||||
interval_control = len(self.single_transformer_blocks) / len(
|
||||
controlnet_single_block_samples
|
||||
)
|
||||
interval_control = int(np.ceil(interval_control))
|
||||
hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
|
||||
hidden_states[:, encoder_hidden_states.shape[1] :, ...]
|
||||
+ controlnet_single_block_samples[index_block // interval_control]
|
||||
)
|
||||
|
||||
hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
|
||||
|
||||
hidden_states = self.norm_out(hidden_states, temb)
|
||||
output = self.proj_out(hidden_states)
|
||||
if return_conditional_latents:
|
||||
condition_latents = (
|
||||
self.norm_out(condition_latents, cond_temb) if use_condition else None
|
||||
)
|
||||
condition_output = self.proj_out(condition_latents) if use_condition else None
|
||||
|
||||
if USE_PEFT_BACKEND:
|
||||
# remove `lora_scale` from each PEFT layer
|
||||
unscale_lora_layers(self, lora_scale)
|
||||
|
||||
if not return_dict:
|
||||
return (
|
||||
(output,) if not return_conditional_latents else (output, condition_output)
|
||||
)
|
||||
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
Reference in New Issue
Block a user