init

.gitignore

@@ -0,0 +1,8 @@
+*.pyc
+*.pyd
+*.pth
+*.pkl
+*_fps15.mp4
+.idea/
+data/results
+pretrained_models/

README.md

@@ -0,0 +1,37 @@
+<h1 align='center'>HelloMeme: Integrating Spatial Knitting Attentions to Embed High-Level and Fidelity-Rich Conditions in Diffusion Models</h1>
+
+<div align='center'>
+    <a href='https://github.com/songkey' target='_blank'>Shengkai Zhang</a>, <a href='https://github.com/RhythmJnh' target='_blank'>Nianhong Jiao</a>, <a href='https://github.com/Shelton0215' target='_blank'>Tian Li</a>, <a href='https://github.com/chaojie12131243' target='_blank'>Chaojie Yang</a>, <a href='https://github.com/xchgit' target='_blank'>Chenhui Xue</a><sup>*</sup>, <a href='https://github.com/boya34' target='_blank'>Boya Niu</a><sup>*</sup>, <a href='https://github.com/HelloVision/HelloMeme' target='_blank'>Jun Gao</a> 
+</div>
+
+<div align='center'>
+    HelloVision | HelloGroup Inc.
+</div>
+
+<div align='center'>
+    <small><sup>*</sup> Intern</small>
+</div>
+
+<br>
+<div align='center'>
+    <a href='https://songkey.github.io/hellomeme/'><img src='https://img.shields.io/badge/Project-HomePage-Green'></a>
+    <a href='https://arxiv.org/pdf/2410.22901'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
+    <a href='https://huggingface.co/songkey'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow'></a>
+</div>
+
+## Introduction
+
+This repository is the official implementation of the [`HelloMeme`](https://arxiv.org/pdf/2410.22901) ComfyUI interface, featuring both image and video generation functionalities. Example workflow files can be found in the `ComfyUI_HelloMeme/workflows` directory. Test images and videos are saved in the `ComfyUI_HelloMeme/examples` directory. Below are screenshots of the interfaces for image and video generation.
+
+### Image Generation Interface
+
+<p align="center">
+  <img src="workflows/hellomeme_image_example.jpg" alt="image_generation_interface">
+</p>
+
+### Video Generation Interface
+
+<p align="center">
+  <img src="workflows/hellomeme_video_example.jpg" alt="video_generation_interface">
+</p>
+

__init__.py

@@ -0,0 +1,3 @@
+from .meme import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

hellomeme/__init__.py

@@ -0,0 +1,12 @@
+# coding: utf-8
+
+"""
+@File   : __init__.py.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/28/2024
+@Desc   : 
+"""
+
+from .piplines import HMVideoPipeline, HMImagePipeline
+from .models import HMDenoising3D, HMDenoisingMotion, HMControlNet, HMReferenceAdapter

hellomeme/models/__init__.py

@@ -0,0 +1,14 @@
+# coding: utf-8
+
+"""
+@File   : __init__.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/14/2024
+@Desc   : 
+"""
+
+from .hm_denoising_motion import HMDenoisingMotion
+from .hm_control import HMControlNet
+from .hm_blocks import HMReferenceAdapter
+from .hm_denoising_3d import HMDenoising3D

hellomeme/models/hm_blocks.py

@@ -0,0 +1,330 @@
+# coding: utf-8
+
+"""
+@File   : models6/hm_blocks.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/14/2024
+@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
+
+from einops import rearrange
+
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention import Attention, FeedForward
+from torchvision.models.resnet import BasicBlock
+from diffusers.models.embeddings import SinusoidalPositionalEmbedding
+from diffusers.utils import is_transformers_available, logging
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.models.attention_processor import (
+    AttnProcessor,
+    AttnProcessor2_0,
+)
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+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,
+                 ):
+        super().__init__()
+
+        self.reference_modules_down = nn.ModuleList([])
+
+        if isinstance(num_attention_heads, int):
+            num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+        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 SKReferenceAttention(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_attention_heads: int = 1,
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        num_positional_embeddings: int = 64*2,
+    ):
+        super().__init__()
+        self.pos_embed = SinusoidalPositionalEmbedding(in_channels, max_seq_length=num_positional_embeddings)
+        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.ff = FeedForward(
+        #     in_channels,
+        #     dropout=0.0,
+        #     activation_fn="geglu",
+        #     final_dropout=False,
+        #     inner_dim=in_channels*2,
+        #     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:]
+        ref_stats = ref_stats.repeat_interleave(num_frames, dim=0)
+        cat_stats = torch.cat([hidden_states, ref_stats], dim=-1)
+
+        cat_stats = rearrange(cat_stats.contiguous(), "b c h w -> (b h) w c")
+        res1 = self.attn1(self.norm(self.pos_embed(cat_stats)))
+        res1 = rearrange(res1[:, :w, :], "(b h) w c -> b c h w", h=h)
+
+        cat_stats2 = torch.cat([res1, ref_stats], dim=-2)
+        cat_stats2 = rearrange(cat_stats2.contiguous(), "b c h w -> (b w) h c")
+        res2 = self.attn2(self.norm(self.pos_embed(cat_stats2)))
+
+        res2 = rearrange(res2[:, :h, :], "(b w) h c -> b c h w", w=w)
+        # res2 = self.ff(res2)
+
+        return hidden_states + self.proj(res2)
+
+class SKCrossAttention(nn.Module):
+    def __init__(
+            self,
+            channel_in,
+            channel_out,
+            heads: int=8,
+            cross_attention_dim: int = 320,
+            norm_elementwise_affine: bool = True,
+            norm_eps: float = 1e-5,
+            num_positional_embeddings: int = 64,
+            num_positional_embeddings_hidden: int = 64,
+    ):
+        super().__init__()
+        self.conv = BasicBlock(
+                    inplanes=channel_in,
+                    planes=channel_out,
+                    stride=2,
+                    downsample=nn.Sequential(
+                        nn.Conv2d(channel_in, channel_out, kernel_size=1, stride=2, bias=False),
+                        nn.InstanceNorm2d(channel_out),
+                        nn.SiLU(),
+                    ),
+                    norm_layer=nn.InstanceNorm2d,
+        )
+
+        self.pos_embed = SinusoidalPositionalEmbedding(channel_out, max_seq_length=num_positional_embeddings)
+        self.pos_embed_hidden = SinusoidalPositionalEmbedding(cross_attention_dim, max_seq_length=num_positional_embeddings_hidden)
+
+        self.norm1 = nn.LayerNorm(channel_out, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
+        self.attn1 = Attention(
+            query_dim=channel_out,
+            heads=heads,
+            dim_head=channel_out // heads,
+            dropout=0.0,
+            bias=False,
+            cross_attention_dim=cross_attention_dim,
+            upcast_attention=False,
+            out_bias=True,
+        )
+
+        self.norm2 = nn.LayerNorm(channel_out, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
+        self.attn2 = Attention(
+            query_dim=channel_out,
+            heads=heads,
+            dim_head=channel_out // heads,
+            dropout=0.0,
+            bias=False,
+            cross_attention_dim=cross_attention_dim,
+            upcast_attention=False,
+            out_bias=True,
+        )
+
+        self.ff = FeedForward(
+            channel_out,
+            dropout=0.0,
+            activation_fn="geglu",
+            final_dropout=False,
+            inner_dim=channel_out*2,
+            bias=True,
+        )
+
+        self.proj = zero_module(nn.Conv2d(channel_out, channel_out, kernel_size=3, padding=1))
+
+    def forward(self, input, hidden_stats):
+        x = self.conv(input)
+        h, w = x.shape[-2:]
+        x = rearrange(x, "b c h w -> (b h) w c")
+        x = self.attn1(self.norm1(self.pos_embed(x)), self.pos_embed_hidden(hidden_stats.repeat_interleave(h, dim=0).contiguous()))
+        x = rearrange(x, "(b h) w c -> (b w) h c", h=h)
+        x = self.ff(self.attn2(self.norm2(self.pos_embed(x)), self.pos_embed_hidden(hidden_stats.repeat_interleave(w, dim=0).contiguous())))
+        x = rearrange(x, "(b w) h c -> b c h w", w=w)
+        x = self.proj(x)
+        return x
+
+# 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
+
+    def insert_reference_adapter(self, adapter: HMReferenceAdapter):
+        self.reference_modules_down = copy.deepcopy(adapter.reference_modules_down)
+        self.reference_modules_mid = copy.deepcopy(adapter.reference_modules_mid)
+        self.reference_modules_up = copy.deepcopy(adapter.reference_modules_up)
+

hellomeme/models/hm_control.py

@@ -0,0 +1,88 @@
+# coding: utf-8
+
+"""
+@File   : models6/hm_control.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/14/2024
+@Desc   : 
+"""
+
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+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
+
+
+class HMControlNet(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+            self,
+            embedding_channels: int,
+            input_channels: int = 3,
+            scale_factor: int = 8,
+            cross_attention_dim: int = 320,
+            block_out_channels: Tuple[int] = (128, 320, 640, 1280),
+    ):
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.embedding_channels = embedding_channels
+        self.cross_attention_dim = cross_attention_dim
+        self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False)
+
+        self.exp_embedding = Timesteps(cross_attention_dim, True, 0)
+        self.exp_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim)
+        self.face_proj = TimestepEmbedding(1024, cross_attention_dim)
+
+        self.blocks_down = nn.ModuleList([])
+        for i in range(1, len(block_out_channels)):
+            channel_in = block_out_channels[i-1]
+            channel_out = block_out_channels[i]
+            self.blocks_down.append(
+                SKCrossAttention(
+                    channel_in=channel_in,
+                    channel_out=channel_out,
+                    cross_attention_dim=cross_attention_dim,
+                    num_positional_embeddings=64,
+                    num_positional_embeddings_hidden=64
+                )
+            )
+
+    def forward(self, condition, drive_coeff, face_parts):
+        bs, _, video_length, h, w = condition.shape
+
+        condition = rearrange(condition, "b c f h w -> (b f) c h w")
+        conditioning = F.interpolate(condition,
+                                   size=(h // self.scale_factor, w // self.scale_factor),
+                                   mode='bilinear',
+                                   align_corners=False)
+        embedding = self.conv_in(conditioning)
+
+        drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)")
+        drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype)
+        drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d",
+                                    b=bs, f=video_length, d=self.cross_attention_dim)
+
+        face_parts = rearrange(face_parts, "b f c d -> (b f) c d")
+        face_embedding = self.face_proj(face_parts)
+
+        face_embedding = rearrange(face_embedding, "(b f) c d -> b f c d", f=video_length)
+
+        drive_embedding = torch.cat([face_embedding, drive_embedding], dim=2)
+        drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d")
+
+        drive_embedding = self.exp_proj(drive_embedding)
+
+        ret_dict = {}
+        for idx, block in enumerate(self.blocks_down):
+            embedding = block(embedding, drive_embedding)
+            ret_dict[f'down_{idx}'] = rearrange(embedding,
+                                                     "(b f) c h w -> b c f h w",
+                                                     f=video_length)
+        return ret_dict

hellomeme/models/hm_denoising_3d.py

@@ -0,0 +1,289 @@
+# coding: utf-8
+
+"""
+@File   : models6/hm_denoising_3d.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/14/2024
+@Desc   : 删除实验代码，精简结构
+adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_2d_condition.py
+"""
+
+import torch
+import torch.utils.checkpoint
+from typing import Any, Dict, Optional, Tuple, Union
+
+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
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class HMDenoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter):
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        reference_hidden_states: Optional[dict] = None,
+        control_hidden_states: Optional[torch.Tensor] = None,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        for dim in sample.shape[-2:]:
+            if dim % default_overall_up_factor != 0:
+                # Forward upsample size to force interpolation output size.
+                forward_upsample_size = True
+                break
+
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        t_emb = self.get_time_embed(sample=sample, timestep=timestep)
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        class_emb = self.get_class_embed(sample=sample, class_labels=class_labels)
+        if class_emb is not None:
+            if self.config.class_embeddings_concat:
+                emb = torch.cat([emb, class_emb], dim=-1)
+            else:
+                emb = emb + class_emb
+
+        aug_emb = self.get_aug_embed(
+            emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+        )
+        if self.config.addition_embed_type == "image_hint":
+            aug_emb, hint = aug_emb
+            sample = torch.cat([sample, hint], dim=1)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        if self.time_embed_act is not None:
+            emb = self.time_embed_act(emb)
+
+        num_frames = sample.shape[2]
+        emb = emb.repeat_interleave(repeats=num_frames, dim=0)
+
+        if len(encoder_hidden_states.shape) == 3:
+            encoder_hidden_states = encoder_hidden_states.repeat_interleave(repeats=num_frames, dim=0)
+        elif len(encoder_hidden_states.shape) == 4:
+            encoder_hidden_states = rearrange(encoder_hidden_states, "b f l d -> (b f) l d")
+
+        if not added_cond_kwargs is None and 'image_embeds' in added_cond_kwargs:
+            if isinstance(added_cond_kwargs['image_embeds'], torch.Tensor):
+                added_cond_kwargs['image_embeds'] = added_cond_kwargs['image_embeds'].repeat_interleave(repeats=num_frames, dim=0)
+            else:
+                added_cond_kwargs['image_embeds'] = [x.repeat_interleave(repeats=num_frames, dim=0) for x in added_cond_kwargs['image_embeds']]
+
+        encoder_hidden_states = self.process_encoder_hidden_states(
+            encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+        )
+
+        # 2. pre-process
+        sample = rearrange(sample, "b c f h w -> (b f) c h w")
+        sample = self.conv_in(sample)
+
+        # 2.5 GLIGEN position net
+        if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            gligen_args = cross_attention_kwargs.pop("gligen")
+            cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
+
+        # 3. down
+        # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated
+        # to the internal blocks and will raise deprecation warnings. this will be confusing for our users.
+        if cross_attention_kwargs is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            lora_scale = cross_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+
+        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+        # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
+        is_adapter = down_intrablock_additional_residuals is not None
+        # maintain backward compatibility for legacy usage, where
+        #       T2I-Adapter and ControlNet both use down_block_additional_residuals arg
+        #       but can only use one or the other
+        if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None:
+            deprecate(
+                "T2I should not use down_block_additional_residuals",
+                "1.3.0",
+                "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
+                       and will be removed in diffusers 1.3.0.  `down_block_additional_residuals` should only be used \
+                       for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
+                standard_warn=False,
+            )
+            down_intrablock_additional_residuals = down_block_additional_residuals
+            is_adapter = True
+
+        res_cache = dict()
+        down_block_res_samples = (sample,)
+        for idx, downsample_block in enumerate(self.down_blocks):
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0)
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+                res_cache[f"down_{idx}"] = sample.clone()
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    sample += down_intrablock_additional_residuals.pop(0)
+
+            if not control_hidden_states is None and f'down_{idx}' in control_hidden_states:
+                sample += rearrange(control_hidden_states[f'down_{idx}'], "b c f h w -> (b f) c h w")
+            if hasattr(self, 'reference_modules_down') and not reference_hidden_states is None and f'down_{idx}' in reference_hidden_states:
+                sample = self.reference_modules_down[idx](sample, reference_hidden_states[f'down_{idx}'], num_frames)
+
+            down_block_res_samples += res_samples
+
+        if is_controlnet:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
+                sample = self.mid_block(
+                    sample,
+                    emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = self.mid_block(sample, emb)
+            if hasattr(self, 'reference_modules_mid') and not reference_hidden_states is None and f'mid' in reference_hidden_states:
+                sample = self.reference_modules_mid(sample, reference_hidden_states[f'mid'], num_frames)
+
+            # To support T2I-Adapter-XL
+            if (
+                is_adapter
+                and len(down_intrablock_additional_residuals) > 0
+                and sample.shape == down_intrablock_additional_residuals[0].shape
+            ):
+                sample += down_intrablock_additional_residuals.pop(0)
+            res_cache[f"mid"] = sample.clone()
+
+        if is_controlnet:
+            sample = sample + mid_block_additional_residual
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                res_cache[f"up_{i}"] = sample.clone()
+                if hasattr(self, "reference_modules_up") and not reference_hidden_states is None and f'up_{i}' in reference_hidden_states:
+                    sample = self.reference_modules_up[i-1](sample, reference_hidden_states[f'up_{i}'], num_frames)
+
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                )
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        # reshape to (batch, channel, framerate, width, height)
+        sample = rearrange(sample, "(b f) c h w -> b c f h w", f=num_frames)
+
+        if not return_dict:
+            return (sample, res_cache)
+
+        return (UNet2DConditionOutput(sample=sample), res_cache)

hellomeme/models/hm_denoising_motion.py

@@ -0,0 +1,241 @@
+# coding: utf-8
+
+"""
+@File   : models6/hm_denoising_motion.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 9/9/2024
+@Desc   :
+adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_motion_model.py
+"""
+
+import torch
+import torch.utils.checkpoint
+from typing import Any, Dict, Optional, Tuple, Union
+
+from einops import rearrange
+
+from diffusers.utils import logging
+from diffusers.models.unets.unet_motion_model import UNetMotionModel, UNetMotionOutput
+from .hm_blocks import InsertReferenceAdapter
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class HMDenoisingMotion(UNetMotionModel, InsertReferenceAdapter):
+    def forward(
+            self,
+            sample: torch.Tensor,
+            timestep: Union[torch.Tensor, float, int],
+            encoder_hidden_states: torch.Tensor,
+            reference_hidden_states: Optional[dict] = None,
+            control_hidden_states: Optional[torch.Tensor] = None,
+            timestep_cond: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+            down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+            mid_block_additional_residual: Optional[torch.Tensor] = None,
+            return_dict: bool = True,
+    ) -> Union[UNetMotionOutput, Tuple[torch.Tensor]]:
+
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2 ** self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        # prepare attention_mask
+        if attention_mask is not None:
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        num_frames = sample.shape[2]
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=self.dtype)
+
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        if self.config.addition_embed_type == "text_time":
+            if "text_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                )
+
+            text_embeds = added_cond_kwargs.get("text_embeds")
+            if "time_ids" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                )
+            time_ids = added_cond_kwargs.get("time_ids")
+            time_embeds = self.add_time_proj(time_ids.flatten())
+            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+            add_embeds = add_embeds.to(emb.dtype)
+            aug_emb = self.add_embedding(add_embeds)
+
+        emb = emb if aug_emb is None else emb + aug_emb
+        emb = emb.repeat_interleave(repeats=num_frames, dim=0)
+        if len(encoder_hidden_states.shape) == 3:
+            encoder_hidden_states = encoder_hidden_states.repeat_interleave(repeats=num_frames, dim=0)
+        elif len(encoder_hidden_states.shape) == 4:
+            encoder_hidden_states = rearrange(encoder_hidden_states, "b f l d -> (b f) l d")
+
+        if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj":
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            image_embeds = self.encoder_hid_proj(image_embeds)
+            image_embeds = [image_embed.repeat_interleave(repeats=num_frames, dim=0) for image_embed in image_embeds]
+            encoder_hidden_states = (encoder_hidden_states, image_embeds)
+
+        # 2. pre-process
+        sample = sample.permute(0, 2, 1, 3, 4).reshape((sample.shape[0] * num_frames, -1) + sample.shape[3:])
+        sample = self.conv_in(sample)
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        for idx, downsample_block in enumerate(self.down_blocks):
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    num_frames=num_frames,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames)
+
+            if not control_hidden_states is None and f'down_{idx}' in control_hidden_states:
+                sample += rearrange(control_hidden_states[f'down_{idx}'], "b c f h w -> (b f) c h w")
+
+
+            if hasattr(self, 'reference_modules_down') and not reference_hidden_states is None and f'down_{idx}' in reference_hidden_states:
+                sample = self.reference_modules_down[idx](sample, reference_hidden_states[f'down_{idx}'], num_frames)
+
+            down_block_res_samples += res_samples
+
+        if down_block_additional_residuals is not None:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                    down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples += (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            # To support older versions of motion modules that don't have a mid_block
+            if hasattr(self.mid_block, "motion_modules"):
+                sample = self.mid_block(
+                    sample,
+                    emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    num_frames=num_frames,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                )
+            else:
+                sample = self.mid_block(
+                    sample,
+                    emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                )
+            if hasattr(self, 'reference_modules_mid') and not reference_hidden_states is None and f'mid' in reference_hidden_states:
+                sample = self.reference_modules_mid(sample, reference_hidden_states[f'mid'], num_frames)
+
+        if mid_block_additional_residual is not None:
+            sample = sample + mid_block_additional_residual
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets):]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    if hasattr(self,
+                               "reference_modules_up") and not reference_hidden_states is None and f'up_{i}' in reference_hidden_states:
+                        sample = self.reference_modules_up[i - 1](sample, reference_hidden_states[f'up_{i}'],
+                                                                  num_frames)
+
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    num_frames=num_frames,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                    num_frames=num_frames,
+                )
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+
+        sample = self.conv_out(sample)
+
+        # reshape to (batch, channel, framerate, width, height)
+        sample = sample[None, :].reshape((-1, num_frames) + sample.shape[1:]).permute(0, 2, 1, 3, 4)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNetMotionOutput(sample=sample)

hellomeme/piplines/__init__.py

@@ -0,0 +1,12 @@
+# coding: utf-8
+
+"""
+@File   : __init__.py.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/29/2024
+@Desc   : 
+"""
+
+from .pipline_hm_video import HMVideoPipeline
+from .pipline_hm_image import HMImagePipeline

hellomeme/piplines/pipline_hm_image.py

@@ -0,0 +1,289 @@
+# coding: utf-8
+
+"""
+@File   : hm_pipline_image.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/29/2024
+@Desc   :
+adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
+"""
+
+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.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
+from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents
+
+from ..models import HMDenoising3D, HMControlNet
+from ..models import HMReferenceAdapter
+
+class HMImagePipeline(StableDiffusionImg2ImgPipeline):
+    def caryomitosis(self,):
+        if not hasattr(self, "unet_ref"):
+            self.unet_ref = HMDenoising3D.from_unet2d(self.unet)
+        if not isinstance(self.unet, HMDenoising3D):
+            unet = HMDenoising3D.from_unet2d(unet=self.unet)
+            # todo: 不够优雅
+            del self.unet
+            self.unet = unet
+
+    def insert_hm_modules(self, dtype, device):
+        if isinstance(self.unet, HMDenoising3D):
+            hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference')
+            self.unet.insert_reference_adapter(hm_adapter)
+            self.unet_ref = self.unet_ref.to(device=device, dtype=dtype).eval()
+            self.unet = self.unet.to(device=device, dtype=dtype).eval()
+
+        if not hasattr(self, "mp_control"):
+            self.mp_control = HMControlNet.from_pretrained('songkey/hm_control')
+            self.mp_control = self.mp_control.to(device=device, 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",
+            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
+        )
+        print('device:', device)
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            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,
+        )
+        # 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])
+
+        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)
+
+        # 5. set timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps, sigmas
+        )
+        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+        ref_latents = [
+            retrieve_latents(self.vae.encode(image[i: i + 1].to(device=device)), generator=generator)
+            for i in range(batch_size)
+        ]
+        ref_latents = torch.cat(ref_latents, dim=0)
+        ref_latents = self.vae.config.scaling_factor * ref_latents
+
+        # 6. Prepare latent variables
+        # latents = self.prepare_latents(
+        #     ref_latents,
+        #     latent_timestep,
+        #     batch_size,
+        #     num_images_per_prompt,
+        #     prompt_embeds.dtype,
+        #     device,
+        #     generator,
+        # )
+
+        # condition, drive_coeff, face_parts
+        if self.do_classifier_free_guidance:
+            drive_params['condition'] = torch.cat([torch.ones_like(drive_params['condition']) * -1, drive_params['condition']], dim=0)
+            drive_params['drive_coeff'] = torch.cat([torch.zeros_like(drive_params['drive_coeff']), drive_params['drive_coeff']], dim=0)
+            drive_params['face_parts'] = torch.cat([torch.zeros_like(drive_params['face_parts']), drive_params['face_parts']], dim=0)
+
+        if hasattr(self, "mp_control"):
+            control_latents = self.mp_control(**drive_params)
+        else:
+            control_latents = None
+
+        # 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
+        )
+
+        latent_model_input = torch.cat([torch.zeros_like(ref_latents), ref_latents]) if self.do_classifier_free_guidance else ref_latents
+        cached_res = self.unet_ref(
+            latent_model_input.unsqueeze(2),
+            0,
+            encoder_hidden_states=prompt_embeds,
+            # added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )[1]
+
+        # 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)
+        base_noise = torch.randn_like(ref_latents, device=device, dtype=prompt_embeds.dtype)
+        latents = self.scheduler.add_noise(ref_latents, base_noise, latent_timestep)
+        # 8. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+        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 = self.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 = self.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) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+        if not output_type == "latent":
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
+                0
+            ]
+        else:
+            image = latents
+
+        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)

hellomeme/piplines/pipline_hm_video.py

@@ -0,0 +1,376 @@
+# coding: utf-8
+
+"""
+@File   : hm_pipline_image.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/29/2024
+@Desc   :
+adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
+"""
+
+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_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents
+from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import EulerDiscreteScheduler
+
+from diffusers import MotionAdapter
+from ..models import HMDenoising3D, HMDenoisingMotion, HMControlNet
+from ..models import HMReferenceAdapter
+
+
+class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
+    def caryomitosis(self,):
+        if not hasattr(self, "unet_ref"):
+            self.unet_ref = HMDenoising3D.from_unet2d(self.unet)
+        if not isinstance(self.unet, HMDenoisingMotion):
+            adapter = MotionAdapter.from_pretrained("songkey/hm_animatediff", torch_dtype=torch.float16)
+            unet = HMDenoisingMotion.from_unet2d(unet=self.unet, motion_adapter=adapter, load_weights=True)
+            # todo: 不够优雅
+            del self.unet
+            self.unet = unet
+
+    def insert_hm_modules(self, dtype, device):
+        if isinstance(self.unet, HMDenoisingMotion):
+            hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference')
+            self.unet.insert_reference_adapter(hm_adapter)
+            self.unet_ref.insert_reference_adapter(hm_adapter)
+            self.unet_ref = self.unet_ref.to(device=device, dtype=dtype).eval()
+            self.unet = self.unet.to(device=device, dtype=dtype).eval()
+
+        if not hasattr(self, "mp_control"):
+            self.mp_control = HMControlNet.from_pretrained('songkey/hm_control')
+            self.mp_control = self.mp_control.to(device=device, dtype=dtype).eval()
+
+    @torch.no_grad()
+    def __call__(
+            self,
+            prompt: Union[str, List[str]] = None,
+            image: PipelineImageInput = None,
+            num_frames: int = 16,
+            chunk_overlap: int = 4,
+            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",
+            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._execution_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
+        )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            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,
+        )
+        # 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])
+
+        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
+        image = self.image_processor.preprocess(image).to(device=device, dtype=prompt_embeds.dtype)
+
+        # 5. set timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps, sigmas
+        )
+        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+        raw_video_len = drive_params['drive_coeff'].size(1)
+        chunk_size = num_frames
+        chunk_overlap = min(max(0, chunk_overlap), chunk_size // 2)
+        chunk_stride = chunk_size - chunk_overlap
+
+        drive_idx_list = list(range(raw_video_len))
+
+        steped_video_len = len(drive_idx_list)
+        post_pad_num = chunk_stride - (steped_video_len - chunk_overlap) % chunk_stride
+
+        if post_pad_num > 0:
+            drive_idx_list = drive_idx_list + [drive_idx_list[-1]] * post_pad_num
+        paded_video_len = len(drive_idx_list)
+
+        drive_idx_chunks = [drive_idx_list[i:i + chunk_size] for i in
+                            range(0, paded_video_len - chunk_stride + 1, chunk_stride)]
+
+        # 6. Prepare latent variables
+        ref_latents = [
+            retrieve_latents(self.vae.encode(image[i: i + 1].to(device=device)), generator=generator)
+            for i in range(batch_size)
+        ]
+        ref_latents = torch.cat(ref_latents, dim=0)
+        ref_latents = self.vae.config.scaling_factor * ref_latents
+        c, h, w = ref_latents.shape[1:]
+
+        # 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
+        )
+
+        # if hasattr(self, "unet_ref"):
+        ref_latents = ref_latents.unsqueeze(2)
+        cached_res = self.unet_ref(
+            torch.cat([torch.zeros_like(ref_latents), ref_latents], dim=0) if
+                self.do_classifier_free_guidance else ref_latents,
+            0,
+            encoder_hidden_states=prompt_embeds,
+            # added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )[1]
+        # else:
+        #     cached_res = None
+
+        # 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)
+
+        if self.do_classifier_free_guidance:
+            drive_params['condition'] = torch.cat(
+                [torch.ones_like(drive_params['condition']) * -1, drive_params['condition']], dim=0)
+            drive_params['drive_coeff'] = torch.cat(
+                [torch.zeros_like(drive_params['drive_coeff']), drive_params['drive_coeff']], dim=0)
+            drive_params['face_parts'] = torch.cat(
+                [torch.zeros_like(drive_params['face_parts']), drive_params['face_parts']], dim=0)
+
+        control_latents_chunks = []
+        for cidx, chunk in enumerate(drive_idx_chunks):
+            condition_chunk = drive_params['condition'][:, :, chunk, :, :]
+            drive_coeff_chunk = drive_params['drive_coeff'][:, chunk, :]
+            face_parts_chunk = drive_params['face_parts'][:, chunk, :, :]
+
+
+            control_latents_chunk = self.mp_control(condition=condition_chunk.to(device=device),
+                                                    drive_coeff=drive_coeff_chunk.to(device=device),
+                                                    face_parts=face_parts_chunk.to(device=device))
+            control_latents_chunks.append(control_latents_chunk)
+        drive_idx_chunks[-1] = list(range(drive_idx_chunks[-1][0], paded_video_len))
+
+        with self.progress_bar(total=len(drive_idx_chunks)) as progress_bar:
+            chunk_latents = torch.zeros([batch_size, c, paded_video_len, h, w], dtype=prompt_embeds.dtype).cpu()
+            chunk_latents_counter = torch.zeros([1, 1, paded_video_len, 1, 1], dtype=prompt_embeds.dtype).cpu()
+            chunk_noise = randn_tensor([batch_size, c, 1, h, w], dtype=prompt_embeds.dtype, generator=generator).to(
+                device=device)
+            chunk_noise = chunk_noise.repeat_interleave(chunk_size, dim=2)
+            for cidx, chunk in enumerate(drive_idx_chunks):
+                control_latents_chunk = control_latents_chunks[cidx]
+
+                scheduler = EulerDiscreteScheduler(
+                    num_train_timesteps=1000,
+                    beta_start=0.00085,
+                    beta_end=0.012,
+                    beta_schedule="scaled_linear",
+                )
+
+                tmp_timesteps, _ = retrieve_timesteps(
+                    scheduler, 8, device, None, sigmas
+                )
+                # chunk_latent = chunk_noise * self.scheduler.init_noise_sigma
+                chunk_latent = scheduler.add_noise(ref_latents, chunk_noise, tmp_timesteps[:1])
+                for i, t in enumerate(tmp_timesteps):
+                    latent_model_input = torch.cat([chunk_latent, chunk_latent], dim=0) if \
+                        self.do_classifier_free_guidance else chunk_latent
+                    latent_model_input = scheduler.scale_model_input(latent_model_input, t)
+                    noise_pred_chunk = self.unet_ref(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=prompt_embeds,
+                        reference_hidden_states=cached_res,
+                        control_hidden_states=control_latents_chunk,
+                        timestep_cond=timestep_cond,
+                        cross_attention_kwargs=self.cross_attention_kwargs,
+                        added_cond_kwargs=added_cond_kwargs,
+                        return_dict=False,
+                    )[0]
+                    if self.do_classifier_free_guidance:
+                        noise_pred_uncond, noise_pred_text = noise_pred_chunk.chunk(2)
+                        noise_pred_chunk = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                    chunk_latent = scheduler.step(noise_pred_chunk, t, chunk_latent,
+                                                        **extra_step_kwargs, return_dict=False)[0]
+
+                chunk_latents[:, :, chunk] = chunk_latents[:, :, chunk] + chunk_latent.cpu()
+                chunk_latents_counter[:, :, chunk] += 1
+                progress_bar.update()
+            chunk_latents = chunk_latents / chunk_latents_counter
+
+        # 5. set timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, None, sigmas
+        )
+        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+        # 8. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        base_noise = randn_tensor([batch_size, c, paded_video_len, h, w], dtype=prompt_embeds.dtype, generator=generator)
+        # ref_latents = ref_latents.repeat_interleave(paded_video_len, dim=2)
+
+        latents = self.scheduler.add_noise(chunk_latents.to(device=device), base_noise.to(device=device), latent_timestep).cpu()
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                noise_pred = torch.zeros([batch_size * 2 if self.do_classifier_free_guidance
+                                          else batch_size, c, paded_video_len, h, w], dtype=prompt_embeds.dtype).cpu()
+                noise_pred_counter = torch.zeros([1, 1, paded_video_len, 1, 1], dtype=prompt_embeds.dtype).cpu()
+
+                latent_model_input_all = torch.cat([latents.clone(), latents.clone()], dim=0) if \
+                    self.do_classifier_free_guidance else latents.clone()
+
+                latent_model_input_all = self.scheduler.scale_model_input(latent_model_input_all, t)
+                for cidx, chunk in enumerate(drive_idx_chunks):
+                    if self.interrupt:
+                        continue
+
+                    control_latents = control_latents_chunks[cidx]
+                    latent_model_input = latent_model_input_all[:, :, chunk].to(device=device)
+
+                    # predict the noise residual
+                    noise_pred_chunk = self.unet(
+                        latent_model_input,
+                        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].cpu()
+
+                    noise_pred[:, :, chunk] = noise_pred[:, :, chunk] + noise_pred_chunk
+                    noise_pred_counter[:, :, chunk] = noise_pred_counter[:, :, chunk] + 1
+
+                noise_pred = noise_pred / noise_pred_counter
+                # 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 = self.scheduler.step(noise_pred.to(device=device), t, latents.to(device=device),
+                                              **extra_step_kwargs, return_dict=False)[0].cpu()
+
+                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) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+        latents = latents / self.vae.config.scaling_factor
+
+        res_frames = [self.vae.decode(latents[:,:,i,:,:].to(device=device), return_dict=False, generator=generator)[0]
+                       for i in range(raw_video_len)]
+        res_frames = [self.image_processor.postprocess(frame, output_type=output_type) for frame in res_frames]
+        return res_frames

hellomeme/tools/__init__.py

@@ -0,0 +1,13 @@
+# coding: utf-8
+
+# @File   : __init__.py
+# @Author : Songkey
+# @Email  : songkey@pku.edu.cn
+# @Date   : 8/28/2024
+# @Desc   :
+
+from .hello_arkit import HelloARKitBSPred
+from .hello_face_det import HelloFaceDet
+from .hello_camera_demo import HelloCameraDemo
+from .hello_3dmm import Hello3DMMPred
+from .hello_face_alignment import HelloFaceAlignment

hellomeme/tools/hello_3dmm.py

@@ -0,0 +1,140 @@
+# coding: utf-8
+
+"""
+@File   : test.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 11/1/2024
+@Desc   : Created by Shengjie Wu (wu.shengjie@immomo.com)
+这可能是一个很强大的模型
+"""
+
+import numpy as np
+import cv2
+
+from huggingface_hub import hf_hub_download
+from .utils import get_warp_mat_bbox_by_gt_pts_float, create_onnx_session
+
+def crop_transl_to_full_transl(crop_trans, crop_center, scale, full_center, focal_length):
+    """
+    :param crop_trans: (3), float
+    :param crop_center: (2), float
+    :param scale: (1), float
+    :param full_center: (2), float
+    :param focal_length: (1), float
+    :return:
+    """
+    crop_c_x, crop_c_y = crop_center
+    full_c_x, full_c_y = full_center
+    bs = 2 * focal_length / scale / crop_trans[2]
+    full_x = crop_trans[0] - 2 * (crop_c_x - full_c_x) / bs
+    full_y = crop_trans[1] + 2 * (crop_c_y - full_c_y) / bs
+    full_z = crop_trans[2] * scale
+
+    full_trans = np.array([full_x, full_y, full_z], dtype=np.float32)
+
+    return full_trans
+
+class Hello3DMMPred(object):
+    def __init__(self, gpu_id=None):
+        self.deep3d_pred_net = create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_3dmm.onnx'), gpu_id=gpu_id)
+        self.deep3d_pred_net_input_name = self.deep3d_pred_net.get_inputs()[0].name
+        self.deep3d_pred_net_output_name = [output.name for output in self.deep3d_pred_net.get_outputs()]
+
+        self.image_size = 224
+        self.camera_init_z = -0.4
+        self.camera_init_focal_len = 386.2879122887948
+        self.used_focal_len = -5.0 / self.camera_init_z * self.camera_init_focal_len
+        self.id_dims = 526
+        self.exp_dims = 203
+        self.tex_dims = 439
+
+    def forward_params(self, src_image, src_pt):
+        align_mat_info = get_warp_mat_bbox_by_gt_pts_float(src_pt, base_angle=0, dst_size=self.image_size, expand_ratio=0.35, return_info=True)
+        align_mat = align_mat_info["M"]
+
+        align_image_rgb_uint8 = cv2.cvtColor(cv2.warpAffine(src_image, align_mat, (self.image_size, self.image_size)), cv2.COLOR_BGR2RGB)
+
+        # cv2.imshow('align_image_rgb_uint8', align_image_rgb_uint8)
+
+        align_image_rgb_fp32 = align_image_rgb_uint8.astype(np.float32) / 255.0
+        align_image_rgb_fp32_onnx_input = align_image_rgb_fp32.copy().transpose((2, 0, 1))[np.newaxis, ...]
+        pred_coeffs = self.deep3d_pred_net.run(self.deep3d_pred_net_output_name,
+                                               {self.deep3d_pred_net_input_name: align_image_rgb_fp32_onnx_input})[0]
+
+        angles = pred_coeffs[:, self.id_dims + self.exp_dims + self.tex_dims:self.id_dims + self.exp_dims + self.tex_dims + 3]
+        translations = pred_coeffs[:, self.id_dims + self.exp_dims + self.tex_dims + 3 + 27:]
+
+        crop_global_transl = crop_transl_to_full_transl(translations[0],
+                                                        crop_center=[align_mat_info["center_x"],
+                                                                     align_mat_info["center_y"]],
+                                                        scale=align_mat_info["scale"],
+                                                        full_center=[src_image.shape[1] * 0.5, src_image.shape[0] * 0.5],
+                                                        focal_length=self.used_focal_len)
+        return angles, crop_global_transl[np.newaxis, :]
+
+    def compute_rotation_matrix(self, angles):
+        n_b = angles.shape[0]
+        sinx = np.sin(angles[:, 0])
+        siny = np.sin(angles[:, 1])
+        sinz = np.sin(angles[:, 2])
+        cosx = np.cos(angles[:, 0])
+        cosy = np.cos(angles[:, 1])
+        cosz = np.cos(angles[:, 2])
+        rotXYZ = np.eye(3).reshape(1, 3, 3).repeat(n_b*3, 0).reshape(3, n_b, 3, 3)
+        rotXYZ[0, :, 1, 1] = cosx
+        rotXYZ[0, :, 1, 2] = -sinx
+        rotXYZ[0, :, 2, 1] = sinx
+        rotXYZ[0, :, 2, 2] = cosx
+        rotXYZ[1, :, 0, 0] = cosy
+        rotXYZ[1, :, 0, 2] = siny
+        rotXYZ[1, :, 2, 0] = -siny
+        rotXYZ[1, :, 2, 2] = cosy
+        rotXYZ[2, :, 0, 0] = cosz
+        rotXYZ[2, :, 0, 1] = -sinz
+        rotXYZ[2, :, 1, 0] = sinz
+        rotXYZ[2, :, 1, 1] = cosz
+        rotation = np.matmul(np.matmul(rotXYZ[2], rotXYZ[1]), rotXYZ[0])
+        return rotation.transpose(0, 2, 1)
+
+    def rigid_transform(self, vs, rot, trans):
+        vs_r = np.matmul(vs, rot)
+        vs_t = vs_r + trans.reshape(-1, 1, 3)
+        return vs_t
+
+    def perspective_projection_points(self, points, image_w, image_h, used_focal_len):
+        batch_size = points.shape[0]
+        K = np.zeros([batch_size, 3, 3])
+        K[:, 0, 0] = used_focal_len
+        K[:, 1, 1] = used_focal_len
+        K[:, 2, 2] = 1.
+        K[:, 0, 2] = image_w * 0.5
+        K[:, 1, 2] = image_h * 0.5
+
+        reverse_z = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]])[np.newaxis, :, :].repeat(batch_size, 0)
+
+        # Transform points
+        aug_projection = np.matmul(points, reverse_z)
+        aug_projection = np.matmul(aug_projection, K.transpose((0, 2, 1)))
+
+        # Apply perspective distortion
+        projected_points = aug_projection[:, :, :2] / aug_projection[:, :, 2:]
+        return projected_points
+
+    def get_project_points_rect(self, angle, trans, image_w, image_h):
+        vs = np.array(
+            [[-1, -1, 0], [-1, 1, 0], [1, 1, 0], [1, -1, 0]],
+        ) * 0.05
+        vs = vs[np.newaxis, :, :]
+
+        rotation = self.compute_rotation_matrix(angle)
+        translation = trans.copy()
+        translation[0, 2] *= 0.05
+
+        vs_t = self.rigid_transform(vs, rotation, translation)
+
+        project_points = self.perspective_projection_points(vs_t, image_w, image_h, self.used_focal_len*0.05)
+        project_points = np.stack([project_points[:, :, 0], image_h - project_points[:, :, 1]], axis=2)
+
+        return project_points[0]
+

hellomeme/tools/hello_arkit.py

@@ -0,0 +1,35 @@
+"""
+@File   : test.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 11/1/2024
+@Desc   : Created by Shengjie Wu (wu.shengjie@immomo.com)
+"""
+
+import numpy as np
+import cv2
+from huggingface_hub import hf_hub_download
+from .utils import create_onnx_session, get_warp_mat_bbox_by_gt_pts_float
+
+class HelloARKitBSPred(object):
+    def __init__(self, gpu_id=0):
+        self.face_rig_net = create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_arkit_blendshape.onnx'), gpu_id=gpu_id)
+        self.onnx_input_name = self.face_rig_net.get_inputs()[0].name
+        self.onnx_output_name = [output.name for output in self.face_rig_net.get_outputs()]
+        self.image_size = 224
+        self.expand_ratio = 0.15
+
+    def forward(self, src_image, src_pt):
+        left_eye_corner = src_pt[74]
+        right_eye_corner = src_pt[96]
+        radian = np.arctan2(right_eye_corner[1] - left_eye_corner[1], right_eye_corner[0] - left_eye_corner[0] + 0.00000001)
+        rotate_angle = np.rad2deg(radian)
+        align_warp_mat = get_warp_mat_bbox_by_gt_pts_float(src_pt, base_angle=rotate_angle, dst_size=self.image_size,
+                                                           expand_ratio=self.expand_ratio)
+        face_rig_input = cv2.warpAffine(src_image, align_warp_mat, (self.image_size, self.image_size))
+
+        face_rig_onnx_input = face_rig_input.transpose((2, 0, 1)).astype(np.float32)[np.newaxis, :, :, :] / 255.0
+        face_rig_params = self.face_rig_net.run(self.onnx_output_name,
+                                                {self.onnx_input_name: face_rig_onnx_input})
+        face_rig_params = face_rig_params[0][0]
+        return face_rig_params

hellomeme/tools/hello_camera_demo.py

@@ -0,0 +1,377 @@
+"""
+@File   : test.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 11/1/2024
+@Desc   : Created by Shengjie Wu (wu.shengjie@immomo.com)
+"""
+
+import numpy as np
+from .utils import OneEuroFilter
+
+def cult_dis(old_kpts, new_kpts):
+    dis = np.sqrt(np.square(new_kpts[:, 0] - old_kpts[:, 0]) + np.square(new_kpts[:, 1] - old_kpts[:, 1]))
+    return dis
+
+class Smoother222(object):
+    def __init__(self):
+        # face config
+        self.face_idx = list(range(0, 33))
+        self.face_down_idx = list(range(9, 24))
+        self.filter_face = OneEuroFilter()
+        # nose config
+        self.nose_idx = list(range(33, 48))
+        self.filter_nose = OneEuroFilter()
+        # eyebrow config
+        self.eyebrow_idx = list(range(48, 74))
+        self.filter_eyebrow = OneEuroFilter()
+        # eye config
+        self.left_eye_idx = list(range(74, 96))
+        self.filter_left_eye = OneEuroFilter()
+        self.right_eye_idx = list(range(96, 118))
+        self.filter_right_eye = OneEuroFilter()
+        # mouth config
+        self.mouth_idx = list(range(118, 182))
+        self.filter_mouth = OneEuroFilter()
+        # pupil config
+        self.left_pupil_idx = list(range(182, 202))
+        self.filter_left_pupil = OneEuroFilter()
+        self.right_pupil_idx = list(range(202, 222))
+        self.filter_right_pupil = OneEuroFilter()
+        self.prev_points = None
+
+    def smooth(self, new_points, face_dis):
+        if self.prev_points is None:
+            self.prev_points = new_points.copy()
+            return new_points
+        dis = cult_dis(self.prev_points, new_points) / face_dis
+        smooth_points = new_points.copy()
+
+        # smooth face
+        if np.mean(dis[self.face_down_idx]) < 0.005:
+            # stable
+            ratio_tmp = np.mean(dis[self.face_down_idx]) / 0.005
+            fcmin_tmp = 0.05 * ratio_tmp
+            beta_tmp = 0.05 * ratio_tmp
+            smooth_points[self.face_idx] = self.filter_face(new_points[self.face_idx], self.prev_points[self.face_idx],
+                                                            fcmin=fcmin_tmp, beta=beta_tmp)
+        elif np.mean(dis[self.face_down_idx]) < 0.02:
+            # filter
+            ratio_tmp = (np.mean(dis[self.face_down_idx]) - 0.005) / (0.02 - 0.005)
+            fcmin_tmp = 0.05 + (0.3 - 0.05) * ratio_tmp
+            beta_tmp = 0.05 + (0.3 - 0.05) * ratio_tmp
+            smooth_points[self.face_idx] = self.filter_face(new_points[self.face_idx], self.prev_points[self.face_idx],
+                                                            fcmin=fcmin_tmp, beta=beta_tmp)
+        else:
+            # filter
+            smooth_points[self.face_idx] = self.filter_face(new_points[self.face_idx], self.prev_points[self.face_idx],
+                                                            fcmin=0.3, beta=0.3)
+        # smooth nose
+        if np.mean(dis[self.nose_idx]) < 0.003:
+            # stable
+            ratio_tmp = np.mean(dis[self.nose_idx]) / 0.003
+            fcmin_tmp = 0.03 * ratio_tmp
+            beta_tmp = 0.03 * ratio_tmp
+            smooth_points[self.nose_idx] = self.filter_nose(new_points[self.nose_idx], self.prev_points[self.nose_idx],
+                                                            fcmin=fcmin_tmp, beta=beta_tmp)
+        elif np.mean(dis[self.nose_idx]) < 0.02:
+            ratio_tmp = (np.mean(dis[self.nose_idx]) - 0.003) / (0.02 - 0.003)
+            fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            smooth_points[self.nose_idx] = self.filter_nose(new_points[self.nose_idx], self.prev_points[self.nose_idx],
+                                                            fcmin=fcmin_tmp, beta=beta_tmp)
+        else:
+            # filter
+            smooth_points[self.nose_idx] = self.filter_nose(new_points[self.nose_idx], self.prev_points[self.nose_idx],
+                                                            fcmin=0.7, beta=0.7)
+        # smooth eyebrow
+        # print(np.mean(dis[self.eyebrow_idx]))
+        if np.mean(dis[self.eyebrow_idx]) < 0.003:
+            # stable
+            ratio_tmp = np.mean(dis[self.eyebrow_idx]) / 0.003
+            fcmin_tmp = 0.02 * ratio_tmp
+            beta_tmp = 0.02 * ratio_tmp
+            smooth_points[self.eyebrow_idx] = self.filter_eyebrow(new_points[self.eyebrow_idx],
+                                                                  self.prev_points[self.eyebrow_idx],
+                                                                  fcmin=fcmin_tmp, beta=beta_tmp)
+        elif np.mean(dis[self.eyebrow_idx]) < 0.02:
+            # filter
+            ratio_tmp = (np.mean(dis[self.eyebrow_idx]) - 0.003) / (0.02 - 0.003)
+            fcmin_tmp = 0.02 + (0.5 - 0.02) * ratio_tmp
+            beta_tmp = 0.02 + (0.5 - 0.02) * ratio_tmp
+            smooth_points[self.eyebrow_idx] = self.filter_eyebrow(new_points[self.eyebrow_idx],
+                                                                  self.prev_points[self.eyebrow_idx],
+                                                                  fcmin=fcmin_tmp, beta=beta_tmp)
+        else:
+            # filter
+            smooth_points[self.eyebrow_idx] = self.filter_eyebrow(new_points[self.eyebrow_idx],
+                                                                  self.prev_points[self.eyebrow_idx],
+                                                                  fcmin=0.5, beta=0.5)
+        # smooth eye
+        if np.mean(dis[self.left_eye_idx]) < 0.003:
+            # stable
+            ratio_tmp = np.mean(dis[self.left_eye_idx]) / 0.003
+            fcmin_tmp = 0.03 * ratio_tmp
+            beta_tmp = 0.03 * ratio_tmp
+            smooth_points[self.left_eye_idx] = self.filter_left_eye(new_points[self.left_eye_idx],
+                                                                    self.prev_points[self.left_eye_idx],
+                                                                    fcmin=fcmin_tmp, beta=beta_tmp)
+        elif np.mean(dis[self.left_eye_idx]) < 0.02:
+            # filter
+            ratio_tmp = (np.mean(dis[self.left_eye_idx]) - 0.003) / (0.02 - 0.003)
+            fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            smooth_points[self.left_eye_idx] = self.filter_left_eye(new_points[self.left_eye_idx],
+                                                                    self.prev_points[self.left_eye_idx],
+                                                                    fcmin=fcmin_tmp, beta=beta_tmp)
+        else:
+            # fast
+            smooth_points[self.left_eye_idx] = self.filter_left_eye(new_points[self.left_eye_idx],
+                                                                    self.prev_points[self.left_eye_idx],
+                                                                    fcmin=0.7, beta=0.7)
+        if np.mean(dis[self.right_eye_idx]) < 0.003:
+            # stable
+            ratio_tmp = np.mean(dis[self.right_eye_idx]) / 0.003
+            fcmin_tmp = 0.03 * ratio_tmp
+            beta_tmp = 0.03 * ratio_tmp
+            smooth_points[self.right_eye_idx] = self.filter_right_eye(new_points[self.right_eye_idx],
+                                                                      self.prev_points[self.right_eye_idx],
+                                                                      fcmin=fcmin_tmp, beta=beta_tmp)
+        elif np.mean(dis[self.right_eye_idx]) < 0.02:
+            # filter
+            ratio_tmp = (np.mean(dis[self.right_eye_idx]) - 0.003) / (0.02 - 0.003)
+            fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            smooth_points[self.right_eye_idx] = self.filter_right_eye(new_points[self.right_eye_idx],
+                                                                      self.prev_points[self.right_eye_idx],
+                                                                      fcmin=fcmin_tmp, beta=beta_tmp)
+        else:
+            # fast
+            smooth_points[self.right_eye_idx] = self.filter_right_eye(new_points[self.right_eye_idx],
+                                                                      self.prev_points[self.right_eye_idx],
+                                                                      fcmin=0.7, beta=0.7)
+
+        # smooth mouth
+        if np.mean(dis[self.mouth_idx]) < 0.003:
+            # stable
+            ratio_tmp = np.mean(dis[self.mouth_idx]) / 0.003
+            fcmin_tmp = 0.05 * ratio_tmp
+            beta_tmp = 0.05 * ratio_tmp
+            smooth_points[self.mouth_idx] = self.filter_mouth(new_points[self.mouth_idx],
+                                                              self.prev_points[self.mouth_idx],
+                                                              fcmin=fcmin_tmp, beta=beta_tmp)
+        elif np.mean(dis[self.mouth_idx]) < 0.02:
+            # filter
+            ratio_tmp = (np.mean(dis[self.mouth_idx]) - 0.003) / (0.02 - 0.003)
+            fcmin_tmp = 0.05 + (0.7 - 0.05) * ratio_tmp
+            beta_tmp = 0.05 + (0.7 - 0.05) * ratio_tmp
+            smooth_points[self.mouth_idx] = self.filter_mouth(new_points[self.mouth_idx],
+                                                              self.prev_points[self.mouth_idx],
+                                                              fcmin=fcmin_tmp, beta=beta_tmp)
+        else:
+            # fast
+            smooth_points[self.mouth_idx] = self.filter_mouth(new_points[self.mouth_idx],
+                                                              self.prev_points[self.mouth_idx],
+                                                              fcmin=0.7, beta=0.7)
+
+        # smooth pupil
+        if np.mean(dis[self.left_pupil_idx]) < 0.003:
+            # stable
+            ratio_tmp = np.mean(dis[self.left_pupil_idx]) / 0.003
+            fcmin_tmp = 0.03 * ratio_tmp
+            beta_tmp = 0.03 * ratio_tmp
+            smooth_points[self.left_pupil_idx] = self.filter_left_pupil(new_points[self.left_pupil_idx],
+                                                                        self.prev_points[self.left_pupil_idx],
+                                                                        fcmin=fcmin_tmp, beta=beta_tmp)
+        elif np.mean(dis[self.left_pupil_idx]) < 0.02:
+            # filter
+            ratio_tmp = (np.mean(dis[self.left_pupil_idx]) - 0.003) / (0.02 - 0.003)
+            fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            smooth_points[self.left_pupil_idx] = self.filter_left_pupil(new_points[self.left_pupil_idx],
+                                                                        self.prev_points[self.left_pupil_idx],
+                                                                        fcmin=fcmin_tmp, beta=beta_tmp)
+        else:
+            # fast
+            smooth_points[self.left_pupil_idx] = self.filter_left_pupil(new_points[self.left_pupil_idx],
+                                                                        self.prev_points[self.left_pupil_idx],
+                                                                        fcmin=0.7, beta=0.7)
+        if np.mean(dis[self.right_pupil_idx]) < 0.003:
+            # stable
+            ratio_tmp = np.mean(dis[self.right_pupil_idx]) / 0.003
+            fcmin_tmp = 0.03 * ratio_tmp
+            beta_tmp = 0.03 * ratio_tmp
+            smooth_points[self.right_pupil_idx] = self.filter_right_pupil(new_points[self.right_pupil_idx],
+                                                                          self.prev_points[self.right_pupil_idx],
+                                                                          fcmin=fcmin_tmp, beta=beta_tmp)
+        elif np.mean(dis[self.right_pupil_idx]) < 0.02:
+            # filter
+            ratio_tmp = (np.mean(dis[self.right_pupil_idx]) - 0.003) / (0.02 - 0.003)
+            fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
+            smooth_points[self.right_pupil_idx] = self.filter_right_pupil(new_points[self.right_pupil_idx],
+                                                                          self.prev_points[self.right_pupil_idx],
+                                                                          fcmin=fcmin_tmp, beta=beta_tmp)
+        else:
+            # fast
+            smooth_points[self.right_pupil_idx] = self.filter_right_pupil(new_points[self.right_pupil_idx],
+                                                                          self.prev_points[self.right_pupil_idx],
+                                                                          fcmin=0.7, beta=0.7)
+
+        # update pre points
+        self.prev_points = smooth_points
+        return smooth_points
+
+class HelloCameraDemo(object):
+    def __init__(self, face_alignment_module, reset=False):
+        self.face_alignment_module = face_alignment_module
+        self.face_prob_th = 0.01
+        self.min_face = 96
+        self.face_image_size = self.face_alignment_module.face_image_size
+        self.trackingFaces = []
+        self.reset = reset
+
+    def reset_track(self):
+        self.trackingFaces = []
+
+    def forward(self, src_image, reset=False, pre_rect=None):
+        if self.reset or reset:
+            self.trackingFaces = []
+
+        if len(self.trackingFaces) == 0:
+            if pre_rect is not None:
+                detected_faces = [pre_rect]
+            else:
+                detected_faces, _, _ = self.face_alignment_module.face_detector.detect(src_image)
+            for face_rect in detected_faces:
+                new_tracking_object = {
+                    'face_rect': face_rect,
+                    'rotate_angle': 0.0,
+                    'pre_kpt_222': None,
+                    'face_dis': np.sqrt(
+                        np.square((face_rect[2] - face_rect[0])) + np.square((face_rect[3] - face_rect[1]))),
+                    'smoother_222': Smoother222(),
+                    'prob': 0
+                }
+                self.trackingFaces.append(new_tracking_object)
+        else:
+            detected_faces, _, _ = self.face_alignment_module.face_detector.detect(src_image)
+            for face_rect in detected_faces:
+                new_tracking_object = {
+                    'face_rect': face_rect,
+                    'rotate_angle': 0.0,
+                    'pre_kpt_222': None,
+                    'face_dis': np.sqrt(
+                        np.square((face_rect[2] - face_rect[0])) + np.square((face_rect[3] - face_rect[1]))),
+                    'smoother_222': Smoother222(),
+                    'prob': 0
+                }
+                self.trackingFaces.append(new_tracking_object)
+
+        delete_idx_list = []
+        for face_idx, tracking_face in enumerate(self.trackingFaces):
+            if tracking_face['pre_kpt_222'] is not None:
+                result_dict = self.face_alignment_module.forward(src_image, pre_pts=tracking_face['pre_kpt_222'],
+                                                                 iterations=3)
+            else:
+                result_dict = self.face_alignment_module.forward(src_image, face_box=tracking_face['face_rect'],
+                                                                 iterations=3)
+
+            if result_dict['prob'] < self.face_prob_th:
+                if not face_idx in delete_idx_list:
+                    delete_idx_list.append(face_idx)
+                continue
+
+            landmarks_final = tracking_face['smoother_222'].smooth(result_dict['pt222'], tracking_face['face_dis'])
+            tracking_face['pre_kpt_222'] = landmarks_final
+
+            left_eye_corner = landmarks_final[74]
+            right_eye_corner = landmarks_final[96]
+
+            radian = np.arctan2(right_eye_corner[1] - left_eye_corner[1],
+                                right_eye_corner[0] - left_eye_corner[0] + 0.00000001)
+            rotate_angle = np.rad2deg(radian)
+            face_x_min, face_x_max = np.min(landmarks_final[:, 0]), np.max(landmarks_final[:, 0])
+            face_y_min, face_y_max = np.min(landmarks_final[:, 1]), np.max(landmarks_final[:, 1])
+            face_bbox = [face_x_min, face_y_min, face_x_max, face_y_max]
+            face_dis = np.linalg.norm(landmarks_final[0] - landmarks_final[32])
+
+            if face_x_max - face_x_min < self.min_face or face_y_max - face_y_min < self.min_face:
+                if not face_idx in delete_idx_list:
+                    delete_idx_list.append(face_idx)
+
+            euler_pred = result_dict["euler_rad"]
+            pitch = np.rad2deg(euler_pred[0])
+            yaw = np.rad2deg(euler_pred[1])
+            roll = np.rad2deg(euler_pred[2])
+            # print("pitch, yaw, roll", pitch, yaw, roll)
+
+            # one filter model
+            max_euler = abs(pitch) + (abs(yaw)*0.6)
+            face_dis *= (1.0 + max_euler / 18.0)
+
+            # two filter model
+            tracking_face['face_rect'] = face_bbox
+            tracking_face['rotate_angle'] = rotate_angle
+            tracking_face['face_dis'] = face_dis
+            tracking_face['prob'] = result_dict['prob']
+            tracking_face['pitch'] = pitch
+            tracking_face['yaw'] = yaw
+            tracking_face['roll'] = roll
+            tracking_face['euler_rad'] = result_dict["euler_rad"]
+
+        if len(self.trackingFaces) > 1:
+            for face_idx, tracking_face_target in enumerate(self.trackingFaces):
+                if face_idx in delete_idx_list:
+                    continue
+                for idx, tracking_face in enumerate(self.trackingFaces):
+                    if idx in delete_idx_list:
+                        continue
+                    if face_idx == idx: continue
+                    iou_temp = self.count_iou(tracking_face_target['face_rect'], tracking_face['face_rect'])
+                    # prog 2
+                    if iou_temp > 0.12:
+                        if self.area(tracking_face_target['face_rect']) - self.area(tracking_face['face_rect']) < 0:
+                            if not face_idx in delete_idx_list:
+                                delete_idx_list.append(face_idx)
+                        else:
+                            if not idx in delete_idx_list:
+                                delete_idx_list.append(idx)
+
+
+        idx_offset = 0
+        for delete_idx in sorted(delete_idx_list):
+            self.trackingFaces.pop(delete_idx - idx_offset)
+            idx_offset += 1
+
+        return self.trackingFaces
+
+    def count_iou(self, boxA, boxB):
+        # determine the (x, y)-coordinates of the intersection rectangle
+        xA = max(boxA[0], boxB[0])
+        yA = max(boxA[1], boxB[1])
+        xB = min(boxA[2], boxB[2])
+        yB = min(boxA[3], boxB[3])
+
+        # compute the area of intersection rectangle
+        interArea = abs(max((xB - xA, 0)) * max((yB - yA), 0))
+        if interArea == 0:
+            return 0
+        # compute the area of both the prediction and ground-truth
+        # rectangles
+        boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))
+        boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))
+
+        # compute the intersection over union by taking the intersection
+        # area and dividing it by the sum of prediction + ground-truth
+        # areas - the interesection area
+        iou = interArea / float(boxAArea + boxBArea - interArea)
+
+        # return the intersection over union value
+        return iou
+
+    def area(self, bbox):
+        w = bbox[3]-bbox[1]
+        h = bbox[2]-bbox[0]
+        return w * h
+

hellomeme/tools/hello_face_alignment.py

@@ -0,0 +1,93 @@
+"""
+@File   : test.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 11/1/2024
+@Desc   : Created by Shengjie Wu (wu.shengjie@immomo.com)
+"""
+
+import cv2
+import numpy as np
+from .hello_face_det import HelloFaceDet
+from .utils import get_warp_mat_bbox, get_warp_mat_bbox_by_gt_pts_float, transform_points
+from .utils import create_onnx_session
+from huggingface_hub import hf_hub_download
+
+class HelloFaceAlignment(object):
+    def __init__(self, gpu_id=None):
+        expand_ratio = 0.15
+        self.face_alignment_net_222 = (
+            create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_landmark.onnx'), gpu_id=gpu_id))
+        self.onnx_input_name_222 = self.face_alignment_net_222.get_inputs()[0].name
+        self.onnx_output_name_222 = [output.name for output in self.face_alignment_net_222.get_outputs()]
+        self.face_image_size = 128
+        self.face_detector = HelloFaceDet(hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_det.onnx'), gpu_id=gpu_id)
+        self.expand_ratio = expand_ratio
+
+    def onnx_infer(self, input_uint8):
+        assert input_uint8.shape[0] == input_uint8.shape[1] == self.face_image_size
+        onnx_input = input_uint8.transpose((2, 0, 1)).astype(np.float32)[np.newaxis, :, :, :] / 255.0
+        landmark, euler, prob = self.face_alignment_net_222.run(self.onnx_output_name_222,
+                                                                {self.onnx_input_name_222: onnx_input})
+
+        landmark = np.reshape(landmark[0], (2, -1)).transpose((1, 0)) * self.face_image_size
+        left_eye_corner = landmark[74]
+        right_eye_corner = landmark[96]
+        radian = np.arctan2(right_eye_corner[1] - left_eye_corner[1],
+                            right_eye_corner[0] - left_eye_corner[0] + 0.00000001)
+        euler_rad = np.array([euler[0, 0], euler[0, 1], radian], dtype=np.float32)
+        prob = prob[0]
+
+        return landmark, euler_rad, prob
+
+    def forward(self, src_image, face_box=None, pre_pts=None, iterations=3):
+        if pre_pts is None:
+            if face_box is None:
+                # Detect max size face
+                bounding_boxes, _, score = self.face_detector.detect(src_image)
+                print("facedet score", score)
+                if len(bounding_boxes) == 0:
+                    return None
+                bbox = np.zeros(4, dtype=np.float32)
+                if len(bounding_boxes) >= 1:
+                    max_area = 0.0
+                    for each_bbox in bounding_boxes:
+                        area = (each_bbox[2] - each_bbox[0]) * (each_bbox[3] - each_bbox[1])
+                        if area > max_area:
+                            bbox[:4] = each_bbox[:4]
+                        max_area = area
+                else:
+                    bbox = bounding_boxes[0, :4]
+            else:
+                bbox = face_box.copy()
+            M_Face = get_warp_mat_bbox(bbox, 0, self.face_image_size, expand_ratio=self.expand_ratio)
+        else:
+            left_eye_corner = pre_pts[74]
+            right_eye_corner = pre_pts[96]
+
+            radian = np.arctan2(right_eye_corner[1] - left_eye_corner[1],
+                                right_eye_corner[0] - left_eye_corner[0] + 0.00000001)
+            M_Face = get_warp_mat_bbox_by_gt_pts_float(pre_pts, np.rad2deg(radian), self.face_image_size,
+                                                       expand_ratio=self.expand_ratio)
+
+        face_input = cv2.warpAffine(src_image, M_Face, (self.face_image_size, self.face_image_size))
+        landmarks, euler, prob = self.onnx_infer(face_input)
+        landmarks = transform_points(landmarks, M_Face, invert=True)
+
+        # Repeat
+        for i in range(iterations - 1):
+            M_Face = get_warp_mat_bbox_by_gt_pts_float(landmarks, np.rad2deg(euler[2]), self.face_image_size,
+                                                       expand_ratio=self.expand_ratio)
+            face_input = cv2.warpAffine(src_image, M_Face, (self.face_image_size, self.face_image_size))
+            landmarks, euler, prob = self.onnx_infer(face_input)
+            landmarks = transform_points(landmarks, M_Face, invert=True)
+
+        return_dict = {
+            "pt222": landmarks,
+            "euler_rad": euler,
+            "prob": prob,
+            "M_Face": M_Face,
+            "face_input": face_input
+        }
+
+        return return_dict

hellomeme/tools/hello_face_det.py

@@ -0,0 +1,326 @@
+"""
+@File   : test.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 11/1/2024
+@Desc   : Created by Zemin An (an.zemin@hellogroup.com)
+"""
+
+from abc import ABCMeta, abstractmethod
+import cv2
+import numpy as np
+from scipy.special import softmax
+import os.path as osp
+from .utils import create_onnx_session
+
+songkey_weights_dir = 'pretrained_models'
+
+_COLORS = (
+    np.array(
+        [
+            0.000,
+            0.447,
+            0.741,
+        ]
+    )
+        .astype(np.float32)
+        .reshape(-1, 3)
+)
+
+def get_resize_matrix(raw_shape, dst_shape, keep_ratio):
+    """
+    Get resize matrix for resizing raw img to input size
+    :param raw_shape: (width, height) of raw image
+    :param dst_shape: (width, height) of input image
+    :param keep_ratio: whether keep original ratio
+    :return: 3x3 Matrix
+    """
+    r_w, r_h = raw_shape
+    d_w, d_h = dst_shape
+    Rs = np.eye(3)
+    if keep_ratio:
+        C = np.eye(3)
+        C[0, 2] = -r_w / 2
+        C[1, 2] = -r_h / 2
+
+        if r_w / r_h < d_w / d_h:
+            ratio = d_h / r_h
+        else:
+            ratio = d_w / r_w
+        Rs[0, 0] *= ratio
+        Rs[1, 1] *= ratio
+
+        T = np.eye(3)
+        T[0, 2] = 0.5 * d_w
+        T[1, 2] = 0.5 * d_h
+        return T @ Rs @ C
+    else:
+        Rs[0, 0] *= d_w / r_w
+        Rs[1, 1] *= d_h / r_h
+        return Rs
+
+def warp_boxes(boxes, M, width, height):
+    """Apply transform to boxes
+    Copy from nanodet/data/transform/warp.py
+    """
+    n = len(boxes)
+    if n:
+        # warp points
+        xy = np.ones((n * 4, 3))
+        xy[:, :2] = boxes[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(
+            n * 4, 2
+        )  # x1y1, x2y2, x1y2, x2y1
+        xy = xy @ M.T  # transform
+        xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8)  # rescale
+        # create new boxes
+        x = xy[:, [0, 2, 4, 6]]
+        y = xy[:, [1, 3, 5, 7]]
+        xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
+        # clip boxes
+        xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)
+        xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)
+        return xy.astype(np.float32)
+    else:
+        return boxes
+
+def overlay_bbox_cv(img, all_box, class_names):
+    """Draw result boxes
+    Copy from nanodet/util/visualization.py
+    """
+    # all_box array of [label, x0, y0, x1, y1, score]
+    all_box.sort(key=lambda v: v[5])
+    for box in all_box:
+        label, x0, y0, x1, y1, score = box
+        # color = self.cmap(i)[:3]
+        color = (_COLORS[label] * 255).astype(np.uint8).tolist()
+        text = "{}:{:.1f}%".format(class_names[label], score * 100)
+        txt_color = (0, 0, 0) if np.mean(_COLORS[label]) > 0.5 else (255, 255, 255)
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        txt_size = cv2.getTextSize(text, font, 0.5, 2)[0]
+        cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)
+
+        cv2.rectangle(
+            img,
+            (x0, y0 - txt_size[1] - 1),
+            (x0 + txt_size[0] + txt_size[1], y0 - 1),
+            color,
+            -1,
+        )
+        cv2.putText(img, text, (x0, y0 - 1), font, 0.5, txt_color, thickness=1)
+    return img
+
+def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):
+    """
+
+    Args:
+        box_scores (N, 5): boxes in corner-form and probabilities.
+        iou_threshold: intersection over union threshold.
+        top_k: keep top_k results. If k <= 0, keep all the results.
+        candidate_size: only consider the candidates with the highest scores.
+    Returns:
+         picked: a list of indexes of the kept boxes
+    """
+    scores = box_scores[:, -1]
+    boxes = box_scores[:, :-1]
+    picked = []
+    # _, indexes = scores.sort(descending=True)
+    indexes = np.argsort(scores)
+    # indexes = indexes[:candidate_size]
+    indexes = indexes[-candidate_size:]
+    while len(indexes) > 0:
+        # current = indexes[0]
+        current = indexes[-1]
+        picked.append(current)
+        if 0 < top_k == len(picked) or len(indexes) == 1:
+            break
+        current_box = boxes[current, :]
+        # indexes = indexes[1:]
+        indexes = indexes[:-1]
+        rest_boxes = boxes[indexes, :]
+        iou = iou_of(
+            rest_boxes,
+            np.expand_dims(current_box, axis=0),
+        )
+        indexes = indexes[iou <= iou_threshold]
+
+    return box_scores[picked, :]
+
+
+def iou_of(boxes0, boxes1, eps=1e-5):
+    """Return intersection-over-union (Jaccard index) of boxes.
+
+    Args:
+        boxes0 (N, 4): ground truth boxes.
+        boxes1 (N or 1, 4): predicted boxes.
+        eps: a small number to avoid 0 as denominator.
+    Returns:
+        iou (N): IoU values.
+    """
+    overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])
+    overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])
+
+    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
+    area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
+    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
+    return overlap_area / (area0 + area1 - overlap_area + eps)
+
+
+def area_of(left_top, right_bottom):
+    """Compute the areas of rectangles given two corners.
+
+    Args:
+        left_top (N, 2): left top corner.
+        right_bottom (N, 2): right bottom corner.
+
+    Returns:
+        area (N): return the area.
+    """
+    hw = np.clip(right_bottom - left_top, 0.0, None)
+    return hw[..., 0] * hw[..., 1]
+
+
+class NanoDetABC(metaclass=ABCMeta):
+    def __init__(
+            self,
+            input_shape=[272, 160],
+            reg_max=7,
+            strides=[8, 16, 32],
+            prob_threshold=0.4,
+            iou_threshold=0.3,
+            num_candidate=1000,
+            top_k=-1,
+            class_names=["face"]
+    ):
+        self.strides = strides
+        self.input_shape = input_shape
+        self.reg_max = reg_max
+        self.prob_threshold = prob_threshold
+        self.iou_threshold = iou_threshold
+        self.num_candidate = num_candidate
+        self.top_k = top_k
+        self.img_mean = [103.53, 116.28, 123.675]
+        self.img_std = [57.375, 57.12, 58.395]
+        self.input_size = (self.input_shape[1], self.input_shape[0])
+        self.class_names = class_names
+        self.num_classes = len(self.class_names)
+
+    def preprocess(self, img):
+        # resize image
+        ResizeM = get_resize_matrix((img.shape[1], img.shape[0]), self.input_size, True)
+        img_resize = cv2.warpPerspective(img, ResizeM, dsize=self.input_size)
+
+        # normalize image
+        img_input = img_resize.astype(np.float32) / 255
+        img_mean = np.array(self.img_mean, dtype=np.float32).reshape(1, 1, 3) / 255
+        img_std = np.array(self.img_std, dtype=np.float32).reshape(1, 1, 3) / 255
+        img_input = (img_input - img_mean) / img_std
+
+        # expand dims
+        img_input = np.transpose(img_input, [2, 0, 1])
+        img_input = np.expand_dims(img_input, axis=0)
+        return img_input, ResizeM
+
+    def postprocess(self, scores, raw_boxes, ResizeM, raw_shape):
+        # generate centers
+        decode_boxes = []
+        select_scores = []
+        for stride, box_distribute, score in zip(self.strides, raw_boxes, scores):
+            # centers
+            fm_h = self.input_shape[0] / stride
+            fm_w = self.input_shape[1] / stride
+
+            h_range = np.arange(fm_h)
+            w_range = np.arange(fm_w)
+            ww, hh = np.meshgrid(w_range, h_range)
+
+            ct_row = hh.flatten() * stride
+            ct_col = ww.flatten() * stride
+
+            center = np.stack((ct_col, ct_row, ct_col, ct_row), axis=1)
+
+            # box distribution to distance
+            reg_range = np.arange(self.reg_max + 1)
+            box_distance = box_distribute.reshape((-1, self.reg_max + 1))
+            box_distance = softmax(box_distance, axis=1)
+            box_distance = box_distance * np.expand_dims(reg_range, axis=0)
+            box_distance = np.sum(box_distance, axis=1).reshape((-1, 4))
+            box_distance = box_distance * stride
+
+            # top K candidate
+            topk_idx = np.argsort(score.max(axis=1))[::-1]
+            topk_idx = topk_idx[: self.num_candidate]
+            center = center[topk_idx]
+            score = score[topk_idx]
+            box_distance = box_distance[topk_idx]
+
+            # decode box
+            decode_box = center + [-1, -1, 1, 1] * box_distance
+
+            select_scores.append(score)
+            decode_boxes.append(decode_box)
+
+        # nms
+        bboxes = np.concatenate(decode_boxes, axis=0)
+        confidences = np.concatenate(select_scores, axis=0)
+        picked_box_probs = []
+        picked_labels = []
+        for class_index in range(0, confidences.shape[1]):
+            probs = confidences[:, class_index]
+            mask = probs > self.prob_threshold
+            probs = probs[mask]
+            if probs.shape[0] == 0:
+                continue
+            subset_boxes = bboxes[mask, :]
+            box_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], axis=1)
+            box_probs = hard_nms(
+                box_probs,
+                iou_threshold=self.iou_threshold,
+                top_k=self.top_k,
+            )
+            picked_box_probs.append(box_probs)
+            picked_labels.extend([class_index] * box_probs.shape[0])
+        if not picked_box_probs:
+            return np.array([]), np.array([]), np.array([])
+        picked_box_probs = np.concatenate(picked_box_probs)
+
+        # resize output boxes
+        picked_box_probs[:, :4] = warp_boxes(
+            picked_box_probs[:, :4], np.linalg.inv(ResizeM), raw_shape[1], raw_shape[0]
+        )
+        return (
+            picked_box_probs[:, :4].astype(np.int32),
+            np.array(picked_labels),
+            picked_box_probs[:, 4],
+        )
+
+    @abstractmethod
+    def infer_image(self, img_input):
+        pass
+
+    def detect(self, img):
+        raw_shape = img.shape
+        img_input, ResizeM = self.preprocess(img)
+        scores, raw_boxes = self.infer_image(img_input)
+        if scores[0].ndim == 1:  # handling num_classes=1 case
+            scores = [x[:, None] for x in scores]
+        bbox, label, score = self.postprocess(scores, raw_boxes, ResizeM, raw_shape)
+
+        return bbox, label, score
+
+class HelloFaceDet(NanoDetABC):
+    def __init__(self, model_path=osp.join(songkey_weights_dir, 'face/nanodet_humandet_320-192_220302_model_20220315_test3.onnx'), gpu_id=None, *args, **kwargs):
+        super(HelloFaceDet, self).__init__(*args, **kwargs)
+        # print("Using ONNX as inference backend")
+        # print(f"Using weight: {model_path}")
+
+        # load model
+        self.model_path = model_path
+        self.ort_session = create_onnx_session(model_path, gpu_id=gpu_id)
+        self.input_name = self.ort_session.get_inputs()[0].name
+
+    def infer_image(self, img_input):
+        inference_results = self.ort_session.run(None, {self.input_name: img_input})
+
+        scores = [np.squeeze(x) for x in inference_results[:3]]
+        raw_boxes = [np.squeeze(x) for x in inference_results[3:]]
+        return scores, raw_boxes

hellomeme/tools/utils.py

@@ -0,0 +1,118 @@
+# coding: utf-8
+
+# @File   : utils.py
+# @Author : Songkey
+# @Email  : songkey@pku.edu.cn
+# @Date   : 8/18/2024
+# @Desc   :
+
+import onnx, onnxruntime
+import time
+import cv2
+import numpy as np
+import math
+
+def create_onnx_session(onnx_path, gpu_id=None)->onnxruntime.InferenceSession:
+    start = time.perf_counter()
+    onnx_model = onnx.load(onnx_path)
+    onnx.checker.check_model(onnx_model)
+    providers = [
+        ('CUDAExecutionProvider', {
+            'device_id': int(gpu_id),
+            'arena_extend_strategy': 'kNextPowerOfTwo',
+            #'cuda_mem_limit': 5 * 1024 * 1024 * 1024,
+            'cudnn_conv_algo_search': 'EXHAUSTIVE',
+            'do_copy_in_default_stream': True,
+        }),
+        'CPUExecutionProvider',
+    ] if (gpu_id is not None and gpu_id >= 0) else ['CPUExecutionProvider']
+
+    sess = onnxruntime.InferenceSession(onnx_path, providers=providers)
+    print('create onnx session cost: {:.3f}s. {}'.format(time.perf_counter() - start, onnx_path))
+    return sess
+
+def smoothing_factor(t_e, cutoff):
+    r = 2 * math.pi * cutoff * t_e
+    return r / (r + 1)
+
+def exponential_smoothing(a, x, x_prev):
+    return a * x + (1 - a) * x_prev
+
+class OneEuroFilter:
+    def __init__(self, dx0=0.0, d_cutoff=1.0):
+        """Initialize the one euro filter."""
+        # self.min_cutoff = float(min_cutoff)
+        # self.beta = float(beta)
+        self.d_cutoff = float(d_cutoff)
+        self.dx_prev = float(dx0)
+        # self.t_e = fcmin
+
+    def __call__(self, x, x_prev, fcmin=1.0, min_cutoff=1.0, beta=0.0):
+        if x_prev is None:
+            return x
+        # t_e = 1
+        a_d = smoothing_factor(fcmin, self.d_cutoff)
+        dx = (x - x_prev) / fcmin
+        dx_hat = exponential_smoothing(a_d, dx, self.dx_prev)
+        cutoff = min_cutoff + beta * abs(dx_hat)
+        a = smoothing_factor(fcmin, cutoff)
+        x_hat = exponential_smoothing(a, x, x_prev)
+        self.dx_prev = dx_hat
+        return x_hat
+
+def get_warp_mat_bbox(face_bbox, base_angle, dst_size=128, expand_ratio=0.15, aug_angle=0.0, aug_scale=1.0):
+    face_x_min, face_y_min, face_x_max, face_y_max = face_bbox
+    face_x_center = (face_x_min + face_x_max) / 2
+    face_y_center = (face_y_min + face_y_max) / 2
+    face_width = face_x_max - face_x_min
+    face_height = face_y_max - face_y_min
+    scale = dst_size / max(face_width, face_height) * (1 - expand_ratio) * aug_scale
+    M = cv2.getRotationMatrix2D((face_x_center, face_y_center), angle=base_angle + aug_angle, scale=scale)
+    offset = [dst_size / 2 - face_x_center, dst_size / 2 - face_y_center]
+    M[:, 2] += offset
+    return M
+
+def transform_points(points, mat, invert=False):
+    if invert:
+        mat = cv2.invertAffineTransform(mat)
+    points = np.expand_dims(points, axis=1)
+    points = cv2.transform(points, mat, points.shape)
+    points = np.squeeze(points)
+    return points
+
+def get_warp_mat_bbox_by_gt_pts_float(gt_pts, base_angle=0.0, dst_size=128, expand_ratio=0.15, return_info=False):
+    # step 1
+    face_x_min, face_x_max = np.min(gt_pts[:, 0]), np.max(gt_pts[:, 0])
+    face_y_min, face_y_max = np.min(gt_pts[:, 1]), np.max(gt_pts[:, 1])
+    face_x_center = (face_x_min + face_x_max) / 2
+    face_y_center = (face_y_min + face_y_max) / 2
+    M_step_1 = cv2.getRotationMatrix2D((face_x_center, face_y_center), angle=base_angle, scale=1.0)
+    pts_step_1 = transform_points(gt_pts, M_step_1)
+    face_x_min_step_1, face_x_max_step_1 = np.min(pts_step_1[:, 0]), np.max(pts_step_1[:, 0])
+    face_y_min_step_1, face_y_max_step_1 = np.min(pts_step_1[:, 1]), np.max(pts_step_1[:, 1])
+    # step 2
+    face_width = face_x_max_step_1 - face_x_min_step_1
+    face_height = face_y_max_step_1 - face_y_min_step_1
+    scale = dst_size / max(face_width, face_height) * (1 - expand_ratio)
+    M_step_2 = cv2.getRotationMatrix2D((face_x_center, face_y_center), angle=base_angle, scale=scale)
+    pts_step_2 = transform_points(gt_pts, M_step_2)
+    face_x_min_step_2, face_x_max_step_2 = np.min(pts_step_2[:, 0]), np.max(pts_step_2[:, 0])
+    face_y_min_step_2, face_y_max_step_2 = np.min(pts_step_2[:, 1]), np.max(pts_step_2[:, 1])
+    face_x_center_step_2 = (face_x_min_step_2 + face_x_max_step_2) / 2
+    face_y_center_step_2 = (face_y_min_step_2 + face_y_max_step_2) / 2
+
+    M = cv2.getRotationMatrix2D((face_x_center, face_y_center), angle=base_angle, scale=scale)
+    offset = [dst_size / 2 - face_x_center_step_2, dst_size / 2 - face_y_center_step_2]
+    M[:, 2] += offset
+
+    if not return_info:
+        return M
+    else:
+        transform_info = {
+            "M": M,
+            "center_x": face_x_center,
+            "center_y": face_y_center,
+            "rotate_angle": base_angle,
+            "scale": scale
+        }
+        return transform_info

hellomeme/utils.py

@@ -0,0 +1,379 @@
+# coding: utf-8
+
+"""
+@File   : utils.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 8/28/2024
+@Desc   : 
+"""
+
+import os
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+import subprocess
+from einops import rearrange
+from .tools.utils import transform_points
+from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_unet_checkpoint
+
+from safetensors import safe_open
+
+def load_safetensors(model_path):
+    tensors = {}
+    with safe_open(model_path, framework="pt", device=0) as f:
+        for k in f.keys():
+            tensors[k] = f.get_tensor(k) # loads the full tensor given a key
+    return tensors
+
+def load_unet_from_safetensors(safetensors_path, unet_config):
+    original_stats = load_safetensors(safetensors_path)
+    return convert_ldm_unet_checkpoint(original_stats, unet_config)
+
+def image_preprocess(np_bgr, size, dtype=torch.float32):
+    img_np = cv2.resize(np_bgr, size)
+    return np_bgr_to_tensor(img_np, dtype)
+
+def np_bgr_to_tensor(img_np, dtype):
+    img_rgb = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB) / 255. * 2 - 1
+    return torch.tensor(img_rgb).permute(2, 0, 1).to(dtype=dtype)
+
+def tensor_to_np_rgb(img_np):
+    img_rgb = np.clip((img_np.permute(1, 2, 0).detach().cpu().numpy() + 1) / 2 * 255, 0, 255)
+    return img_rgb.astype(np.uint8)
+
+def clip_preprocess_from_bgr(image_bgr, dtype=torch.float32):
+    temp_image = cv2.resize(image_bgr, (224, 224), interpolation=cv2.INTER_CUBIC)
+    temp_image = cv2.cvtColor(temp_image, cv2.COLOR_BGR2RGB)
+    temp_image_np = temp_image * 0.00392156862745098
+    temp_image_np = (temp_image_np - [0.48145466, 0.4578275, 0.40821073]) / [0.26862954, 0.26130258, 0.27577711]
+    temp_image_np = np.transpose(temp_image_np, (2, 0, 1))
+    return torch.tensor(temp_image_np, dtype=dtype)
+
+def clip_preprocess_to_bgr(image_tensor):
+    image_tensor_np = image_tensor.detach().cpu().numpy().astype(float)
+    image_tensor_np = np.transpose(image_tensor_np, (1, 2, 0))
+    image_tensor_np = (image_tensor_np * [0.26862954, 0.26130258, 0.27577711]) + [0.48145466, 0.4578275, 0.40821073]
+    image_tensor_np = np.clip(image_tensor_np / 0.00392156862745098, 0, 255).astype(np.uint8)
+    return cv2.cvtColor(image_tensor_np, cv2.COLOR_BGR2RGB)
+
+def clip_preprocess_from_pil(image_pil, dtype):
+    temp_image = image_pil.resize((224, 224), resample=Image.BICUBIC)
+    temp_image_np = np.asarray(temp_image) * 0.00392156862745098
+    temp_image_np = (temp_image_np - [0.48145466, 0.4578275, 0.40821073]) / [0.26862954, 0.26130258, 0.27577711]
+    temp_image_np = np.transpose(temp_image_np, (2, 0, 1))
+    return torch.tensor(temp_image_np, dtype=dtype)
+
+def draw_skl_by_rect(save_size, rect):
+    rect = rect.astype(np.int32)
+    ret_img = np.zeros((save_size, save_size, 3), dtype=np.uint8)
+    # cv2.rectangle(ret_img, rect[0], rect[2], (255, 255, 255), -1)
+    cv2.line(ret_img, rect[0], rect[1], (0, 255, 0), 15)
+    cv2.line(ret_img, rect[1], rect[2], (255, 0, 0), 15)
+    cv2.line(ret_img, rect[2], rect[3], (0, 0, 255), 15)
+    cv2.line(ret_img, rect[3], rect[0], (255, 0, 255), 15)
+    return ret_img
+
+def get_face_part_1frame(image, parts, save_size=256):
+    ret_list = []
+    for part_landmarks in parts:
+        part_tls, part_brs = np.min(part_landmarks, axis=0), np.max(part_landmarks, axis=0)
+        center = (part_tls + part_brs) * 0.5
+        wh = max(part_brs - part_tls) * 1.2
+
+        scale = save_size / wh
+
+        M = cv2.getRotationMatrix2D(center, 0, scale)
+        M[:, 2] += save_size * 0.5 - center
+
+        part = cv2.warpAffine(image, M, (save_size, save_size))
+        ret_list.append(part)
+    return np.stack(ret_list, axis=0)
+
+face_parts_regions = [
+    [48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94],
+    [61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116],
+    [118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148],
+]
+face_region = [33, 74, 96, 133, 149]
+face_temp = np.array([[0.50951539, 0.40128625], [0.29259476, 0.40203411], [0.72196212, 0.40857479], [0.63420714, 0.71678643], [0.38437035, 0.71723224]])
+
+def get_face_parts(images, landmarks, save_size=256):
+    H, W = images[0].shape[:2]
+    warped_imgs = []
+    warped_landmarks = []
+
+    for landmark, frame in zip(landmarks, images):
+        src_landmarks = landmark[face_region, :]
+        dst_landmarks = face_temp * (W, H)
+
+        M = umeyama(src_landmarks, dst_landmarks, True)[0:2]
+        warped_img = cv2.warpAffine(frame, M, (W, H))
+
+        warped_landmark = transform_points(landmark.astype(np.float32), M)[:,:2]
+
+        # for x, y in warped_landmark:
+        #     warped_img = cv2.circle(warped_img, (int(x), int(y)), 1, (255, 255, 0), -1)
+        # for x, y in dst_landmarks:
+        #     warped_img = cv2.circle(warped_img, (int(x), int(y)), 1, (0, 255, 0), -1)
+        # cv2.imshow('warped_img', warped_img)
+        # cv2.waitKey()
+
+        warped_imgs.append(warped_img)
+        warped_landmarks.append(warped_landmark)
+    warped_imgs = np.stack(warped_imgs, axis=0)
+    warped_landmarks = np.stack(warped_landmarks, axis=0)
+
+    ret_list = []
+    for region in face_parts_regions:
+        part_landmarks = warped_landmarks[:, region, :]
+        part_tls, part_brs = np.min(np.min(part_landmarks, axis=1), axis=0), np.max(np.max(part_landmarks, axis=1), axis=0)
+        center = (part_tls + part_brs) * 0.5
+        wh = max(part_brs - part_tls) * 1.2
+
+        scale = save_size / wh
+
+        M = cv2.getRotationMatrix2D(center, 0, scale)
+        M[:, 2] += save_size * 0.5 - center
+        tmp_part_list = []
+        for warped_img in warped_imgs:
+            part = cv2.warpAffine(warped_img, M, (save_size, save_size))
+            # cv2.imshow('warped_img', warped_img)
+            # cv2.imshow('part', part)
+            # cv2.waitKey()
+            tmp_part_list.append(part)
+        ret_list.append(np.stack(tmp_part_list, axis=0))
+    return np.stack(ret_list, axis=0).transpose(1, 0, 2, 3, 4)
+
+def get_face_params(h3dmm, harkit_bs, frames, landmarks, save_size=512, align=True):
+    face_parts = []
+    for tmp_face_parts_list in get_face_parts(frames, landmarks, save_size=256):
+        tmp_face_parts_list = [cv2.GaussianBlur(x, (9, 9), 0) for x in tmp_face_parts_list]
+        face_parts.append(tmp_face_parts_list)
+
+    if align:
+        new_frames, new_landmarks = crop_and_resize(np.stack(frames, axis=0), np.stack(landmarks, axis=0),
+                                                    save_size=save_size, crop=False)
+    else:
+        new_frames, new_landmarks = frames, landmarks
+
+    rot_list = []
+    trans_list = []
+    arkit_bs_list = []
+    for landmark, frame in zip(new_landmarks, new_frames):
+        drive_rot, drive_trans = h3dmm.forward_params(frame, landmark)
+        rot_list.append(drive_rot)
+        trans_list.append(drive_trans)
+        arkit_bs_list.append(harkit_bs.forward(frame, landmark))
+    drive_coeff = torch.from_numpy(np.stack(arkit_bs_list, axis=0))
+    return face_parts, drive_coeff, rot_list, trans_list
+
+def face_params_to_tensor(clip_encoder, h3dmm, face_parts, drive_rot, drive_trans, ref_trans, save_size=512, trans_ratio=0.0):
+    face_parts_list = []
+    for tmp_face_parts_list in face_parts:
+        tmp_face_parts_list = [clip_preprocess_from_bgr(x).unsqueeze(1) for x in tmp_face_parts_list]
+        face_parts_list.append(torch.cat(tmp_face_parts_list, dim=1))
+    face_parts = torch.stack(face_parts_list, dim=2)
+
+    face_parts_tensor = rearrange(face_parts, "c f p h w -> (f p) c h w")
+    face_parts_embedding = clip_encoder(face_parts_tensor.to(device=clip_encoder.device, dtype=clip_encoder.dtype)).image_embeds
+    face_parts_embedding = rearrange(face_parts_embedding,
+                                     "(f p) c -> f p c",
+                                     f=face_parts.size(2)).cpu()
+
+    control_list = []
+    control_show_list = []
+    for rot, trans in zip(drive_rot, drive_trans):
+        rect = h3dmm.get_project_points_rect(rot, ref_trans + (trans - drive_trans[0]) * trans_ratio, save_size, save_size)
+        control_heatmap = draw_skl_by_rect(save_size, rect)
+
+        control_show_list.append(control_heatmap.copy())
+        control_heatmap = image_preprocess(control_heatmap, (save_size, save_size))
+        control_list.append(control_heatmap)
+    control_heatmaps = torch.stack(control_list, dim=1)
+
+    return face_parts_embedding, control_heatmaps
+
+def det_landmarks(face_aligner, frame_list, save_size=(512, 512), reset=False):
+    rect_list = []
+    new_frame_list = []
+    for frame in frame_list:
+        frame = cv2.resize(frame, save_size)
+        faces = face_aligner.forward(frame, reset=reset)
+        if len(faces) > 0:
+            face = sorted(faces, key=lambda x: (x['face_rect'][2] - x['face_rect'][0]) * (
+                    x['face_rect'][3] - x['face_rect'][1]))[-1]
+            rect_list.append(face['face_rect'])
+            new_frame_list.append(frame)
+
+    face_aligner.reset_track()
+    save_frame_list = []
+    save_landmark_list = []
+    for frame, rect in zip(new_frame_list, rect_list):
+        faces = face_aligner.forward(frame, pre_rect=rect, reset=reset)
+        if len(faces) > 0:
+            face = sorted(faces, key=lambda x: (x['face_rect'][2] - x['face_rect'][0]) * (
+                    x['face_rect'][3] - x['face_rect'][1]))[-1]
+            landmarks = face['pre_kpt_222']
+            save_frame_list.append(frame)
+            save_landmark_list.append(landmarks)
+    save_landmark_list = np.stack(save_landmark_list, axis=0).astype(np.float16)
+
+    face_aligner.reset_track()
+    return save_frame_list, save_landmark_list
+
+def get_drive_params(face_aligner, h3dmm, harkit_bs, frame_list, save_size, align):
+    frame_num = len(frame_list)
+    frame_list, landmark_list = det_landmarks(face_aligner, frame_list, save_size=(512, 512), reset=False)
+    assert len(frame_list) == frame_num
+    return get_face_params(h3dmm, harkit_bs, frame_list, landmark_list, save_size=save_size, align=align)
+
+def load_data_list(data_dir, post_fix='.pickle;.txt'):
+    post_fixs = post_fix.split(';')
+    ret_list = []
+    for root, dirnames, filenames in os.walk(data_dir):
+            for name in filenames:
+                if os.path.splitext(name)[1].lower() in post_fixs and not name.startswith('.'):
+                    data_path = os.path.join(root, name)
+                    ret_list.append(data_path)
+    return ret_list
+
+def crop_and_resize(frames, landmarks, save_size=512, crop=True):
+    H, W = frames[0].shape[:2]
+    if crop:
+        all_tl, all_br = np.min(landmarks, axis=1), np.max(landmarks, axis=1)
+        mean_wh = np.mean(all_br - all_tl, axis=0)
+        tl, br = np.min(all_tl, axis=0), np.max(all_br, axis=0)
+        new_size = min(max(mean_wh) * 2.2, min(H, W) - 1)
+        fcenter = (tl + br) * 0.5
+        ftl = fcenter - new_size * 0.5
+        ftl[1] -= mean_wh[1] * 0.2
+        fbr = fcenter + new_size * 0.5
+        fbr[1] -= mean_wh[1] * 0.2
+
+        if ftl[0] < 0:
+            fbr[0] -= ftl[0]
+            ftl[0] = 0
+        if ftl[1] < 0:
+            fbr[1] -= ftl[1]
+            ftl[1] = 0
+
+        if fbr[0] >= W:
+            ftl[0] -= fbr[0] - W + 1
+            fbr[0] = W - 1
+        if fbr[1] >= H:
+            ftl[1] -= fbr[1] - H + 1
+            fbr[1] = H - 1
+        ftl = ftl.astype(int)
+        fbr = (fbr - ftl)[0].astype(int) + ftl
+
+        frames = frames[:, ftl[1]:fbr[1], ftl[0]:fbr[0], :].copy()
+        landmarks = landmarks - ftl
+        ratio = save_size / (fbr - ftl)
+    else:
+        ratio = (save_size / W, save_size / H)
+    landmarks = landmarks * ratio
+    frames = np.stack([cv2.resize(frame, (save_size, save_size), interpolation=cv2.INTER_CUBIC) for frame in frames], axis=0)
+    return frames, landmarks
+
+def umeyama(src, dst, estimate_scale):
+    """Estimate N-D similarity transformation with or without scaling.
+    Parameters
+    ----------
+    src : (M, N) array
+        Source coordinates.
+    dst : (M, N) array
+        Destination coordinates.
+    estimate_scale : bool
+        Whether to estimate scaling factor.
+    Returns
+    -------
+    T : (N + 1, N + 1)
+        The homogeneous similarity transformation matrix. The matrix contains
+        NaN values only if the problem is not well-conditioned.
+    References
+    ----------
+    .. [1] "Least-squares estimation of transformation parameters between two
+            point patterns", Shinji Umeyama, PAMI 1991, DOI: 10.1109/34.88573
+    """
+
+    num = src.shape[0]
+    dim = src.shape[1]
+
+    # Compute mean of src and dst.
+    src_mean = src.mean(axis=0)
+    dst_mean = dst.mean(axis=0)
+
+    # Subtract mean from src and dst.
+    src_demean = src - src_mean
+    dst_demean = dst - dst_mean
+
+    # Eq. (38).
+    A = np.dot(dst_demean.T, src_demean) / num
+
+    # Eq. (39).
+    d = np.ones((dim,), dtype=np.double)
+    if np.linalg.det(A) < 0:
+        d[dim - 1] = -1
+
+    T = np.eye(dim + 1, dtype=np.double)
+
+    U, S, V = np.linalg.svd(A)
+
+    # Eq. (40) and (43).
+    rank = np.linalg.matrix_rank(A)
+    if rank == 0:
+        return np.nan * T
+    elif rank == dim - 1:
+        if np.linalg.det(U) * np.linalg.det(V) > 0:
+            T[:dim, :dim] = np.dot(U, V)
+        else:
+            s = d[dim - 1]
+            d[dim - 1] = -1
+            T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V))
+            d[dim - 1] = s
+    else:
+        T[:dim, :dim] = np.dot(U, np.dot(np.diag(d), V.T))
+
+    if estimate_scale:
+        # Eq. (41) and (42).
+        scale = 1.0 / src_demean.var(axis=0).sum() * np.dot(S, d)
+    else:
+        scale = 1.0
+
+    T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T)
+    T[:dim, :dim] *= scale
+
+    return T
+
+def ff_cat_video_and_audio(video_path, audio_path, output_path):
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-v", "quiet", "-stats",
+        "-i", video_path,
+        "-i", audio_path,
+        "-c:v", "copy",
+        "-map", "0:v:0",
+        "-map", "1:a:0",
+        "-shortest",
+        output_path
+    ]
+    subprocess.run(cmd)
+
+def ff_change_fps(input_video, output_video, fps=15):
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-v", "quiet", "-stats",
+        # "-hwaccel", "cuda",
+        "-i", input_video,
+        "-c:v", "libx264",
+        "-r", f"{fps}",
+        "-c:a", "copy",
+        "-crf", "18",
+        output_video
+    ]
+
+    subprocess.run(cmd)

meme.py

@@ -0,0 +1,368 @@
+import os
+import os.path as osp
+import folder_paths
+import torch
+import numpy as np
+import cv2
+import sys
+import copy
+
+
+from PIL import Image
+import subprocess
+
+import importlib.metadata
+
+cur_dir = osp.dirname(osp.abspath(__file__))
+
+installed_packages = [package.name for package in importlib.metadata.distributions()]
+
+REQUIRED = {
+  'diffusers', 'transformers', 'einops', 'opencv-python', 'tqdm', 'pillow', 'onnxruntime', 'onnx', 'safetensors', 'accelerate', 'peft'
+}
+
+missing = [name for name in REQUIRED if name not in installed_packages]
+print("missing pkgs", missing)
+
+if missing:
+    python = sys.executable
+    subprocess.check_call([python, '-m', 'pip', 'install', *missing], stdout=subprocess.DEVNULL)
+
+from .hellomeme.utils import (face_params_to_tensor,
+                             get_drive_params,
+                             get_face_params,
+                             load_data_list,
+                             load_unet_from_safetensors)
+
+from .hellomeme.tools import Hello3DMMPred, HelloARKitBSPred, HelloFaceAlignment, HelloCameraDemo
+from .hellomeme import HMImagePipeline, HMVideoPipeline
+from transformers import CLIPVisionModelWithProjection
+
+class HMImagePipelineLoader:
+    @classmethod
+    def INPUT_TYPES(s):
+        checkpoint_files = sorted(load_data_list(osp.join(cur_dir, '../../models/checkpoints'), '.pt;.pth;.ckpt;.safetensors'))
+        lora_files = sorted(load_data_list(osp.join(cur_dir, '../../models/loras'), '.safetensors'))
+
+        return {
+            "optional": {
+                "checkpoint_path": (['None'] + checkpoint_files, ),
+                "lora_path": (['None'] + lora_files, ),
+                "gpu_id": ("INT", {"default": 0}),
+            }
+        }
+    RETURN_TYPES = ("HMIMAGEPIPELINE", )
+    FUNCTION = "load_pipeline"
+    CATEGORY = "hellomeme"
+    def load_pipeline(self, checkpoint_path=None, lora_path=None, gpu_id=0):
+        dtype = torch.float16
+        if gpu_id >= 0:
+            device = torch.device("cuda:{}".format(gpu_id))
+        else:
+            device = torch.device("cpu")
+        pipeline = HMImagePipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5").to(dtype=dtype, device=device)
+        pipeline.caryomitosis()
+
+        ### load customized checkpoint or lora here:
+        ## checkpoints
+
+        if checkpoint_path and osp.isfile(checkpoint_path):
+            if checkpoint_path.endswith('.safetensors'):
+                state_dict = load_unet_from_safetensors(checkpoint_path, pipeline.unet_ref.config)
+                pipeline.unet.load_state_dict(state_dict, strict=False)
+            elif osp.splitext(checkpoint_path)[-1] in ['.pt', '.pth', '.ckpt']:
+                state_dict = torch.load(checkpoint_path)
+                pipeline.unet.load_state_dict(state_dict, strict=False)
+            else:
+                print("Invalid checkpoint path", checkpoint_path)
+
+        ### lora
+        if lora_path and osp.isfile(lora_path):
+            pipeline.load_lora_weights(osp.dirname(lora_path), weight_name=osp.basename(lora_path), adapter_name="lora")
+
+        pipeline.insert_hm_modules(dtype=dtype, device=device)
+        
+        return (pipeline, )
+
+class HMVideoPipelineLoader:
+    @classmethod
+    def INPUT_TYPES(s):
+        checkpoint_files = sorted(load_data_list(osp.join(cur_dir, '../../models/checkpoints'), '.pt;.pth;.ckpt;.safetensors'))
+        lora_files = sorted(load_data_list(osp.join(cur_dir, '../../models/loras'), '.safetensors'))
+
+        return {
+            "optional": {
+                "checkpoint_path": (['None'] + checkpoint_files, ),
+                "lora_path": (['None'] + lora_files, ),
+                "gpu_id": ("INT", {"default": 0}),
+            }
+        }
+
+    RETURN_TYPES = ("HMVIDEOPIPELINE",)
+    FUNCTION = "load_pipeline"
+    CATEGORY = "hellomeme"
+
+    def load_pipeline(self, checkpoint_path=None, lora_path=None, gpu_id=0):
+        dtype = torch.float16
+        if gpu_id >= 0:
+            device = torch.device("cuda:{}".format(gpu_id))
+        else:
+            device = torch.device("cpu")
+        pipeline = HMVideoPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5").to(dtype=dtype,
+                                                                                                     device=device)
+        pipeline.caryomitosis()
+
+        ### load customized checkpoint or lora here:
+        ## checkpoints
+
+        if checkpoint_path and osp.isfile(checkpoint_path):
+            if checkpoint_path.endswith('.safetensors'):
+                state_dict = load_unet_from_safetensors(checkpoint_path, pipeline.unet_ref.config)
+                pipeline.unet.load_state_dict(state_dict, strict=False)
+            elif osp.splitext(checkpoint_path)[-1] in ['.pt', '.pth', '.ckpt']:
+                state_dict = torch.load(checkpoint_path)
+                pipeline.unet.load_state_dict(state_dict, strict=False)
+            else:
+                print("Invalid checkpoint path", checkpoint_path)
+
+        ### lora
+        if lora_path and osp.isfile(lora_path):
+            pipeline.load_lora_weights(osp.dirname(lora_path), weight_name=osp.basename(lora_path), adapter_name="lora")
+
+        pipeline.insert_hm_modules(dtype=dtype, device=device)
+
+        return (pipeline,)
+
+class HMFaceToolkitsLoader:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "gpu_id": ("INT", {"default": 0}),
+            }
+        }
+
+    RETURN_TYPES = ("FACE_TOOLKITS",)
+    FUNCTION = "load_face_toolkits"
+    CATEGORY = "hellomeme"
+    def load_face_toolkits(self, gpu_id):
+        dtype = torch.float16
+        device = torch.device(f'cuda:{gpu_id}')
+
+        face_aligner = HelloCameraDemo(face_alignment_module=HelloFaceAlignment(gpu_id=gpu_id), reset=True)
+        harkit_bs = HelloARKitBSPred(gpu_id=gpu_id)
+        h3dmm = Hello3DMMPred(gpu_id=gpu_id)
+        clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+            'h94/IP-Adapter',subfolder='models/image_encoder').to(dtype=dtype, device=device)
+        return (dict(face_aligner=face_aligner, harkit_bs=harkit_bs, h3dmm=h3dmm, image_encoder=clip_image_encoder), )
+
+class GetReferenceImageRT:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "face_toolkits": ("FACE_TOOLKITS",),
+                "image": ("IMAGE",),
+            }
+        }
+
+    RETURN_TYPES = ("REFRT",)
+    FUNCTION = "get_reference_image_rt"
+    CATEGORY = "hellomeme"
+
+    def get_reference_image_rt(self, face_toolkits, image):
+        image_np = cv2.cvtColor((image[0] * 255).cpu().numpy().astype(np.uint8), cv2.COLOR_BGR2RGB)
+        image_np = cv2.resize(image_np, (512, 512))
+        print(image_np.shape)
+        face_toolkits['face_aligner'].reset_track()
+        faces = face_toolkits['face_aligner'].forward(image_np)
+        assert len(faces) > 0
+        face = sorted(faces, key=lambda x: (x['face_rect'][2] - x['face_rect'][0]) * (
+                x['face_rect'][3] - x['face_rect'][1]))[-1]
+        ref_landmark = face['pre_kpt_222']
+
+        ref_rot, ref_trans = face_toolkits['h3dmm'].forward_params(image_np, ref_landmark)
+        return (dict(rot=ref_rot, trans=ref_trans), )
+
+class GetImageDriveParams:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "face_toolkits": ("FACE_TOOLKITS",),
+                "image": ("IMAGE",),
+                "ref_rt": ("REFRT",),
+                "gpu_id": ("INT", {"default": 0})
+            }
+        }
+
+    RETURN_TYPES = ("DRIVE_IMAGE_PARAMS",)
+    RETURN_NAMES = ("drive_image_params",)
+    FUNCTION = "get_face_params"
+    CATEGORY = "hellomeme"
+    def get_face_params(self, face_toolkits, image, ref_rt, gpu_id):
+        dtype = torch.float16
+        device = torch.device(f'cuda:{gpu_id}')
+
+        image_np = cv2.cvtColor((image[0] * 255).cpu().numpy().astype(np.uint8), cv2.COLOR_BGR2RGB)
+
+        face_toolkits['face_aligner'].reset_track()
+        faces = face_toolkits['face_aligner'].forward(image_np)
+        assert len(faces) > 0
+
+        face = sorted(faces, key=lambda x: (x['face_rect'][2] - x['face_rect'][0]) * (
+                x['face_rect'][3] - x['face_rect'][1]))[-1]
+        drive_landmark = face['pre_kpt_222']
+
+        drive_face_parts, drive_coeff, drive_rot, drive_trans = get_face_params(face_toolkits['h3dmm'],
+                                                                                face_toolkits['harkit_bs'],
+                                                                                [image_np],
+                                                                                [drive_landmark],
+                                                                                save_size=(512, 512),
+                                                                                align=False)
+
+        face_parts_embedding, control_heatmaps = face_params_to_tensor(face_toolkits['image_encoder'], face_toolkits['h3dmm'],
+                                                                       drive_face_parts,
+                                                                       drive_rot, drive_trans, ref_rt['trans'],
+                                                                       save_size=512, trans_ratio=0.0)
+
+        drive_params = dict(
+            face_parts=face_parts_embedding.unsqueeze(0).to(dtype=dtype, device=device),
+            drive_coeff=drive_coeff.unsqueeze(0).to(dtype=dtype, device=device),
+            condition=control_heatmaps.unsqueeze(0).to(dtype=dtype, device=device),
+        )
+        return (drive_params, )
+
+class HMPipelineImage:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "pipeline": ("HMIMAGEPIPELINE",),
+                "image": ("IMAGE",),
+                "drive_image_params": ("DRIVE_IMAGE_PARAMS",),
+                "prompt": ("STRING", {"default": '(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture'}),
+                "negative_prompt": ("STRING", {"default": ''}),
+                "steps": ("INT", {"default": 25, "min": 1, "max": 10000}),
+                "guidance_scale": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "sample"
+    CATEGORY = "hellomeme"
+
+    def sample(self, pipeline, image, drive_image_params,  prompt, negative_prompt, steps=25, guidance_scale=2.0):
+        image_np = (image[0] * 255).cpu().numpy().astype(np.uint8)
+        image_np = cv2.resize(image_np, (512, 512))
+        image_pil = Image.fromarray(image_np)
+        result_img = pipeline(
+            prompt=[prompt],
+            strength=1.0,
+            image=image_pil,
+            drive_params=copy.deepcopy(drive_image_params),
+            num_inference_steps=steps,
+            negative_prompt=[negative_prompt],
+            guidance_scale=guidance_scale,
+            output_type='np'
+        )
+        return (torch.from_numpy(np.clip(result_img[0], 0, 1)), )
+
+class GetVideoDriveParams:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "face_toolkits": ("FACE_TOOLKITS",),
+                "images": ("IMAGE",),
+                "ref_rt": ("REFRT",),
+                "trans_ratio": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+            }
+        }
+    RETURN_TYPES = ("DRIVE_VIDEO_PARAMS",)
+    RETURN_NAMES = ("drive_video_params",)
+    FUNCTION = "get_face_params"
+    CATEGORY = "hellomeme"
+    def get_face_params(self, face_toolkits, images, ref_rt, trans_ratio):
+        dtype = torch.float16
+
+        frame_list = [cv2.cvtColor((frame * 255).cpu().numpy().astype(np.uint8), cv2.COLOR_BGR2RGB) for frame in images]
+
+        face_toolkits['face_aligner'].reset_track()
+        (drive_face_parts, drive_coeff, drive_rot, drive_trans) = get_drive_params(face_toolkits['face_aligner'],
+                                                                                   face_toolkits['h3dmm'],
+                                                                                   face_toolkits['harkit_bs'],
+                                                                                   frame_list=frame_list,
+                                                                                   save_size=512,
+                                                                                   align=True)
+        face_toolkits['face_aligner'].reset_track()
+        face_parts_embedding, control_heatmaps = face_params_to_tensor(
+            face_toolkits['image_encoder'], face_toolkits['h3dmm'],
+            drive_face_parts,
+            drive_rot, drive_trans, ref_rt['trans'],
+            save_size=512, trans_ratio=trans_ratio)
+
+        drive_params = dict(
+            face_parts=face_parts_embedding.unsqueeze(0).to(dtype=dtype),
+            drive_coeff=drive_coeff.unsqueeze(0).to(dtype=dtype),
+            condition=control_heatmaps.unsqueeze(0).to(dtype=dtype),
+        )
+        return (drive_params, )
+
+class HMPipelineVideo:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {
+                        "pipeline": ("HMVIDEOPIPELINE",),
+                        "image": ("IMAGE",),
+                        "drive_video_params": ("DRIVE_VIDEO_PARAMS",),
+                        "prompt": ("STRING", {"default": '(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture'}),
+                        "negative_prompt": ("STRING", {"default": ''}),
+                        "steps": ("INT", {"default": 25, "min": 1, "max": 10000}),
+                        "guidance_scale": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+                     }
+                }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "sample"
+    CATEGORY = "hellomeme"
+
+    def sample(self, pipeline, image, drive_video_params,  prompt, negative_prompt, steps=25, guidance_scale=2.0):
+        image_np = (image[0] * 255).cpu().numpy().astype(np.uint8)
+        image_np = cv2.resize(image_np, (512, 512))
+        image_pil = Image.fromarray(image_np)
+        res_frames = pipeline(
+            prompt=[prompt],
+            strength=1.0,
+            image=image_pil,
+            drive_params=copy.deepcopy(drive_video_params),
+            num_inference_steps=steps,
+            negative_prompt=[negative_prompt],
+            guidance_scale=guidance_scale,
+            output_type='np'
+        )
+        res_frames = [np.clip(x[0], 0, 1) for x in res_frames]
+        return (torch.from_numpy(np.array(res_frames)), )
+
+NODE_CLASS_MAPPINGS = {
+    "HMImagePipelineLoader": HMImagePipelineLoader,
+    "HMVideoPipelineLoader": HMVideoPipelineLoader,
+    "HMFaceToolkitsLoader": HMFaceToolkitsLoader,
+    "GetReferenceImageRT": GetReferenceImageRT,
+    "GetImageDriveParams": GetImageDriveParams,
+    "GetVideoDriveParams": GetVideoDriveParams,
+    "HMPipelineImage": HMPipelineImage,
+    "HMPipelineVideo": HMPipelineVideo,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "HMImagePipelineLoader": "Load HelloMemeImage Pipeline",
+    "HMVideoPipelineLoader": "Load HelloMemeVideo Pipeline",
+    "HMFaceToolkitsLoader": "Load Face Toolkits",
+    "GetReferenceImageT": "Get Reference Image Translation",
+    "GetImageDriveParams": "Get Drive Image Parameters",
+    "GetVideoDriveParams": "Get Drive Video Parameters",
+    "HMPipelineImage": "HelloMeme Image Pipeline",
+    "HMPipelineVideo": "HelloMeme Video Pipeline",
+}

workflows/hellomeme_image_workflow.json

@@ -0,0 +1,403 @@
+{
+  "last_node_id": 14,
+  "last_link_id": 18,
+  "nodes": [
+    {
+      "id": 11,
+      "type": "GetReferenceImageRT",
+      "pos": {
+        "0": 621,
+        "1": 262
+      },
+      "size": {
+        "0": 241.79998779296875,
+        "1": 46
+      },
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "face_toolkits",
+          "type": "FACE_TOOLKITS",
+          "link": 13
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 12
+        }
+      ],
+      "outputs": [
+        {
+          "name": "REFRT",
+          "type": "REFRT",
+          "links": [
+            17
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "GetReferenceImageRT"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 14,
+      "type": "GetImageDriveParams",
+      "pos": {
+        "0": 987,
+        "1": 226
+      },
+      "size": {
+        "0": 330,
+        "1": 98
+      },
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "face_toolkits",
+          "type": "FACE_TOOLKITS",
+          "link": 16
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 18
+        },
+        {
+          "name": "ref_rt",
+          "type": "REFRT",
+          "link": 17
+        }
+      ],
+      "outputs": [
+        {
+          "name": "drive_image_params",
+          "type": "DRIVE_IMAGE_PARAMS",
+          "links": [
+            15
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "GetImageDriveParams"
+      },
+      "widgets_values": [
+        0
+      ]
+    },
+    {
+      "id": 6,
+      "type": "HMPipelineImage",
+      "pos": {
+        "0": 1325,
+        "1": 654
+      },
+      "size": {
+        "0": 315,
+        "1": 170
+      },
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "pipeline",
+          "type": "HMIMAGEPIPELINE",
+          "link": 7
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 9
+        },
+        {
+          "name": "drive_image_params",
+          "type": "DRIVE_IMAGE_PARAMS",
+          "link": 15
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            11
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "HMPipelineImage"
+      },
+      "widgets_values": [
+        "(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture",
+        "",
+        25,
+        2
+      ]
+    },
+    {
+      "id": 10,
+      "type": "PreviewImage",
+      "pos": {
+        "0": 1695,
+        "1": 652
+      },
+      "size": {
+        "0": 210,
+        "1": 246
+      },
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 11
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 2,
+      "type": "HMFaceToolkitsLoader",
+      "pos": {
+        "0": 211,
+        "1": 180
+      },
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "FACE_TOOLKITS",
+          "type": "FACE_TOOLKITS",
+          "links": [
+            13,
+            16
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "HMFaceToolkitsLoader"
+      },
+      "widgets_values": [
+        0
+      ]
+    },
+    {
+      "id": 1,
+      "type": "HMImagePipelineLoader",
+      "pos": {
+        "0": 226,
+        "1": 356
+      },
+      "size": {
+        "0": 315,
+        "1": 106
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "HMIMAGEPIPELINE",
+          "type": "HMIMAGEPIPELINE",
+          "links": [
+            7
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "HMImagePipelineLoader"
+      },
+      "widgets_values": [
+        "None",
+        "None",
+        0
+      ]
+    },
+    {
+      "id": 8,
+      "type": "LoadImage",
+      "pos": {
+        "0": 233,
+        "1": 536
+      },
+      "size": {
+        "0": 315,
+        "1": 314
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            9,
+            12
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "majicmix2.jpg",
+        "image"
+      ]
+    },
+    {
+      "id": 9,
+      "type": "LoadImage",
+      "pos": {
+        "0": 227,
+        "1": 915
+      },
+      "size": {
+        "0": 315,
+        "1": 314
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            18
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "yao.jpg",
+        "image"
+      ]
+    }
+  ],
+  "links": [
+    [
+      7,
+      1,
+      0,
+      6,
+      0,
+      "HMIMAGEPIPELINE"
+    ],
+    [
+      9,
+      8,
+      0,
+      6,
+      1,
+      "IMAGE"
+    ],
+    [
+      11,
+      6,
+      0,
+      10,
+      0,
+      "IMAGE"
+    ],
+    [
+      12,
+      8,
+      0,
+      11,
+      1,
+      "IMAGE"
+    ],
+    [
+      13,
+      2,
+      0,
+      11,
+      0,
+      "FACE_TOOLKITS"
+    ],
+    [
+      15,
+      14,
+      0,
+      6,
+      2,
+      "DRIVE_PARAMS"
+    ],
+    [
+      16,
+      2,
+      0,
+      14,
+      0,
+      "FACE_TOOLKITS"
+    ],
+    [
+      17,
+      11,
+      0,
+      14,
+      2,
+      "REFRT"
+    ],
+    [
+      18,
+      9,
+      0,
+      14,
+      1,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        150,
+        -42.66668701171875
+      ]
+    }
+  },
+  "version": 0.4
+}

workflows/hellomeme_video_workflow.json

@@ -0,0 +1,502 @@
+{
+  "last_node_id": 24,
+  "last_link_id": 30,
+  "nodes": [
+    {
+      "id": 15,
+      "type": "HMPipelineVideo",
+      "pos": {
+        "0": 1408,
+        "1": 390
+      },
+      "size": {
+        "0": 315,
+        "1": 170
+      },
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "pipeline",
+          "type": "HMVIDEOPIPELINE",
+          "link": 16
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 17
+        },
+        {
+          "name": "drive_video_params",
+          "type": "DRIVE_VIDEO_PARAMS",
+          "link": 29
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            24
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "HMPipelineVideo"
+      },
+      "widgets_values": [
+        "(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture",
+        "",
+        25,
+        2
+      ]
+    },
+    {
+      "id": 2,
+      "type": "HMFaceToolkitsLoader",
+      "pos": {
+        "0": 124,
+        "1": 276
+      },
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "FACE_TOOLKITS",
+          "type": "FACE_TOOLKITS",
+          "links": [
+            13,
+            27
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "HMFaceToolkitsLoader"
+      },
+      "widgets_values": [
+        0
+      ]
+    },
+    {
+      "id": 11,
+      "type": "GetReferenceImageRT",
+      "pos": {
+        "0": 536,
+        "1": 516
+      },
+      "size": {
+        "0": 241.79998779296875,
+        "1": 46
+      },
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "face_toolkits",
+          "type": "FACE_TOOLKITS",
+          "link": 13
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 12
+        }
+      ],
+      "outputs": [
+        {
+          "name": "REFRT",
+          "type": "REFRT",
+          "links": [
+            26
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "GetReferenceImageRT"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 21,
+      "type": "VHS_VideoCombine",
+      "pos": {
+        "0": 1927,
+        "1": 604
+      },
+      "size": [
+        214.7587890625,
+        518.7587890625
+      ],
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 24
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 25,
+          "shape": 7
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null,
+          "shape": 7
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null,
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": []
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 15,
+        "loop_count": 0,
+        "filename_prefix": "HelloMeme",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 19,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": true,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "HelloMeme_00001-audio.mp4",
+            "subfolder": "",
+            "type": "output",
+            "format": "video/h264-mp4",
+            "frame_rate": 15
+          },
+          "muted": false
+        }
+      }
+    },
+    {
+      "id": 22,
+      "type": "GetVideoDriveParams",
+      "pos": {
+        "0": 932,
+        "1": 503
+      },
+      "size": {
+        "0": 330,
+        "1": 98
+      },
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "face_toolkits",
+          "type": "FACE_TOOLKITS",
+          "link": 27
+        },
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 28
+        },
+        {
+          "name": "ref_rt",
+          "type": "REFRT",
+          "link": 26
+        }
+      ],
+      "outputs": [
+        {
+          "name": "drive_video_params",
+          "type": "DRIVE_VIDEO_PARAMS",
+          "links": [
+            29
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "GetVideoDriveParams"
+      },
+      "widgets_values": [
+        0
+      ]
+    },
+    {
+      "id": 14,
+      "type": "HMVideoPipelineLoader",
+      "pos": {
+        "0": 116,
+        "1": 105
+      },
+      "size": {
+        "0": 352.79998779296875,
+        "1": 106
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "HMVIDEOPIPELINE",
+          "type": "HMVIDEOPIPELINE",
+          "links": [
+            16
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "HMVideoPipelineLoader"
+      },
+      "widgets_values": [
+        "None",
+        "None",
+        0
+      ]
+    },
+    {
+      "id": 8,
+      "type": "LoadImage",
+      "pos": {
+        "0": 120,
+        "1": 414
+      },
+      "size": {
+        "0": 315,
+        "1": 314
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            12,
+            17
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "majicmix2.jpg",
+        "image"
+      ]
+    },
+    {
+      "id": 20,
+      "type": "VHS_LoadVideoFFmpeg",
+      "pos": {
+        "0": 139,
+        "1": 793
+      },
+      "size": [
+        303.46875,
+        535.46875
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null,
+          "shape": 7
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null,
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            28
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "mask",
+          "type": "MASK",
+          "links": null
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            25
+          ],
+          "slot_index": 2
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadVideoFFmpeg"
+      },
+      "widgets_values": {
+        "video": "jgz.mp4",
+        "force_rate": 15,
+        "force_size": "Disabled",
+        "custom_width": 512,
+        "custom_height": 512,
+        "frame_load_cap": 0,
+        "start_time": 0,
+        "choose video to upload": "image",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "force_rate": 15,
+            "frame_load_cap": 0,
+            "start_time": 0,
+            "filename": "jgz.mp4",
+            "type": "input",
+            "format": "video/mp4"
+          },
+          "muted": false
+        }
+      }
+    }
+  ],
+  "links": [
+    [
+      12,
+      8,
+      0,
+      11,
+      1,
+      "IMAGE"
+    ],
+    [
+      13,
+      2,
+      0,
+      11,
+      0,
+      "FACE_TOOLKITS"
+    ],
+    [
+      16,
+      14,
+      0,
+      15,
+      0,
+      "HMVIDEOPIPELINE"
+    ],
+    [
+      17,
+      8,
+      0,
+      15,
+      1,
+      "IMAGE"
+    ],
+    [
+      24,
+      15,
+      0,
+      21,
+      0,
+      "IMAGE"
+    ],
+    [
+      25,
+      20,
+      2,
+      21,
+      1,
+      "AUDIO"
+    ],
+    [
+      26,
+      11,
+      0,
+      22,
+      2,
+      "REFRT"
+    ],
+    [
+      27,
+      2,
+      0,
+      22,
+      0,
+      "FACE_TOOLKITS"
+    ],
+    [
+      28,
+      20,
+      0,
+      22,
+      1,
+      "IMAGE"
+    ],
+    [
+      29,
+      22,
+      0,
+      15,
+      2,
+      "DRIVE_PARAMS"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        0,
+        0
+      ]
+    }
+  },
+  "version": 0.4
+}