464 lines
19 KiB
Python
464 lines
19 KiB
Python
import torch
|
|
|
|
from .utils.config import Config
|
|
from .tools.inferences import inference_unianimate_entrance
|
|
from .tools.inferences import inference_unianimate_long_entrance
|
|
from .tools.inferences import inference_animate_x_entrance
|
|
from .tools.inferences import inference_animate_x_long_entrance
|
|
from .tools.inferences import inference_animate_x_entrance_v2
|
|
from . import run_align_pose
|
|
from . import run_align_posev2
|
|
from . import run_align_pose_Animate_X
|
|
from . import original_run_align_pose_animate_X
|
|
from nodes import MAX_RESOLUTION
|
|
|
|
# from tools import *
|
|
|
|
|
|
|
|
class UniAnimateImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": 11, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"steps": ("INT", {"default": 30, "min": 10, "max": 50, "step": 1}),
|
|
"useFirstFrame": ("BOOLEAN", { "default": False }),
|
|
"reference_image": ("IMAGE",), # single image
|
|
"ref_pose": ("IMAGE",), # single image
|
|
"pose_sequence": ("IMAGE",), # Batch of pose images
|
|
"frame_interval": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1}),
|
|
"max_frames": ("INT", {"default": 32, "min": 1, "max": 1024, "step": 1}),
|
|
"resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("video", "empty_mask")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, seed, steps, useFirstFrame, reference_image, ref_pose, pose_sequence, frame_interval, max_frames, resolution_x):
|
|
cfg_update = Config('configs/UniAnimate_infer.yaml', load=True)
|
|
resolution_y = 768
|
|
if resolution_x == 768:
|
|
resolution_y = 1216
|
|
resolution = [resolution_x, resolution_y]
|
|
print("Ready for inference.")
|
|
|
|
# print(f"image is: {reference_image}")
|
|
|
|
frames = inference_unianimate_entrance(seed, steps, useFirstFrame, reference_image, ref_pose, pose_sequence, frame_interval, max_frames, resolution, cfg_update=cfg_update.cfg_dict)
|
|
mask_template = torch.zeros((1, resolution_y, resolution_x), dtype=torch.float32)
|
|
masks = [mask_template.clone() for _ in range(len(pose_sequence))]
|
|
masks = torch.cat(masks, dim=0)
|
|
return (frames, masks)
|
|
|
|
class Gen_align_pose:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"reference_image": ("IMAGE",), # single image
|
|
"video": ("IMAGE",), # video
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("ref_pose", "pose_seq")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, reference_image, video):
|
|
if torch.cuda.is_available():
|
|
print(f"CUDA version: {torch.version.cuda}")
|
|
print(f"CUDNN version: {torch.backends.cudnn.version()}")
|
|
print(f"Device name: {torch.cuda.get_device_name(0)}")
|
|
else:
|
|
print("CUDA is not available")
|
|
poses, refPose = run_align_pose.mp_main(reference_image, video)
|
|
return (refPose, poses)
|
|
|
|
|
|
class Gen_align_pose2:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"reference_image": ("IMAGE",), # single image
|
|
"video": ("IMAGE",), # video
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("ref_pose", "pose_seq")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, reference_image, video):
|
|
if torch.cuda.is_available():
|
|
print(f"CUDA version: {torch.version.cuda}")
|
|
print(f"CUDNN version: {torch.backends.cudnn.version()}")
|
|
print(f"Device name: {torch.cuda.get_device_name(0)}")
|
|
else:
|
|
print("CUDA is not available")
|
|
poses, refPose = original_run_align_pose_animate_X.mp_main(reference_image, video)
|
|
return (refPose, poses)
|
|
|
|
|
|
|
|
class UniAnimateImageLong:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": 7, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"steps": ("INT", {"default": 30, "min": 10, "max": 50, "step": 1}),
|
|
"useFirstFrame": ("BOOLEAN", { "default": False }),
|
|
"dontAlignPose": ("BOOLEAN", { "default": False }),
|
|
"image": ("IMAGE",), # single image
|
|
"video": ("IMAGE",), # Batch of pose images
|
|
"frame_interval": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}),
|
|
"context_size": ("INT", {"default": 32, "min": 16, "max": 64, "step": 16}),
|
|
"context_stride": ("INT", {"default": 1, "min": 1, "max": 4, "step": 1}),
|
|
"context_overlap": ("INT", {"default": 8, "min": 4, "max": 16, "step": 4}),
|
|
"max_frames": ("INT", {"default": 1024000, "min": 16, "max": 1024000, "step": 1}),
|
|
"resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("video", "poses")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, seed, steps, useFirstFrame, dontAlignPose, image, video, frame_interval, context_size, context_stride, context_overlap, max_frames, resolution_x):
|
|
cfg_update = Config('configs/UniAnimate_infer_long.yaml', load=True)
|
|
resolution_y = 768
|
|
if resolution_x == 768:
|
|
resolution_y = 1216
|
|
resolution = [resolution_x, resolution_y]
|
|
|
|
context_size = context_size
|
|
context_overlap = context_overlap
|
|
max_frames = max_frames
|
|
if context_size == context_overlap:
|
|
context_overlap = 8
|
|
print("context_size equal to context_overlap; context_overlap changed to default.")
|
|
|
|
if context_size > max_frames or context_size > len(video) :
|
|
context_size = 32
|
|
print("context_size greater than max_frames; context_size changed to default.")
|
|
|
|
if max_frames < 32 or len(video) < 32:
|
|
context_size = 16
|
|
context_overlap = 4
|
|
print("Video frames less than 32; context_size changed to 16 and context_overlap changed to 4.")
|
|
|
|
pose_sequence, refPose = run_align_posev2.mp_main(dontAlignPose, image, video)
|
|
|
|
print("Ready for inference.")
|
|
|
|
frames = inference_unianimate_long_entrance(seed, steps, useFirstFrame, image, refPose, pose_sequence, frame_interval, context_size, context_stride, context_overlap, max_frames, resolution, cfg_update=cfg_update.cfg_dict)
|
|
|
|
return (frames, pose_sequence)
|
|
|
|
|
|
|
|
class ReposeImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": 11, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"steps": ("INT", {"default": 30, "min": 10, "max": 50, "step": 1}),
|
|
"dontAlignPose": ("BOOLEAN", { "default": False }),
|
|
"image": ("IMAGE",),
|
|
"pose": ("IMAGE",),
|
|
"resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("newPose", "pose")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, seed, steps, dontAlignPose, image, pose, resolution_x):
|
|
cfg_update = Config('configs/UniAnimate_infer.yaml', load=True)
|
|
resolution_y = 768
|
|
if resolution_x == 768:
|
|
resolution_y = 1216
|
|
resolution = [resolution_x, resolution_y]
|
|
|
|
pose_i, refPose = run_align_posev2.mp_main(dontAlignPose, image, pose)
|
|
|
|
print(f"shape of image: {image.shape}")
|
|
print(f"shape of pose: {pose.shape}")
|
|
|
|
|
|
print("Ready for inference.")
|
|
frame = inference_unianimate_entrance(seed, steps, False, image, refPose, pose_i, 1, 1, resolution, cfg_update=cfg_update.cfg_dict)
|
|
# mask_template = torch.zeros((1, resolution_y, resolution_x), dtype=torch.float32)
|
|
# masks = [mask_template.clone() for _ in range(len(pose_i))]
|
|
# masks = torch.cat(masks, dim=0)
|
|
|
|
return (frame, pose_i)
|
|
|
|
|
|
|
|
|
|
|
|
class Animate_X_ReposeImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": 13, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"steps": ("INT", {"default": 25, "min": 10, "max": 50, "step": 1}),
|
|
"dontAlignPose": ("BOOLEAN", { "default": False }),
|
|
"image": ("IMAGE",),
|
|
"pose": ("IMAGE",),
|
|
"resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("newPose", "pose")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, seed, steps, dontAlignPose, image, pose, resolution_x):
|
|
cfg_update = Config('configs/Animate_X_infer.yaml', load=True)
|
|
resolution_y = 768
|
|
if resolution_x == 768:
|
|
resolution_y = 1216
|
|
resolution = [resolution_x, resolution_y]
|
|
|
|
# pose_i, refPose, pose_embeding, refpose_embeding = run_align_pose_Animate_X.mp_main(dontAlignPose, image, pose)
|
|
pose_i, refPose = run_align_posev2.mp_main(dontAlignPose, image, pose)
|
|
|
|
# print(f"shape of image: {image.shape}")
|
|
# print(f"shape of pose: {pose.shape}")
|
|
|
|
|
|
print("Ready for inference.")
|
|
# frame = inference_animate_x_entrance(seed, steps, False, image, refPose, pose_i, pose, refpose_embeding, pose_embeding, 1, 1, resolution, cfg_update=cfg_update.cfg_dict)
|
|
frame = inference_animate_x_entrance(seed, steps, False, image, refPose, pose_i, 1, 1, resolution, cfg_update=cfg_update.cfg_dict)
|
|
# mask_template = torch.zeros((1, resolution_y, resolution_x), dtype=torch.float32)
|
|
# masks = [mask_template.clone() for _ in range(len(pose_i))]
|
|
# masks = torch.cat(masks, dim=0)
|
|
|
|
return (frame, pose_i)
|
|
|
|
|
|
|
|
|
|
class Animate_X_Image:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": 13, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"steps": ("INT", {"default": 30, "min": 10, "max": 50, "step": 1}),
|
|
"useFirstFrame": ("BOOLEAN", { "default": False }),
|
|
"reference_image": ("IMAGE",), # single image
|
|
"ref_pose": ("IMAGE",), # single image
|
|
"pose_sequence": ("IMAGE",), # Batch of pose images
|
|
"frame_interval": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1}),
|
|
"max_frames": ("INT", {"default": 32, "min": 2, "max": 1024, "step": 1}),
|
|
"resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("video", "empty_mask")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, seed, steps, useFirstFrame, reference_image, ref_pose, pose_sequence, frame_interval, max_frames, resolution_x):
|
|
cfg_update = Config('configs/Animate_X_infer.yaml', load=True)
|
|
resolution_y = 768
|
|
if resolution_x == 768:
|
|
resolution_y = 1216
|
|
resolution = [resolution_x, resolution_y]
|
|
print("Ready for inference.")
|
|
|
|
# print(f"image is: {reference_image}")
|
|
|
|
frames = inference_animate_x_entrance(seed, steps, useFirstFrame, reference_image, ref_pose, pose_sequence, frame_interval, max_frames, resolution, cfg_update=cfg_update.cfg_dict)
|
|
mask_template = torch.zeros((1, resolution_y, resolution_x), dtype=torch.float32)
|
|
masks = [mask_template.clone() for _ in range(len(pose_sequence))]
|
|
masks = torch.cat(masks, dim=0)
|
|
return (frames, masks)
|
|
|
|
|
|
|
|
class Animate_X_Image_Long:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": 13, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"steps": ("INT", {"default": 30, "min": 10, "max": 50, "step": 1}),
|
|
"useFirstFrame": ("BOOLEAN", { "default": False }),
|
|
"dontAlignPose": ("BOOLEAN", { "default": False }),
|
|
"image": ("IMAGE",), # single image
|
|
"video": ("IMAGE",), # Batch of pose images
|
|
"frame_interval": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}),
|
|
"context_size": ("INT", {"default": 32, "min": 16, "max": 64, "step": 16}),
|
|
"context_stride": ("INT", {"default": 1, "min": 1, "max": 4, "step": 1}),
|
|
"context_overlap": ("INT", {"default": 8, "min": 4, "max": 16, "step": 4}),
|
|
"max_frames": ("INT", {"default": 1024000, "min": 16, "max": 1024000, "step": 1}),
|
|
"resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("video", "poses")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, seed, steps, useFirstFrame, dontAlignPose, image, video, frame_interval, context_size, context_stride, context_overlap, max_frames, resolution_x):
|
|
cfg_update = Config('configs/Animate_X_infer_long.yaml', load=True)
|
|
resolution_y = 768
|
|
if resolution_x == 768:
|
|
resolution_y = 1216
|
|
resolution = [resolution_x, resolution_y]
|
|
|
|
context_size = context_size
|
|
context_overlap = context_overlap
|
|
max_frames = max_frames
|
|
if context_size == context_overlap:
|
|
context_overlap = 8
|
|
print("context_size equal to context_overlap; context_overlap changed to default.")
|
|
|
|
if context_size > max_frames or context_size > len(video) :
|
|
context_size = 32
|
|
print("context_size greater than max_frames; context_size changed to default.")
|
|
|
|
if max_frames < 32 or len(video) < 32:
|
|
context_size = 16
|
|
context_overlap = 4
|
|
print("Video frames less than 32; context_size changed to 16 and context_overlap changed to 4.")
|
|
|
|
pose_sequence, refPose = run_align_posev2.mp_main(dontAlignPose, image, video)
|
|
|
|
print("Ready for inference.")
|
|
|
|
frames = inference_animate_x_long_entrance(seed, steps, useFirstFrame, image, refPose, pose_sequence, frame_interval, context_size, context_stride, context_overlap, max_frames, resolution, cfg_update=cfg_update.cfg_dict)
|
|
|
|
return (frames, pose_sequence)
|
|
|
|
|
|
|
|
class Animate_X_ReposeImage_v2:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": 13, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"steps": ("INT", {"default": 25, "min": 10, "max": 50, "step": 1}),
|
|
"dontAlignPose": ("BOOLEAN", { "default": True }),
|
|
"image": ("IMAGE",),
|
|
"pose": ("IMAGE",),
|
|
"resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("newPose", "pose")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, seed, steps, dontAlignPose, image, pose, resolution_x):
|
|
cfg_update = Config('configs/Animate_X_infer.yaml', load=True)
|
|
resolution_y = 768
|
|
if resolution_x == 768:
|
|
resolution_y = 1216
|
|
resolution = [resolution_x, resolution_y]
|
|
|
|
# pose_i, refPose, pose_embeding, refpose_embeding = run_align_pose_Animate_X.mp_main(dontAlignPose, image, pose)
|
|
pose_i, refPose, pose_embeding, refpose_embeding = run_align_pose_Animate_X.mp_main(dontAlignPose, image, pose)
|
|
|
|
# print(f"shape of image: {image.shape}")
|
|
# print(f"shape of pose: {pose.shape}")
|
|
|
|
|
|
print("Ready for inference.")
|
|
frame = inference_animate_x_entrance_v2(seed, steps, False, image, refPose, pose_i, pose, refpose_embeding, pose_embeding, 1, 1, resolution, cfg_update=cfg_update.cfg_dict)
|
|
|
|
return (frame, pose_i)
|
|
|
|
|
|
|
|
class Animate_X_Image_v2:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"seed": ("INT", {"default": 13, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
|
"steps": ("INT", {"default": 25, "min": 10, "max": 50, "step": 1}),
|
|
"useFirstFrame": ("BOOLEAN", { "default": True }),
|
|
"dontAlignPose": ("BOOLEAN", { "default": True }),
|
|
"image": ("IMAGE",),
|
|
"video": ("IMAGE",),
|
|
"frame_interval": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1}),
|
|
"max_frames": ("INT", {"default": 32, "min": 2, "max": 1024, "step": 1}),
|
|
"resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("video", "poses")
|
|
FUNCTION = "process"
|
|
CATEGORY = "image"
|
|
|
|
def process(self, seed, steps, useFirstFrame, dontAlignPose, image, video, frame_interval, max_frames, resolution_x):
|
|
cfg_update = Config('configs/Animate_X_infer.yaml', load=True)
|
|
resolution_y = 768
|
|
if resolution_x == 768:
|
|
resolution_y = 1216
|
|
resolution = [resolution_x, resolution_y]
|
|
|
|
# pose_i, refPose, pose_embeding, refpose_embeding = run_align_pose_Animate_X.mp_main(dontAlignPose, image, pose)
|
|
pose_i, refPose, pose_embeding, refpose_embeding = run_align_pose_Animate_X.mp_main(dontAlignPose, image, video)
|
|
|
|
# print(f"shape of image: {image.shape}")
|
|
# print(f"shape of pose: {pose.shape}")
|
|
|
|
|
|
print("Ready for inference.")
|
|
frame = inference_animate_x_entrance_v2(seed, steps, useFirstFrame, image, refPose, pose_i, video, refpose_embeding, pose_embeding, frame_interval, max_frames, resolution, cfg_update=cfg_update.cfg_dict)
|
|
|
|
return (frame, pose_i)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UniAnimateImage" : UniAnimateImage,
|
|
"Gen_align_pose" : Gen_align_pose,
|
|
"Gen_align_pose2" : Gen_align_pose2,
|
|
"UniAnimateImageLong" : UniAnimateImageLong,
|
|
"ReposeImage" : ReposeImage,
|
|
"Animate_X_ReposeImage" : Animate_X_ReposeImage,
|
|
"Animate_X_ReposeImage_v2" : Animate_X_ReposeImage_v2,
|
|
"Animate_X_Image" : Animate_X_Image,
|
|
"Animate_X_Image_Long" : Animate_X_Image_Long,
|
|
"Animate_X_Image_v2" : Animate_X_Image_v2,
|
|
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"UniAnimateImage" :"Animate image with UniAnimate",
|
|
"Gen_align_pose" :"Align & Generate poses for UniAnimate",
|
|
"Gen_align_pose2" :"Generate dwpose",
|
|
"UniAnimateImageLong" :"Animate image with UniAnimate_Long",
|
|
"ReposeImage" :"Repose image with UniAnimate",
|
|
"Animate_X_ReposeImage" :"Repose image with Animate_X",
|
|
"Animate_X_ReposeImage_v2" :"Repose image with Animate_X_v2",
|
|
"Animate_X_Image" : "Animate image with Animate_X",
|
|
"Animate_X_Image_Long" : "Animate image with Animate_X_Long",
|
|
"Animate_X_Image_v2" : "Animate image with Animate_X_v2",
|
|
|
|
}
|