518 lines
17 KiB
Python
518 lines
17 KiB
Python
import math
|
|
import os
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
import torch
|
|
|
|
import comfy.utils
|
|
import folder_paths
|
|
import node_helpers
|
|
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, _value):
|
|
return False
|
|
|
|
|
|
ANY_TYPE = AnyType("*")
|
|
VIDEO_OUTPUT_EXTENSIONS = {".mp4", ".webm", ".mov", ".mkv", ".avi", ".gif", ".webp", ".m4v"}
|
|
|
|
|
|
SIZE_TABLE = {
|
|
"480p": {
|
|
"1:1": (480, 480),
|
|
"3:4": (480, 640),
|
|
"4:3": (640, 480),
|
|
"9:16": (480, 832),
|
|
"16:9": (832, 480),
|
|
"21:9": (1120, 480),
|
|
},
|
|
"720p": {
|
|
"1:1": (720, 720),
|
|
"3:4": (720, 960),
|
|
"4:3": (960, 720),
|
|
"9:16": (720, 1280),
|
|
"16:9": (1280, 720),
|
|
"21:9": (1680, 720),
|
|
},
|
|
"1K": {
|
|
"1:1": (1024, 1024),
|
|
"3:4": (896, 1152),
|
|
"4:3": (1152, 896),
|
|
"9:16": (768, 1344),
|
|
"16:9": (1344, 768),
|
|
"21:9": (1536, 640),
|
|
},
|
|
"2K": {
|
|
"1:1": (2048, 2048),
|
|
"3:4": (1536, 2048),
|
|
"4:3": (2048, 1536),
|
|
"9:16": (1152, 2048),
|
|
"16:9": (2048, 1152),
|
|
"21:9": (2560, 1088),
|
|
},
|
|
"4K": {
|
|
"1:1": (4096, 4096),
|
|
"3:4": (3072, 4096),
|
|
"4:3": (4096, 3072),
|
|
"9:16": (2304, 4096),
|
|
"16:9": (4096, 2304),
|
|
"21:9": (4096, 1792),
|
|
},
|
|
}
|
|
|
|
VIDEO_SIZE_TABLE = {
|
|
"480p": {
|
|
"1:1": (480, 480),
|
|
"3:4": (480, 640),
|
|
"4:3": (640, 480),
|
|
"9:16": (480, 832),
|
|
"16:9": (832, 480),
|
|
"21:9": (1120, 480),
|
|
},
|
|
"720p": {
|
|
"1:1": (720, 720),
|
|
"3:4": (720, 960),
|
|
"4:3": (960, 720),
|
|
"9:16": (720, 1280),
|
|
"16:9": (1280, 720),
|
|
"21:9": (1680, 720),
|
|
},
|
|
"1080p": {
|
|
"1:1": (1088, 1088),
|
|
"3:4": (1088, 1456),
|
|
"4:3": (1456, 1088),
|
|
"9:16": (1088, 1920),
|
|
"16:9": (1920, 1088),
|
|
"21:9": (2544, 1088),
|
|
},
|
|
}
|
|
|
|
MAX_RESOLUTION = 16384
|
|
|
|
|
|
def resolve_dimensions(aspect_ratio, image_size):
|
|
size_key = image_size if image_size in SIZE_TABLE else "1K"
|
|
aspect_key = aspect_ratio if aspect_ratio in SIZE_TABLE[size_key] else "1:1"
|
|
return SIZE_TABLE[size_key][aspect_key]
|
|
|
|
|
|
def resolve_video_dimensions(aspect_ratio, resolution_tier):
|
|
size_key = resolution_tier if resolution_tier in VIDEO_SIZE_TABLE else "720p"
|
|
aspect_key = aspect_ratio if aspect_ratio in VIDEO_SIZE_TABLE[size_key] else "16:9"
|
|
return VIDEO_SIZE_TABLE[size_key][aspect_key]
|
|
|
|
|
|
def resolve_input_path(filename):
|
|
input_dir = os.path.abspath(folder_paths.get_input_directory())
|
|
relative_name = filename.strip()
|
|
candidate_path = os.path.abspath(os.path.join(input_dir, os.path.normpath(relative_name)))
|
|
|
|
if os.path.commonpath((input_dir, candidate_path)) != input_dir:
|
|
raise ValueError(f"Reference image path escapes the ComfyUI input directory: {filename}")
|
|
|
|
if not os.path.exists(candidate_path):
|
|
raise FileNotFoundError(f"Reference image was not found in the ComfyUI input directory: {filename}")
|
|
|
|
return candidate_path
|
|
|
|
|
|
def resolve_output_path(path_or_filename):
|
|
output_dir = os.path.abspath(folder_paths.get_output_directory())
|
|
candidate_value = str(path_or_filename or "").strip()
|
|
if not candidate_value:
|
|
raise ValueError("RenderZero video output is missing a filename or path from the upstream save/combine node.")
|
|
|
|
candidate_path = os.path.abspath(
|
|
candidate_value if os.path.isabs(candidate_value)
|
|
else os.path.join(output_dir, os.path.normpath(candidate_value))
|
|
)
|
|
|
|
try:
|
|
if os.path.commonpath((output_dir, candidate_path)) != output_dir:
|
|
raise ValueError(
|
|
"RenderZero video output only accepts files saved inside the ComfyUI output directory."
|
|
)
|
|
except ValueError as error:
|
|
raise ValueError(
|
|
"RenderZero video output only accepts files saved inside the ComfyUI output directory."
|
|
) from error
|
|
|
|
if not os.path.exists(candidate_path):
|
|
raise FileNotFoundError(f"Video file was not found in the ComfyUI output directory: {candidate_value}")
|
|
|
|
relative_path = os.path.relpath(candidate_path, output_dir)
|
|
subfolder = os.path.dirname(relative_path)
|
|
if subfolder == ".":
|
|
subfolder = ""
|
|
|
|
return {
|
|
"filename": os.path.basename(candidate_path),
|
|
"subfolder": subfolder.replace("\\", "/"),
|
|
"type": "output",
|
|
"fullpath": candidate_path,
|
|
}
|
|
|
|
|
|
def infer_video_format(filename):
|
|
extension = os.path.splitext(filename)[1].lower()
|
|
if extension == ".gif":
|
|
return "image/gif"
|
|
if extension == ".webp":
|
|
return "image/webp"
|
|
if extension == ".webm":
|
|
return "video/webm"
|
|
if extension == ".mov":
|
|
return "video/quicktime"
|
|
if extension == ".mkv":
|
|
return "video/x-matroska"
|
|
if extension == ".avi":
|
|
return "video/x-msvideo"
|
|
return "video/mp4"
|
|
|
|
|
|
def build_video_preview_payload(reference):
|
|
return {
|
|
"filename": reference["filename"],
|
|
"subfolder": reference.get("subfolder", ""),
|
|
"type": reference.get("type", "output"),
|
|
"format": infer_video_format(reference["filename"]),
|
|
"fullpath": reference.get("fullpath"),
|
|
}
|
|
|
|
|
|
def normalize_video_output_candidate(value):
|
|
if isinstance(value, dict):
|
|
if isinstance(value.get("filename"), str) and value["filename"].strip():
|
|
relative_parts = []
|
|
subfolder = value.get("subfolder")
|
|
if isinstance(subfolder, str) and subfolder.strip():
|
|
relative_parts.append(subfolder.strip())
|
|
relative_parts.append(value["filename"].strip())
|
|
return resolve_output_path(os.path.join(*relative_parts))
|
|
|
|
for key in ("fullpath", "path"):
|
|
path_value = value.get(key)
|
|
if isinstance(path_value, str) and path_value.strip():
|
|
return resolve_output_path(path_value)
|
|
|
|
raise TypeError("Unsupported dictionary payload.")
|
|
|
|
if isinstance(value, str) and value.strip():
|
|
return resolve_output_path(value)
|
|
|
|
raise TypeError("Unsupported payload type.")
|
|
|
|
|
|
def select_best_video_reference(values):
|
|
normalized_values = []
|
|
last_error = None
|
|
|
|
for value in values:
|
|
try:
|
|
normalized_values.append(normalize_video_output_reference(value))
|
|
except (TypeError, ValueError, FileNotFoundError) as error:
|
|
last_error = error
|
|
|
|
if normalized_values:
|
|
video_matches = [
|
|
reference for reference in normalized_values
|
|
if os.path.splitext(reference["filename"])[1].lower() in VIDEO_OUTPUT_EXTENSIONS
|
|
]
|
|
if video_matches:
|
|
return video_matches[-1]
|
|
return normalized_values[0]
|
|
|
|
if last_error is not None:
|
|
raise last_error
|
|
|
|
raise TypeError("No usable filename or path was found.")
|
|
|
|
|
|
def normalize_video_output_reference(value):
|
|
try:
|
|
return normalize_video_output_candidate(value)
|
|
except TypeError:
|
|
pass
|
|
|
|
if isinstance(value, (list, tuple)):
|
|
if len(value) >= 2 and isinstance(value[0], bool) and isinstance(value[1], (list, tuple)):
|
|
return select_best_video_reference(value[1])
|
|
|
|
return select_best_video_reference(value)
|
|
|
|
raise TypeError("Unsupported payload type.")
|
|
|
|
|
|
def load_reference_image_tensor(filename):
|
|
image_path = resolve_input_path(filename)
|
|
|
|
with Image.open(image_path) as image:
|
|
image = ImageOps.exif_transpose(image)
|
|
image = image.convert("RGB")
|
|
image_array = np.array(image).astype(np.float32) / 255.0
|
|
|
|
return torch.from_numpy(image_array)[None,]
|
|
|
|
|
|
def encode_reference_latent(vae, image):
|
|
samples = image.movedim(-1, 1)
|
|
total_pixels = int(1024 * 1024)
|
|
scale_by = math.sqrt(total_pixels / (samples.shape[3] * samples.shape[2]))
|
|
width = max(8, round(samples.shape[3] * scale_by / 8.0) * 8)
|
|
height = max(8, round(samples.shape[2] * scale_by / 8.0) * 8)
|
|
upscaled = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
|
|
return vae.encode(upscaled.movedim(1, -1)[:, :, :, :3])
|
|
|
|
|
|
class RenderZeroPromptBridge:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"positive_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"aspect_ratio": (["1:1", "3:4", "4:3", "9:16", "16:9", "21:9"], {"default": "16:9"}),
|
|
"image_size": (["480p", "720p", "1K", "2K", "4K"], {"default": "1K"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "INT", "INT", "INT")
|
|
RETURN_NAMES = ("positive_prompt", "width", "height", "seed")
|
|
FUNCTION = "build"
|
|
CATEGORY = "RenderZero"
|
|
|
|
def build(self, positive_prompt, aspect_ratio, image_size, seed=0):
|
|
width, height = resolve_dimensions(aspect_ratio, image_size)
|
|
return (positive_prompt, width, height, seed)
|
|
|
|
|
|
class RenderZeroVideoPromptBridge:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"positive_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"aspect_ratio": (["1:1", "3:4", "4:3", "9:16", "16:9", "21:9"], {"default": "16:9"}),
|
|
"resolution_tier": (["480p", "720p", "1080p"], {"default": "720p"}),
|
|
"frame_count": ("INT", {"default": 81, "min": 1, "max": 999, "step": 1}),
|
|
"frame_rate": ("FLOAT", {"default": 24.0, "min": 0.1, "max": 240.0, "step": 0.1}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "INT", "INT", "INT", "INT", "FLOAT", "INT")
|
|
RETURN_NAMES = ("positive_prompt", "width", "height", "frame_count", "frame_rate_int", "frame_rate_float", "seed")
|
|
FUNCTION = "build"
|
|
CATEGORY = "RenderZero"
|
|
|
|
def build(self, positive_prompt, aspect_ratio, resolution_tier, frame_count, frame_rate, seed=0):
|
|
width, height = resolve_video_dimensions(aspect_ratio, resolution_tier)
|
|
safe_frame_count = max(1, int(frame_count))
|
|
safe_frame_rate_float = max(0.1, float(frame_rate))
|
|
safe_frame_rate_int = max(1, int(round(safe_frame_rate_float)))
|
|
return (
|
|
positive_prompt,
|
|
width,
|
|
height,
|
|
safe_frame_count,
|
|
safe_frame_rate_int,
|
|
safe_frame_rate_float,
|
|
seed,
|
|
)
|
|
|
|
|
|
class RenderZeroImageInput:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"slot_index": ("INT", {"default": 1, "min": 1, "max": 999, "step": 1}),
|
|
"image_slot": ("STRING", {"default": ""}),
|
|
"target_width": ("INT", {"default": 1024, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
|
|
"target_height": ("INT", {"default": 1024, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "load_image"
|
|
CATEGORY = "RenderZero"
|
|
|
|
def load_image(self, slot_index, image_slot, target_width, target_height):
|
|
del slot_index, target_width, target_height
|
|
if not isinstance(image_slot, str) or not image_slot.strip():
|
|
raise ValueError("RenderZero image input is missing image_slot. Upload an Additional Reference Image before queueing this workflow.")
|
|
|
|
image = load_reference_image_tensor(image_slot)
|
|
return (image,)
|
|
|
|
|
|
class RenderZeroVideoStartFrameBridge:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_slot": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "load_image"
|
|
CATEGORY = "RenderZero"
|
|
|
|
def load_image(self, image_slot):
|
|
if not isinstance(image_slot, str) or not image_slot.strip():
|
|
raise ValueError("RenderZero video start frame is missing image_slot. Select or upload a start frame before queueing this workflow.")
|
|
|
|
image = load_reference_image_tensor(image_slot)
|
|
return (image,)
|
|
|
|
|
|
class RenderZeroImageOutput:
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.prefix_append = ""
|
|
self.compress_level = 4
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"filename_prefix": ("STRING", {"default": "RenderZero"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "RenderZero"
|
|
|
|
def save_images(self, images, filename_prefix="RenderZero"):
|
|
filename_prefix += self.prefix_append
|
|
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
|
filename_prefix,
|
|
self.output_dir,
|
|
images[0].shape[1],
|
|
images[0].shape[0],
|
|
)
|
|
|
|
results = []
|
|
for image in images:
|
|
array = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
|
|
file_name = f"{filename}_{counter:05}_.png"
|
|
img.save(os.path.join(full_output_folder, file_name), compress_level=self.compress_level)
|
|
results.append({
|
|
"filename": file_name,
|
|
"subfolder": subfolder,
|
|
"type": self.type,
|
|
})
|
|
counter += 1
|
|
|
|
return {"ui": {"images": results}}
|
|
|
|
|
|
class RenderZeroVideoOutput:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"video_payload": (ANY_TYPE, {}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "collect_video_output"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "RenderZero"
|
|
|
|
def collect_video_output(self, video_payload):
|
|
try:
|
|
reference = normalize_video_output_reference(video_payload)
|
|
except (TypeError, ValueError, FileNotFoundError) as error:
|
|
raise ValueError(
|
|
"RenderZero video output requires a connectable filename or path-like payload from an upstream save/combine node. "
|
|
"Make sure the connected node exposes a filename/path output and that the saved file lives inside the ComfyUI output directory."
|
|
) from error
|
|
|
|
return {"ui": {"gifs": [build_video_preview_payload(reference)]}}
|
|
|
|
|
|
class RenderZeroReferenceBridge:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"vae": ("VAE",),
|
|
"character1_face": ("STRING", {"default": ""}),
|
|
"character1_outfit": ("STRING", {"default": ""}),
|
|
"character1_object": ("STRING", {"default": ""}),
|
|
"global_reference": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("positive", "negative")
|
|
FUNCTION = "apply"
|
|
CATEGORY = "RenderZero"
|
|
|
|
def apply(
|
|
self,
|
|
positive,
|
|
negative,
|
|
vae,
|
|
character1_face="",
|
|
character1_outfit="",
|
|
character1_object="",
|
|
global_reference="",
|
|
):
|
|
filenames = [
|
|
character1_face,
|
|
character1_outfit,
|
|
character1_object,
|
|
global_reference,
|
|
]
|
|
|
|
reference_latents = []
|
|
for filename in filenames:
|
|
if not isinstance(filename, str) or not filename.strip():
|
|
continue
|
|
|
|
image = load_reference_image_tensor(filename)
|
|
reference_latents.append(encode_reference_latent(vae, image))
|
|
|
|
if reference_latents:
|
|
values = {"reference_latents": reference_latents}
|
|
positive = node_helpers.conditioning_set_values(positive, values, append=True)
|
|
negative = node_helpers.conditioning_set_values(negative, values, append=True)
|
|
|
|
return (positive, negative)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RenderZeroPromptBridge": RenderZeroPromptBridge,
|
|
"RenderZeroVideoPromptBridge": RenderZeroVideoPromptBridge,
|
|
"RenderZeroImageInput": RenderZeroImageInput,
|
|
"RenderZeroVideoStartFrameBridge": RenderZeroVideoStartFrameBridge,
|
|
"RenderZeroReferenceBridge": RenderZeroReferenceBridge,
|
|
"RenderZeroImageOutput": RenderZeroImageOutput,
|
|
"RenderZeroVideoOutput": RenderZeroVideoOutput,
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RenderZeroPromptBridge": "RenderZero Prompt Bridge",
|
|
"RenderZeroVideoPromptBridge": "RenderZero Video Prompt Bridge",
|
|
"RenderZeroImageInput": "RenderZero Image Input",
|
|
"RenderZeroVideoStartFrameBridge": "RenderZero Video Start Frame Bridge",
|
|
"RenderZeroReferenceBridge": "RenderZero Reference Bridge",
|
|
"RenderZeroImageOutput": "RenderZero Image Output",
|
|
"RenderZeroVideoOutput": "RenderZero Video Output",
|
|
}
|