nodes/MrForExample-ComfyUI-AnimateAnyone-Evolved
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
ea56d60cded760be34c0fe3c93caa858f74af891
MrForExample-ComfyUI-Animat…/nodes.py
T
MrForExample 5c7a6a3f9b Update After RCFG experiment 
- Update readme to log the RCFG experiment
- Refactor code
- Add delta parameter on Animate_Anyone_Sampler node
- Add ComfyUI Sampler progress bar
2024-01-21 18:02:27 +01:00

331 lines
12 KiB
Python
Raw Blame History

import os
import inspect
from collections import OrderedDict
from omegaconf import OmegaConf
from comfy import latent_formats
import folder_paths as comfy_paths
import torch
from diffusers.image_processor import VaeImageProcessor
from diffusers.schedulers import (
DDIMScheduler,
DPMSolverMultistepScheduler,
LCMScheduler,
EulerDiscreteScheduler,
EulerAncestralDiscreteScheduler,
LMSDiscreteScheduler,
PNDMScheduler,
)
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 AADiffusion
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "./ComfyUI-AnimateAnyone-Evolved")
DEFAULT_CONFIG_PATH = os.path.join(ROOT_PATH, "./configs/default.yaml")
CONFIG = OmegaConf.load(DEFAULT_CONFIG_PATH)
DEVICE = 'cuda'
WEIGHT_DETYPE = torch.float16
SCHEDULER_DICT = OrderedDict([
("DDIM", DDIMScheduler),
("DPM++ 2M Karras", DPMSolverMultistepScheduler),
("LCM", LCMScheduler),
("Euler", EulerDiscreteScheduler),
("Euler Ancestral", EulerAncestralDiscreteScheduler),
("LMS", LMSDiscreteScheduler),
("PNDM", PNDMScheduler),
])
class Animate_Anyone_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 = "animate_anyone"
CATEGORY = "AnimateAnyone-Evolved"
def animate_anyone(
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 = AADiffusion(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}, )
class Load_UNet2D_ConditionModel:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pretrained_base_unet_folder_path": ("STRING", {"default": "./pretrained_weights/stable-diffusion-v1-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 = "AnimateAnyone-Evolved/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_UNet3D_ConditionModel:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pretrained_base_unet_folder_path": ("STRING", {"default": "./pretrained_weights/stable-diffusion-v1-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 = "AnimateAnyone-Evolved/loaders"
def load_unet3d(self, 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)
unet3d = UNet3DConditionModel.from_pretrained_2d(
pretrained_base_unet_folder_path,
motion_module_path,
unet_additional_kwargs=CONFIG.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 = "AnimateAnyone-Evolved/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",),
},
}
RETURN_TYPES = (
"POSE_LATENT",
)
RETURN_NAMES = (
"pose_latent",
)
FUNCTION = "pose_guider_encode"
CATEGORY = "AnimateAnyone-Evolved/processors"
def pose_guider_encode(self, pose_guider, pose_images):
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=CONFIG.input_img_height, width=CONFIG.input_img_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,)
Reference in New Issue View Git Blame Copy Permalink