117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
import os
|
|
import math
|
|
import json
|
|
from typing import List
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
from nodes import SaveImage
|
|
|
|
class GridImage(SaveImage):
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
'required': {
|
|
'images': ('IMAGE',),
|
|
'filename_prefix': ('STRING', {'default': 'ComfyUI-Grid'}),
|
|
'x': ('INT', {
|
|
'default': 1,
|
|
'min': 1,
|
|
'max': 64,
|
|
'step': 1
|
|
}),
|
|
'gap': ('INT', {
|
|
'default': 0,
|
|
'min': 0,
|
|
'max': 32,
|
|
'step': 1
|
|
}),
|
|
},
|
|
'hidden': {
|
|
'prompt': 'PROMPT',
|
|
'extra_pnginfo': 'EXTRA_PNGINFO'
|
|
},
|
|
}
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
RETURN_TYPES = ()
|
|
|
|
FUNCTION = 'execute'
|
|
|
|
CATEGORY = 'image'
|
|
|
|
def execute(self, images: List[torch.Tensor], filename_prefix: str = 'ComfyUI-Grid', x: int = 1, gap: int = 0, prompt=None, extra_pnginfo=None):
|
|
y = max([math.ceil(len(images) / x), 1])
|
|
|
|
def map_filename(filename):
|
|
prefix_len = len(os.path.basename(filename_prefix))
|
|
prefix = filename[:prefix_len + 1]
|
|
try:
|
|
digits = int(filename[prefix_len + 1:].split('_')[0])
|
|
except:
|
|
digits = 0
|
|
return (digits, prefix)
|
|
|
|
subfolder = os.path.dirname(os.path.normpath(filename_prefix))
|
|
filename = os.path.basename(os.path.normpath(filename_prefix))
|
|
|
|
full_output_folder = os.path.join(self.output_dir, subfolder)
|
|
|
|
if os.path.commonpath((self.output_dir, os.path.realpath(full_output_folder))) != self.output_dir:
|
|
print("Saving image outside the output folder is not allowed.")
|
|
return {}
|
|
|
|
try:
|
|
counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1
|
|
except ValueError:
|
|
counter = 1
|
|
except FileNotFoundError:
|
|
os.makedirs(full_output_folder, exist_ok=True)
|
|
counter = 1
|
|
|
|
if not os.path.exists(self.output_dir):
|
|
os.makedirs(self.output_dir)
|
|
|
|
results = list()
|
|
|
|
canvas = self.grid_image(images, x, y, gap)
|
|
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"
|
|
canvas.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=4)
|
|
results.append({
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": 'output'
|
|
})
|
|
counter += 1
|
|
|
|
return { "ui": { "images": results } }
|
|
|
|
def grid_image(self, images: List[torch.Tensor], x: int, y: int, gap: int):
|
|
width, height, _ = images[0].shape
|
|
canvas = Image.new('RGB', (x*(width+gap)-gap, y*(height+gap)-gap), color='black')
|
|
|
|
for Y in range(y):
|
|
for X in range(x):
|
|
idx = Y * x + X
|
|
if len(images) <= idx:
|
|
return canvas
|
|
|
|
image = images[idx]
|
|
|
|
i = 255. * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
canvas.paste(img, (X*(width+gap), Y*(height+gap)))
|
|
|
|
return canvas
|