Files
Sxela-ComfyWarp/frame_nodes.py
T
2025-02-08 23:50:52 +01:00

581 lines
21 KiB
Python

import glob
import os
from PIL import Image, ImageOps
import torch
import numpy as np
import folder_paths
from .frame_utils import FrameDataset, StylizedFrameDataset, get_scheduled_arg, get_size, save_video
class ApplyMask:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"destination": ("IMAGE",),
"source": ("IMAGE",),
},
"optional": {
"mask": ("MASK",),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "composite"
CATEGORY = "WarpFusion"
def composite(self, destination, source, mask = None):
mask = mask[..., None].repeat(1,1,1,destination.shape[-1])
res = destination*(1-mask) + source*(mask)
return (res,)
class ApplyMaskConditional:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"destination": ("IMAGE",),
"source": ("IMAGE",),
"current_frame_number": ("INT",),
"apply_at_frames": ("STRING",),
"don_not_apply_at_frames": ("BOOLEAN",),
},
"optional": {
"mask": ("MASK",),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "composite"
CATEGORY = "WarpFusion"
def composite(self, destination, source, current_frame_number, apply_at_frames, don_not_apply_at_frames, mask = None):
idx_list = [int(i) for i in apply_at_frames.split(',')]
if (current_frame_number not in idx_list) if don_not_apply_at_frames else (current_frame_number in idx_list):
# Convert mask to correct format for interpolation [b,c,h,w]
mask = mask[None,...]
# Resize mask to destination size using explicit dimensions
mask = torch.nn.functional.interpolate(mask, size=(destination.shape[1], destination.shape[2]), mode='bilinear')
# Convert back to [b,h,w,1] format
mask = mask[0,...,None].repeat(1,1,1,destination.shape[-1])
source = source.permute(0,3,1,2)
source = torch.nn.functional.interpolate(source, size=(destination.shape[1], destination.shape[2]), mode='bilinear')
source = source.permute(0,2,3,1)
res = destination*(1-mask) + source*(mask)
return (res,)
else:
return (destination,)
class ApplyMaskLatent:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"destination": ("LATENT",),
"source": ("LATENT",),
},
"optional": {
"mask": ("MASK",),
}
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "composite"
CATEGORY = "WarpFusion"
def composite(self, destination, source, mask = None):
destination = destination['samples']
source = source['samples']
mask = mask[None, ...]
mask = torch.nn.functional.interpolate(mask, size=(destination.shape[2], destination.shape[3]))
res = destination*(1-mask) + source*(mask)
return ({"samples":res}, )
class ApplyMaskLatentConditional:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"destination": ("LATENT",),
"source": ("LATENT",),
"current_frame_number": ("INT",),
"apply_at_frames": ("STRING",),
"don_not_apply_at_frames": ("BOOLEAN",),
},
"optional": {
"mask": ("MASK",),
}
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "composite"
CATEGORY = "WarpFusion"
def composite(self, destination, source, current_frame_number, apply_at_frames, don_not_apply_at_frames, mask = None):
destination = destination['samples']
source = source['samples']
idx_list = [int(i) for i in apply_at_frames.split(',')]
if (current_frame_number not in idx_list) if don_not_apply_at_frames else (current_frame_number in idx_list):
mask = mask[None, ...]
mask = torch.nn.functional.interpolate(mask, size=(destination.shape[2], destination.shape[3]))
res = destination*(1-mask) + source*(mask)
return ({"samples":res}, )
else:
return ({"samples":destination}, )
class LoadFrameSequence:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"file_path": ("STRING", {"multiline": True,
"default":"C:\\code\\warp\\19_cn_venv\\images_out\\stable_warpfusion_0.20.0\\videoFrames\\650571deef_0_0_1"})
}
}
CATEGORY = "WarpFusion"
OUTPUT_IS_LIST = (True,False)
RETURN_TYPES = ("FRAMES", "INT")
RETURN_NAMES = ("Frames", "Total_frames")
FUNCTION = "get_frames"
def get_frames(self, file_path):
print(file_path)
self.frames = glob.glob(file_path+'/**/*.*', recursive=True)
self.max_frames = len(self.frames)
print(f'Found {len(self.frames)} frames.')
out = [{'image':frame, 'max_frames':self.max_frames}for frame in self.frames]
return (out,self.max_frames)
@classmethod
def IS_CHANGED(self, file_path):
self.get_frames(self, file_path)
@classmethod
def VALIDATE_INPUTS(self, file_path):
self.get_frames(self, file_path)
if len(self.frames)==0:
return f"Found 0 frames in path {file_path}"
return True
class LoadFrame:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"file_paths": ("FRAMES",),
"seed": ("INT", {"default": 0, "min": 0, "max": 9999999999}),
"total_frames":("INT", {"default": 0, "min": 0, "max": 9999999999})
}
}
CATEGORY = "WarpFusion"
INPUT_IS_LIST = True
RETURN_TYPES = ("IMAGE","INT")
RETURN_NAMES = ("Image","Frame number")
FUNCTION = "load_frame"
#validation fails here for some reason
def load_frame(self, file_paths, seed, total_frames):
frame_number = seed
print(file_paths[:10], frame_number, total_frames)
frame_number = frame_number[0]
total_frames = total_frames[0]
frame_number = min(min(frame_number, total_frames), len(file_paths)-1)
print(frame_number)
image_path = file_paths[frame_number]['image']
i = Image.open(image_path)
i = ImageOps.exif_transpose(i)
image = i.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
print(image.shape, frame_number, total_frames)
return (image, frame_number)
class MakeFrameDataset:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"file_path": ("STRING", {"multiline": True,
"default":"C:\\code\\warp\\19_cn_venv\\images_out\\stable_warpfusion_0.20.0\\videoFrames\\650571deef_0_0_1"}),
"update_on_frame_load": ("BOOLEAN", {"default": True}),
"start_frame":("INT", {"default": 0, "min": 0, "max": 9999999999}),
"end_frame":("INT", {"default": -1, "min": -1, "max": 9999999999}),
"nth_frame":("INT", {"default": 1, "min": 1, "max": 9999999999}),
"overwrite":("BOOLEAN", {"default": False})
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("FRAME_DATASET", "INT")
RETURN_NAMES = ("FRAME_DATASET", "Total_frames")
FUNCTION = "get_frames"
def get_frames(self, file_path, update_on_frame_load, start_frame, end_frame, nth_frame, overwrite):
ds = FrameDataset(file_path, outdir_prefix='', videoframes_root=folder_paths.get_output_directory(),
update_on_getitem=update_on_frame_load, start_frame=start_frame, end_frame=end_frame, nth_frame=nth_frame, overwrite=overwrite)
if len(ds)==0:
raise Exception(f"Found 0 frames in path {file_path}") #thanks to https://github.com/Aljnk
return (ds,len(ds))
class LoadFrameFromFolder:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"file_path": ("STRING", {"multiline": True,
"default":"C:\\code\\warp\\19_cn_venv\\images_out\\stable_warpfusion_0.20.0\\videoFrames\\650571deef_0_0_1"}),
"init_image":("IMAGE",) ,
"frame_number":("INT", {"default": 0, "min": 0, "max": 9999999999}),
"fit_into": ("INT", {"default": 1280, "min": 0, "max": 8196*2}),
},
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("IMAGE",)
FUNCTION = "get_frames"
def load_image(self, image_path, fit_into):
i = Image.open(image_path)
i = ImageOps.exif_transpose(i)
size = get_size(i.size, max_size=fit_into, divisible_by=8)
image = i.convert("RGB").resize(size)
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
return image
def get_frames(self, file_path, init_image, frame_number, fit_into):
if frame_number == -1: return (init_image,)
if not os.path.exists(file_path):
os.makedirs(file_path, exist_ok=True)
ds = StylizedFrameDataset(file_path)
frame_number = min(frame_number, len(ds)-1)
frame_number = max(0, frame_number)
if len(ds) == 0: return (init_image,)
return (self.load_image(ds[frame_number], fit_into),)
class LoadFrameFromDataset:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"frame_dataset": ("FRAME_DATASET",),
"seed": ("INT", {"default": 0, "min": 0, "max": 9999999999}),
"total_frames":("INT", {"default": 0, "min": 0, "max": 9999999999})
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("IMAGE","INT")
RETURN_NAMES = ("Image","Frame number")
FUNCTION = "load_frame"
def load_frame(self, frame_dataset, seed, total_frames):
frame_number = seed
frame_number = min(min(frame_number, total_frames), len(frame_dataset)-1)
frame_number = max(0, frame_number)
image_path = frame_dataset[frame_number]
i = Image.open(image_path)
i = ImageOps.exif_transpose(i)
image = i.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
return (image, frame_number)
class LoadFramePairFromDataset:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"frame_dataset": ("FRAME_DATASET",),
"seed": ("INT", {"default": 0, "min": 0, "max": 9999999999}),
"total_frames":("INT", {"default": 0, "min": 0, "max": 9999999999}),
"fit_into": ("INT", {"default": 1280, "min": 0, "max": 8196*2}),
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("IMAGE","IMAGE","INT")
RETURN_NAMES = ("Current frame","Previous Frame","Frame number")
FUNCTION = "load_frames"
def load_frame(self, frame_dataset, seed, total_frames, fit_into):
frame_number = seed
frame_number = min(min(frame_number, total_frames), len(frame_dataset)-1)
frame_number = max(0, frame_number)
image_path = frame_dataset[frame_number]
i = Image.open(image_path)
size = get_size(i.size, fit_into, divisible_by=8)
i = ImageOps.exif_transpose(i).resize(size)
image = i.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
return image
def load_frames(self, frame_dataset, seed, total_frames, fit_into):
current_frame = self.load_frame(frame_dataset, seed, total_frames, fit_into)
previous_frame = self.load_frame(frame_dataset, seed-1, total_frames, fit_into)
return (current_frame, previous_frame, seed)
class ResizeToFit:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"image": ("IMAGE",),
"max_size": ("INT", {"default": 1280, "min": 0, "max": 9999999999}),
"divisible_by": ("INT", {"default": 64, "min": 2, "max": 2048}),
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("Image",)
FUNCTION = "resize"
def resize(self, image, max_size, divisible_by):
image = image.transpose(1,-1)
size = image.shape[2:]
size = get_size(size, max_size, divisible_by)
image = torch.nn.functional.interpolate(image, size)
image = image.transpose(1,-1)
return (image, )
class SaveFrame:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"image": ("IMAGE",),
"output_dir": ("STRING", {"multiline": True,
"default":''}),
"batch_name": ("STRING",{"default": "ComfyWarp"}),
"frame_number":("INT",{"default": 0, "min": 0, "max": 9999999999}),
}
}
CATEGORY = "WarpFusion"
FUNCTION = "save_img"
RETURN_TYPES = ()
OUTPUT_NODE = True
def save_img(self, image, output_dir, batch_name, frame_number):
os.makedirs(output_dir, exist_ok=True)
fname = f'{batch_name}_{frame_number:06}.png'
out_fname = os.path.join(output_dir, fname)
print('image.shape', image.shape, image.max(), image.min())
image = (image[0].clip(0,1)*255.).cpu().numpy().astype('uint8')
image = Image.fromarray(image)
image.save(out_fname)
print('fname', out_fname)
return ()
class RenderVideo:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"output_dir": ("STRING", {"multiline": True,
"default":''}),
"frames_input_dir": ("STRING", {"multiline": True,
"default":''}),
"batch_name": ("STRING", {"default":'ComfyWarp'}),
"first_frame":("INT",{"default": 0, "min": 0, "max": 9999999999}),
"last_frame":("INT",{"default": -1, "min": -1, "max": 9999999999}),
"render_at_frame":("INT",{"default": 0, "min": 0, "max": 9999999999}),
"current_frame":("INT",{"default": 0, "min": 0, "max": 9999999999}),
"fps":("FLOAT",{"default": 24, "min": 0, "max": 9999999999}),
"output_format":(["h264_mp4", "qtrle_mov", "prores_mov"],),
"use_deflicker": ("BOOLEAN", {"default": False})
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ()
FUNCTION = "export_video"
OUTPUT_NODE = True
def export_video(self, output_dir, frames_input_dir, batch_name, first_frame=1, last_frame=-1,
render_at_frame=999999, current_frame=0, fps=30, output_format='h264_mp4', use_deflicker=False):
if current_frame>=render_at_frame:
print('Exporting video.')
save_video(indir=frames_input_dir, video_out=output_dir, batch_name=batch_name, start_frame=first_frame,
last_frame=last_frame, fps=fps, output_format=output_format, use_deflicker=use_deflicker)
# raise Exception(f'Exported video successfully. This exception is raised to just stop the endless cycle :D.\n you can find your video at {output_dir}')
return ()
class SchedulerInt:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"schedule": ("STRING", {"multiline": True,
"default":''}),
"frame_number":("INT",{"default": 0, "min": 0, "max": 9999999999}),
"blend_json": ("BOOLEAN", {"default": True})
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("INT", )
FUNCTION = "get_value"
def get_value(self, schedule, frame_number, blend_json=True):
value = get_scheduled_arg(frame_num=frame_number, schedule=schedule, blend_json_schedules=blend_json)
return (int(value),)
class SchedulerFloat:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"schedule": ("STRING", {"multiline": True,
"default":''}),
"frame_number":("INT",{"default": 0, "min": 0, "max": 9999999999}),
"blend_json": ("BOOLEAN", {"default": True})
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("FLOAT", )
FUNCTION = "get_value"
def get_value(self, schedule, frame_number, blend_json=True):
value = get_scheduled_arg(frame_num=frame_number, schedule=schedule, blend_json_schedules=blend_json)
return (float(value),)
class SchedulerString:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"schedule": ("STRING", {"multiline": True,
"default":''}),
"frame_number":("INT",{"default": 0, "min": 0, "max": 9999999999}),
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("STRING", )
FUNCTION = "get_value"
def get_value(self, schedule, frame_number):
value = get_scheduled_arg(frame_num=frame_number, schedule=schedule, blend_json_schedules=False)
return (str(value),)
"""
inspired by https://github.com/Fannovel16/ComfyUI-Loopchain
"""
class FixedQueue:
@classmethod
def INPUT_TYPES(self):
return {"required":
{
"start": ("INT",{"default": 0, "min": 0, "max": 9999999999}),
"end":("INT",{"default": 1, "min": 0, "max": 9999999999}),
"current_number":("INT",{"default": 0, "min": 0, "max": 9999999999}),
}
}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("INT", "INT", "INT",)
RETURN_NAMES = ("current", "start", "end",)
FUNCTION = "get_value"
def get_value(self, start, end, current_number):
return (current_number, start, end)
class MakePaths:
@classmethod
def INPUT_TYPES(self):
return {"required": {
"root_path": ("STRING", {"multiline": True, "default": "./"}),
"experiment": ("STRING", {"default": "experiment"}),
"video": ("STRING", {"default": "video"}),
"frames": ("STRING", {"default": "frames"}),
"smoothed": ("STRING", {"default": "smoothed"}),
}}
CATEGORY = "WarpFusion"
RETURN_TYPES = ("STRING", "STRING", "STRING")
RETURN_NAMES = ("video_path", "frames_path", "smoothed_frames_path")
FUNCTION = "build_paths"
def build_paths(self, root_path, experiment, video, frames, smoothed):
base_path = os.path.join(root_path, experiment)
video_path = os.path.join(base_path, video)
frames_path = os.path.join(base_path, frames)
smoothed_frames_path = os.path.join(base_path, smoothed)
return (video_path, frames_path, smoothed_frames_path)
NODE_CLASS_MAPPINGS = {
"LoadFrameSequence": LoadFrameSequence,
"LoadFrame": LoadFrame,
"LoadFrameFromDataset":LoadFrameFromDataset,
"MakeFrameDataset":MakeFrameDataset,
"LoadFramePairFromDataset":LoadFramePairFromDataset,
"LoadFrameFromFolder":LoadFrameFromFolder,
"ResizeToFit":ResizeToFit,
"SaveFrame":SaveFrame,
"RenderVideo": RenderVideo,
"SchedulerString":SchedulerString,
"SchedulerFloat":SchedulerFloat,
"SchedulerInt":SchedulerInt,
"FixedQueue":FixedQueue,
"ApplyMask":ApplyMask,
"ApplyMaskConditional":ApplyMaskConditional,
"ApplyMaskLatent":ApplyMaskLatent,
"ApplyMaskLatentConditional":ApplyMaskLatentConditional,
"MakePaths": MakePaths,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LoadFrameSequence": "Load Frame Sequence",
"LoadFrame":"Load Frame",
"LoadFrameFromDataset":"Load Frame From Dataset",
"MakeFrameDataset":"Make Frame Dataset",
"LoadFramePairFromDataset":"Load Frame Pair From Dataset",
"LoadFrameFromFolder": "Maybe Load Frame From Folder",
"ResizeToFit":"Resize To Fit",
"SaveFrame":"SaveFrame",
"RenderVideo": "RenderVideo",
"SchedulerString":"SchedulerString",
"SchedulerFloat":"SchedulerFloat",
"SchedulerInt":"SchedulerInt",
"FixedQueue":"FixedQueue",
"ApplyMask":"ApplyMask",
"ApplyMaskConditional":"ApplyMaskConditional",
"ApplyMaskLatent":"ApplyMaskLatent",
"ApplyMaskLatentConditional":"ApplyMaskLatentConditional",
"MakePaths": "Make Paths",
}