489 lines
20 KiB
Python
489 lines
20 KiB
Python
# coding: utf-8
|
|
|
|
"""
|
|
@File : hm_adapters.py
|
|
@Author : Songkey
|
|
@Email : songkey@pku.edu.cn
|
|
@Date : 5/19/2025
|
|
@Desc :
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.utils.checkpoint
|
|
|
|
from typing import Optional, Union, Dict, List, Tuple
|
|
import copy
|
|
import torch.nn.functional as F
|
|
from diffusers import StableDiffusionImg2ImgPipeline
|
|
from diffusers.models.modeling_utils import ModelMixin
|
|
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
|
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
|
|
|
|
from diffusers.configuration_utils import FrozenDict
|
|
from diffusers.models.attention_processor import (
|
|
AttnProcessor,
|
|
AttnProcessor2_0,
|
|
)
|
|
|
|
from diffusers.utils import (USE_PEFT_BACKEND,
|
|
scale_lora_layers,
|
|
unscale_lora_layers,
|
|
logging, is_torch_version,
|
|
is_peft_available,
|
|
is_peft_version,
|
|
is_transformers_available,
|
|
is_transformers_version
|
|
)
|
|
|
|
from diffusers.loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin
|
|
from diffusers.models.lora import adjust_lora_scale_text_encoder
|
|
|
|
from .hm_blocks import (SKReferenceAttentionV3,
|
|
SKReferenceAttentionV5,
|
|
SKReferenceAttention,
|
|
SKMotionModule)
|
|
|
|
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
|
|
|
_LOW_CPU_MEM_USAGE_DEFAULT_LORA = False
|
|
if is_torch_version(">=", "1.9.0"):
|
|
if (
|
|
is_peft_available()
|
|
and is_peft_version(">=", "0.13.1")
|
|
and is_transformers_available()
|
|
and is_transformers_version(">", "4.45.2")
|
|
):
|
|
_LOW_CPU_MEM_USAGE_DEFAULT_LORA = True
|
|
|
|
class HMReferenceAdapter(ModelMixin, ConfigMixin):
|
|
@register_to_config
|
|
def __init__(self,
|
|
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
|
|
down_block_types: Tuple[str] = (
|
|
"CrossAttnDownBlock2D",
|
|
"CrossAttnDownBlock2D",
|
|
"CrossAttnDownBlock2D",
|
|
"DownBlock2D",
|
|
),
|
|
up_block_types: Tuple[str] = (
|
|
"UpBlock2D",
|
|
"CrossAttnUpBlock2D",
|
|
"CrossAttnUpBlock2D",
|
|
"CrossAttnUpBlock2D"),
|
|
num_attention_heads: Optional[Union[int, Tuple[int]]] = 8,
|
|
version='v1'
|
|
):
|
|
super().__init__()
|
|
|
|
if isinstance(num_attention_heads, int):
|
|
num_attention_heads = (num_attention_heads,) * len(down_block_types)
|
|
|
|
if version == 'v1':
|
|
self.reference_modules_down = nn.ModuleList([])
|
|
for i, down_block_type in enumerate(down_block_types):
|
|
output_channel = block_out_channels[i]
|
|
|
|
self.reference_modules_down.append(
|
|
SKReferenceAttention(
|
|
in_channels=output_channel,
|
|
num_attention_heads=num_attention_heads[i],
|
|
)
|
|
)
|
|
|
|
self.reference_modules_mid = SKReferenceAttention(
|
|
in_channels=block_out_channels[-1],
|
|
num_attention_heads=num_attention_heads[-1],
|
|
)
|
|
|
|
self.reference_modules_up = nn.ModuleList([])
|
|
|
|
reversed_block_out_channels = list(reversed(block_out_channels))
|
|
reversed_num_attention_heads = list(reversed(num_attention_heads))
|
|
|
|
output_channel = reversed_block_out_channels[0]
|
|
for i, up_block_type in enumerate(up_block_types):
|
|
prev_output_channel = output_channel
|
|
output_channel = reversed_block_out_channels[i]
|
|
|
|
if i > 0:
|
|
self.reference_modules_up.append(
|
|
SKReferenceAttention(
|
|
in_channels=prev_output_channel,
|
|
num_attention_heads=reversed_num_attention_heads[i],
|
|
num_positional_embeddings=64 * 2
|
|
)
|
|
)
|
|
|
|
|
|
class HM3ReferenceAdapter(ModelMixin, ConfigMixin):
|
|
@register_to_config
|
|
def __init__(self, block_down_channels: Tuple[int] = (320, 640, 1280, 1280),
|
|
block_up_channels: Tuple[int] = (1280, 1280, 1280, 640),
|
|
num_attention_heads: int = 8,
|
|
use_3d: bool = False):
|
|
super().__init__()
|
|
|
|
self.reference_modules_up = nn.ModuleList([])
|
|
for i, in_channels in enumerate(block_up_channels):
|
|
self.reference_modules_up.append(
|
|
SKReferenceAttentionV3(
|
|
in_channels=in_channels,
|
|
num_attention_heads=num_attention_heads,
|
|
num_positional_embeddings=64*2
|
|
)
|
|
)
|
|
|
|
|
|
class HM5ReferenceAdapter(ModelMixin, ConfigMixin):
|
|
@register_to_config
|
|
def __init__(self, block_down_channels: Tuple[int] = (320, 640, 1280, 1280),
|
|
block_up_channels: Tuple[int] = (1280, 1280, 1280, 640),
|
|
num_attention_heads: int = 8,
|
|
use_3d: bool = False):
|
|
super().__init__()
|
|
|
|
self.reference_modules_up = nn.ModuleList([])
|
|
for i, in_channels in enumerate(block_up_channels):
|
|
self.reference_modules_up.append(
|
|
SKReferenceAttentionV5(
|
|
in_channels=in_channels,
|
|
num_attention_heads=num_attention_heads,
|
|
)
|
|
)
|
|
|
|
|
|
class HM3MotionAdapter(ModelMixin, ConfigMixin):
|
|
@register_to_config
|
|
def __init__(self, block_down_channels: Tuple[int] = (320, 640, 1280, 1280),
|
|
block_up_channels: Tuple[int] = (1280, 1280, 1280, 640),
|
|
num_attention_heads: int = 8,
|
|
use_3d: bool = True):
|
|
super().__init__()
|
|
blocks_time_embed_dim = 1280
|
|
self.motion_down = nn.ModuleList([])
|
|
|
|
for i, in_channels in enumerate(block_down_channels):
|
|
self.motion_down.append(
|
|
SKMotionModule(
|
|
in_channels=in_channels,
|
|
num_attention_heads=num_attention_heads,
|
|
blocks_time_embed_dim=blocks_time_embed_dim,
|
|
)
|
|
)
|
|
|
|
self.motion_up = nn.ModuleList([])
|
|
for i, in_channels in enumerate(block_up_channels):
|
|
self.motion_up.append(
|
|
SKMotionModule(
|
|
in_channels=in_channels,
|
|
num_attention_heads=num_attention_heads,
|
|
blocks_time_embed_dim=blocks_time_embed_dim,
|
|
)
|
|
)
|
|
|
|
|
|
class HMPipeline(StableDiffusionImg2ImgPipeline):
|
|
@torch.no_grad()
|
|
def load_lora_weights_sk(self, unet, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name=None, text_encoder=None, lora_scale=1.0, **kwargs):
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
state_dict, network_alphas = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_unet(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
unet=unet,
|
|
adapter_name=adapter_name,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
)
|
|
|
|
if not text_encoder is None:
|
|
self.load_lora_into_text_encoder(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
text_encoder=text_encoder,
|
|
lora_scale=lora_scale,
|
|
adapter_name=adapter_name,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
)
|
|
|
|
def encode_prompt_sk(
|
|
self,
|
|
text_encoder,
|
|
prompt,
|
|
device,
|
|
num_images_per_prompt,
|
|
do_classifier_free_guidance,
|
|
negative_prompt=None,
|
|
prompt_embeds: Optional[torch.Tensor] = None,
|
|
negative_prompt_embeds: Optional[torch.Tensor] = None,
|
|
lora_scale: Optional[float] = None,
|
|
clip_skip: Optional[int] = None,
|
|
):
|
|
|
|
# set lora scale so that monkey patched LoRA
|
|
# function of text encoder can correctly access it
|
|
if lora_scale is not None and isinstance(self, StableDiffusionLoraLoaderMixin):
|
|
self._lora_scale = lora_scale
|
|
|
|
# dynamically adjust the LoRA scale
|
|
if not USE_PEFT_BACKEND:
|
|
adjust_lora_scale_text_encoder(text_encoder , lora_scale)
|
|
else:
|
|
scale_lora_layers(text_encoder , lora_scale)
|
|
|
|
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]
|
|
|
|
if prompt_embeds is None:
|
|
# textual inversion: process multi-vector tokens if necessary
|
|
if isinstance(self, TextualInversionLoaderMixin):
|
|
prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
|
|
|
|
text_inputs = self.tokenizer(
|
|
prompt,
|
|
padding="max_length",
|
|
max_length=self.tokenizer.model_max_length,
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
)
|
|
text_input_ids = text_inputs.input_ids
|
|
untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
|
|
|
|
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
|
|
text_input_ids, untruncated_ids
|
|
):
|
|
removed_text = self.tokenizer.batch_decode(
|
|
untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
|
|
)
|
|
logger.warning(
|
|
"The following part of your input was truncated because CLIP can only handle sequences up to"
|
|
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
|
|
)
|
|
|
|
if hasattr(text_encoder .config, "use_attention_mask") and text_encoder .config.use_attention_mask:
|
|
attention_mask = text_inputs.attention_mask.to(device)
|
|
else:
|
|
attention_mask = None
|
|
|
|
if clip_skip is None:
|
|
prompt_embeds = text_encoder (text_input_ids.to(device), attention_mask=attention_mask)
|
|
prompt_embeds = prompt_embeds[0]
|
|
else:
|
|
prompt_embeds = text_encoder (
|
|
text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True
|
|
)
|
|
# Access the `hidden_states` first, that contains a tuple of
|
|
# all the hidden states from the encoder layers. Then index into
|
|
# the tuple to access the hidden states from the desired layer.
|
|
prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]
|
|
# We also need to apply the final LayerNorm here to not mess with the
|
|
# representations. The `last_hidden_states` that we typically use for
|
|
# obtaining the final prompt representations passes through the LayerNorm
|
|
# layer.
|
|
prompt_embeds = text_encoder .text_model.final_layer_norm(prompt_embeds)
|
|
|
|
if text_encoder is not None:
|
|
prompt_embeds_dtype = text_encoder .dtype
|
|
elif self.unet is not None:
|
|
prompt_embeds_dtype = self.unet.dtype
|
|
else:
|
|
prompt_embeds_dtype = prompt_embeds.dtype
|
|
|
|
prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
|
|
|
|
bs_embed, seq_len, _ = prompt_embeds.shape
|
|
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
|
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
|
prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
|
|
|
|
# get unconditional embeddings for classifier free guidance
|
|
if do_classifier_free_guidance and negative_prompt_embeds is None:
|
|
uncond_tokens: List[str]
|
|
if negative_prompt is None:
|
|
uncond_tokens = [""] * batch_size
|
|
elif prompt is not None and type(prompt) is not type(negative_prompt):
|
|
raise TypeError(
|
|
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
|
|
f" {type(prompt)}."
|
|
)
|
|
elif isinstance(negative_prompt, str):
|
|
uncond_tokens = [negative_prompt]
|
|
elif batch_size != len(negative_prompt):
|
|
raise ValueError(
|
|
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
|
|
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
|
|
" the batch size of `prompt`."
|
|
)
|
|
else:
|
|
uncond_tokens = negative_prompt
|
|
|
|
# textual inversion: process multi-vector tokens if necessary
|
|
if isinstance(self, TextualInversionLoaderMixin):
|
|
uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
|
|
|
|
max_length = prompt_embeds.shape[1]
|
|
uncond_input = self.tokenizer(
|
|
uncond_tokens,
|
|
padding="max_length",
|
|
max_length=max_length,
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
)
|
|
|
|
if hasattr(text_encoder .config, "use_attention_mask") and text_encoder .config.use_attention_mask:
|
|
attention_mask = uncond_input.attention_mask.to(device)
|
|
else:
|
|
attention_mask = None
|
|
|
|
negative_prompt_embeds = text_encoder (
|
|
uncond_input.input_ids.to(device),
|
|
attention_mask=attention_mask,
|
|
)
|
|
negative_prompt_embeds = negative_prompt_embeds[0]
|
|
|
|
if do_classifier_free_guidance:
|
|
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
|
|
seq_len = negative_prompt_embeds.shape[1]
|
|
|
|
negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
|
|
|
|
negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
|
negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
|
|
|
|
if text_encoder is not None:
|
|
if isinstance(self, StableDiffusionLoraLoaderMixin) and USE_PEFT_BACKEND:
|
|
# Retrieve the original scale by scaling back the LoRA layers
|
|
unscale_lora_layers(text_encoder , lora_scale)
|
|
|
|
return prompt_embeds, negative_prompt_embeds
|
|
|
|
|
|
# https://github.com/huggingface/diffusers/blob/82058a5413ca09561cc5cc236c4abc5eeda7b209/src/diffusers/loaders/ip_adapter.py
|
|
|
|
if is_transformers_available():
|
|
from diffusers.models.attention_processor import (
|
|
IPAdapterAttnProcessor,
|
|
IPAdapterAttnProcessor2_0,
|
|
)
|
|
|
|
class CopyWeights(object):
|
|
@classmethod
|
|
def from_unet2d(cls, unet: UNet2DConditionModel):
|
|
# adapted from :https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_motion_model.py
|
|
|
|
config = dict(unet.config)
|
|
|
|
# Need this for backwards compatibility with UNet2DConditionModel checkpoints
|
|
if not config.get("num_attention_heads"):
|
|
config["num_attention_heads"] = config["attention_head_dim"]
|
|
|
|
config = FrozenDict(config)
|
|
model = cls.from_config(config)
|
|
|
|
model.conv_in.load_state_dict(unet.conv_in.state_dict())
|
|
|
|
model.time_proj.load_state_dict(unet.time_proj.state_dict())
|
|
model.time_embedding.load_state_dict(unet.time_embedding.state_dict())
|
|
|
|
if any(
|
|
isinstance(proc, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0))
|
|
for proc in unet.attn_processors.values()
|
|
):
|
|
attn_procs = {}
|
|
for name, processor in unet.attn_processors.items():
|
|
if name.endswith("attn1.processor"):
|
|
attn_processor_class = (
|
|
AttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else AttnProcessor
|
|
)
|
|
attn_procs[name] = attn_processor_class()
|
|
else:
|
|
attn_processor_class = (
|
|
IPAdapterAttnProcessor2_0
|
|
if hasattr(F, "scaled_dot_product_attention")
|
|
else IPAdapterAttnProcessor
|
|
)
|
|
attn_procs[name] = attn_processor_class(
|
|
hidden_size=processor.hidden_size,
|
|
cross_attention_dim=processor.cross_attention_dim,
|
|
scale=processor.scale,
|
|
num_tokens=processor.num_tokens,
|
|
)
|
|
for name, processor in model.attn_processors.items():
|
|
if name not in attn_procs:
|
|
attn_procs[name] = processor.__class__()
|
|
model.set_attn_processor(attn_procs)
|
|
model.config.encoder_hid_dim_type = "ip_image_proj"
|
|
model.encoder_hid_proj = unet.encoder_hid_proj
|
|
|
|
for i, down_block in enumerate(unet.down_blocks):
|
|
model.down_blocks[i].resnets.load_state_dict(down_block.resnets.state_dict())
|
|
if hasattr(model.down_blocks[i], "attentions"):
|
|
model.down_blocks[i].attentions.load_state_dict(down_block.attentions.state_dict())
|
|
if model.down_blocks[i].downsamplers:
|
|
model.down_blocks[i].downsamplers.load_state_dict(down_block.downsamplers.state_dict())
|
|
|
|
for i, up_block in enumerate(unet.up_blocks):
|
|
model.up_blocks[i].resnets.load_state_dict(up_block.resnets.state_dict())
|
|
if hasattr(model.up_blocks[i], "attentions"):
|
|
model.up_blocks[i].attentions.load_state_dict(up_block.attentions.state_dict())
|
|
if model.up_blocks[i].upsamplers:
|
|
model.up_blocks[i].upsamplers.load_state_dict(up_block.upsamplers.state_dict())
|
|
|
|
model.mid_block.resnets.load_state_dict(unet.mid_block.resnets.state_dict())
|
|
model.mid_block.attentions.load_state_dict(unet.mid_block.attentions.state_dict())
|
|
|
|
if unet.conv_norm_out is not None:
|
|
model.conv_norm_out.load_state_dict(unet.conv_norm_out.state_dict())
|
|
if unet.conv_act is not None:
|
|
model.conv_act.load_state_dict(unet.conv_act.state_dict())
|
|
model.conv_out.load_state_dict(unet.conv_out.state_dict())
|
|
|
|
# ensure that the Motion UNet is the same dtype as the UNet2DConditionModel
|
|
model.to(unet.dtype)
|
|
|
|
return model
|
|
|
|
class InsertReferenceAdapter(object):
|
|
def __init__(self):
|
|
self.reference_modules_down = None
|
|
self.reference_modules_mid = None
|
|
self.reference_modules_up = None
|
|
self.motion_down = None
|
|
self.motion_up = None
|
|
|
|
def insert_reference_adapter(self, adapter):
|
|
if hasattr(adapter, "reference_modules_down"):
|
|
self.reference_modules_down = copy.deepcopy(adapter.reference_modules_down)
|
|
if hasattr(adapter, "reference_modules_mid"):
|
|
self.reference_modules_mid = copy.deepcopy(adapter.reference_modules_mid)
|
|
if hasattr(adapter, "reference_modules_up"):
|
|
self.reference_modules_up = copy.deepcopy(adapter.reference_modules_up)
|
|
if hasattr(adapter, "motion_down"):
|
|
self.motion_down = copy.deepcopy(adapter.motion_down)
|
|
if hasattr(adapter, "motion_up"):
|
|
self.motion_up = copy.deepcopy(adapter.motion_up) |