diff --git a/.gitignore b/.gitignore index b931f2e..dc6d818 100644 --- a/.gitignore +++ b/.gitignore @@ -2,7 +2,10 @@ *.pyd *.pth *.pkl +*.pyc* +*.pyd* *_fps15.mp4 +*__pycache__/* .idea/ data/results pretrained_models/ diff --git a/hellomeme/__init__.py b/hellomeme/__init__.py index 6b1739a..4fa767d 100644 --- a/hellomeme/__init__.py +++ b/hellomeme/__init__.py @@ -8,4 +8,6 @@ @Desc : """ -from .pipelines import HMImagePipeline, HMVideoPipeline, HM3ImagePipeline, HM3VideoPipeline \ No newline at end of file +from .pipelines import (HMImagePipeline, HMVideoPipeline, + HM3ImagePipeline, HM3VideoPipeline, + HM5ImagePipeline) \ No newline at end of file diff --git a/hellomeme/models/__init__.py b/hellomeme/models/__init__.py index 5f06d5e..fe72eb7 100644 --- a/hellomeme/models/__init__.py +++ b/hellomeme/models/__init__.py @@ -9,8 +9,11 @@ """ from .hm_denoising_motion import HMDenoisingMotion -from .hm_control import HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2, HMV3ControlNet, HMControlNetBase, HM4SD15ControlProj -from .hm_blocks import HMReferenceAdapter, HM3ReferenceAdapter, HM3MotionAdapter, HMPipeline +from .hm_control import (HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2, + HMV3ControlNet, HMControlNetBase, HM5ControlNetBase, + HM4SD15ControlProj, HM5SD15ControlProj) +from .hm_adapters import (HMReferenceAdapter, HM3ReferenceAdapter, HM5ReferenceAdapter, + HM3MotionAdapter, HMPipeline) from .hm_denoising_3d import HMDenoising3D from .hm3_denoising_3d import HM3Denoising3D from .hm3_denoising_motion import HM3DenoisingMotion diff --git a/hellomeme/models/hm3_denoising_3d.py b/hellomeme/models/hm3_denoising_3d.py index e15d8cf..e69050c 100644 --- a/hellomeme/models/hm3_denoising_3d.py +++ b/hellomeme/models/hm3_denoising_3d.py @@ -17,7 +17,7 @@ from einops import rearrange from diffusers.utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel, UNet2DConditionOutput -from .hm_blocks import CopyWeights, InsertReferenceAdapter +from .hm_adapters import CopyWeights, InsertReferenceAdapter logger = logging.get_logger(__name__) # pylint: disable=invalid-name diff --git a/hellomeme/models/hm3_denoising_motion.py b/hellomeme/models/hm3_denoising_motion.py index b3ef351..bd16144 100644 --- a/hellomeme/models/hm3_denoising_motion.py +++ b/hellomeme/models/hm3_denoising_motion.py @@ -17,7 +17,7 @@ from einops import rearrange from diffusers.utils import logging from diffusers.models.unets.unet_motion_model import UNetMotionModel, UNetMotionOutput -from .hm_blocks import InsertReferenceAdapter +from .hm_adapters import InsertReferenceAdapter logger = logging.get_logger(__name__) # pylint: disable=invalid-name diff --git a/hellomeme/models/hm_adapters.py b/hellomeme/models/hm_adapters.py new file mode 100644 index 0000000..064d3a6 --- /dev/null +++ b/hellomeme/models/hm_adapters.py @@ -0,0 +1,489 @@ +# 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) \ No newline at end of file diff --git a/hellomeme/models/hm_blocks.py b/hellomeme/models/hm_blocks.py index f11eabd..bacb6e4 100644 --- a/hellomeme/models/hm_blocks.py +++ b/hellomeme/models/hm_blocks.py @@ -8,159 +8,20 @@ @Desc : """ -import copy import torch import torch.nn as nn import torch.utils.checkpoint -import torch.nn.functional as F -from typing import Optional, Tuple, Union, Dict, List +from typing import Tuple from einops import rearrange -from diffusers import StableDiffusionImg2ImgPipeline -from diffusers.models.modeling_utils import ModelMixin -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.models.attention import Attention, FeedForward +from diffusers.models.attention import Attention, FeedForward, BasicTransformerBlock from torchvision.models.resnet import BasicBlock -from diffusers.models.embeddings import SinusoidalPositionalEmbedding -from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel +from diffusers.models.embeddings import SinusoidalPositionalEmbedding, get_1d_sincos_pos_embed_from_grid, TimestepEmbedding -from diffusers.utils import (logging, is_torch_version, - is_peft_available, is_peft_version, - is_transformers_available, is_transformers_version, - USE_PEFT_BACKEND, - scale_lora_layers, - unscale_lora_layers, - ) - -from diffusers.configuration_utils import FrozenDict -from diffusers.models.attention_processor import ( - AttnProcessor, - AttnProcessor2_0, -) - -from diffusers.loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin -from diffusers.models.lora import adjust_lora_scale_text_encoder +from diffusers.utils import logging 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 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 SKReferenceAttention(nn.Module): def __init__( @@ -271,6 +132,64 @@ class SKReferenceAttentionV3(nn.Module): return hidden_states + self.proj(res2.contiguous()) +class SKReferenceAttentionV5(nn.Module): + def __init__( + self, + in_channels: int, + num_attention_heads: int = 1, + blocks_time_embed_dim: int = 1280, + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-5, + ): + super().__init__() + self.time_emb_proj = nn.Linear(blocks_time_embed_dim, in_channels) + self.attn1 = Attention( + query_dim=in_channels, + heads=num_attention_heads, + dim_head=in_channels // num_attention_heads, + dropout=0.0, + bias=False, + cross_attention_dim=None, + upcast_attention=False, + out_bias=True, + ) + self.attn2 = Attention( + query_dim=in_channels, + heads=num_attention_heads, + dim_head=in_channels // num_attention_heads, + dropout=0.0, + bias=False, + cross_attention_dim=None, + upcast_attention=False, + out_bias=True, + ) + self.norm = nn.LayerNorm(in_channels, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.proj = zero_module(nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)) + + def forward(self, hidden_states, ref_stats, num_frames): + h, w = hidden_states.shape[-2:] + + if ref_stats.shape[0] != hidden_states.shape[0]: + ref_stats = ref_stats.repeat_interleave(num_frames, dim=0) + cat_stats = torch.cat([hidden_states, ref_stats], dim=-1).contiguous() + + cat_stats = rearrange(cat_stats, "b c h w -> (b h) w c") + pose_embed1 = get_posembed_linear(cat_stats.shape[1], cat_stats.shape[2], + dtype=cat_stats.dtype, device=cat_stats.device) + res1 = self.attn1(self.norm(cat_stats+pose_embed1[None,:,:])) + res1 = rearrange(res1, "(b h) w c -> b c h w", h=h, w=w+w)[:, :, :, :w] + + cat_stats2 = torch.cat([res1.contiguous(), ref_stats], dim=-2).contiguous() + cat_stats2 = rearrange(cat_stats2, "b c h w -> (b w) h c") + pose_embed2 = get_posembed_linear(cat_stats.shape[1], cat_stats.shape[2], + dtype=cat_stats.dtype, device=cat_stats.device) + res2 = self.attn2(self.norm(cat_stats2+pose_embed2[None,:,:])) + + res2 = rearrange(res2, "(b w) h c -> b c h w", h=h+h, w=w)[:, :, :h, :] + + return hidden_states + self.proj(res2.contiguous()) + + class SmallUnet(nn.Module): def __init__(self, in_channels: int = 3, out_channels: int = 16, @@ -332,6 +251,264 @@ class SmallUnet(nn.Module): ret_dict[f'feat_{i+2}'] = rearrange(x, "(b f) c h w -> b c f h w", f=f) return ret_dict +class SmallUnetV5(nn.Module): + def __init__(self, in_channels: int = 3, + out_channels: int = 16, + cross_attention_dim: int = 1024, + mid_channels: Tuple[int] = (320, 640, 1280), + temporal_attn: bool = False): + super().__init__() + + down_channels = [in_channels] + mid_channels + up_channels = mid_channels[::-1] + [out_channels] + self.down_blocks = nn.ModuleList([ + BasicBlock( + inplanes=down_channels[i-1], + planes=down_channels[i], + stride=2, + downsample=nn.Sequential( + nn.Conv2d(down_channels[i-1], down_channels[i], kernel_size=1, stride=2, bias=False), + nn.InstanceNorm2d(down_channels[i]), + nn.SiLU(), + ), + norm_layer=nn.InstanceNorm2d, + ) for i in range(1, len(down_channels)) + ]) + self.up_blocks = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(up_channels[i-1], up_channels[i], kernel_size=3, padding=1), + nn.InstanceNorm2d(up_channels[i]), + nn.SiLU(), + nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), + ) for i in range(1, len(up_channels)) + ]) + self.attn_blocks = nn.ModuleList([ + STKCrossAttentionV5( + channel_in=c, + channel_mid=c, + cross_attention_dim=cross_attention_dim, + temporal_attn=temporal_attn, + ) for c in mid_channels + ]) + + def forward(self, x, condition): + f = x.size(2) + x = rearrange(x, "b c f h w -> (b f) c h w").contiguous() + skips = [] + ret_dict = {} + for down_block, attn_block in zip(self.down_blocks, self.attn_blocks): + x = rearrange(down_block(x), "(b f) c h w -> b f c h w", f=f).contiguous() + x = attn_block(x, condition) + x = rearrange(x, "b f c h w -> (b f) c h w").contiguous() + skips.append(x) + skips = skips[::-1][1:] + ret_dict['feat_0'] = rearrange(x, "(b f) c h w -> b c f h w", f=f).contiguous() + for i, block in enumerate(self.up_blocks[:-1]): + x = block(x) + skips[i] + ret_dict[f'feat_{i+1}'] = rearrange(x, "(b f) c h w -> b c f h w", f=f).contiguous() + x = self.up_blocks[-1](x) + ret_dict[f'feat_{i+2}'] = rearrange(x, "(b f) c h w -> b c f h w", f=f).contiguous() + return ret_dict + + +def get_posembed_linear(length, dim, dtype, device): + return get_1d_sincos_pos_embed_from_grid(dim, torch.arange(length), output_type='pt').to(dtype=dtype, device=device) + + +class STKAttentionV5(nn.Module): + def __init__( + self, + in_channels: int = 16, + mid_channels: int = 16, + num_attention_heads: int = 8, + time_embed_dim: int = 512, + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-5, + temporal_attn: bool = False, + ): + super().__init__() + self.conv_in = nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1) + + if not time_embed_dim is None: + self.emb_proj = TimestepEmbedding(time_embed_dim, mid_channels) + + self.attnx = Attention( + query_dim=mid_channels, + heads=num_attention_heads, + dim_head=mid_channels // num_attention_heads, + dropout=0.0, + bias=False, + cross_attention_dim=None, + upcast_attention=False, + out_bias=True, + ) + self.attny = Attention( + query_dim=mid_channels, + heads=num_attention_heads, + dim_head=mid_channels // num_attention_heads, + dropout=0.0, + bias=False, + cross_attention_dim=None, + upcast_attention=False, + out_bias=True, + ) + if temporal_attn: + self.attnt = Attention( + query_dim=mid_channels, + heads=num_attention_heads, + dim_head=mid_channels // num_attention_heads, + dropout=0.0, + bias=False, + cross_attention_dim=None, + upcast_attention=False, + out_bias=True, + ) + self.norm = nn.LayerNorm(mid_channels, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.proj = zero_module(nn.Conv2d(mid_channels, in_channels, kernel_size=3, padding=1)) + + def forward(self, hidden_states, ref_stats=None, emb=None): + b, f, h, w, _ = hidden_states.shape + + hidden_states_in = rearrange(hidden_states, "b f h w c -> (b f) c h w") + hidden_states_in = self.conv_in(hidden_states_in) + + if ref_stats is not None: + ref_stats_in = rearrange(ref_stats, "b f h w c -> (b f) c h w") + ref_stats_in = self.conv_in(ref_stats_in) + if emb is not None and hasattr(self, 'emb_proj'): + emb = self.emb_proj(emb)[:,:,None,None] + emb = emb.repeat_interleave(f, dim=0) + ref_stats_in = ref_stats_in + emb + hidden_states_in = rearrange(hidden_states_in, "(b f) c h w -> b f h w c", f=f) + + if ref_stats is not None: + ref_stats_in = rearrange(ref_stats_in, "(b f) c h w -> b f h w c", f=f) + cat_stats = torch.cat([hidden_states_in, ref_stats_in], dim=3).contiguous() + else: + cat_stats = hidden_states_in + + cat_stats = rearrange(cat_stats, "b f h w c -> (b f h) w c") + + pose_embedx = get_posembed_linear(cat_stats.shape[1], cat_stats.shape[2], + dtype=cat_stats.dtype, device=cat_stats.device) + resx = self.attnx(self.norm(cat_stats+pose_embedx[None,:,:])) + resx = rearrange(resx[:,:w], "(b f h) w c -> b f h w c", h=h, f=f) + + if ref_stats is not None: + cat_stats = torch.cat([resx, ref_stats_in], dim=2).contiguous() + else: + cat_stats = resx + cat_stats = rearrange(cat_stats, "b f h w c -> (b f w) h c") + + pose_embedy = get_posembed_linear(cat_stats.shape[1], cat_stats.shape[2], + dtype=cat_stats.dtype, device=cat_stats.device) + resy = self.attny(self.norm(cat_stats+pose_embedy[None,:,:])) + resy = rearrange(resy[:,:h], "(b f w) h c -> b f h w c", f=f, w=w) + + + if hasattr(self, 'attnt'): + if ref_stats is not None: + cat_stats = torch.cat([resy, ref_stats_in], dim=1).contiguous() + else: + cat_stats = resy + cat_stats = rearrange(cat_stats, "b f h w c -> (b h w) f c") + pose_embedt = get_posembed_linear(cat_stats.shape[1], cat_stats.shape[2], + dtype=cat_stats.dtype, device=cat_stats.device) + rest = self.attnt(self.norm(cat_stats+pose_embedt[None,:,:])) + rest = rearrange(rest[:,:f], "(b h w) f c -> (b f) c h w", h=h, w=w) + else: + rest = rearrange(resy, "b f h w c -> (b f) c h w") + + return hidden_states + rearrange(self.proj(rest), "(b f) c h w -> b f h w c", f=f) + +class STKCrossAttentionV5(nn.Module): + def __init__( + self, + channel_in: int, + channel_mid: int, + heads: int = 8, + cross_attention_dim: int = 320, + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-5, + temporal_attn: bool = False, + ): + super().__init__() + + self.conv_in = nn.Conv2d(channel_in, channel_mid, kernel_size=3, padding=1) + + self.norm = nn.LayerNorm(channel_mid, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.attn1 = Attention( + query_dim=channel_mid, + heads=heads, + dim_head=channel_mid // heads, + dropout=0.0, + bias=False, + cross_attention_dim=cross_attention_dim, + upcast_attention=False, + out_bias=True, + ) + + self.attn2 = Attention( + query_dim=channel_mid, + heads=heads, + dim_head=channel_mid // heads, + dropout=0.0, + bias=False, + cross_attention_dim=cross_attention_dim, + upcast_attention=False, + out_bias=True, + ) + + if temporal_attn: + self.attn3 = Attention( + query_dim=channel_mid, + heads=heads, + dim_head=channel_mid // heads, + dropout=0.0, + bias=False, + cross_attention_dim=cross_attention_dim, + upcast_attention=False, + out_bias=True, + ) + + self.ff = FeedForward( + channel_mid, + dropout=0.0, + activation_fn="geglu", + final_dropout=False, + inner_dim=channel_mid*16, + bias=True, + ) + + self.proj = nn.Conv2d(channel_mid, channel_in, kernel_size=3, padding=1) + + def forward(self, input, hidden_stats): + b, f, _, h, w = input.shape + x = rearrange(input, "b f c h w -> (b f) c h w") + x = self.conv_in(x) + + hidden_stats = rearrange(hidden_stats, "b f c d -> (b f) c d") + + pose_embed_hidden = get_posembed_linear(hidden_stats.shape[1], hidden_stats.shape[2], + dtype=hidden_stats.dtype, device=hidden_stats.device)[None,:,:] + + x = rearrange(x, "b c h w -> (b h) w c") + pose_embed1 = get_posembed_linear(x.shape[1], x.shape[2], dtype=x.dtype, device=x.device)[None,:,:] + x = self.attn1(self.norm(x + pose_embed1), hidden_stats.repeat_interleave(h, dim=0)+pose_embed_hidden) + x = rearrange(x, "(b h) w c -> (b w) h c", h=h) + pose_embed2 = get_posembed_linear(x.shape[1], x.shape[2], dtype=x.dtype, device=x.device)[None,:,:] + x = self.attn2(self.norm(x + pose_embed2), hidden_stats.repeat_interleave(w, dim=0)+pose_embed_hidden) + if hasattr(self, 'attn3'): + x = rearrange(x, "(b f w) h c -> (b w h) f c", w=w, f=f) + pose_embed3 = get_posembed_linear(x.shape[1], x.shape[2], dtype=x.dtype, device=x.device)[None,:,:] + x = self.attn3(self.norm(x+pose_embed3), hidden_stats.repeat_interleave(w*h, dim=0)+pose_embed_hidden) + x = rearrange(x, "(b w h) f c -> (b f) (h w) c", w=w, h=h) + else: + x = rearrange(x, "(b f w) h c -> (b f) (h w) c", w=w, f=f) + x = self.norm(self.ff(x)) + x = rearrange(x, "(b f) (h w) c -> (b f) c h w", w=w, f=f) + x = self.proj(x) + return rearrange(x, "(b f) c h w -> b f c h w", f=f) class STKCrossAttention(nn.Module): def __init__( @@ -547,313 +724,9 @@ class SKMotionModule(nn.Module): return hidden_states + self.proj(res_temp) + # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/controlnet.py def zero_module(module): for p in module.parameters(): nn.init.zeros_(p) - return module - -# 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) - - -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 + return module \ No newline at end of file diff --git a/hellomeme/models/hm_control.py b/hellomeme/models/hm_control.py index c647e6f..083eb68 100644 --- a/hellomeme/models/hm_control.py +++ b/hellomeme/models/hm_control.py @@ -17,7 +17,8 @@ from diffusers.models.modeling_utils import ModelMixin from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.embeddings import TimestepEmbedding, Timesteps from einops import rearrange -from .hm_blocks import SKCrossAttention, SmallUnet +from .hm_blocks import SKCrossAttention, STKAttentionV5, SmallUnet, SmallUnetV5 + class HMControlNetBase(ModelMixin, ConfigMixin): @register_to_config @@ -73,6 +74,60 @@ class HMControlNetBase(ModelMixin, ConfigMixin): return self.ctrl_attn(condition, drive_embedding) +class HM5ControlNetBase(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + in_channels: int = 3, + cross_attention_dim: int = 2048, + ): + super().__init__() + + self.emo_embedding = Timesteps(1024, True, 0) + self.emo_pre_proj = TimestepEmbedding(548, 16) + self.emo_proj = TimestepEmbedding(1024, cross_attention_dim) + + self.exp_embedding = Timesteps(1024, True, 0) + self.exp_pre_proj = TimestepEmbedding(51, 16) + self.face_pre_proj = TimestepEmbedding(3, 16) + self.exp_proj = TimestepEmbedding(1024, cross_attention_dim) + self.face_proj = TimestepEmbedding(1024, cross_attention_dim) + + self.ctrl_attn = SmallUnetV5(in_channels=in_channels, out_channels=320, + mid_channels=[320, 640, 1280], + cross_attention_dim=cross_attention_dim, + temporal_attn=False) + + def forward(self, condition, drive_coeff=None, face_parts=None, emo_embedding=None): + bs, video_length = condition.size(0), condition.size(2) + condition = rearrange(condition, "b c f h w -> (b f) c h w") + condition = F.interpolate(condition, size=(condition.size(2)//8, condition.size(3)//8), mode='bilinear', align_corners=False) + condition = rearrange(condition, "(b f) c h w -> b c f h w", f=video_length) + + if drive_coeff is None: + emo_coeff = rearrange(emo_embedding * 20., "b f c -> (b f c)") + emo_embedding = self.emo_embedding(emo_coeff) + emo_embedding = rearrange(emo_embedding, "(b f c) d -> (b f) d c", b=bs, f=video_length) + emo_embedding = self.emo_pre_proj(emo_embedding.to(dtype=condition.dtype)) + emo_embedding = rearrange(emo_embedding, "(b f) d c -> (b f) c d", f=video_length) + drive_embedding = self.emo_proj(emo_embedding) + else: + drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)") + exp_embedding = self.exp_embedding(drive_coeff) + exp_embedding = rearrange(exp_embedding, "(b f c) d -> (b f) d c", b=bs, f=video_length) + exp_embedding = self.exp_pre_proj(exp_embedding.to(dtype=condition.dtype)) + exp_embedding = rearrange(exp_embedding, "(b f) d c -> (b f) c d", f=video_length) + + face_parts = rearrange(face_parts, "b f c d -> (b f) d c") + face_embedding = self.face_pre_proj(face_parts) + face_embedding = rearrange(face_embedding, "(b f) d c -> (b f) c d", f=video_length) + + drive_embedding = self.exp_proj(exp_embedding) + self.face_proj(face_embedding) + + drive_embedding = rearrange(drive_embedding, "(b f) c d -> b f c d", f=video_length) + return self.ctrl_attn(condition, drive_embedding) + + class HM4SD15ControlProj(ModelMixin, ConfigMixin): @register_to_config def __init__( @@ -126,6 +181,55 @@ class HM4SD15ControlProj(ModelMixin, ConfigMixin): ) return ret_dict + +class HMProjector(nn.Module): + def __init__( + self, + in_channels: int = 3, + min_dim: int = 2048, + ): + super().__init__() + self.proj = nn.Conv2d(in_channels, min_dim, kernel_size=3, padding=1, bias=False) + self.attn = STKAttentionV5(in_channels=min_dim, mid_channels=min_dim // 4, temporal_attn=False, time_embed_dim=None) + + def forward(self, x): + f = x.size(2) + x = rearrange(x, "b c f h w -> (b f) c h w") + x = self.proj(x) + x = rearrange(x, "(b f) c h w -> b f h w c", f=f) + x = self.attn(x) + return rearrange(x, "b f h w c -> b c f h w") + + +class HM5SD15ControlProj(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + ): + super().__init__() + self.map640_to_640 = HMProjector(640, 640) # nn.Conv2d(640, 640, kernel_size=3, padding=1, bias=False) + self.map640_to_1280 = HMProjector(640, 1280) # nn.Conv2d(640, 1280, kernel_size=3, padding=1, bias=False) + self.map320_to_320 = HMProjector(320, 320) # nn.Conv2d(320, 320, kernel_size=3, padding=1, bias=False) + self.map320_to_640 = HMProjector(320, 640) # nn.Conv2d(320, 640, kernel_size=3, padding=1, bias=False) + self.map320_to_1280 = HMProjector(320, 1280) # nn.Conv2d(320, 1280, kernel_size=3, padding=1, bias=False) + self.map1280_to_1280 = HMProjector(1280, 1280) # nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) + self.map1280_to_1280_2 = HMProjector(1280, 1280) # nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) + self.map1280_to_1280_3 = HMProjector(1280, 1280) # nn.Conv2d(1280, 1280, kernel_size=3, padding=1, bias=False) + + def forward(self, control_dict): + ret_dict = dict( + up3_0=self.map1280_to_1280(control_dict['feat_0']), + up3_1=self.map640_to_1280(control_dict['feat_1']), + up3_2=self.map320_to_1280(control_dict['feat_2']), + up3_3=self.map320_to_640(control_dict['feat_3']), + down3_0=self.map320_to_320(control_dict['feat_2']), + down3_1=self.map640_to_640(control_dict['feat_1']), + down3_2=self.map1280_to_1280_2(control_dict['feat_0']), + down3_3=self.map1280_to_1280_3(control_dict['feat_0']), + ) + return ret_dict + + class HMControlNet(ModelMixin, ConfigMixin): @register_to_config def __init__( diff --git a/hellomeme/models/hm_denoising_3d.py b/hellomeme/models/hm_denoising_3d.py index d5312d0..29fb4fe 100644 --- a/hellomeme/models/hm_denoising_3d.py +++ b/hellomeme/models/hm_denoising_3d.py @@ -17,7 +17,7 @@ from einops import rearrange from diffusers.utils import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel, UNet2DConditionOutput -from .hm_blocks import CopyWeights, InsertReferenceAdapter +from .hm_adapters import CopyWeights, InsertReferenceAdapter logger = logging.get_logger(__name__) # pylint: disable=invalid-name diff --git a/hellomeme/models/hm_denoising_motion.py b/hellomeme/models/hm_denoising_motion.py index ada2450..ede9645 100644 --- a/hellomeme/models/hm_denoising_motion.py +++ b/hellomeme/models/hm_denoising_motion.py @@ -17,7 +17,7 @@ from einops import rearrange from diffusers.utils import logging from diffusers.models.unets.unet_motion_model import UNetMotionModel, UNetMotionOutput -from .hm_blocks import InsertReferenceAdapter +from .hm_adapters import InsertReferenceAdapter logger = logging.get_logger(__name__) # pylint: disable=invalid-name diff --git a/hellomeme/pipelines/__init__.py b/hellomeme/pipelines/__init__.py index dfc6a26..ae89f78 100644 --- a/hellomeme/pipelines/__init__.py +++ b/hellomeme/pipelines/__init__.py @@ -11,4 +11,5 @@ from .pipline_hm_image import HMImagePipeline from .pipline_hm_video import HMVideoPipeline from .pipline_hm3_image import HM3ImagePipeline -from .pipline_hm3_video import HM3VideoPipeline \ No newline at end of file +from .pipline_hm3_video import HM3VideoPipeline +from .pipline_hm5_image import HM5ImagePipeline \ No newline at end of file diff --git a/hellomeme/pipelines/pipline_hm5_image.py b/hellomeme/pipelines/pipline_hm5_image.py new file mode 100644 index 0000000..3dd4554 --- /dev/null +++ b/hellomeme/pipelines/pipline_hm5_image.py @@ -0,0 +1,370 @@ +# coding: utf-8 + +""" +@File : hm_pipline_image.py +@Author : Songkey +@Email : songkey@pku.edu.cn +@Date : 1/3/2025 +@Desc : +adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +""" + +import copy +from typing import Any, Callable, Dict, List, Optional, Union +import torch + +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.image_processor import PipelineImageInput +from diffusers.utils import deprecate +from diffusers.utils.torch_utils import randn_tensor +from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput +from diffusers import DPMSolverMultistepScheduler +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents +from custom_nodes.ComfyUI_HelloMeme.hellomeme.models import (HM3Denoising3D, + HMPipeline, HM5ReferenceAdapter, + HM5ControlNetBase, + HM5SD15ControlProj) + +class HM5ImagePipeline(HMPipeline): + def caryomitosis(self, **kwargs): + if hasattr(self, "unet_ref"): + del self.unet_ref + self.unet_ref = HM3Denoising3D.from_unet2d(self.unet) + self.unet_ref.cpu() + + if not isinstance(self.unet, HM3Denoising3D): + unet = HM3Denoising3D.from_unet2d(unet=self.unet) + # todo: 不够优雅 + del self.unet + self.unet = unet + self.unet.cpu() + + self.vae.cpu() + self.vae_decode = copy.deepcopy(self.vae) + self.text_encoder.cpu() + self.text_encoder_ref = copy.deepcopy(self.text_encoder) + self.safety_checker.cpu() + + def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False): + + self.version = version + if modelscope: + from modelscope import snapshot_download + hm_reference_dir = snapshot_download('songkey/hm5_reference') + hm_control_dir = snapshot_download('songkey/hm5_control_base') + hm_control_proj_dir = snapshot_download('songkey/hm5_control_proj') + else: + hm_reference_dir = 'songkey/hm5_reference' + hm_control_dir = 'songkey/hm5_control_base' + hm_control_proj_dir = 'songkey/hm5_control_proj' + + if isinstance(self.unet, HM3Denoising3D): + hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir) + + self.unet.insert_reference_adapter(hm_adapter) + self.unet.to(device='cpu', dtype=dtype).eval() + + if hasattr(self, "unet_ref"): + self.unet_ref.to(device='cpu', dtype=dtype).eval() + + if hasattr(self, "mp_control"): + del self.mp_control + + if hasattr(self, "mp_control_proj"): + del self.mp_control_proj + + self.mp_control = HM5ControlNetBase.from_pretrained(hm_control_dir) + self.mp_control_proj = HM5SD15ControlProj.from_pretrained(hm_control_proj_dir) + + self.mp_control.to(device='cpu', dtype=dtype).eval() + self.mp_control_proj.to(device='cpu', dtype=dtype).eval() + + self.vae.to(device='cpu', dtype=dtype).eval() + self.vae_decode.to(device='cpu', dtype=dtype).eval() + self.text_encoder.to(device='cpu', dtype=dtype).eval() + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + image: PipelineImageInput = None, + drive_params: Dict[str, Any] = None, + strength: float = 0.8, + num_inference_steps: Optional[int] = 50, + timesteps: List[int] = None, + sigmas: List[float] = None, + guidance_scale: Optional[float] = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + eta: Optional[float] = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + output_type: Optional[str] = "pil", + device: Optional[str] = "cpu", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + clip_skip: int = None, + callback_on_step_end: Optional[ + Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + **kwargs, + ): + callback = kwargs.pop("callback", None) + callback_steps = kwargs.pop("callback_steps", None) + num_images_per_prompt = 1 + + if callback is not None: + deprecate( + "callback", + "1.0.0", + "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", + ) + if callback_steps is not None: + deprecate( + "callback_steps", + "1.0.0", + "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", + ) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + strength, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + ip_adapter_image, + ip_adapter_image_embeds, + callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._clip_skip = clip_skip + self._cross_attention_kwargs = cross_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.device + + # 3. Encode input prompt + text_encoder_lora_scale = ( + self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None + ) + + self.text_encoder_ref.to(device=device) + prompt_embeds_ref, negative_prompt_embeds_ref = self.encode_prompt_sk( + self.text_encoder_ref, + prompt, + device, + num_images_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=text_encoder_lora_scale, + clip_skip=self.clip_skip, + ) + self.text_encoder_ref.cpu() + + self.text_encoder.to(device=device) + prompt_embeds, negative_prompt_embeds = self.encode_prompt_sk( + self.text_encoder, + prompt, + device, + num_images_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=text_encoder_lora_scale, + clip_skip=self.clip_skip, + ) + self.text_encoder.cpu() + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + prompt_embeds_ref = torch.cat([negative_prompt_embeds_ref, prompt_embeds_ref]) + + if ip_adapter_image is not None or ip_adapter_image_embeds is not None: + image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image, + ip_adapter_image_embeds, + device, + batch_size * num_images_per_prompt, + self.do_classifier_free_guidance, + ) + + # 4. Preprocess + image = self.image_processor.preprocess(image).to(device=device, dtype=prompt_embeds.dtype) + + scheduler = DPMSolverMultistepScheduler( + num_train_timesteps=1000, + beta_start=0.00085, + beta_end=0.012, + beta_schedule="scaled_linear", + # use_karras_sigmas=True, + algorithm_type="sde-dpmsolver++", + ) + + # 5. set timesteps + timesteps, num_inference_steps = retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas) + + # 6. Prepare reference latents + self.vae.to(device=device) + ref_latents = [ + retrieve_latents(self.vae.encode(image[i: i + 1].to(device=device)), generator=generator) + for i in range(batch_size) + ] + self.vae.cpu() + + ref_latents = torch.cat(ref_latents, dim=0) + ref_latents = self.vae.config.scaling_factor * ref_latents + c, h, w = ref_latents.shape[1:] + + condition = drive_params['condition'].clone().to(device=device) + if self.do_classifier_free_guidance: + condition = torch.cat([torch.ones_like(condition) * -1, condition], dim=0) + + control_latents = {} + self.mp_control.to(device=device) + self.mp_control_proj.to(device=device) + if 'drive_coeff' in drive_params: + drive_coeff = drive_params['drive_coeff'].clone().to(device=device) + face_parts = drive_params['face_parts'].clone().to(device=device) + if self.do_classifier_free_guidance: + drive_coeff = torch.cat([torch.zeros_like(drive_coeff), drive_coeff], dim=0) + face_parts = torch.cat([torch.zeros_like(face_parts), face_parts], dim=0) + control_latents1 = self.mp_control(condition=condition, drive_coeff=drive_coeff, face_parts=face_parts) + control_latents1 = self.mp_control_proj(control_latents1) + control_latents.update(control_latents1) + elif 'pd_fpg' in drive_params: + pd_fpg = drive_params['pd_fpg'].clone().to(device=device) + if self.do_classifier_free_guidance: + pd_fpg = torch.cat([torch.zeros_like(pd_fpg), pd_fpg], dim=0) + control_latents2 = self.mp_control(condition=condition, emo_embedding=pd_fpg) + control_latents2 = self.mp_control_proj(control_latents2) + control_latents.update(control_latents2) + self.mp_control.cpu() + self.mp_control_proj.cpu() + + # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7.1 Add image embeds for IP-Adapter + added_cond_kwargs = ( + {"image_embeds": image_embeds} + if ip_adapter_image is not None or ip_adapter_image_embeds is not None + else None + ) + base_noise = randn_tensor([batch_size, c, h, w], dtype=prompt_embeds.dtype, generator=generator).to(device=device) + + latent_model_input = torch.cat([torch.zeros_like(ref_latents), ref_latents]) if ( + self.do_classifier_free_guidance) else ref_latents + # latent_model_input = torch.cat([ref_latents_neg, ref_latents], dim=0) + self.unet_ref.to(device=device) + cached_res = self.unet_ref( + latent_model_input.unsqueeze(2), + 0, + encoder_hidden_states=prompt_embeds_ref, + return_dict=False, + )[1] + self.unet_ref.cpu() + + # 7.2 Optionally get Guidance Scale Embedding + timestep_cond = None + if self.unet.config.time_cond_proj_dim is not None: + guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) + timestep_cond = self.get_guidance_scale_embedding( + guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim + ).to(device=device, dtype=prompt_embeds.dtype) + + latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) + + latents = base_noise * scheduler.init_noise_sigma + # 8. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * scheduler.order + self._num_timesteps = len(timesteps) + self.unet.to(device=device) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = scheduler.scale_model_input(latent_model_input, t) + + # predict the noise residual + noise_pred = self.unet( + latent_model_input.unsqueeze(2), + t, + encoder_hidden_states=prompt_embeds, + reference_hidden_states=cached_res, + control_hidden_states=control_latents, + timestep_cond=timestep_cond, + cross_attention_kwargs=self.cross_attention_kwargs, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[0][:,:,0,:,:] + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + 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) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(scheduler, "order", 1) + callback(step_idx, t, latents) + + self.unet.cpu() + + self.vae_decode.to(device=device) + if not output_type == "latent": + image = self.vae_decode.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ + 0 + ] + else: + image = latents + self.vae_decode.cpu() + + do_denormalize = [True] * image.shape[0] + + image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + + # Offload all models + self.maybe_free_model_hooks() + + return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=None), latents.detach().cpu() / self.vae.config.scaling_factor diff --git a/meme.py b/meme.py index 5f57a9d..41892e8 100644 --- a/meme.py +++ b/meme.py @@ -52,7 +52,9 @@ from .hellomeme.utils import (get_drive_expression, append_pipline_weights, load_face_toolkits ) -from .hellomeme import HMImagePipeline, HMVideoPipeline, HM3ImagePipeline, HM3VideoPipeline +from .hellomeme import (HMImagePipeline, HMVideoPipeline, + HM3ImagePipeline, HM3VideoPipeline, + HM5ImagePipeline) config_path = osp.join(cur_dir, 'hellomeme', 'model_config.json') with open(config_path, 'r') as f: @@ -123,7 +125,7 @@ class HMImagePipelineLoader: "checkpoint": (checkpoint_files, ), "lora": (lora_files, ), "vae": (vae_files, ), - "version": (['v1', 'v2', 'v3', 'v4'], ), + "version": (['v1', 'v2', 'v3', 'v4', 'v5'], ), "stylize": (['x1', 'x2'], ), "deployment": (['huggingface', 'modelscope'], ), "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}), @@ -143,6 +145,8 @@ class HMImagePipelineLoader: sd1_5_dir = "songkey/stable-diffusion-v1-5" if version == 'v3' or version == 'v4': pipeline = HM3ImagePipeline.from_pretrained(sd1_5_dir) + elif version == 'v5': + pipeline = HM5ImagePipeline.from_pretrained(sd1_5_dir) else: pipeline = HMImagePipeline.from_pretrained(sd1_5_dir) pipeline.to(dtype=dtype) diff --git a/pyproject.toml b/pyproject.toml index 11ab53f..13e6bb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_hellomeme" description = "This repository is the official implementation of the [a/HelloMeme](https://arxiv.org/pdf/2410.22901) ComfyUI interface" -version = "1.2.3" +version = "1.3.0" license = {file = "LICENSE"} [project.urls]