261 lines
9.1 KiB
Python
261 lines
9.1 KiB
Python
import re
|
|
from copy import deepcopy
|
|
import folder_paths
|
|
from PIL import Image
|
|
import numpy as np
|
|
from comfy.cli_args import args
|
|
import json
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import os.path
|
|
|
|
static_x = 1
|
|
static_y = 1
|
|
|
|
reset_registry = {}
|
|
|
|
class GridFloats:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"index": ( "INT", {"default": 1, "min": 1, "max": 6 } ),
|
|
"float1": ("FLOAT", {"default": 1.0}),
|
|
"float2": ("FLOAT", {"default": 1.0}),
|
|
"float3": ("FLOAT", {"default": 1.0}),
|
|
"float4": ("FLOAT", {"default": 1.0}),
|
|
"float5": ("FLOAT", {"default": 1.0}),
|
|
"float6": ("FLOAT", {"default": 1.0}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
FUNCTION = "ReturnFloat"
|
|
CATEGORY = "EasyGrids"
|
|
|
|
def ReturnFloat( self, index: int, float1 : float, float2 : float, float3 : float, float4 : float, float5 : float, float6: float ):
|
|
#TODO: probably a more pythonic way to do this
|
|
ret_list = [float1, float2, float3, float4, float5, float6]
|
|
if ( index > len(ret_list) ):
|
|
return ( ret_list[len(ret_list) - 1], )
|
|
return (ret_list[ index - 1 ], )
|
|
|
|
class GridFloatList:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"index": ( "INT", {"default": 1, "min": 1, "max": 100 } ),
|
|
"float_list": ("STRING", {"multiline": True}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
FUNCTION = "ParseAndReturnFloat"
|
|
CATEGORY = "EasyGrids"
|
|
|
|
def __init__(self):
|
|
self.static_text = ""
|
|
self.static_out_arr = []
|
|
|
|
def ParseAndReturnFloat( self, index: int, float_list: str ):
|
|
if float_list != self.static_text:
|
|
split_str = re.split( ",|;|\s|:", float_list )
|
|
out_arr = []
|
|
for val in split_str:
|
|
# let the exception happen if invalid
|
|
out_arr.append(float(val))
|
|
self.static_text = float_list
|
|
self.static_out_arr = deepcopy( out_arr )
|
|
#import pdb; pdb.set_trace()
|
|
if ( index > len(self.static_out_arr) ):
|
|
return ( self.static_out_arr[len(self.static_out_arr) - 1], )
|
|
return (self.static_out_arr[ index - 1 ],)
|
|
|
|
class ImageGridCommander:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": { "x_count": ("INT", {"default": 1, "min": 1, "max": 12, "step": 1}),
|
|
"y_count": ("INT", {"default": 1, "min": 1, "max": 12, "step": 1}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT","INT","INT","INT",)
|
|
RETURN_NAMES = ("x_index", "y_index", "x_size", "y_size",)
|
|
FUNCTION = "queue_batch"
|
|
CATEGORY = "EasyGrids"
|
|
|
|
def __init__(self):
|
|
self.curr_x_index = 1
|
|
self.curr_y_index = 1
|
|
self.last_x_count = 0
|
|
self.last_y_count = 0
|
|
|
|
def queue_batch(self, x_count, y_count, unique_id ):
|
|
#wish we could do this on init but there doesn't seem to be a way to get the unique_id at that point
|
|
#there shouldn't be any need to reset before the first run in any case
|
|
if unique_id not in reset_registry:
|
|
reset_registry[unique_id] = self
|
|
if x_count != self.last_x_count or y_count != self.last_y_count:
|
|
self.last_x_count = x_count
|
|
self.last_y_count = y_count
|
|
self.curr_x_index = 1
|
|
self.curr_y_index = 1
|
|
last_x_index = self.curr_x_index
|
|
last_y_index = self.curr_y_index
|
|
self.curr_x_index += 1
|
|
if self.curr_x_index > x_count:
|
|
self.curr_x_index = 1
|
|
self.curr_y_index += 1
|
|
if self.curr_y_index > y_count:
|
|
self.curr_y_index = 1
|
|
return (last_x_index, last_y_index, x_count, y_count,)
|
|
|
|
# This node will always be run
|
|
@classmethod
|
|
def IS_CHANGED( s, x_count, y_count ):
|
|
return float("NaN")
|
|
|
|
|
|
class SimpleLooper:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return { "required": { "loop_count": ("INT", {"default":1, "min": 1, "max": 10, "step": 1}),}}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = None
|
|
CATEGORY = "EasyGrids"
|
|
|
|
class TextConcatenator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return { "required" : { "text_1": ("STRING", {"multiline": True}),
|
|
"text_2": ("STRING", {"multiline": True}), } }
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "concat_text"
|
|
CATEGORY = "EasyGrids"
|
|
|
|
def concat_text( self, text_1, text_2 ):
|
|
#simple as!
|
|
return ((text_1 + text_2), )
|
|
|
|
class FloatToText:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return { "required": { "float_input": ("FLOAT", {"default": 1.0}),
|
|
"decimal_places": ("INT", {"default": 3, "min": 1, "max": 10 }), }}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "convert_to_str"
|
|
CATEGORY = "EasyGrids"
|
|
|
|
def convert_to_str(self, float_input, decimal_places):
|
|
# if this doesn't work, blame Copilot
|
|
formatted_float = "{:.{}f}".format(float_input, decimal_places)
|
|
return (formatted_float,)
|
|
|
|
class SaveImageGrid:
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.prefix_append = ""
|
|
self.image_grid = [[]]
|
|
self.curr_x_size = 1
|
|
self.curr_x_idx = 0
|
|
self.curr_y_size = 1
|
|
self.curr_y_idx = 0
|
|
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"images": ("IMAGE", ),
|
|
"x_size": ("INT", {"default": 1, "min": 1, "max": 12, "step": 1}),
|
|
"y_size": ("INT", {"default": 1, "min": 1, "max": 12, "step": 1}),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"})},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "accumulate_images"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "EasyGrids"
|
|
|
|
def accumulate_images(self, images, x_size, y_size, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
|
|
filename_prefix += self.prefix_append
|
|
if x_size != self.curr_x_size or y_size != self.curr_y_size:
|
|
self.curr_x_size = x_size
|
|
self.curr_y_size = y_size
|
|
self.curr_x_idx = 0
|
|
self.curr_y_idx = 0
|
|
self.image_grid = []
|
|
for image in images:
|
|
self.image_grid.append(image)
|
|
self.curr_x_idx += 1
|
|
if self.curr_x_idx >= self.curr_x_size:
|
|
self.curr_y_idx += 1
|
|
self.curr_x_idx = 0
|
|
|
|
if len( self.image_grid ) >= self.curr_x_size * self.curr_y_size:
|
|
#complete grid
|
|
grid_image = self.assemble_grid()
|
|
return self.save_grid( grid_image, filename_prefix, prompt, extra_pnginfo )
|
|
return { "ui": { "images": [] } }
|
|
|
|
def assemble_grid(self):
|
|
space_height = max( [ len(image) for image in self.image_grid ] )
|
|
space_width = max( [ len(image[0]) for image in self.image_grid ] )
|
|
grid_canvas = Image.new("RGBA", (space_width * self.curr_x_size, space_height * self.curr_y_size), color="#ffffff00")
|
|
for y_idx in range( self.curr_y_size ):
|
|
for x_idx in range( self.curr_x_size ):
|
|
pil_image = Image.fromarray( np.clip( ( self.image_grid[ ( y_idx * self.curr_x_size ) + x_idx].cpu().numpy() * 255. ), 0, 255 ).astype( np.uint8 ) )
|
|
grid_canvas.paste(pil_image, (x_idx * space_width, y_idx * space_height ))
|
|
return grid_canvas
|
|
|
|
|
|
def save_grid( self, grid_image, filename_prefix, prompt=None, extra_pnginfo=None ):
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, grid_image.width, grid_image.height)
|
|
results = list()
|
|
metadata = None
|
|
if not args.disable_metadata:
|
|
metadata = PngInfo()
|
|
if prompt is not None:
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
|
|
|
file = f"{filename}_{counter:05}_.png"
|
|
grid_image.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4)
|
|
results.append({
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": self.type
|
|
})
|
|
|
|
return { "ui": { "images": results } }
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageGridCommander": ImageGridCommander,
|
|
"GridFloatList": GridFloatList,
|
|
"GridFloats": GridFloats,
|
|
"SimpleLooper": SimpleLooper,
|
|
"TextConcatenator": TextConcatenator,
|
|
"FloatToText": FloatToText,
|
|
"SaveImageGrid": SaveImageGrid,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageGridCommander": "Create Image Grid",
|
|
"GridFloatList": "Float List fom Text",
|
|
"GridFloats" : "Float List",
|
|
"SimpleLooper": "Simple Looper",
|
|
"TextConcatenator": "Text Concatenator",
|
|
"FloatToText": "Float to Text",
|
|
"SaveImageGrid": "Save Image Grid",
|
|
} |