diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..b931f2e
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,8 @@
+*.pyc
+*.pyd
+*.pth
+*.pkl
+*_fps15.mp4
+.idea/
+data/results
+pretrained_models/
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..e66f6b0
--- /dev/null
+++ b/README.md
@@ -0,0 +1,37 @@
+HelloMeme: Integrating Spatial Knitting Attentions to Embed High-Level and Fidelity-Rich Conditions in Diffusion Models
+
+
+
+
+ HelloVision | HelloGroup Inc.
+
+
+
+ * Intern
+
+
+
+
+
+## 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
+
+
+ 
+
+
+### Video Generation Interface
+
+
+ 
+
+
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..e315d4e
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,3 @@
+from .meme import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
\ No newline at end of file
diff --git a/examples/jgz.mp4 b/examples/jgz.mp4
new file mode 100644
index 0000000..3b15ab8
Binary files /dev/null and b/examples/jgz.mp4 differ
diff --git a/examples/majicmix2.jpg b/examples/majicmix2.jpg
new file mode 100644
index 0000000..cd46460
Binary files /dev/null and b/examples/majicmix2.jpg differ
diff --git a/examples/yao.jpg b/examples/yao.jpg
new file mode 100644
index 0000000..491c852
Binary files /dev/null and b/examples/yao.jpg differ
diff --git a/hellomeme/__init__.py b/hellomeme/__init__.py
new file mode 100644
index 0000000..632a0f7
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/hellomeme/models/__init__.py b/hellomeme/models/__init__.py
new file mode 100644
index 0000000..45e2de5
--- /dev/null
+++ b/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
diff --git a/hellomeme/models/hm_blocks.py b/hellomeme/models/hm_blocks.py
new file mode 100644
index 0000000..31e2175
--- /dev/null
+++ b/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)
+
diff --git a/hellomeme/models/hm_control.py b/hellomeme/models/hm_control.py
new file mode 100644
index 0000000..d077214
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/hellomeme/models/hm_denoising_3d.py b/hellomeme/models/hm_denoising_3d.py
new file mode 100644
index 0000000..e989d81
--- /dev/null
+++ b/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)
diff --git a/hellomeme/models/hm_denoising_motion.py b/hellomeme/models/hm_denoising_motion.py
new file mode 100644
index 0000000..49451af
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/hellomeme/piplines/__init__.py b/hellomeme/piplines/__init__.py
new file mode 100644
index 0000000..a5d36f8
--- /dev/null
+++ b/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
diff --git a/hellomeme/piplines/pipline_hm_image.py b/hellomeme/piplines/pipline_hm_image.py
new file mode 100644
index 0000000..062a524
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/hellomeme/piplines/pipline_hm_video.py b/hellomeme/piplines/pipline_hm_video.py
new file mode 100644
index 0000000..344932c
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/hellomeme/tools/__init__.py b/hellomeme/tools/__init__.py
new file mode 100644
index 0000000..d0ba181
--- /dev/null
+++ b/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
diff --git a/hellomeme/tools/hello_3dmm.py b/hellomeme/tools/hello_3dmm.py
new file mode 100644
index 0000000..4be6d8a
--- /dev/null
+++ b/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]
+
diff --git a/hellomeme/tools/hello_arkit.py b/hellomeme/tools/hello_arkit.py
new file mode 100644
index 0000000..2c8b2f8
--- /dev/null
+++ b/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
diff --git a/hellomeme/tools/hello_camera_demo.py b/hellomeme/tools/hello_camera_demo.py
new file mode 100644
index 0000000..105d3ae
--- /dev/null
+++ b/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
+
diff --git a/hellomeme/tools/hello_face_alignment.py b/hellomeme/tools/hello_face_alignment.py
new file mode 100644
index 0000000..bd9b527
--- /dev/null
+++ b/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
diff --git a/hellomeme/tools/hello_face_det.py b/hellomeme/tools/hello_face_det.py
new file mode 100644
index 0000000..f11f708
--- /dev/null
+++ b/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
diff --git a/hellomeme/tools/utils.py b/hellomeme/tools/utils.py
new file mode 100644
index 0000000..4281d9d
--- /dev/null
+++ b/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
diff --git a/hellomeme/utils.py b/hellomeme/utils.py
new file mode 100644
index 0000000..20f3ca8
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/meme.py b/meme.py
new file mode 100644
index 0000000..ac276cc
--- /dev/null
+++ b/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",
+}
\ No newline at end of file
diff --git a/workflows/hellomeme_image_example.jpg b/workflows/hellomeme_image_example.jpg
new file mode 100644
index 0000000..75ef420
Binary files /dev/null and b/workflows/hellomeme_image_example.jpg differ
diff --git a/workflows/hellomeme_image_workflow.json b/workflows/hellomeme_image_workflow.json
new file mode 100644
index 0000000..eb7af40
--- /dev/null
+++ b/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
+}
\ No newline at end of file
diff --git a/workflows/hellomeme_video_example.jpg b/workflows/hellomeme_video_example.jpg
new file mode 100644
index 0000000..932a8fc
Binary files /dev/null and b/workflows/hellomeme_video_example.jpg differ
diff --git a/workflows/hellomeme_video_workflow.json b/workflows/hellomeme_video_workflow.json
new file mode 100644
index 0000000..748362b
--- /dev/null
+++ b/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
+}
\ No newline at end of file
+ 
+ 
+