nodes/HelloVision-ComfyUI_HelloMeme
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

arc opt that do not affect functionality

Browse Source
This commit is contained in:
songkey-win
2024-12-12 11:52:18 +08:00
parent 0d46bb867f
commit 0924445042
3 changed files with 143 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hellomeme/pipelines/pipline_hm_video.py
+30 -1
View File
@@ -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
hellomeme/utils.py
+72 -23
View File
@@ -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 = {}
meme.py
+41 -74
View File
@@ -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