diff --git a/hellomeme/__init__.py b/hellomeme/__init__.py index 2aefa89..ec3c7ff 100644 --- a/hellomeme/__init__.py +++ b/hellomeme/__init__.py @@ -10,4 +10,5 @@ from .pipelines import (HMImagePipeline, HMVideoPipeline, HM3ImagePipeline, HM3VideoPipeline, - HM5ImagePipeline, HM5VideoPipeline) \ No newline at end of file + HM5ImagePipeline, HM5VideoPipeline) +from .tools.utils import download_file_from_cloud, creat_model_from_cloud \ No newline at end of file diff --git a/hellomeme/pipelines/pipline_hm5_video.py b/hellomeme/pipelines/pipline_hm5_video.py index 510f1c2..c8da26a 100644 --- a/hellomeme/pipelines/pipline_hm5_video.py +++ b/hellomeme/pipelines/pipline_hm5_video.py @@ -24,7 +24,8 @@ from diffusers import DPMSolverMultistepScheduler from diffusers.utils.torch_utils import randn_tensor from ..models import (HM3Denoising3D, HM5MotionAdapter, - HM5ReferenceAdapter, HMPipeline, HM5ControlNetBase, HM5SD15ControlProj) + HM5ReferenceAdapter, HM5bReferenceAdapter, + HMPipeline, HM5ControlNetBase, HM5SD15ControlProj) class HM5VideoPipeline(HMPipeline): def caryomitosis(self, version='v5', modelscope=False, **kwargs): @@ -53,19 +54,33 @@ class HM5VideoPipeline(HMPipeline): if modelscope: from modelscope import snapshot_download - hm_reference_dir = snapshot_download('songkey/hm5_reference') + if version == 'v5b': + hm_reference_dir = snapshot_download('songkey/hm5b_reference') + else: + hm_reference_dir = snapshot_download('songkey/hm5_reference') hm_control_dir = snapshot_download('songkey/hm5_control_base') hm_control_proj_dir = snapshot_download('songkey/hm5_control_proj') hm_motion_dir = snapshot_download('songkey/hm5_motion') else: - hm_reference_dir = 'songkey/hm5_reference' + if version == 'v5b': + hm_reference_dir = 'songkey/hm5b_reference' + else: + hm_reference_dir = 'songkey/hm5_reference' hm_control_dir = 'songkey/hm5_control_base' hm_control_proj_dir = 'songkey/hm5_control_proj' hm_motion_dir = 'songkey/hm5_motion' - hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir) + if version == 'v5b': + hm_adapter = HM5bReferenceAdapter.from_pretrained(hm_reference_dir) + else: + hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir) + + # hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir) motion_adapter = HM5MotionAdapter.from_pretrained(hm_motion_dir) + stats = torch.load(r"E:\upload2\weights\a100-0617mpsd15vff\checkpoint-41000.pth", map_location='cpu') + motion_adapter.load_state_dict(stats['motion'], strict=False) + if isinstance(self.unet, HM3Denoising3D): self.unet.insert_reference_adapter(hm_adapter) self.unet.insert_reference_adapter(motion_adapter) diff --git a/hellomeme/tools/hello_3dmm.py b/hellomeme/tools/hello_3dmm.py index 4584855..90c2a67 100644 --- a/hellomeme/tools/hello_3dmm.py +++ b/hellomeme/tools/hello_3dmm.py @@ -13,7 +13,7 @@ import numpy as np import cv2 import os.path as osp -from .utils import get_warp_mat_bbox_by_gt_pts_float, create_onnx_session +from .utils import get_warp_mat_bbox_by_gt_pts_float, create_onnx_session, download_file_from_cloud def crop_transl_to_full_transl(crop_trans, crop_center, scale, full_center, focal_length): """ @@ -37,12 +37,9 @@ def crop_transl_to_full_transl(crop_trans, crop_center, scale, full_center, foca class Hello3DMMPred(object): def __init__(self, gpu_id=None, modelscope=False): - if modelscope: - from modelscope import snapshot_download - model_path = osp.join(snapshot_download('songkey/hello_group_facemodel'), 'hello_3dmm.onnx') - else: - from huggingface_hub import hf_hub_download - model_path = hf_hub_download('songkey/hello_group_facemodel', filename='hello_3dmm.onnx') + model_path = download_file_from_cloud(model_id='songkey/hello_group_facemodel', + file_name='hello_3dmm.onnx', + modelscope=modelscope) self.deep3d_pred_net = create_onnx_session(model_path, 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()] diff --git a/hellomeme/tools/hello_arkit.py b/hellomeme/tools/hello_arkit.py index b92cc81..69fc219 100644 --- a/hellomeme/tools/hello_arkit.py +++ b/hellomeme/tools/hello_arkit.py @@ -9,16 +9,13 @@ import numpy as np import cv2 import os.path as osp -from .utils import create_onnx_session, get_warp_mat_bbox_by_gt_pts_float +from .utils import create_onnx_session, get_warp_mat_bbox_by_gt_pts_float, download_file_from_cloud class HelloARKitBSPred(object): def __init__(self, gpu_id=0, modelscope=False): - if modelscope: - from modelscope import snapshot_download - model_path = osp.join(snapshot_download('songkey/hello_group_facemodel'), 'hello_arkit_blendshape.onnx') - else: - from huggingface_hub import hf_hub_download - model_path = hf_hub_download('songkey/hello_group_facemodel', filename='hello_arkit_blendshape.onnx') + model_path = download_file_from_cloud(model_id='songkey/hello_group_facemodel', + file_name='hello_arkit_blendshape.onnx', + modelscope=modelscope) self.face_rig_net = create_onnx_session(model_path, gpu_id=gpu_id) self.onnx_input_name = self.face_rig_net.get_inputs()[0].name diff --git a/hellomeme/tools/hello_face_alignment.py b/hellomeme/tools/hello_face_alignment.py index 607622b..dc48603 100644 --- a/hellomeme/tools/hello_face_alignment.py +++ b/hellomeme/tools/hello_face_alignment.py @@ -11,20 +11,19 @@ import os.path as osp 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 .utils import create_onnx_session, download_file_from_cloud class HelloFaceAlignment(object): def __init__(self, gpu_id=None, modelscope=False): expand_ratio = 0.15 - if modelscope: - from modelscope import snapshot_download - alignment_model_path = osp.join(snapshot_download('songkey/hello_group_facemodel'), 'hello_face_landmark.onnx') - det_model_path = osp.join(snapshot_download('songkey/hello_group_facemodel'), 'hello_face_det.onnx') - else: - from huggingface_hub import hf_hub_download - alignment_model_path = hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_landmark.onnx') - det_model_path = hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_det.onnx') + alignment_model_path = download_file_from_cloud(model_id='songkey/hello_group_facemodel', + file_name='hello_face_landmark.onnx', + modelscope=modelscope) + + det_model_path = download_file_from_cloud(model_id='songkey/hello_group_facemodel', + file_name='hello_face_det.onnx', + modelscope=modelscope) self.face_alignment_net_222 = ( create_onnx_session(alignment_model_path, gpu_id=gpu_id)) self.onnx_input_name_222 = self.face_alignment_net_222.get_inputs()[0].name diff --git a/hellomeme/tools/utils.py b/hellomeme/tools/utils.py index 4281d9d..aedbebd 100644 --- a/hellomeme/tools/utils.py +++ b/hellomeme/tools/utils.py @@ -11,6 +11,7 @@ import time import cv2 import numpy as np import math +import os.path as osp def create_onnx_session(onnx_path, gpu_id=None)->onnxruntime.InferenceSession: start = time.perf_counter() @@ -116,3 +117,47 @@ def get_warp_mat_bbox_by_gt_pts_float(gt_pts, base_angle=0.0, dst_size=128, expa "scale": scale } return transform_info + + +def download_file_from_cloud(model_id, + file_name, + modelscope=False, + cache_dir=None, + hf_token='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',): + if modelscope: + from modelscope import snapshot_download + try: + model_path = osp.join(snapshot_download(model_id, cache_dir=cache_dir), file_name) + except Exception as e: + print(e) + assert False, "@@ Failed to download model from modelscope" + else: + from huggingface_hub import hf_hub_download + try: + model_path = hf_hub_download(model_id, filename=file_name, cache_dir=cache_dir, token=hf_token) + except Exception as e: + print(e) + assert False, "@@ Failed to download model from huggingface hub" + return model_path + +def creat_model_from_cloud(model_cls, + model_id, + modelscope=False, + cache_dir=None, + hf_token='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'): + if modelscope: + from modelscope import snapshot_download + try: + model_path = snapshot_download(model_id, cache_dir=cache_dir) + except Exception as e: + print(e) + assert False, "@@ Failed to download model from modelscope" + + model = model_cls.from_pretrained(model_path) + else: + try: + model = model_cls.from_pretrained(model_id, cache_dir=cache_dir, token=hf_token) + except Exception as e: + print(e) + assert False, "@@ Failed to download model from huggingface hub" + return model \ No newline at end of file diff --git a/hellomeme/utils.py b/hellomeme/utils.py index 7543ce8..a070d5d 100644 --- a/hellomeme/utils.py +++ b/hellomeme/utils.py @@ -18,7 +18,7 @@ import torch from PIL import Image import subprocess from einops import rearrange -from .tools.utils import transform_points +from .tools.utils import transform_points, creat_model_from_cloud from .tools.hello_3dmm import get_project_points_rect from safetensors import safe_open @@ -64,7 +64,9 @@ def load_face_toolkits(dtype=torch.float16, gpu_id=-1, modelscope=False): image_encoder=image_encoder ) -def append_pipline_weights(pipeline, checkpoint_path=None, lora_path=None, vae_path=None, stylize='x1', lora_scale=1.0): +def append_pipline_weights(pipeline, checkpoint_path=None, + lora_path=None, vae_path=None, + stylize='x1', lora_scale=1.0, modelscope=False): ### load customized checkpoint or lora here: # print("## checkpoint_path", checkpoint_path) @@ -84,7 +86,7 @@ def append_pipline_weights(pipeline, checkpoint_path=None, lora_path=None, vae_p unet_stats = convert_ldm_unet_checkpoint(raw_stats, pipeline.unet_ref.config) vae_stats = convert_ldm_vae_checkpoint(raw_stats, pipeline.vae.config) else: - tmp_pipeline = HMImagePipeline.from_pretrained(checkpoint_path) + tmp_pipeline = creat_model_from_cloud(HMImagePipeline, checkpoint_path, modelscope=modelscope) unet_stats = tmp_pipeline.unet.state_dict() vae_stats = tmp_pipeline.vae.state_dict() diff --git a/meme.py b/meme.py index 6c813e6..f826e34 100644 --- a/meme.py +++ b/meme.py @@ -23,13 +23,6 @@ if missing: os.system(f'{sys.executable} -m pip install {missing_params}') # subprocess.check_call([python, '-m', 'pip', 'install', missing_params], stdout=subprocess.DEVNULL) -try: - from huggingface_hub import login - # Model read permission - login(token='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', new_session=False, write_permission=False) -except: - print("login error!") - import numpy as np import cv2 @@ -52,7 +45,8 @@ from .hellomeme.utils import (get_drive_expression, ) from .hellomeme import (HMImagePipeline, HMVideoPipeline, HM3ImagePipeline, HM3VideoPipeline, - HM5ImagePipeline, HM5VideoPipeline) + HM5ImagePipeline, HM5VideoPipeline, + download_file_from_cloud) config_path = osp.join(cur_dir, 'hellomeme', 'model_config.json') with open(config_path, 'r') as f: @@ -78,12 +72,7 @@ def get_models_files(): def format_model_path(pipeline, config, checkpoint, vae, lora, stylize, lora_scale, deployment): if checkpoint and not checkpoint.startswith('SD1.5'): if checkpoint in config['sd15']['checkpoints']: - tmp_checkpoint_name = config['sd15']['checkpoints'][checkpoint] - if deployment == 'modelscope': - from modelscope import snapshot_download - checkpoint_path = snapshot_download(tmp_checkpoint_name) - else: - checkpoint_path = tmp_checkpoint_name + checkpoint_path = config['sd15']['checkpoints'][checkpoint] else: checkpoint_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint) else: @@ -99,19 +88,14 @@ def format_model_path(pipeline, config, checkpoint, vae, lora, stylize, lora_sca if lora and not lora.startswith('None'): if lora in config['sd15']['loras']: tmp_lora_info = config['sd15']['loras'][lora] - if deployment == 'modelscope': - from modelscope import snapshot_download - lora_path = osp.join(snapshot_download(tmp_lora_info[0]), tmp_lora_info[1]) - else: - from huggingface_hub import hf_hub_download - lora_path = hf_hub_download(tmp_lora_info[0], filename=tmp_lora_info[1]) + lora_path = download_file_from_cloud(tmp_lora_info[0], tmp_lora_info[1], modelscope=deployment == 'modelscope') else: 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, lora_scale=lora_scale) + stylize=stylize, lora_scale=lora_scale, modelscope=deployment == 'modelscope') class HMImagePipelineLoader: @@ -169,7 +153,7 @@ class HMVideoPipelineLoader: "checkpoint": (checkpoint_files, ), "lora": (lora_files, ), "vae": (vae_files, ), - "version": (['v5', 'v4', 'v3', 'v2', 'v1'], ), + "version": (['v5b', 'v5', 'v4', 'v3', 'v2', 'v1'], ), "stylize": (['x1', 'x2'], ), "deployment": (['huggingface', 'modelscope'], ), "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}), @@ -193,7 +177,7 @@ class HMVideoPipelineLoader: if version == 'v3' or version == 'v4': pipeline = HM3VideoPipeline.from_pretrained(sd1_5_dir) - elif version == 'v5': + elif version == 'v5' or version == 'v5b': pipeline = HM5VideoPipeline.from_pretrained(sd1_5_dir) else: pipeline = HMVideoPipeline.from_pretrained(sd1_5_dir) diff --git a/pyproject.toml b/pyproject.toml index 063a556..54ffe01 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.3.16" +version = "1.3.17" license = {file = "LICENSE"} [project.urls]