Files

569 lines
20 KiB
Python

import math
from datetime import datetime
import os
import shutil
import random
from typing import Optional, List
class EbuScalingResolution:
aspect_ratios = {
"16:9": [
"640x360", "960x540", "1024x576", "1088x612", "1152x648", "1216x684", "1280x720", "1344x756", "1408x792", "1472x828", "1536x864", "1600x900", "1792x1024", "1856x1044", "1920x1080", "2048x1152", "2240x1260", "2304x1296", "2560x1440"
],
"16:10": [
"640x400", "704x440", "768x480", "832x520", "896x560", "960x600", "1024x640", "1088x680", "1152x720", "1216x760", "1280x800", "1344x840", "1408x880", "1472x920", "1536x960", "1600x1000", "1792x1120", "1856x1160", "1920x1200", "2048x1280", "2240x1400", "2560x1600"
],
"3:2": [
"768x512", "832x554", "896x597", "960x640", "1024x682", "1088x725", "1152x768", "1216x810", "1280x853", "1344x896", "1408x938", "1472x981", "1536x1024", "1600x1066"
],
"4:3": [
"640x480", "704x528", "768x576", "832x624", "896x672", "960x720", "1024x768", "1088x816", "1152x864", "1216x912", "1280x960", "1344x1008", "1408x1056", "1472x1104", "1536x1152", "1600x1200"
],
"5:4": [
"640x512", "704x564", "768x615", "832x667", "896x718", "960x770", "1024x821", "1088x872", "1152x924", "1216x975", "1280x1024", "1344x1076", "1408x1128", "1472x1179", "1536x1230", "1600x1280"
],
"6:5": [
"768x640", "832x693", "896x746", "960x800", "1024x853", "1088x906", "1152x960", "1216x1013", "1280x1066", "1344x1120", "1408x1173", "1472x1226", "1536x1280", "1600x1333"
],
"1:1": [
"512x512", "576x576", "640x640", "704x704", "768x768", "832x832", "896x896", "960x960", "1024x1024", "1088x1088", "1152x1152", "1216x1216", "1280x1280", "1344x1344", "1408x1408", "1472x1472", "1536x1536", "1600x1600"
]
}
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"active_aspect_ratio": (list(cls.aspect_ratios.keys()) + ["Other"],),
**{key: (value,) for key, value in cls.aspect_ratios.items()},
"other_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8, "display": "number"}),
"other_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8, "display": "number"}),
"mode": (["Landscape", "Profile"], {"default": "Landscape"}),
"upscale_by": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 10.0, "step": 0.05, "display": "number"})
}
}
RETURN_TYPES = ("INT", "INT", "INT", "INT", "FLOAT", "STRING",)
RETURN_NAMES = ("width", "height", "upscaled_width", "upscaled_height", "upscale_by", "upscaled_resolution_string",)
FUNCTION = "compute_resolution"
CATEGORY = "Resolution"
def compute_resolution(self, active_aspect_ratio, **kwargs):
if active_aspect_ratio == "Other":
width = kwargs["other_width"]
height = kwargs["other_height"]
else:
resolution = kwargs[active_aspect_ratio]
width, height = [int(x) for x in resolution.split('x')]
def round_up_to_multiple_of_eight(value):
return math.ceil(value / 8) * 8
upscale_by = kwargs["upscale_by"]
scaled_width = round_up_to_multiple_of_eight(width * upscale_by)
scaled_height = round_up_to_multiple_of_eight(height * upscale_by)
if kwargs["mode"] == "Profile":
width, height = height, width
scaled_width, scaled_height = scaled_height, scaled_width
upscaled_resolution_string = f"{scaled_width}x{scaled_height}"
print(f"Original: {width}x{height}, Scaled: {upscaled_resolution_string}, Upscale factor: {upscale_by}")
return width, height, scaled_width, scaled_height, upscale_by, upscaled_resolution_string
class EbuScalingTile:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"upscaled_image_width": ("INT", {"default": 768, "step": 1, "display": "number"}),
"upscaled_image_height": ("INT", {"default": 1024, "step": 1, "display": "number"}),
"profile_width_div_by": ("FLOAT", {"default": 1, "min": 1.0, "max": 20.0, "step": 0.05, "display": "number"}),
"profile_height_div_by": ("FLOAT", {"default": 3, "min": 1.0, "max": 20.0, "step": 0.05, "display": "number"}),
"landscape_width_div_by": ("FLOAT", {"default": 2, "min": 1.0, "max": 20.0, "step": 0.05, "display": "number"}),
"landscape_height_div_by": ("FLOAT", {"default": 2, "min": 1.0, "max": 20.0, "step": 0.05, "display": "number"}),
"tile_width_padding": ("INT", {"default": 0, "step": 1, "display": "number"}),
"tile_height_padding": ("INT", {"default": 0, "step": 1, "display": "number"}),
}
}
RETURN_TYPES = ("INT", "INT",)
RETURN_NAMES = ("tile_width", "tile_height",)
FUNCTION = "tile"
CATEGORY = "Resolution"
def tile(self, upscaled_image_width, upscaled_image_height, profile_width_div_by, profile_height_div_by, landscape_width_div_by, landscape_height_div_by, tile_width_padding, tile_height_padding):
if upscaled_image_height > upscaled_image_width:
new_width = int(upscaled_image_width / profile_width_div_by) + tile_width_padding
new_height = int(upscaled_image_height / profile_height_div_by) + tile_height_padding
else:
new_width = int(upscaled_image_width / landscape_width_div_by) + tile_width_padding
new_height = int(upscaled_image_height / landscape_height_div_by) + tile_height_padding
return new_width, new_height
class EbuGetImageAspectRatio:
ASPECT_RATIOS = {
"1:1": 1.0,
"6:5": 1.2,
"5:4": 1.25,
"4:3": 1.333,
"3:2": 1.5,
"2:1": 2.0,
"16:10": 1.6,
"16:9": 1.777,
"5:6": 0.833,
"4:5": 0.8,
"3:4": 0.75,
"2:3": 0.666,
"10:16": 0.625,
"9:16": 0.5625
}
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",)
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("aspect_ratio",)
FUNCTION = "get_aspect_ratio"
CATEGORY = "Resolution"
def get_aspect_ratio(self, image):
img = image[0] # Access the first image tensor
height, width = img.shape[:2] # Extract height and width
ratio = width / height
tolerance = 0.08 # 8% tolerance
closest = min(self.ASPECT_RATIOS.items(), key=lambda x: abs(x[1] - ratio))
label, value = closest
diff = abs(value - ratio)
print(f"DEBUG: width={width}, height={height}, ratio={ratio:.4f}, closest={label} ({value}), diff={diff:.4f}")
if diff <= tolerance:
return (label,)
else:
return ("Unknown",)
class EbuGetImageAspectRatioFromImage:
ASPECT_RATIOS = {
"1:1": 1.0,
"6:5": 1.2,
"5:4": 1.25,
"4:3": 1.333,
"3:2": 1.5,
"2:1": 2.0,
"16:10": 1.6,
"16:9": 1.777,
"5:6": 0.833,
"4:5": 0.8,
"3:4": 0.75,
"2:3": 0.666,
"10:16": 0.625,
"9:16": 0.5625
}
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"delimiter": ("STRING", {"default": ":"}),
"tolerance": ("FLOAT", {"default": 0.08, "min": 0.0, "max": 1.0, "step": 0.01}),
}
}
RETURN_TYPES = ("STRING", "STRING", "INT", "INT",)
RETURN_NAMES = ("aspect_ratio", "resolution", "width", "height",)
FUNCTION = "get_aspect_ratio_from_image"
CATEGORY = "Resolution"
def get_aspect_ratio_from_image(self, image, delimiter, tolerance):
# Access the first image tensor and get its dimensions
img = image[0]
height, width = img.shape[:2]
# Calculate the aspect ratio
ratio = width / height
# Find the closest aspect ratio from the predefined list
closest = min(self.ASPECT_RATIOS.items(), key=lambda x: abs(x[1] - ratio))
label, value = closest
diff = abs(value - ratio)
print(f"DEBUG: width={width}, height={height}, ratio={ratio:.4f}, closest={label} ({value}), diff={diff:.4f}")
resolution_str = f"{width}{delimiter}{height}"
# Check if the closest match is within the specified tolerance
if diff <= tolerance:
# Replace the default delimiter with the user-specified one
if delimiter != ":":
label = label.replace(":", delimiter)
return (label, resolution_str, width, height,)
else:
return ("custom", resolution_str, width, height,)
class EbuUniqueFileName:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"str": ("STRING", {"default": "file"}),
"join_str": ("STRING", {"default": "-"}),
"seed": ("INT", {"default": 0, "max": 0xffffffffffffffff})
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("unique_filename",)
FUNCTION = "generate_filename"
CATEGORY = "Utility"
def generate_filename(self, str, join_str, seed):
now = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
return (f"{str}{join_str}{now}",)
class EbuAppendToFile:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"string_to_append": ("STRING", {"multiline": True}),
"directory_name": ("STRING", {"default": "store"}),
"file_name": ("STRING", {"default": "output.txt"}),
"overwrite": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES = ()
FUNCTION = "append_to_file"
OUTPUT_NODE = True
CATEGORY = "Utility"
def append_to_file(self, string_to_append, directory_name, file_name, overwrite):
os.makedirs(directory_name, exist_ok=True)
full_path = os.path.join(directory_name, file_name)
mode = 'w' if overwrite else 'a'
with open(full_path, mode) as file:
file.write(string_to_append + "\n")
print(f"{'Overwritten' if overwrite else 'Appended to'} file: {full_path}")
return ()
class EbuReadFromFile:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"directory_name": ("STRING", {"default": "store"}),
"file_name": ("STRING", {"default": "output.txt"})
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("file_contents",)
FUNCTION = "read_from_file"
OUTPUT_NODE = True
CATEGORY = "Utility"
def read_from_file(self, directory_name, file_name):
full_path = os.path.join(directory_name, file_name)
if not os.path.exists(full_path):
print(f"File not found: {full_path}")
return ("",)
with open(full_path, 'r') as file:
contents = file.read()
print(f"Read from file: {full_path}")
return (contents,)
class EbuFileListCache:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"directory_name": ("STRING", {"default": "store"}),
"file_name": ("STRING", {"default": "output.txt"}),
"num_return_items": ("INT", {"default": 5, "min": 0, "max": 1000}),
"input_items": ("STRING", {"multiline": True}),
"limit_list_size": ("INT", {"default": 100, "min": 1, "max": 10000}),
"seed": ("INT", {"default": 0, "max": 0xffffffffffffffff})
}
}
RETURN_TYPES = ("STRING", "STRING", "STRING")
RETURN_NAMES = ("selected_items", "input_items", "combined_items")
FUNCTION = "process_file_list_cache"
OUTPUT_NODE = True
CATEGORY = "Utility"
def process_file_list_cache(self,
directory_name: str,
file_name: str,
num_return_items: int,
input_items: str,
limit_list_size: int,
seed: Optional[int] = None
) -> List[str]:
"""
Reads directory_name/file_name, merges any new lines from input_items,
shuffles & trims to limit_list_size, optionally rewrites the file,
then returns a random sample of num_return_items.
"""
# Ensure storage directory exists
os.makedirs(directory_name, exist_ok=True)
full_path = os.path.join(directory_name, file_name)
backup_path = full_path + ".bk"
# 1. Load existing lines
current_lines = set()
if os.path.exists(full_path):
with open(full_path, "r", encoding="utf-8") as f:
current_lines = {line.strip() for line in f if line.strip()}
# 2. If no new input → just sample & return, no file changes
if not input_items.strip():
selected = random.sample(
list(current_lines),
min(num_return_items, len(current_lines))
)
return (
"\n".join(selected),
"", # no input_items
"\n".join(selected) # combined == selected
)
# 3. Parse & merge new input lines
input_lines = [
line.strip()
for line in input_items.splitlines()
if line.strip()
]
input_set = set(input_lines)
# 4. Combine, optionally reseed, then shuffle
combined = current_lines.union(input_set)
combined_list = list(combined)
if seed is not None:
random.seed(seed)
random.shuffle(combined_list)
# 5. Trim to limit_list_size
trimmed = combined_list[:limit_list_size]
# 6. Backup old file & rewrite with trimmed list
if os.path.exists(full_path):
shutil.copy2(full_path, backup_path)
with open(full_path, "w", encoding="utf-8") as f:
f.write("\n".join(trimmed) + "\n")
# 7. From the trimmed list, pick new random sample (excluding the inputs)
available = list(set(trimmed) - input_set)
selected = random.sample(
available,
min(num_return_items, len(available))
)
# 8. Build the combined output (inputs first, then selected)
combined_output = input_lines + selected
return (
"\n".join(selected),
"\n".join(input_lines),
"\n".join(combined_output)
)
class EbuEncodeNewLines:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING", {"multiline": True}),
"new_line_encoding": ("STRING", {"default": "|||"})
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("encoded_text",)
FUNCTION = "encode"
OUTPUT_NODE = True
CATEGORY = "Utility"
def encode(self, text, new_line_encoding):
encoded = text.replace("\n", new_line_encoding)
return (encoded,)
class EbuDecodeNewLines:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"encoded_text": ("STRING",),
"new_line_encoding": ("STRING", {"default": "|||"})
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("decoded_text",)
FUNCTION = "decode"
OUTPUT_NODE = True
CATEGORY = "Utility"
def decode(self, encoded_text, new_line_encoding):
decoded = encoded_text.replace(new_line_encoding, "\n")
return (decoded,)
class EbuStringWaitForImage:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"input_string": ("STRING", {"multiline": True}),
"wait_for_image": ("IMAGE",),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("output_string",)
FUNCTION = "passthrough"
CATEGORY = "Utility"
def passthrough(self, input_string, wait_for_image):
if wait_for_image is None or len(wait_for_image) == 0:
raise ValueError("Image not loaded. Cannot proceed.")
return (input_string,)
class EbuImageWaitForImage:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"input_image": ("IMAGE",),
"wait_for_image": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("output_image",)
FUNCTION = "passthrough"
CATEGORY = "Utility"
def passthrough(self, input_image, wait_for_image):
if wait_for_image is None or len(wait_for_image) == 0:
raise ValueError("Image not loaded. Cannot proceed.")
return (input_image,)
class EbuModelWaitForImage:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"input_model": ("MODEL",),
"wait_for_image": ("IMAGE",),
}
}
RETURN_TYPES = ("MODEL",)
RETURN_NAMES = ("output_model",)
FUNCTION = "passthrough"
CATEGORY = "Utility"
def passthrough(self, input_model, wait_for_image):
if wait_for_image is None or len(wait_for_image) == 0:
raise ValueError("Image not loaded. Cannot proceed.")
return (input_model,)
class EbuComputeImageUpscale:
"""
Computes the upscale factor and new dimensions for an image
to meet minimum width and height requirements.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"min_width": ("INT", {"default": 2560, "min": 0, "max": 16384, "step": 1, "display": "number"}),
"min_height": ("INT", {"default": 1440, "min": 0, "max": 16384, "step": 1, "display": "number"}),
}
}
RETURN_TYPES = ("FLOAT", "INT", "INT",)
RETURN_NAMES = ("upscale_by", "upscaled_width", "upscaled_height",)
FUNCTION = "compute_upscale"
CATEGORY = "Resolution"
def compute_upscale(self, image, min_width, min_height):
# Access the first image tensor and get its dimensions
img = image[0]
original_height, original_width = img.shape[:2]
# Calculate upscale factors for width and height
upscale_by_width = 1.0
if min_width > 0:
upscale_by_width = min_width / original_width
upscale_by_height = 1.0
if min_height > 0:
upscale_by_height = min_height / original_height
# The final upscale factor is the maximum of the two to satisfy both conditions
upscale_by = max(upscale_by_width, upscale_by_height, 1.0)
# Calculate the upscaled dimensions, rounded up to the nearest integer
upscaled_width = math.ceil(original_width * upscale_by)
upscaled_height = math.ceil(original_height * upscale_by)
return (upscale_by, upscaled_width, upscaled_height,)
NODE_CLASS_MAPPINGS = {
"EbuGetImageAspectRatio": EbuGetImageAspectRatio,
"EbuScalingResolution": EbuScalingResolution,
"EbuScalingTile": EbuScalingTile,
"EbuUniqueFileName": EbuUniqueFileName,
"EbuAppendToFile": EbuAppendToFile,
"EbuReadFromFile": EbuReadFromFile,
"EbuFileListCache": EbuFileListCache,
"EbuEncodeNewLines": EbuEncodeNewLines,
"EbuDecodeNewLines": EbuDecodeNewLines,
"EbuStringWaitForImage": EbuStringWaitForImage,
"EbuImageWaitForImage": EbuImageWaitForImage,
"EbuModelWaitForImage": EbuModelWaitForImage,
"EbuGetImageAspectRatioFromImage": EbuGetImageAspectRatioFromImage,
"EbuComputeImageUpscale": EbuComputeImageUpscale,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"EbuGetImageAspectRatio": "EBU Get Image Aspect Ratio",
"EbuScalingResolution": "EBU Scaling Resolution",
"EbuScalingTile": "EBU Scaling Tile",
"EbuUniqueFileName": "EBU Unique File Name",
"EbuAppendToFile": "EBU Append To File",
"EbuReadFromFile": "EBU Read From File",
"EbuFileListCache": "EBU File List Cache",
"EbuEncodeNewLines": "EBU Encode New Lines",
"EbuDecodeNewLines": "EBU Decode New Lines",
"EbuStringWaitForImage": "EBU String Wait For Image",
"EbuImageWaitForImage": "EBU Image Wait For Image",
"EbuModelWaitForImage": "EBU Model Wait For Image",
"EbuGetImageAspectRatioFromImage": "EBU Get Image Aspect Ratio From Image",
"EbuComputeImageUpscale": "EBU Compute Image Upscale",
}