nodes/Isi-dev-ComfyUI-UniAnimate-W
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update uniAnimate_Inference.py

Browse Source 
Added two nodes
This commit is contained in:
Isi
2024-09-07 19:01:16 +01:00
committed by GitHub
parent 56a033eed6
commit 439ecf7436
1 changed files with 104 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
uniAnimate_Inference.py
+104 -2
View File
@@ -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",
}