diff --git a/hellomeme/pipelines/pipline_hm3_image.py b/hellomeme/pipelines/pipline_hm3_image.py index 8b30086..1a266ea 100644 --- a/hellomeme/pipelines/pipline_hm3_image.py +++ b/hellomeme/pipelines/pipline_hm3_image.py @@ -21,6 +21,8 @@ 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, HMPipeline, HM3ReferenceAdapter, HMControlNetBase, HM4SD15ControlProj +from ..tools.utils import creat_model_from_cloud + class HM3ImagePipeline(HMPipeline): def caryomitosis(self, **kwargs): @@ -44,26 +46,17 @@ class HM3ImagePipeline(HMPipeline): def insert_hm_modules(self, version='v3', dtype=torch.float16, modelscope=False): self.version = version - if modelscope: - from modelscope import snapshot_download - if version == 'v3': - hm_reference_dir = snapshot_download('songkey/hm3_reference') - hm_control_dir = snapshot_download('songkey/hm3_control_mix') - else: - hm_reference_dir = snapshot_download('songkey/hm4_reference') - hm_control_dir = snapshot_download('songkey/hm_control_base') - hm_control_proj_dir = snapshot_download('songkey/hm4_control_proj') + + if version == 'v3': + hm_reference_dir = 'songkey/hm3_reference' + hm_control_dir = 'songkey/hm3_control_mix' else: - if version == 'v3': - hm_reference_dir = 'songkey/hm3_reference' - hm_control_dir = 'songkey/hm3_control_mix' - else: - hm_reference_dir = 'songkey/hm4_reference' - hm_control_dir = 'songkey/hm_control_base' - hm_control_proj_dir = 'songkey/hm4_control_proj' + hm_reference_dir = 'songkey/hm4_reference' + hm_control_dir = 'songkey/hm_control_base' + hm_control_proj_dir = 'songkey/hm4_control_proj' if isinstance(self.unet, HM3Denoising3D): - hm_adapter = HM3ReferenceAdapter.from_pretrained(hm_reference_dir) + hm_adapter = creat_model_from_cloud(HM3ReferenceAdapter, hm_reference_dir, modelscope=modelscope) self.unet.insert_reference_adapter(hm_adapter) self.unet.to(device='cpu', dtype=dtype).eval() @@ -77,10 +70,10 @@ class HM3ImagePipeline(HMPipeline): del self.mp_control_proj if version == 'v3': - self.mp_control = HMV3ControlNet.from_pretrained(hm_control_dir) + self.mp_control = creat_model_from_cloud(HMV3ControlNet, hm_control_dir, modelscope=modelscope) else: - self.mp_control = HMControlNetBase.from_pretrained(hm_control_dir) - self.mp_control_proj = HM4SD15ControlProj.from_pretrained(hm_control_proj_dir) + self.mp_control = creat_model_from_cloud(HMControlNetBase, hm_control_dir, modelscope=modelscope) + self.mp_control_proj = creat_model_from_cloud(HM4SD15ControlProj, hm_control_proj_dir, modelscope=modelscope) self.mp_control_proj.to(device='cpu', dtype=dtype).eval() diff --git a/hellomeme/pipelines/pipline_hm3_video.py b/hellomeme/pipelines/pipline_hm3_video.py index 6cfef6b..0c42bde 100644 --- a/hellomeme/pipelines/pipline_hm3_video.py +++ b/hellomeme/pipelines/pipline_hm3_video.py @@ -24,7 +24,10 @@ from diffusers import MotionAdapter, EulerDiscreteScheduler # DPMSolverMultistep from diffusers.utils.torch_utils import randn_tensor from ..models import (HM3Denoising3D, HM3DenoisingMotion, HM3MotionAdapter, - HMV3ControlNet, HM3ReferenceAdapter, HMPipeline, HMControlNetBase, HM4SD15ControlProj) + HMV3ControlNet, HM3ReferenceAdapter, HMPipeline, + HMControlNetBase, HM4SD15ControlProj) +from ..tools.utils import creat_model_from_cloud + class HM3VideoPipeline(HMPipeline): def caryomitosis(self, version='v3', modelscope=False, **kwargs): @@ -38,21 +41,14 @@ class HM3VideoPipeline(HMPipeline): self.unet_pre = HM3Denoising3D.from_unet2d(self.unet) self.unet_pre.cpu() - if modelscope: - from modelscope import snapshot_download - if version == 'v3': - hm_animatediff_dir = snapshot_download('songkey/hm3_animatediff') - else: - hm_animatediff_dir = snapshot_download('songkey/hm4_animatediff') + if version == 'v3': + hm_animatediff_dir = 'songkey/hm3_animatediff' else: - if version == 'v3': - hm_animatediff_dir = 'songkey/hm3_animatediff' - else: - hm_animatediff_dir = 'songkey/hm4_animatediff' + hm_animatediff_dir = 'songkey/hm4_animatediff' self.num_frames = 12 - adapter = MotionAdapter.from_pretrained(hm_animatediff_dir, torch_dtype=torch.float16) + adapter = creat_model_from_cloud(MotionAdapter, hm_animatediff_dir, modelscope=modelscope) unet = HM3DenoisingMotion.from_unet2d(unet=self.unet, motion_adapter=adapter, load_weights=True) del self.unet @@ -67,30 +63,18 @@ class HM3VideoPipeline(HMPipeline): def insert_hm_modules(self, version='v3', dtype=torch.float16, modelscope=False): self.version = version - if modelscope: - from modelscope import snapshot_download - if version == 'v3': - hm_reference_dir = snapshot_download('songkey/hm3_reference') - hm_control_dir = snapshot_download('songkey/hm3_control_mix') - hm_motion_dir = snapshot_download('songkey/hm3_motion') - else: - hm_reference_dir = snapshot_download('songkey/hm4_reference') - hm_control_dir = snapshot_download('songkey/hm_control_base') - hm_control_proj_dir = snapshot_download('songkey/hm4_control_proj') - hm_motion_dir = snapshot_download('songkey/hm4_motion') + if version == 'v3': + hm_reference_dir = 'songkey/hm3_reference' + hm_control_dir = 'songkey/hm3_control_mix' + hm_motion_dir = 'songkey/hm3_motion' else: - if version == 'v3': - hm_reference_dir = 'songkey/hm3_reference' - hm_control_dir = 'songkey/hm3_control_mix' - hm_motion_dir = 'songkey/hm3_motion' - else: - hm_reference_dir = 'songkey/hm4_reference' - hm_control_dir = 'songkey/hm_control_base' - hm_control_proj_dir = 'songkey/hm4_control_proj' - hm_motion_dir = 'songkey/hm4_motion' + hm_reference_dir = 'songkey/hm4_reference' + hm_control_dir = 'songkey/hm_control_base' + hm_control_proj_dir = 'songkey/hm4_control_proj' + hm_motion_dir = 'songkey/hm4_motion' - hm_adapter = HM3ReferenceAdapter.from_pretrained(hm_reference_dir) - motion_adapter = HM3MotionAdapter.from_pretrained(hm_motion_dir) + hm_adapter = creat_model_from_cloud(HM3ReferenceAdapter, hm_reference_dir, modelscope=modelscope) + motion_adapter = creat_model_from_cloud(HM3MotionAdapter, hm_motion_dir, modelscope=modelscope) if isinstance(self.unet, HM3DenoisingMotion): self.unet.insert_reference_adapter(hm_adapter) @@ -111,10 +95,10 @@ class HM3VideoPipeline(HMPipeline): del self.mp_control_proj if version == 'v3': - self.mp_control = HMV3ControlNet.from_pretrained(hm_control_dir) + self.mp_control = creat_model_from_cloud(HMV3ControlNet, hm_control_dir, modelscope=modelscope) else: - self.mp_control = HMControlNetBase.from_pretrained(hm_control_dir) - self.mp_control_proj = HM4SD15ControlProj.from_pretrained(hm_control_proj_dir) + self.mp_control = creat_model_from_cloud(HMControlNetBase, hm_control_dir, modelscope=modelscope) + self.mp_control_proj = creat_model_from_cloud(HM4SD15ControlProj, hm_control_proj_dir, modelscope=modelscope) self.mp_control_proj.to(device='cpu', dtype=dtype).eval() self.mp_control.to(device='cpu', dtype=dtype).eval() diff --git a/hellomeme/pipelines/pipline_hm5_image.py b/hellomeme/pipelines/pipline_hm5_image.py index 880baae..2c105dc 100644 --- a/hellomeme/pipelines/pipline_hm5_image.py +++ b/hellomeme/pipelines/pipline_hm5_image.py @@ -22,6 +22,9 @@ from diffusers import DPMSolverMultistepScheduler from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents from ..models import (HM3Denoising3D, HMPipeline, HM5ReferenceAdapter, HM5bReferenceAdapter, HM5ControlNetBase, HM5SD15ControlProj) +from ..tools.utils import creat_model_from_cloud + + class HM5ImagePipeline(HMPipeline): def caryomitosis(self, **kwargs): @@ -45,27 +48,19 @@ class HM5ImagePipeline(HMPipeline): def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False): self.version = version - if modelscope: - from modelscope import snapshot_download - 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') + + if version == 'v5b': + hm_reference_dir = 'songkey/hm5b_reference' else: - 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_reference_dir = 'songkey/hm5_reference' + hm_control_dir = 'songkey/hm5_control_base' + hm_control_proj_dir = 'songkey/hm5_control_proj' if isinstance(self.unet, HM3Denoising3D): if version == 'v5b': - hm_adapter = HM5bReferenceAdapter.from_pretrained(hm_reference_dir) + hm_adapter = creat_model_from_cloud(HM5bReferenceAdapter, hm_reference_dir, modelscope=modelscope) else: - hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir) + hm_adapter = creat_model_from_cloud(HM5ReferenceAdapter, hm_reference_dir, modelscope=modelscope) self.unet.insert_reference_adapter(hm_adapter) self.unet.to(device='cpu', dtype=dtype).eval() @@ -79,8 +74,8 @@ class HM5ImagePipeline(HMPipeline): if hasattr(self, "mp_control_proj"): del self.mp_control_proj - self.mp_control = HM5ControlNetBase.from_pretrained(hm_control_dir) - self.mp_control_proj = HM5SD15ControlProj.from_pretrained(hm_control_proj_dir) + self.mp_control = creat_model_from_cloud(HM5ControlNetBase, hm_control_dir, modelscope=modelscope) + self.mp_control_proj = creat_model_from_cloud(HM5SD15ControlProj, hm_control_proj_dir, modelscope=modelscope) self.mp_control.to(device='cpu', dtype=dtype).eval() self.mp_control_proj.to(device='cpu', dtype=dtype).eval() diff --git a/hellomeme/pipelines/pipline_hm5_video.py b/hellomeme/pipelines/pipline_hm5_video.py index 30ed9b1..12d4b9e 100644 --- a/hellomeme/pipelines/pipline_hm5_video.py +++ b/hellomeme/pipelines/pipline_hm5_video.py @@ -26,6 +26,8 @@ from diffusers.utils.torch_utils import randn_tensor from ..models import (HM3Denoising3D, HM5MotionAdapter, HM5ReferenceAdapter, HM5bReferenceAdapter, HMPipeline, HM5ControlNetBase, HM5SD15ControlProj) +from ..tools.utils import creat_model_from_cloud + class HM5VideoPipeline(HMPipeline): def caryomitosis(self, version='v5', modelscope=False, **kwargs): @@ -52,30 +54,19 @@ class HM5VideoPipeline(HMPipeline): def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False): self.version = version - if modelscope: - from modelscope import snapshot_download - 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') + if version == 'v5b': + hm_reference_dir = 'songkey/hm5b_reference' else: - 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_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' if version == 'v5b': - hm_adapter = HM5bReferenceAdapter.from_pretrained(hm_reference_dir) + hm_adapter = creat_model_from_cloud(HM5bReferenceAdapter, hm_reference_dir, modelscope=modelscope) else: - hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir) - - motion_adapter = HM5MotionAdapter.from_pretrained(hm_motion_dir) + hm_adapter = creat_model_from_cloud(HM5ReferenceAdapter, hm_reference_dir, modelscope=modelscope) + motion_adapter = creat_model_from_cloud(HM5MotionAdapter, hm_motion_dir, modelscope=modelscope) if isinstance(self.unet, HM3Denoising3D): self.unet.insert_reference_adapter(hm_adapter) @@ -90,8 +81,8 @@ class HM5VideoPipeline(HMPipeline): if hasattr(self, "mp_control_proj"): del self.mp_control_proj - self.mp_control = HM5ControlNetBase.from_pretrained(hm_control_dir) - self.mp_control_proj = HM5SD15ControlProj.from_pretrained(hm_control_proj_dir) + self.mp_control = creat_model_from_cloud(HM5ControlNetBase, hm_control_dir, modelscope=modelscope) + self.mp_control_proj = creat_model_from_cloud(HM5SD15ControlProj, hm_control_proj_dir, modelscope=modelscope) self.mp_control_proj.to(device='cpu', dtype=dtype).eval() self.mp_control.to(device='cpu', dtype=dtype).eval() diff --git a/hellomeme/pipelines/pipline_hm_image.py b/hellomeme/pipelines/pipline_hm_image.py index 6143da9..899eb39 100644 --- a/hellomeme/pipelines/pipline_hm_image.py +++ b/hellomeme/pipelines/pipline_hm_image.py @@ -23,6 +23,8 @@ from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img impo from ..models import HMDenoising3D, HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2, HMPipeline from ..models import HMReferenceAdapter +from ..tools.utils import creat_model_from_cloud + class HMImagePipeline(HMPipeline): def caryomitosis(self, **kwargs): @@ -46,27 +48,18 @@ class HMImagePipeline(HMPipeline): def insert_hm_modules(self, version, dtype, modelscope=False): self.version = version - if modelscope: - from modelscope import snapshot_download - hm_reference_dir = snapshot_download('songkey/hm_reference') - hm2_reference_dir = snapshot_download('songkey/hm2_reference') - hm_control_dir = snapshot_download('songkey/hm_control') - hm_control2_dir = snapshot_download('songkey/hm_control2') - hm2_control_dir = snapshot_download('songkey/hm2_control') - hm2_control2_dir = snapshot_download('songkey/hm2_control2') - else: + + if version == 'v1': hm_reference_dir = 'songkey/hm_reference' - hm2_reference_dir = 'songkey/hm2_reference' hm_control_dir = 'songkey/hm_control' hm_control2_dir = 'songkey/hm_control2' - hm2_control_dir = 'songkey/hm2_control' - hm2_control2_dir = 'songkey/hm2_control2' + else: + hm_reference_dir = 'songkey/hm2_reference' + hm_control_dir = 'songkey/hm2_control' + hm_control2_dir = 'songkey/hm2_control2' if isinstance(self.unet, HMDenoising3D): - if version == 'v1': - hm_adapter = HMReferenceAdapter.from_pretrained(hm_reference_dir) - else: - hm_adapter = HMReferenceAdapter.from_pretrained(hm2_reference_dir) + hm_adapter = creat_model_from_cloud(HMReferenceAdapter, hm_reference_dir, modelscope=modelscope) self.unet.insert_reference_adapter(hm_adapter) self.unet.to(device='cpu', dtype=dtype).eval() @@ -76,17 +69,17 @@ class HMImagePipeline(HMPipeline): if hasattr(self, "mp_control"): del self.mp_control if version == 'v1': - self.mp_control = HMControlNet.from_pretrained(hm_control_dir) + self.mp_control = creat_model_from_cloud(HMControlNet, hm_control_dir, modelscope=modelscope) else: - self.mp_control = HMV2ControlNet.from_pretrained(hm2_control_dir) + self.mp_control = creat_model_from_cloud(HMV2ControlNet, hm_control_dir, modelscope=modelscope) self.mp_control.to(device='cpu', dtype=dtype).eval() if hasattr(self, "mp_control2"): del self.mp_control2 if version == 'v1': - self.mp_control2 = HMControlNet2.from_pretrained(hm_control2_dir) + self.mp_control2 = creat_model_from_cloud(HMControlNet2, hm_control2_dir, modelscope=modelscope) else: - self.mp_control2 = HMV2ControlNet2.from_pretrained(hm2_control2_dir) + self.mp_control2 = creat_model_from_cloud(HMV2ControlNet2, hm_control2_dir, modelscope=modelscope) self.mp_control2.to(device='cpu', dtype=dtype).eval() self.vae.to(device='cpu', dtype=dtype).eval() diff --git a/hellomeme/pipelines/pipline_hm_video.py b/hellomeme/pipelines/pipline_hm_video.py index 35d9177..f106acb 100644 --- a/hellomeme/pipelines/pipline_hm_video.py +++ b/hellomeme/pipelines/pipline_hm_video.py @@ -25,6 +25,8 @@ from diffusers import MotionAdapter, EulerDiscreteScheduler from ..models import (HMDenoising3D, HMDenoisingMotion, HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2, HMPipeline) from ..models import HMReferenceAdapter +from ..tools.utils import creat_model_from_cloud + class HMVideoPipeline(HMPipeline): def caryomitosis(self, version, modelscope=False, **kwargs): @@ -38,19 +40,13 @@ class HMVideoPipeline(HMPipeline): self.unet_pre = HMDenoising3D.from_unet2d(self.unet) self.unet_pre.cpu() - if modelscope: - from modelscope import snapshot_download - hm_animatediff_dir = snapshot_download('songkey/hm_animatediff_frame12') - hm2_animatediff_dir = snapshot_download('songkey/hm2_animatediff_frame12') - else: + if version == 'v1': hm_animatediff_dir = 'songkey/hm_animatediff_frame12' - hm2_animatediff_dir = 'songkey/hm2_animatediff_frame12' + else: + hm_animatediff_dir = 'songkey/hm2_animatediff_frame12' self.num_frames = 12 - if version == 'v1': - adapter = MotionAdapter.from_pretrained(hm_animatediff_dir, torch_dtype=torch.float16) - else: - adapter = MotionAdapter.from_pretrained(hm2_animatediff_dir, torch_dtype=torch.float16) + adapter = creat_model_from_cloud(MotionAdapter, hm_animatediff_dir, modelscope=modelscope) unet = HMDenoisingMotion.from_unet2d(unet=self.unet, motion_adapter=adapter, load_weights=True) # todo: 不够优雅 del self.unet @@ -63,26 +59,17 @@ class HMVideoPipeline(HMPipeline): self.safety_checker.cpu() def insert_hm_modules(self, version, dtype, modelscope=False): - if modelscope: - from modelscope import snapshot_download - hm_reference_dir = snapshot_download('songkey/hm_reference') - hm2_reference_dir = snapshot_download('songkey/hm2_reference') - hm_control_dir = snapshot_download('songkey/hm_control') - hm_control2_dir = snapshot_download('songkey/hm_control2') - hm2_control_dir = snapshot_download('songkey/hm2_control') - hm2_control2_dir = snapshot_download('songkey/hm2_control2') - else: + if version == 'v1': hm_reference_dir = 'songkey/hm_reference' - hm2_reference_dir = 'songkey/hm2_reference' hm_control_dir = 'songkey/hm_control' hm_control2_dir = 'songkey/hm_control2' - hm2_control_dir = 'songkey/hm2_control' - hm2_control2_dir = 'songkey/hm2_control2' - - if version == 'v1': - hm_adapter = HMReferenceAdapter.from_pretrained(hm_reference_dir) else: - hm_adapter = HMReferenceAdapter.from_pretrained(hm2_reference_dir) + hm_reference_dir = 'songkey/hm2_reference' + hm_control_dir = 'songkey/hm2_control' + hm_control2_dir = 'songkey/hm2_control2' + + hm_adapter = creat_model_from_cloud(HMReferenceAdapter, hm_reference_dir, modelscope=modelscope) + if isinstance(self.unet, HMDenoisingMotion): self.unet.insert_reference_adapter(hm_adapter) self.unet.to(device='cpu', dtype=dtype).eval() @@ -97,18 +84,18 @@ class HMVideoPipeline(HMPipeline): if hasattr(self, "mp_control"): del self.mp_control if version == 'v1': - self.mp_control = HMControlNet.from_pretrained(hm_control_dir) + self.mp_control = creat_model_from_cloud(HMControlNet, hm_control_dir, modelscope=modelscope) else: - self.mp_control = HMV2ControlNet.from_pretrained(hm2_control_dir) + self.mp_control = creat_model_from_cloud(HMV2ControlNet, hm_control_dir, modelscope=modelscope) self.mp_control.to(device='cpu', dtype=dtype).eval() if hasattr(self, "mp_control2"): del self.mp_control2 if version == 'v1': - self.mp_control2 = HMControlNet2.from_pretrained(hm_control2_dir) + self.mp_control2 = creat_model_from_cloud(HMControlNet2, hm_control2_dir, modelscope=modelscope) else: - self.mp_control2 = HMV2ControlNet2.from_pretrained(hm2_control2_dir) + self.mp_control2 = creat_model_from_cloud(HMV2ControlNet2, hm2_control2_dir, modelscope=modelscope) self.mp_control2.to(device='cpu', dtype=dtype).eval() self.vae.to(device='cpu', dtype=dtype).eval() diff --git a/hellomeme/tools/sr.py b/hellomeme/tools/sr.py index 702483e..b84069f 100644 --- a/hellomeme/tools/sr.py +++ b/hellomeme/tools/sr.py @@ -15,6 +15,7 @@ import cv2 import numpy as np import math import os.path as osp +from .utils import download_file_from_cloud def pixel_unshuffle(x, scale): """ Pixel unshuffle. @@ -183,13 +184,7 @@ class RealESRGANer(): else: self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device - if modelscope: - from modelscope import snapshot_download - model_path = osp.join(snapshot_download('songkey/ESRGAN'), 'RealESRGAN_x2plus.pth') - else: - from huggingface_hub import hf_hub_download - model_path = hf_hub_download('songkey/ESRGAN', filename='RealESRGAN_x2plus.pth') - + model_path = download_file_from_cloud(model_id='songkey/ESRGAN', file_name='RealESRGAN_x2plus.pth', modelscope=modelscope) loadnet = torch.load(model_path, map_location=torch.device('cpu')) # prefer to use params_ema diff --git a/hellomeme/tools/utils.py b/hellomeme/tools/utils.py index aedbebd..5786218 100644 --- a/hellomeme/tools/utils.py +++ b/hellomeme/tools/utils.py @@ -144,6 +144,7 @@ def creat_model_from_cloud(model_cls, model_id, modelscope=False, cache_dir=None, + subfolder=None, hf_token='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'): if modelscope: from modelscope import snapshot_download @@ -153,10 +154,10 @@ def creat_model_from_cloud(model_cls, print(e) assert False, "@@ Failed to download model from modelscope" - model = model_cls.from_pretrained(model_path) + model = model_cls.from_pretrained(model_path, subfolder=subfolder) else: try: - model = model_cls.from_pretrained(model_id, cache_dir=cache_dir, token=hf_token) + model = model_cls.from_pretrained(model_id, subfolder=subfolder, cache_dir=cache_dir, token=hf_token) except Exception as e: print(e) assert False, "@@ Failed to download model from huggingface hub" diff --git a/hellomeme/utils.py b/hellomeme/utils.py index a070d5d..e6eced5 100644 --- a/hellomeme/utils.py +++ b/hellomeme/utils.py @@ -41,16 +41,8 @@ def get_torch_device(gpu_id=-1): return torch.device(f"cuda:{gpu_id}") def load_face_toolkits(dtype=torch.float16, gpu_id=-1, modelscope=False): - if modelscope: - from modelscope import snapshot_download - pd_fpg_dir = snapshot_download('songkey/pd_fgc_motion') - clip_image_dir = snapshot_download('songkey/IP-Adapter') - else: - pd_fpg_dir = 'songkey/pd_fgc_motion' - clip_image_dir = 'songkey/IP-Adapter' - - pd_fpg_motion = FanEncoder.from_pretrained(pd_fpg_dir) - image_encoder = CLIPVisionModelWithProjection.from_pretrained(clip_image_dir, subfolder='models/image_encoder') + pd_fpg_motion = creat_model_from_cloud(FanEncoder, 'songkey/pd_fgc_motion') + image_encoder = creat_model_from_cloud(CLIPVisionModelWithProjection, 'songkey/IP-Adapter', subfolder='models/image_encoder') image_encoder.to(dtype=dtype).cpu() pd_fpg_motion.to(dtype=dtype).cpu() diff --git a/pyproject.toml b/pyproject.toml index fc4b0ba..0ab8d3c 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.18" +version = "1.3.19" license = {file = "LICENSE"} [project.urls]