hm3:merge control modules

This commit is contained in:
songkey-win
2025-02-14 18:22:28 +08:00
parent 3b54076292
commit 136fab456f
9 changed files with 61 additions and 164 deletions
+1 -1
View File
@@ -9,7 +9,7 @@
"""
from .hm_denoising_motion import HMDenoisingMotion
from .hm_control import HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2, HMV3ControlNet2, HMV3ControlNet
from .hm_control import HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2, HMV3ControlNet
from .hm_blocks import HMReferenceAdapter, HM3ReferenceAdapter, HM3MotionAdapter, HMPipeline
from .hm_denoising_3d import HMDenoising3D
from .hm3_denoising_3d import HM3Denoising3D
+4 -8
View File
@@ -190,10 +190,8 @@ class HM3Denoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter):
sample += down_intrablock_additional_residuals.pop(0)
res_cache[f"down_{idx}"] = sample.clone()
if not control_hidden_states is None and f'down3_v1_{idx}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'down3_v1_{idx}'], "b c f h w -> (b f) c h w")
if not control_hidden_states is None and f'down3_v2_{idx}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'down3_v2_{idx}'], "b c f h w -> (b f) c h w")
if not control_hidden_states is None and f'down3_{idx}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'down3_{idx}'], "b c f h w -> (b f) c h w")
if hasattr(self, 'motion_down') and use_motion:
sample = self.motion_down[idx](sample, motion_pad_hidden_states[f'down_{idx}'], emb, num_frames)
@@ -248,10 +246,8 @@ class HM3Denoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter):
upsample_size = down_block_res_samples[-1].shape[2:]
res_cache[f"up_{i}"] = sample.clone()
if not control_hidden_states is None and f'up3_v1_{i}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'up3_v1_{i}'], "b c f h w -> (b f) c h w")
if not control_hidden_states is None and f'up3_v2_{i}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'up3_v2_{i}'], "b c f h w -> (b f) c h w")
if not control_hidden_states is None and f'up3_{i}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'up3_{i}'], "b c f h w -> (b f) c h w")
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](sample, reference_hidden_states[f'up_{i}'], num_frames)
if hasattr(self, 'motion_up') and use_motion:
+4 -8
View File
@@ -144,10 +144,8 @@ class HM3DenoisingMotion(UNetMotionModel, InsertReferenceAdapter):
else:
sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames)
if not control_hidden_states is None and f'down3_v1_{idx}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'down3_v1_{idx}'], "b c f h w -> (b f) c h w")
if not control_hidden_states is None and f'down3_v2_{idx}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'down3_v2_{idx}'], "b c f h w -> (b f) c h w")
if not control_hidden_states is None and f'down3_{idx}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'down3_{idx}'], "b c f h w -> (b f) c h w")
if hasattr(self, 'motion_down') and use_motion:
sample = self.motion_down[idx](sample, motion_pad_hidden_states[f'down_{idx}'], emb, num_frames)
@@ -200,10 +198,8 @@ class HM3DenoisingMotion(UNetMotionModel, InsertReferenceAdapter):
if not is_final_block and forward_upsample_size:
upsample_size = down_block_res_samples[-1].shape[2:]
if not control_hidden_states is None and f'up3_v1_{i}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'up3_v1_{i}'], "b c f h w -> (b f) c h w")
if not control_hidden_states is None and f'up3_v2_{i}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'up3_v2_{i}'], "b c f h w -> (b f) c h w")
if not control_hidden_states is None and f'up3_{i}' in control_hidden_states:
sample += rearrange(control_hidden_states[f'up3_{i}'], "b c f h w -> (b f) c h w")
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](sample, reference_hidden_states[f'up_{i}'], num_frames)
if hasattr(self, 'motion_up') and use_motion:
+32 -94
View File
@@ -227,13 +227,16 @@ class HMV3ControlNet(ModelMixin, ConfigMixin):
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.conv_up2_down0 = nn.Conv2d(1280, 320, kernel_size=1, padding=0, bias=False)
self.conv_up1_down1 = nn.Conv2d(1280, 640, kernel_size=1, padding=0, bias=False)
self.emo_embedding = Timesteps(cross_attention_dim, True, 0)
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.emo_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim)
self.blocks_down = nn.ModuleList([])
for i in range(1, len(block_out_channels)):
channel_in = block_out_channels[i-1]
@@ -244,11 +247,11 @@ class HMV3ControlNet(ModelMixin, ConfigMixin):
channel_out=channel_out,
cross_attention_dim=cross_attention_dim,
num_positional_embeddings=64,
num_positional_embeddings_hidden=64
num_positional_embeddings_hidden=1024
)
)
def forward(self, condition, drive_coeff, face_parts):
def forward(self, condition, drive_coeff=None, face_parts=None, emo_embedding=None):
bs, _, video_length, h, w = condition.shape
conditioning = rearrange(condition, "b c f h w -> (b f) c h w")
@@ -258,39 +261,47 @@ class HMV3ControlNet(ModelMixin, ConfigMixin):
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)
if drive_coeff is None or face_parts is None:
emo_coeff = rearrange(emo_embedding * 20., "b f c -> (b f c)")
emo_embedding = self.emo_embedding(emo_coeff).to(dtype=embedding.dtype)
emo_embedding = rearrange(emo_embedding, "(b f c) d -> b f c d",
b=bs, f=video_length, d=self.cross_attention_dim)
emo_embedding = rearrange(emo_embedding, "b f c d -> (b f) c d")
drive_embedding = self.emo_proj(emo_embedding)
else:
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_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)
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 = 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)
drive_embedding = self.exp_proj(drive_embedding)
ret_dict = {}
for idx, block in enumerate(self.blocks_down):
embedding = block(embedding, drive_embedding)
up_idx = len(self.blocks_down)-idx-1
ret_dict[f'up3_v1_{up_idx}'] = rearrange(embedding,
ret_dict[f'up3_{up_idx}'] = rearrange(embedding,
"(b f) c h w -> b c f h w",
f=video_length)
if up_idx == 2:
ret_dict[f'down3_v1_0'] = rearrange(self.conv_up2_down0(embedding),
ret_dict[f'down3_0'] = rearrange(self.conv_up2_down0(embedding),
"(b f) c h w -> b c f h w",
f=video_length)
if up_idx == 1:
ret_dict[f'down3_v1_1'] = rearrange(self.conv_up1_down1(embedding),
ret_dict[f'down3_1'] = rearrange(self.conv_up1_down1(embedding),
"(b f) c h w -> b c f h w",
f=video_length)
if up_idx == 0:
ret_dict[f'down3_v1_2'] = ret_dict[f'up3_v1_0']
ret_dict[f'down3_v1_3'] = ret_dict[f'up3_v1_0']
ret_dict[f'down3_2'] = ret_dict[f'up3_0']
ret_dict[f'down3_3'] = ret_dict[f'up3_0']
return ret_dict
@@ -350,77 +361,4 @@ class HMV2ControlNet2(ModelMixin, ConfigMixin):
ret_dict[f'up2_v2_{len(self.blocks_down)-idx-1}'] = rearrange(embedding,
"(b f) c h w -> b c f h w",
f=video_length)
return ret_dict
class HMV3ControlNet2(ModelMixin, ConfigMixin):
def __init__(
self,
embedding_channels: int = 1280,
input_channels: int = 3,
scale_factor: int = 4,
cross_attention_dim: int = 320,
block_out_channels: Tuple[int] = (128, 640, 1280, 1280, 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.conv_up2_down0 = nn.Conv2d(1280, 320, kernel_size=1, padding=0, bias=False)
self.conv_up1_down1 = nn.Conv2d(1280, 640, kernel_size=1, padding=0, bias=False)
self.exp_embedding = Timesteps(cross_attention_dim, True, 0)
self.emo_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=1024
)
)
def forward(self, condition, emo_embedding):
bs, _, video_length, h, w = condition.shape
conditioning = rearrange(condition, "b c f h w -> (b f) c h w")
conditioning = F.interpolate(conditioning,
size=(h // self.scale_factor, w // self.scale_factor),
mode='bilinear',
align_corners=False)
embedding = self.conv_in(conditioning)
drive_coeff = rearrange(emo_embedding * 20., "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)
drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d")
drive_embedding = self.emo_proj(drive_embedding)
ret_dict = {}
for idx, block in enumerate(self.blocks_down):
embedding = block(embedding, drive_embedding)
up_idx = len(self.blocks_down)-idx-1
ret_dict[f'up3_v2_{up_idx}'] = rearrange(embedding,
"(b f) c h w -> b c f h w",
f=video_length)
if up_idx == 2:
ret_dict[f'down3_v2_0'] = rearrange(self.conv_up2_down0(embedding),
"(b f) c h w -> b c f h w",
f=video_length)
if up_idx == 1:
ret_dict[f'down3_v2_1'] = rearrange(self.conv_up1_down1(embedding),
"(b f) c h w -> b c f h w",
f=video_length)
if up_idx == 0:
ret_dict[f'down3_v2_2'] = ret_dict[f'up3_v2_0']
ret_dict[f'down3_v2_3'] = ret_dict[f'up3_v2_0']
return ret_dict
+8 -27
View File
@@ -21,7 +21,7 @@ from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusion
from diffusers import DPMSolverMultistepScheduler
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents
from ..models import HM3Denoising3D, HMV3ControlNet, HMV3ControlNet2, HMPipeline, HM3ReferenceAdapter
from ..models import HM3Denoising3D, HMV3ControlNet, HMPipeline, HM3ReferenceAdapter
class HM3ImagePipeline(HMPipeline):
@@ -42,18 +42,17 @@ class HM3ImagePipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
def insert_hm_modules(self, version='v3', dtype=torch.float16, modelscope=False):
if modelscope:
from modelscope import snapshot_download
hm_reference_dir = snapshot_download('songkey/hm3_reference')
hm_control_dir = snapshot_download('songkey/hm3_control')
hm_control2_dir = snapshot_download('songkey/hm3_control2')
hm_control_dir = snapshot_download('songkey/hm3_control_mix')
else:
hm_reference_dir = 'songkey/hm3_reference'
hm_control_dir = 'songkey/hm3_control'
hm_control2_dir = 'songkey/hm3_control2'
hm_control_dir = 'songkey/hm3_control_mix'
if isinstance(self.unet, HM3Denoising3D):
hm_adapter = HM3ReferenceAdapter.from_pretrained(hm_reference_dir)
@@ -69,11 +68,6 @@ class HM3ImagePipeline(HMPipeline):
self.mp_control = HMV3ControlNet.from_pretrained(hm_control_dir)
self.mp_control.to(device='cpu', dtype=dtype).eval()
if hasattr(self, "mp_control2"):
del self.mp_control2
self.mp_control2 = HMV3ControlNet2.from_pretrained(hm_control2_dir)
self.mp_control2.to(device='cpu', dtype=dtype).eval()
self.vae.to(device='cpu', dtype=dtype).eval()
self.vae_decode.to(device='cpu', dtype=dtype).eval()
self.text_encoder.to(device='cpu', dtype=dtype).eval()
@@ -238,8 +232,8 @@ class HM3ImagePipeline(HMPipeline):
condition = torch.cat([torch.ones_like(condition) * -1, condition], dim=0)
control_latents = {}
self.mp_control.to(device=device)
if 'drive_coeff' in drive_params:
self.mp_control.to(device=device)
drive_coeff = drive_params['drive_coeff'].clone().to(device=device)
face_parts = drive_params['face_parts'].clone().to(device=device)
if self.do_classifier_free_guidance:
@@ -247,26 +241,13 @@ class HM3ImagePipeline(HMPipeline):
face_parts = torch.cat([torch.zeros_like(face_parts), face_parts], dim=0)
control_latents1 = self.mp_control(condition=condition, drive_coeff=drive_coeff, face_parts=face_parts)
control_latents.update(control_latents1)
self.mp_control.cpu()
# debug_dict = torch.load(r'C:\Users\songkey-win\Downloads\debug\debug_dict.pth')
# # latents_cntrl = debug_dict['latents_cntrl']
# ref_latents = debug_dict['ref_latents']
# ref_latents_neg = debug_dict['ref_latents_neg']
# # cached_res = debug_dict['cached_res']
# tgt_skl_tensor = debug_dict['tgt_skl_tensor']
# emo_embedding = debug_dict['emo_embedding']
if 'pd_fpg' in drive_params:
self.mp_control2.to(device=device)
elif 'pd_fpg' in drive_params:
pd_fpg = drive_params['pd_fpg'].clone().to(device=device)
# pd_fpg = emo_embedding[:,:1]
# condition = tgt_skl_tensor[:,:,:1]
if self.do_classifier_free_guidance:
pd_fpg = torch.cat([torch.zeros_like(pd_fpg), pd_fpg], dim=0)
control_latents2 = self.mp_control2(condition=condition, emo_embedding=pd_fpg)
control_latents2 = self.mp_control(condition=condition, emo_embedding=pd_fpg)
control_latents.update(control_latents2)
self.mp_control2.cpu()
self.mp_control.cpu()
# 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+9 -25
View File
@@ -12,7 +12,6 @@ adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/p
import copy
from typing import Any, Callable, Dict, List, Optional, Union
import cv2
import torch
from einops import rearrange
@@ -24,9 +23,8 @@ from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img impo
from diffusers import MotionAdapter, DPMSolverMultistepScheduler
from diffusers.utils.torch_utils import randn_tensor
from ..models import (HM3Denoising3D, HM3DenoisingMotion, HMV3ControlNet2, HM3MotionAdapter,
from ..models import (HM3Denoising3D, HM3DenoisingMotion, HM3MotionAdapter,
HMV3ControlNet, HM3ReferenceAdapter, HMPipeline)
from ..utils import np_rgb_to_tensor, tensor_to_np_rgb
class HM3VideoPipeline(HMPipeline):
def caryomitosis(self, version='v3', modelscope=False, **kwargs):
@@ -57,22 +55,23 @@ class HM3VideoPipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
def insert_hm_modules(self, version='v3', dtype=torch.float16, modelscope=False):
if modelscope:
from modelscope import snapshot_download
hm_reference_dir = snapshot_download('songkey/hm3_reference')
hm_control_dir = snapshot_download('songkey/hm3_control')
hm_control2_dir = snapshot_download('songkey/hm3_control2')
hm_control_dir = snapshot_download('songkey/hm3_control_mix')
hm_motion_dir = snapshot_download('songkey/hm3_motion')
else:
hm_reference_dir = 'songkey/hm3_reference'
hm_control_dir = 'songkey/hm3_control'
hm_control2_dir = 'songkey/hm3_control2'
hm_control_dir = 'songkey/hm3_control_mix'
hm_motion_dir = 'songkey/hm3_motion'
hm_adapter = HM3ReferenceAdapter.from_pretrained(hm_reference_dir)
motion_adapter = HM3MotionAdapter.from_pretrained(hm_motion_dir)
motion_adapter = motion_adapter.to(device='cpu', dtype=dtype).eval()
motion_adapter.push_to_hub('songkey/hm3_motion')
if isinstance(self.unet, HM3DenoisingMotion):
self.unet.insert_reference_adapter(hm_adapter)
@@ -93,12 +92,6 @@ class HM3VideoPipeline(HMPipeline):
self.mp_control = HMV3ControlNet.from_pretrained(hm_control_dir)
self.mp_control.to(device='cpu', dtype=dtype).eval()
if hasattr(self, "mp_control2"):
del self.mp_control2
self.mp_control2 = HMV3ControlNet2.from_pretrained(hm_control2_dir)
self.mp_control2.to(device='cpu', dtype=dtype).eval()
self.vae.to(device='cpu', dtype=dtype).eval()
self.vae_decode.to(device='cpu', dtype=dtype).eval()
self.text_encoder.to(device='cpu', dtype=dtype).eval()
@@ -284,12 +277,6 @@ class HM3VideoPipeline(HMPipeline):
**extra_step_kwargs, return_dict=False)[0]
self.unet.cpu()
# images = self.vae_decode_latents(self.vae_decode, latents / self.vae_decode.config.scaling_factor, device)
# for idx, image in enumerate(images):
# image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# cv2.imwrite(f"C:/Users/songkey-win/Downloads/debug0203/output_d{depth:04d}_{idx:04d}.png", image)
if depth == 0:
return latents[:,:,:num_indexes]
else:
@@ -470,7 +457,6 @@ class HM3VideoPipeline(HMPipeline):
control_latents = []
control_latent1_neg, control_latent2_neg = None, None
self.mp_control.to(device=device)
self.mp_control2.to(device=device)
with self.progress_bar(total=raw_video_len) as progress_bar:
for idx in range(raw_video_len):
progress_bar.set_description(f"GEN stage1")
@@ -495,22 +481,20 @@ class HM3VideoPipeline(HMPipeline):
control_latent1 = cat_dicts([control_latent1_neg, control_latent1], dim=0)
control_latent_input.update(control_latent1)
if 'pd_fpg' in drive_params:
else:
pd_fpg = drive_params['pd_fpg'][:, idx:idx+1].clone().to(device=device)
if control_latent2_neg is None and self.do_classifier_free_guidance:
control_latent2_neg = self.mp_control2(condition=torch.ones_like(condition)*-1,
control_latent2_neg = self.mp_control(condition=torch.ones_like(condition)*-1,
emo_embedding=torch.zeros_like(pd_fpg))
control_latent2 = self.mp_control2(condition=condition, emo_embedding=pd_fpg)
control_latent2 = self.mp_control(condition=condition, emo_embedding=pd_fpg)
if self.do_classifier_free_guidance:
control_latent2 = cat_dicts([control_latent2_neg, control_latent2], dim=0)
control_latent_input.update(control_latent2)
control_latents.append(dicts_to_device([control_latent_input], device='cpu')[0])
self.mp_control2.cpu()
self.mp_control.cpu()
more_params = dict(
+1
View File
@@ -42,6 +42,7 @@ class HMImagePipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
def insert_hm_modules(self, version, dtype, modelscope=False):
if modelscope:
+1
View File
@@ -60,6 +60,7 @@ class HMVideoPipeline(HMPipeline):
self.vae_decode = copy.deepcopy(self.vae)
self.text_encoder.cpu()
self.text_encoder_ref = copy.deepcopy(self.text_encoder)
self.safety_checker.cpu()
def insert_hm_modules(self, version, dtype, modelscope=False):
if modelscope:
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "comfyui_hellomeme"
description = "This repository is the official implementation of the [a/HelloMeme](https://arxiv.org/pdf/2410.22901) ComfyUI interface"
version = "1.1.1"
version = "1.1.2"
license = {file = "LICENSE"}
[project.urls]