From 439ecf74367e8632aab6cf1a80587c97f110cb1e Mon Sep 17 00:00:00 2001 From: Isi Date: Sat, 7 Sep 2024 19:01:16 +0100 Subject: [PATCH] Update uniAnimate_Inference.py Added two nodes --- uniAnimate_Inference.py | 106 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 104 insertions(+), 2 deletions(-) diff --git a/uniAnimate_Inference.py b/uniAnimate_Inference.py index 31fdf12..e03ffb6 100644 --- a/uniAnimate_Inference.py +++ b/uniAnimate_Inference.py @@ -2,7 +2,9 @@ import torch from .utils.config import Config from .tools.inferences import inference_unianimate_entrance +from .tools.inferences import inference_unianimate_long_entrance from . import run_align_pose +from . import run_align_posev2 # from tools import * @@ -13,7 +15,7 @@ class UniAnimateImage: def INPUT_TYPES(cls): return { "required": { - "seed": ("INT", {"default": 11, "min": 0, "max": 1000000000, "step": 1}), + "seed": ("INT", {"default": 11, "min": 0, "max": 1000000000000, "step": 1}), "steps": ("INT", {"default": 30, "min": 25, "max": 50, "step": 1}), "useFirstFrame": ("BOOLEAN", { "default": False }), "reference_image": ("IMAGE",), # single image @@ -30,7 +32,7 @@ class UniAnimateImage: CATEGORY = "image" def process(self, seed, steps, useFirstFrame, reference_image, ref_pose, pose_sequence, frame_interval, max_frames, resolution_x): - cfg_update = Config(load=True) + cfg_update = Config('configs/UniAnimate_infer.yaml', load=True) resolution_y = 768 if resolution_x == 768: resolution_y = 1216 @@ -68,17 +70,117 @@ class Gen_align_pose: print("CUDA is not available") poses, refPose = run_align_pose.mp_main(reference_image, video) return (refPose, poses) + + + +class UniAnimateImageLong: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "seed": ("INT", {"default": 7, "min": 0, "max": 1000000000000, "step": 1}), + "steps": ("INT", {"default": 30, "min": 20, "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": 1000000000000, "step": 1}), + "steps": ("INT", {"default": 30, "min": 20, "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("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) NODE_CLASS_MAPPINGS = { "UniAnimateImage" : UniAnimateImage, "Gen_align_pose" : Gen_align_pose, + "UniAnimateImageLong" : UniAnimateImageLong, + "ReposeImage" : ReposeImage, } NODE_DISPLAY_NAME_MAPPINGS = { "UniAnimateImage" :"Animate image with UniAnimate", "Gen_align_pose" :"Align & Generate poses for UniAnimate", + "UniAnimateImageLong" :"Animate image with UniAnimate_Long", + "ReposeImage" :"Repose image with UniAnimate", }