hm3:merge control modules
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user