500 lines
18 KiB
Python
500 lines
18 KiB
Python
import os
|
|
import folder_paths as comfy_paths
|
|
import torch
|
|
import inspect
|
|
from tqdm import tqdm
|
|
from omegaconf import OmegaConf
|
|
from collections import OrderedDict
|
|
from mmpose.apis import MMPoseInferencer
|
|
|
|
from diffusers.image_processor import VaeImageProcessor
|
|
from diffusers.schedulers import (
|
|
DDIMScheduler,
|
|
DPMSolverMultistepScheduler,
|
|
LCMScheduler,
|
|
EulerDiscreteScheduler,
|
|
EulerAncestralDiscreteScheduler,
|
|
LMSDiscreteScheduler,
|
|
PNDMScheduler,
|
|
)
|
|
|
|
from comfy import latent_formats
|
|
from .src.pose.dwpose import DWposeDetector
|
|
from .src.pose.draw_pose import get_info, draw_pose
|
|
from .src.models.pose_guider import PoseGuider
|
|
from .src.models.unet_2d_condition import UNet2DConditionModel
|
|
from .src.models.unet_3d import UNet3DConditionModel
|
|
from .src.models.main_diffuser import DPDiffusion
|
|
from .src.utils.align import align_process
|
|
|
|
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "dynamicpose")
|
|
DEVICE = 'cuda'
|
|
WEIGHT_DETYPE = torch.float16
|
|
DEFAULT_CONFIG1_PATH = os.path.join(ROOT_PATH, "./configs/inference/inference_v1.yaml")
|
|
DEFAULT_CONFIG2_PATH = os.path.join(ROOT_PATH, "./configs/inference/inference_v2.yaml")
|
|
CONFIG1 = OmegaConf.load(DEFAULT_CONFIG1_PATH)
|
|
CONFIG2 = OmegaConf.load(DEFAULT_CONFIG2_PATH)
|
|
|
|
SCHEDULER_DICT = OrderedDict([
|
|
("DDIM", DDIMScheduler),
|
|
("DPM++ 2M Karras", DPMSolverMultistepScheduler),
|
|
("LCM", LCMScheduler),
|
|
("Euler", EulerDiscreteScheduler),
|
|
("Euler Ancestral", EulerAncestralDiscreteScheduler),
|
|
("LMS", LMSDiscreteScheduler),
|
|
("PNDM", PNDMScheduler),
|
|
])
|
|
|
|
class align:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return{
|
|
"required": {
|
|
"ref_image": ('IMAGE',),
|
|
"ref_info": ('POSEINFOS',),
|
|
"pose_images" :('IMAGE',),
|
|
"pose_info": ('POSEINFOS',),
|
|
},
|
|
"optional":{
|
|
"height": ("INT",{"default": 784, "min": 100, "max": 4000}),
|
|
"width": ("INT",{"default": 512, "min": 100, "max": 4000}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","IMAGE")
|
|
RETURN_NAMES = ('aligend_ref_image',"aligned_pose_images",)
|
|
FUNCTION = "align_pose"
|
|
CATEGORY = "DynamicPose"
|
|
|
|
def align_pose(self,ref_image,ref_info,pose_images,pose_info,height,width):
|
|
ref_img_align, pose_imaegs_align = align_process(ref_image,ref_info,pose_images,pose_info,(width,height))
|
|
return (torch.from_numpy(ref_img_align).unsqueeze(0), torch.from_numpy(pose_imaegs_align),)
|
|
|
|
|
|
class load_pose_model:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return{
|
|
"optional": {
|
|
"body_model": (["rtmpose-x", "dwpose"],),
|
|
"hand_model": (["rtmw-x", "dwpose"],),
|
|
"foot_model": (["dwpose","rtmw-x"],),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("POSEMODELS",)
|
|
RETURN_NAMES = ("pose_models",)
|
|
FUNCTION = "load_pose"
|
|
CATEGORY = "DynamicPose/loaders"
|
|
|
|
def load_pose(self,body_model,hand_model,foot_model):
|
|
model_dir = ROOT_PATH
|
|
if foot_model == 'dwpose':
|
|
model_path = model_dir + '/pretrained_weights/dwpose'
|
|
inferencer_foot = DWposeDetector(foot = True, model_path = model_path).to("cuda")
|
|
else:
|
|
print('not support yet')
|
|
|
|
if body_model == "rtmpose-x":
|
|
inferencer_body = MMPoseInferencer(
|
|
pose2d = model_dir + '/pretrained_weights/rtmpose/rtmpose-x_8xb256-700e_coco-384x288.py',
|
|
pose2d_weights= model_dir + '/pretrained_weights/rtmpose/rtmpose-x_simcc-body7_pt-body7_700e-384x288-71d7b7e9_20230629.pth',
|
|
det_model= model_dir + '/pretrained_weights/rtmpose/rtmdet_m_640-8xb32_coco-person.py',
|
|
det_weights = model_dir + '/pretrained_weights/rtmpose/rtmdet_m_8xb32-100e_coco-obj365-person-235e8209.pth',
|
|
det_cat_ids=[0],
|
|
)
|
|
else:
|
|
print('not support yet')
|
|
|
|
if hand_model == "rtmw-x":
|
|
inferencer_hand = MMPoseInferencer(
|
|
pose2d = model_dir + '/pretrained_weights/rtmpose/rtmw-x_8xb320-270e_cocktail14-384x288.py',
|
|
pose2d_weights = model_dir + '/pretrained_weights/rtmpose/rtmw-x_simcc-cocktail14_pt-ucoco_270e-384x288-f840f204_20231122.pth',
|
|
det_model = model_dir + '/pretrained_weights/rtmpose/rtmdet_m_640-8xb32_coco-person.py',
|
|
det_weights = model_dir + '/pretrained_weights/rtmpose/rtmdet_m_8xb32-100e_coco-obj365-person-235e8209.pth',
|
|
det_cat_ids=[0],
|
|
)
|
|
else:
|
|
print('not support yet')
|
|
|
|
return ([inferencer_body,inferencer_hand,inferencer_foot],)
|
|
|
|
|
|
class pose_extraction:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return{
|
|
"required":{
|
|
"images": ("IMAGE",),
|
|
"models": ("POSEMODELS",)
|
|
},
|
|
"optional": {
|
|
"show_org" : ("BOOLEAN",{"default":False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","POSEINFOS" )
|
|
RETURN_NAMES = ("pose_images","pose_infos")
|
|
FUNCTION = "process"
|
|
CATEGORY = "DynamicPose"
|
|
|
|
def process(self,images,models,show_org):
|
|
inferencer_body = models[0]
|
|
inferencer_hand = models[1]
|
|
inferencer_foot = models[2]
|
|
pose_frames = []
|
|
pose_infos = []
|
|
for image in tqdm(images):
|
|
frame = image.numpy() * 255
|
|
H,W = frame.shape[0:2]
|
|
body = inferencer_body(frame)
|
|
body = next(body)['predictions'][0][0]
|
|
body_w = inferencer_hand(frame)
|
|
hand = next(body_w)['predictions'][0][0]
|
|
feet, feet_scores = inferencer_foot(frame)
|
|
feet = feet.tolist()[0]
|
|
for points in feet:
|
|
points[0] = points[0] * W
|
|
points[1] = points[1] * H
|
|
foot = {
|
|
'keypoints':feet,
|
|
'keypoint_scores':feet_scores.tolist()[0],
|
|
}
|
|
info = get_info(body,hand,foot, W, H)
|
|
pose_infos.append(info)
|
|
|
|
if not show_org:
|
|
detected_map = draw_pose(info)
|
|
else:
|
|
detected_map = draw_pose(info, frame)
|
|
|
|
pose_frames.append(detected_map)
|
|
|
|
|
|
pose_images = torch.cat(pose_frames, dim=0)
|
|
|
|
return (pose_images,pose_infos,)
|
|
|
|
class Load_reference_unet:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"optional": {
|
|
"pretrained_base_unet_folder_path": ("STRING", {"default": ".../models/diffusers/SD1.5/unet/", "multiline": False}),
|
|
"unet2d_model_path": ("STRING", {"default": "./pretrained_weights/reference_unet.pth", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"UNET2D",
|
|
)
|
|
RETURN_NAMES = (
|
|
"unet2d",
|
|
)
|
|
FUNCTION = "load_unet2d"
|
|
|
|
CATEGORY = "DynamicPose/loaders"
|
|
|
|
def load_unet2d(self, pretrained_base_unet_folder_path, unet2d_model_path):
|
|
if not os.path.isabs(pretrained_base_unet_folder_path):
|
|
pretrained_base_unet_folder_path = os.path.join(ROOT_PATH, pretrained_base_unet_folder_path)
|
|
if not os.path.isabs(unet2d_model_path):
|
|
unet2d_model_path = os.path.join(ROOT_PATH, unet2d_model_path)
|
|
|
|
|
|
unet2d = UNet2DConditionModel.from_pretrained(
|
|
pretrained_base_unet_folder_path,
|
|
).to(dtype=WEIGHT_DETYPE, device=DEVICE)
|
|
|
|
unet2d.load_state_dict(
|
|
torch.load(unet2d_model_path, map_location="cpu"),
|
|
)
|
|
|
|
return (unet2d,)
|
|
|
|
class Load_denoising_unet:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"optional": {
|
|
"type": (['image', 'video'],),
|
|
"pretrained_base_unet_folder_path": ("STRING", {"default": ".../models/diffusers/SD1.5/unet/", "multiline": False}),
|
|
"unet3d_model_path": ("STRING", {"default": "./pretrained_weights/denoising_unet.pth", "multiline": False}),
|
|
"motion_module_path": ("STRING", {"default": "./pretrained_weights/motion_module.pth", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"UNET3D",
|
|
)
|
|
RETURN_NAMES = (
|
|
"unet3d",
|
|
)
|
|
FUNCTION = "load_unet3d"
|
|
|
|
CATEGORY = "DynamicPose/loaders"
|
|
|
|
def load_unet3d(self, type, pretrained_base_unet_folder_path, unet3d_model_path, motion_module_path):
|
|
|
|
if not os.path.isabs(pretrained_base_unet_folder_path):
|
|
pretrained_base_unet_folder_path = os.path.join(ROOT_PATH, pretrained_base_unet_folder_path)
|
|
if not os.path.isabs(unet3d_model_path):
|
|
unet3d_model_path = os.path.join(ROOT_PATH, unet3d_model_path)
|
|
if not os.path.isabs(motion_module_path):
|
|
motion_module_path = os.path.join(ROOT_PATH, motion_module_path)
|
|
|
|
if type == 'video':
|
|
unet3d = UNet3DConditionModel.from_pretrained_2d(
|
|
pretrained_base_unet_folder_path,
|
|
motion_module_path,
|
|
unet_additional_kwargs=CONFIG2.unet_additional_kwargs,
|
|
).to(dtype=WEIGHT_DETYPE, device=DEVICE)
|
|
else:
|
|
unet3d = UNet3DConditionModel.from_pretrained_2d(
|
|
pretrained_base_unet_folder_path,
|
|
"",
|
|
unet_additional_kwargs=CONFIG1.unet_additional_kwargs,
|
|
).to(dtype=WEIGHT_DETYPE, device=DEVICE)
|
|
|
|
unet3d.load_state_dict(
|
|
torch.load(unet3d_model_path, map_location="cpu"),
|
|
strict=False,
|
|
)
|
|
|
|
return (unet3d,)
|
|
|
|
class Load_Pose_Guider:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pose_guider_model_path": ("STRING", {"default": "./pretrained_weights/pose_guider.pth", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"POSE_GUIDER",
|
|
)
|
|
RETURN_NAMES = (
|
|
"pose_guider",
|
|
)
|
|
FUNCTION = "load_pose_guider"
|
|
|
|
CATEGORY = "DynamicPose/loaders"
|
|
|
|
def load_pose_guider(self, pose_guider_model_path):
|
|
if not os.path.isabs(pose_guider_model_path):
|
|
pose_guider_model_path = os.path.join(ROOT_PATH, pose_guider_model_path)
|
|
|
|
pose_guider = PoseGuider(320, block_out_channels=(16, 32, 96, 256)).to(
|
|
dtype=WEIGHT_DETYPE, device=DEVICE
|
|
)
|
|
pose_guider.load_state_dict(
|
|
torch.load(pose_guider_model_path, map_location="cpu"),
|
|
)
|
|
|
|
return (pose_guider,)
|
|
|
|
class Pose_Guider_Encode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pose_guider": ("POSE_GUIDER",),
|
|
"pose_images": ("IMAGE",),
|
|
},
|
|
"optional":{
|
|
"height": ("INT",{"default": 784, "min": 100, "max": 4000}),
|
|
"width": ("INT",{"default": 512, "min": 100, "max": 4000}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"POSE_LATENT",
|
|
)
|
|
RETURN_NAMES = (
|
|
"pose_latent",
|
|
)
|
|
FUNCTION = "pose_guider_encode"
|
|
|
|
CATEGORY = "DynamicPose/processors"
|
|
|
|
def pose_guider_encode(self, pose_guider, pose_images,height, width):
|
|
|
|
cond_image_processor = VaeImageProcessor(
|
|
do_convert_rgb=True,
|
|
do_normalize=False,
|
|
)
|
|
|
|
# (b, h, w, c) -> (b, c, h, w)
|
|
pose_images = pose_images.permute(0, 3, 1, 2).to(DEVICE, dtype=WEIGHT_DETYPE)
|
|
|
|
# Prepare a list of pose condition images
|
|
pose_cond_tensor_list = []
|
|
for pose_image in pose_images:
|
|
pose_cond_tensor = cond_image_processor.preprocess(
|
|
pose_image, height=height, width=width
|
|
)
|
|
pose_cond_tensor = pose_cond_tensor.unsqueeze(2) # (b, c, 1, h, w)
|
|
pose_cond_tensor_list.append(pose_cond_tensor)
|
|
|
|
pose_cond_tensor = torch.cat(pose_cond_tensor_list, dim=2) # (b, c, f, h, w)
|
|
pose_cond_tensor = pose_cond_tensor.to(
|
|
device=DEVICE, dtype=pose_guider.dtype
|
|
)
|
|
pose_latent = pose_guider(pose_cond_tensor)
|
|
#print(f"pose_cond_tensor.shape: {pose_cond_tensor.shape}\pose_latent.shape: {pose_latent.shape}")
|
|
#pose_cond_tensor.shape: torch.Size([1, 3, 24, 768, 512]) pose_latent.shape: torch.Size([1, 320, 24, 96, 64])
|
|
|
|
return (pose_latent,)
|
|
|
|
class DynamicPose_Sampler:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"reference_unet": ("UNET2D",),
|
|
"denoising_unet": ("UNET3D",),
|
|
"ref_image_latent": ("LATENT",),
|
|
"clip_image_embeds": ("CLIP_VISION_OUTPUT",),
|
|
"pose_latent": ("POSE_LATENT",),
|
|
"seed": ("INT", {"default": 999999999, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
|
|
"delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
|
|
"context_frames": ("INT", {"default": 24, "min": 1}),
|
|
"context_stride": ("INT", {"default": 1, "min": 1}),
|
|
"context_overlap": ("INT", {"default": 4, "min": 0}),
|
|
"context_batch_size": ("INT", {"default": 1, "min": 1}),
|
|
"interpolation_factor": ("INT", {"default": 1, "min": 0}),
|
|
"sampler_scheduler_pairs": (list(SCHEDULER_DICT.keys()),),
|
|
"beta_start": ("FLOAT", {"default": 0.00085, "min": 0.0, "step":0.00001}),
|
|
"beta_end": ("FLOAT", {"default": 0.012, "min": 0.0, "step":0.00001}),
|
|
"beta_schedule": (["linear", "scaled_linear", "squaredcos_cap_v2"],),
|
|
"prediction_type": (["v_prediction", "epsilon", "sample"],),
|
|
"timestep_spacing": (["trailing", "linspace", "leading"],),
|
|
"steps_offset": ("INT", {"default": 1, "min": 0, "max": 10000}),
|
|
},
|
|
"optional": {
|
|
"clip_sample": ("BOOLEAN", {"default": False},),
|
|
"rescale_betas_zero_snr": ("BOOLEAN", {"default": True},),
|
|
"use_lora": ("BOOLEAN", {"default": False},),
|
|
"lora_name": (comfy_paths.get_filename_list("loras"),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"LATENT",
|
|
)
|
|
RETURN_NAMES = (
|
|
"latent",
|
|
)
|
|
FUNCTION = "inferrence"
|
|
|
|
CATEGORY = "DynamicPose"
|
|
|
|
def inferrence(
|
|
self,
|
|
reference_unet,
|
|
denoising_unet,
|
|
ref_image_latent,
|
|
clip_image_embeds,
|
|
pose_latent,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
delta,
|
|
context_frames,
|
|
context_stride,
|
|
context_overlap,
|
|
context_batch_size,
|
|
interpolation_factor,
|
|
sampler_scheduler_pairs,
|
|
beta_start,
|
|
beta_end,
|
|
beta_schedule,
|
|
prediction_type,
|
|
timestep_spacing,
|
|
steps_offset,
|
|
clip_sample=False,
|
|
rescale_betas_zero_snr=True,
|
|
use_lora=False,
|
|
lora_name=None
|
|
):
|
|
|
|
latent_format = latent_formats.SD15()
|
|
|
|
# encoder_hidden_states.shape: torch.Size([1, 1, 768]) clip_image_embeds.shape: torch.Size([1, 768])
|
|
encoder_hidden_states = clip_image_embeds["image_embeds"].unsqueeze(1).to(DEVICE, dtype=WEIGHT_DETYPE)
|
|
|
|
# forward reference image latent with shape (1, 4, 96, 64) to reference net
|
|
ref_image_latent = latent_format.process_in(ref_image_latent["samples"]).to(DEVICE, dtype=WEIGHT_DETYPE)
|
|
|
|
# setup scheduler from user inputs
|
|
scheduler_class = SCHEDULER_DICT[sampler_scheduler_pairs]
|
|
|
|
sched_kwargs = {
|
|
"beta_start": beta_start,
|
|
"beta_end": beta_end,
|
|
"beta_schedule": beta_schedule,
|
|
"steps_offset": steps_offset,
|
|
"prediction_type": prediction_type,
|
|
"timestep_spacing": timestep_spacing,
|
|
}
|
|
|
|
if "clip_sample" in inspect.signature(scheduler_class.__init__).parameters:
|
|
sched_kwargs["clip_sample"] = clip_sample
|
|
if "rescale_betas_zero_snr" in inspect.signature(scheduler_class.__init__).parameters:
|
|
sched_kwargs["rescale_betas_zero_snr"] = rescale_betas_zero_snr
|
|
|
|
scheduler = scheduler_class(**sched_kwargs)
|
|
|
|
# setup diffuser and then denoise
|
|
diffuser = DPDiffusion(reference_unet, denoising_unet, scheduler)
|
|
|
|
if use_lora:
|
|
lora_path = comfy_paths.get_full_path("loras", lora_name)
|
|
diffuser.load_lora(lora_path)
|
|
|
|
samples = diffuser(
|
|
steps,
|
|
cfg,
|
|
delta,
|
|
ref_image_latent,
|
|
pose_latent,
|
|
encoder_hidden_states,
|
|
seed,
|
|
context_frames=context_frames,
|
|
context_stride=context_stride,
|
|
context_overlap=context_overlap,
|
|
context_batch_size=context_batch_size,
|
|
interpolation_factor=interpolation_factor,
|
|
)
|
|
samples = latent_format.process_out(samples)
|
|
|
|
# (1, 4, f, h, w) -> (f, 4, h, w)
|
|
samples = torch.squeeze(samples, 0).permute(1, 0, 2, 3)
|
|
|
|
return ({"samples":samples}, )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"pose_extraction" : pose_extraction,
|
|
"Load_reference_unet" : Load_reference_unet,
|
|
"Load_denoising_unet": Load_denoising_unet,
|
|
"Load_Pose_Guider": Load_Pose_Guider,
|
|
"Pose_Guider_Encode": Pose_Guider_Encode,
|
|
"DynamicPose_Sampler": DynamicPose_Sampler,
|
|
"load_pose_model": load_pose_model,
|
|
"align": align,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"pose_extraction" : "pose_extraction",
|
|
"Load_reference_unet" : 'Load_reference_unet',
|
|
"Load_denoising_unet": 'Load_denoising_unet',
|
|
"Load_Pose_Guider": 'Load_Pose_Guider',
|
|
"Pose_Guider_Encode": 'Pose_Guider_Encode',
|
|
"DynamicPose_Sampler": 'DynamicPose_Sampler',
|
|
"load_pose_model": 'load_pose_model',
|
|
"align":"align",
|
|
}
|