diff --git a/hellomeme/pipelines/pipline_hm_video.py b/hellomeme/pipelines/pipline_hm_video.py index d6dc9f1..09d1ee6 100644 --- a/hellomeme/pipelines/pipline_hm_video.py +++ b/hellomeme/pipelines/pipline_hm_video.py @@ -12,6 +12,7 @@ adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/p import copy from typing import Any, Callable, Dict, List, Optional, Union import torch +from einops import rearrange from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback from diffusers.image_processor import PipelineImageInput @@ -23,7 +24,6 @@ from diffusers import StableDiffusionImg2ImgPipeline, MotionAdapter, EulerDiscre from ..models import HMDenoising3D, HMDenoisingMotion, HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2 from ..models import HMReferenceAdapter -from ..utils import dicts_to_device, cat_dicts, cat_dicts_ref class HMVideoPipeline(StableDiffusionImg2ImgPipeline): def caryomitosis(self, version, **kwargs): @@ -446,3 +446,32 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): progress_bar.update() self.vae_decode.cpu() return res_frames, latents_res + + +def merge_dicts(dictl, dictr, wl=0.5): + res = {} + for k in dictl.keys(): + res[k] = dictl[k] * wl + dictr[k] * (1-wl) + return res + + +def cat_dicts(dicts, dim=0): + res = {} + for k in dicts[0].keys(): + res[k] = torch.cat([d[k].clone() for d in dicts], dim=dim) + return res + +def cat_dicts_ref(dicts): + res = {} + for k in dicts[0].keys(): + res[k] = rearrange(torch.cat([d[k].clone().unsqueeze(1) for d in dicts], dim=1), "b f c h w -> (b f) c h w ") + return res + +def dicts_to_device(dicts, device): + ret = [] + for d in dicts: + tmpd = {} + for k in d.keys(): + tmpd[k] = d[k].clone().to(device) + ret.append(tmpd) + return ret \ No newline at end of file diff --git a/hellomeme/utils.py b/hellomeme/utils.py index ec06703..91d2e6f 100644 --- a/hellomeme/utils.py +++ b/hellomeme/utils.py @@ -10,6 +10,7 @@ import os import cv2 +import os.path as osp import numpy as np from tqdm import tqdm import torch @@ -20,40 +21,88 @@ from .tools.utils import transform_points from .tools.hello_3dmm import get_project_points_rect from safetensors import safe_open +from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (convert_ldm_unet_checkpoint, + convert_ldm_vae_checkpoint) +from .tools import Hello3DMMPred, HelloARKitBSPred, HelloFaceAlignment, HelloCameraDemo, FanEncoder +from transformers import CLIPVisionModelWithProjection def get_torch_device(gpu_id=-1): if gpu_id < 0: return torch.device("cpu") gpu_id = min(max(-1, gpu_id), torch.cuda.device_count()) return torch.device(f"cuda:{gpu_id}") -def merge_dicts(dictl, dictr, wl=0.5): - res = {} - for k in dictl.keys(): - res[k] = dictl[k] * wl + dictr[k] * (1-wl) - return res +def load_face_toolkits(dtype=torch.float16, gpu_id=-1): + image_encoder = CLIPVisionModelWithProjection.from_pretrained( + 'h94/IP-Adapter', subfolder='models/image_encoder') + image_encoder.to(dtype=dtype).cpu() + pd_fpg_motion = FanEncoder.from_pretrained("songkey/pd_fgc_motion") + pd_fpg_motion.to(dtype=dtype).cpu() + return dict( + device=get_torch_device(gpu_id), + dtype=dtype, + pd_fpg_motion=pd_fpg_motion, + face_aligner=HelloCameraDemo(face_alignment_module=HelloFaceAlignment(gpu_id=gpu_id), reset=False), + harkit_bs=HelloARKitBSPred(gpu_id=gpu_id), + h3dmm=Hello3DMMPred(gpu_id=gpu_id), + image_encoder=image_encoder + ) +def append_pipline_weights(pipeline, checkpoint_path=None, lora_path=None, vae_path=None, stylize='x1'): + ### load customized checkpoint or lora here: -def cat_dicts(dicts, dim=0): - res = {} - for k in dicts[0].keys(): - res[k] = torch.cat([d[k].clone() for d in dicts], dim=dim) - return res + # print("## checkpoint_path", checkpoint_path) + # print("## lora_path", lora_path) + # print("## vae_path", vae_path) -def cat_dicts_ref(dicts): - res = {} - for k in dicts[0].keys(): - res[k] = rearrange(torch.cat([d[k].clone().unsqueeze(1) for d in dicts], dim=1), "b f c h w -> (b f) c h w ") - return res + raw_stats = None + if checkpoint_path and not checkpoint_path.startswith('SD1.5'): + if osp.exists(checkpoint_path): + print("Loading checkpoint from", checkpoint_path) + if checkpoint_path.endswith('.safetensors'): + raw_stats = load_safetensors(checkpoint_path) + else: + raw_stats = torch.load(checkpoint_path) -def dicts_to_device(dicts, device): - ret = [] - for d in dicts: - tmpd = {} - for k in d.keys(): - tmpd[k] = d[k].clone().to(device) - ret.append(tmpd) - return ret + if raw_stats: + try: + state_dict = convert_ldm_unet_checkpoint(raw_stats, pipeline.unet_ref.config) + if hasattr(pipeline, 'unet_pre'): + pipeline.unet_pre.load_state_dict(state_dict, strict=False) + pipeline.unet.load_state_dict(state_dict, strict=False) + if stylize == 'x2' and hasattr(pipeline, 'unet_ref'): + pipeline.unet_ref.load_state_dict(state_dict, strict=False) + except: + raise ValueError("Failed to load checkpoint from", checkpoint_path) + + if vae_path and not vae_path.startswith('SD1.5 default vae'): + raw_vae_stats = raw_stats + + if osp.isfile(vae_path): + print("Loading vae from", vae_path) + if vae_path.endswith('.safetensors'): + raw_vae_stats = load_safetensors(vae_path) + else: + raw_vae_stats = torch.load(vae_path) + + if raw_vae_stats: + try: + vae_state_dict = convert_ldm_vae_checkpoint(raw_vae_stats, pipeline.vae.config) + if hasattr(pipeline, 'vae_decode'): + pipeline.vae_decode.load_state_dict(vae_state_dict, strict=True) + if stylize == 'x2': + pipeline.vae.load_state_dict(vae_state_dict, strict=True) + except: + raise ValueError("Failed to load vae from", vae_path) + + ### lora + if lora_path and not lora_path.startswith('None'): + if osp.exists(lora_path): + print("Loading lora from", lora_path) + try: + pipeline.load_lora_weights(osp.dirname(lora_path), weight_name=osp.basename(lora_path), adapter_name="lora") + except: + raise ValueError("Failed to load lora from", lora_path) def load_safetensors(model_path): tensors = {} diff --git a/meme.py b/meme.py index eead5bb..f508417 100644 --- a/meme.py +++ b/meme.py @@ -13,7 +13,6 @@ from einops import rearrange import importlib.metadata import folder_paths -import logging cur_dir = osp.dirname(osp.abspath(__file__)) @@ -37,13 +36,10 @@ from .hellomeme.utils import (get_drive_expression, crop_and_resize, det_landmarks, get_torch_device, - load_safetensors, + append_pipline_weights, + load_face_toolkits ) -from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (convert_ldm_unet_checkpoint, - convert_ldm_vae_checkpoint) -from .hellomeme.tools import Hello3DMMPred, HelloARKitBSPred, HelloFaceAlignment, HelloCameraDemo, FanEncoder from .hellomeme import HMImagePipeline, HMVideoPipeline -from transformers import CLIPVisionModelWithProjection DEFAULT_PROMPT = '(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture' @@ -58,56 +54,6 @@ def get_models_files(): return ['SD1.5'] + checkpoint_files, ['same as checkpoint', 'SD1.5 default vae'] + vae_files, ['None'] + lora_files -def append_pipline_weights(pipeline, checkpoint=None, lora=None, vae=None, stylize='x1'): - ### load customized checkpoint or lora here: - ## checkpoints - - raw_stats = None - if checkpoint and not checkpoint.startswith('SD1.5'): - checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint) - if osp.exists(checkpoint_path): - print("Loading checkpoint from", checkpoint_path) - if checkpoint_path.endswith('.safetensors'): - raw_stats = load_safetensors(checkpoint_path) - else: - raw_stats = torch.load(checkpoint_path) - - if raw_stats: - state_dict = convert_ldm_unet_checkpoint(raw_stats, pipeline.unet_ref.config) - if hasattr(pipeline, 'unet_pre'): - pipeline.unet_pre.load_state_dict(state_dict, strict=False) - pipeline.unet.load_state_dict(state_dict, strict=False) - - if stylize == 'x2' and hasattr(pipeline, 'unet_ref'): - pipeline.unet_ref.load_state_dict(state_dict, strict=False) - - if vae and not vae.startswith('SD1.5 default vae'): - raw_vae_stats = raw_stats - real_vae_path = '' - if vae.startswith("[checkpoint] "): - real_vae_path = folder_paths.get_full_path_or_raise("checkpoints", vae.replace("[checkpoint] ", "")) - elif vae.startswith("[vae] "): - real_vae_path = folder_paths.get_full_path_or_raise("vae", vae.replace("[vae] ", "")) - if osp.isfile(real_vae_path): - print("Loading vae from", real_vae_path) - if real_vae_path.endswith('.safetensors'): - raw_vae_stats = load_safetensors(real_vae_path) - else: - raw_vae_stats = torch.load(real_vae_path) - - if raw_vae_stats: - vae_state_dict = convert_ldm_vae_checkpoint(raw_vae_stats, pipeline.vae.config) - if hasattr(pipeline, 'vae_decode'): - pipeline.vae_decode.load_state_dict(vae_state_dict, strict=True) - if stylize == 'x2': - pipeline.vae.load_state_dict(vae_state_dict, strict=True) - - ### lora - if lora and not lora.startswith('None'): - lora_path = folder_paths.get_full_path_or_raise("loras", lora) - if osp.exists(lora_path): - print("Loading lora from", lora_path) - pipeline.load_lora_weights(osp.dirname(lora_path), weight_name=osp.basename(lora_path), adapter_name="lora") class HMImagePipelineLoader: @classmethod @@ -133,7 +79,24 @@ class HMImagePipelineLoader: pipeline.to(dtype=dtype) pipeline.caryomitosis(version=version) - append_pipline_weights(pipeline, checkpoint=checkpoint, lora=lora, vae=vae, stylize=stylize) + if checkpoint and not checkpoint.startswith('SD1.5'): + checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint) + else: + checkpoint_path = checkpoint + if vae and vae.startswith("[checkpoint] "): + vae_path = folder_paths.get_full_path_or_raise("checkpoints", vae.replace("[checkpoint] ", "")) + elif vae and vae.startswith("[vae] "): + vae_path = folder_paths.get_full_path_or_raise("vae", vae.replace("[vae] ", "")) + else: + vae_path = vae + + if lora and not lora.startswith('None'): + lora_path = folder_paths.get_full_path_or_raise("loras", lora) + else: + lora_path = lora + + append_pipline_weights(pipeline, checkpoint_path=checkpoint_path, lora_path=lora_path, vae_path=vae_path, + stylize=stylize) pipeline.insert_hm_modules(version=version, dtype=dtype) @@ -166,7 +129,23 @@ class HMVideoPipelineLoader: pipeline.to(dtype=dtype) pipeline.caryomitosis(version=version) - append_pipline_weights(pipeline, checkpoint=checkpoint, lora=lora, vae=vae, stylize=stylize) + if checkpoint and not checkpoint.startswith('SD1.5'): + checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint) + else: + checkpoint_path = checkpoint + if vae and vae.startswith("[checkpoint] "): + vae_path = folder_paths.get_full_path_or_raise("checkpoints", vae.replace("[checkpoint] ", "")) + elif vae and vae.startswith("[vae] "): + vae_path = folder_paths.get_full_path_or_raise("vae", vae.replace("[vae] ", "")) + else: + vae_path = vae + + if lora and not lora.startswith('None'): + lora_path = folder_paths.get_full_path_or_raise("loras", lora) + else: + lora_path = lora + + append_pipline_weights(pipeline, checkpoint_path=checkpoint_path, lora_path=lora_path, vae_path=vae_path, stylize=stylize) pipeline.insert_hm_modules(version=version, dtype=dtype) @@ -186,24 +165,12 @@ class HMFaceToolkitsLoader: RETURN_NAMES = ("face_toolkits",) FUNCTION = "load_face_toolkits" CATEGORY = "hellomeme" + def load_face_toolkits(self, gpu_id): dtype = torch.float16 - image_encoder = CLIPVisionModelWithProjection.from_pretrained( - 'h94/IP-Adapter', subfolder='models/image_encoder') - image_encoder.to(dtype=dtype).cpu() - pd_fpg_motion = FanEncoder.from_pretrained("songkey/pd_fgc_motion") - pd_fpg_motion.to(dtype=dtype).cpu() - return ( - dict( - device=get_torch_device(gpu_id), - dtype=dtype, - pd_fpg_motion=pd_fpg_motion, - face_aligner=HelloCameraDemo(face_alignment_module=HelloFaceAlignment(gpu_id=gpu_id), reset=False), - harkit_bs=HelloARKitBSPred(gpu_id=gpu_id), - h3dmm=Hello3DMMPred(gpu_id=gpu_id), - image_encoder=image_encoder - ), - ) + face_toolkits = load_face_toolkits(dtype=dtype, gpu_id=gpu_id) + return (face_toolkits, ) + class CropPortrait: @classmethod