|
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 |
|
|
