little change

.gitignore

@@ -2,7 +2,10 @@
 *.pyd
 *.pth
 *.pkl
+*.pyc*
+*.pyd*
 *_fps15.mp4
+*__pycache__/*
 .idea/
 data/results
 pretrained_models/

hellomeme/__init__.py

@@ -8,4 +8,6 @@
 @Desc   : 
 """
 
-from .pipelines import HMImagePipeline, HMVideoPipeline, HM3ImagePipeline, HM3VideoPipeline
+from .pipelines import (HMImagePipeline, HMVideoPipeline,
+                        HM3ImagePipeline, HM3VideoPipeline,
+                        HM5ImagePipeline)

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

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
 

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
 

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)

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

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__(

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
 

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
 

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
+from .pipline_hm3_video import HM3VideoPipeline
+from .pipline_hm5_image import HM5ImagePipeline

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

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)

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]