From 136fab456f51df1038a071e53fd9a0b70e7088b1 Mon Sep 17 00:00:00 2001 From: songkey-win Date: Fri, 14 Feb 2025 18:22:28 +0800 Subject: [PATCH] hm3:merge control modules --- hellomeme/models/__init__.py | 2 +- hellomeme/models/hm3_denoising_3d.py | 12 +-- hellomeme/models/hm3_denoising_motion.py | 12 +-- hellomeme/models/hm_control.py | 126 ++++++----------------- hellomeme/pipelines/pipline_hm3_image.py | 35 ++----- hellomeme/pipelines/pipline_hm3_video.py | 34 ++---- hellomeme/pipelines/pipline_hm_image.py | 1 + hellomeme/pipelines/pipline_hm_video.py | 1 + pyproject.toml | 2 +- 9 files changed, 61 insertions(+), 164 deletions(-) diff --git a/hellomeme/models/__init__.py b/hellomeme/models/__init__.py index f225943..608997c 100644 --- a/hellomeme/models/__init__.py +++ b/hellomeme/models/__init__.py @@ -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 diff --git a/hellomeme/models/hm3_denoising_3d.py b/hellomeme/models/hm3_denoising_3d.py index e5120d5..e15d8cf 100644 --- a/hellomeme/models/hm3_denoising_3d.py +++ b/hellomeme/models/hm3_denoising_3d.py @@ -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: diff --git a/hellomeme/models/hm3_denoising_motion.py b/hellomeme/models/hm3_denoising_motion.py index 90f099f..b3ef351 100644 --- a/hellomeme/models/hm3_denoising_motion.py +++ b/hellomeme/models/hm3_denoising_motion.py @@ -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: diff --git a/hellomeme/models/hm_control.py b/hellomeme/models/hm_control.py index 45c3409..4259524 100644 --- a/hellomeme/models/hm_control.py +++ b/hellomeme/models/hm_control.py @@ -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 \ No newline at end of file diff --git a/hellomeme/pipelines/pipline_hm3_image.py b/hellomeme/pipelines/pipline_hm3_image.py index d7efa08..0c62a5f 100644 --- a/hellomeme/pipelines/pipline_hm3_image.py +++ b/hellomeme/pipelines/pipline_hm3_image.py @@ -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) diff --git a/hellomeme/pipelines/pipline_hm3_video.py b/hellomeme/pipelines/pipline_hm3_video.py index 5100ff5..5adc64f 100644 --- a/hellomeme/pipelines/pipline_hm3_video.py +++ b/hellomeme/pipelines/pipline_hm3_video.py @@ -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( diff --git a/hellomeme/pipelines/pipline_hm_image.py b/hellomeme/pipelines/pipline_hm_image.py index e86a8b5..146c38f 100644 --- a/hellomeme/pipelines/pipline_hm_image.py +++ b/hellomeme/pipelines/pipline_hm_image.py @@ -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: diff --git a/hellomeme/pipelines/pipline_hm_video.py b/hellomeme/pipelines/pipline_hm_video.py index 4c0e10a..35d9177 100644 --- a/hellomeme/pipelines/pipline_hm_video.py +++ b/hellomeme/pipelines/pipline_hm_video.py @@ -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: diff --git a/pyproject.toml b/pyproject.toml index 42d9123..f49a774 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_hellomeme" description = "This repository is the official implementation of the [a/HelloMeme](https://arxiv.org/pdf/2410.22901) ComfyUI interface" -version = "1.1.1" +version = "1.1.2" license = {file = "LICENSE"} [project.urls]