auto_x.py

import torch
from PIL import Image, ImageDraw, ImageFont
import numpy as np
from comfy.cli_args import args
import os

FONT_PATH = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"

def flatten_dict(d):
    items = {}
    for node_id, node_data in d.items():
        #class_type = node_data.get("class_type", "UnknownType") # 長いとはみでる・・・
        if "inputs" in node_data:
            for key, value in node_data["inputs"].items():
                new_key = f"{key}"
                items[new_key] = value
    return items

def find_differences(dicts):
    flat_dicts = [flatten_dict(d) for d in dicts]
    all_keys = set(k for d in flat_dicts for k in d)
    differences = {}

    for key in all_keys:
        values = [d.get(key) for d in flat_dicts]
        if not all(v == values[0] for v in values):
            differences[key] = values

    return differences

def create_labeled_image(label, img, font_path=None, font_size=32):
    width, height = img.size

    # フォント読み込み（指定なければデフォルト）
    if font_path and os.path.exists(font_path):
        font = ImageFont.truetype(font_path, font_size)
    else:
        font = ImageFont.load_default()

    dummy_img = Image.new("RGB", (1, 1))
    draw = ImageDraw.Draw(dummy_img)
    bbox = draw.textbbox((0, 0), label, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]

    padding = 10
    new_img = Image.new("RGB", (width, height + text_height + padding), "white")
    draw = ImageDraw.Draw(new_img)
    draw.text(((width - text_width) / 2, padding / 2), label, fill="black", font=font)
    new_img.paste(img, (0, text_height + padding))
    return new_img

def generate_labeled_image_strip(pairs, font_path=None, font_size=32):
    labeled_images = [
        create_labeled_image(label, img, font_path, font_size)
        for label, img in pairs
    ]

    total_width = sum(img.width for img in labeled_images)
    max_height = max(img.height for img in labeled_images)

    combined = Image.new("RGB", (total_width, max_height), "white")

    x_offset = 0
    for img in labeled_images:
        combined.paste(img, (x_offset, 0))
        x_offset += img.width

    return combined

def tensor_to_pil(tensor):
    image = 255. * tensor.cpu().numpy()
    image = np.clip(image, 0, 255).astype(np.uint8)
    return Image.fromarray(image)

def pil_to_tensor(pil_image):
    image = np.array(pil_image).astype(np.float32) / 255.0
    tensor = torch.from_numpy(image)
    return tensor.unsqueeze(0)  # Add batch dimension

class AutoXImage:
    def __init__(self):
        self.images = None
        self.prompts = []

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE", ),
            },
            "hidden": {
                "prompt": "PROMPT",
            },
        }

    RETURN_TYPES = ("IMAGE", )
    FUNCTION = "x_images"

    OUTPUT_NODE = True

    CATEGORY = "image"
    DESCRIPTION = "Saves the input images to your ComfyUI output directory."

    def x_images(self, images, prompt=None):
        b, h, w, c = images.shape
        images_vertical = images.reshape(b * h, w, c).unsqueeze(0)
        self.images = images_vertical if self.images is None else torch.cat((self.images, images_vertical), dim=0)
        assert not args.disable_metadata, "metadata is disabled, but this node requires it to be enabled."
        self.prompts.append(prompt)
        diffs = find_differences(self.prompts)
        labels = [""] * len(self.prompts)
        for key, values in diffs.items():
            for i, value in enumerate(values):
                if value is not None:
                    labels[i] += f"{key}: {value}\n"
            
        if len(diffs) > 0:
            pil_images = [tensor_to_pil(img) for img in self.images]
            labeled_images = generate_labeled_image_strip([(label, img) for label, img in zip(labels, pil_images)], FONT_PATH, 64)
            labeled_images = pil_to_tensor(labeled_images)
        else:
            labeled_images = self.images
        
        return (labeled_images, )
    
NODE_CLASS_MAPPINGS = {
    "AutoXImage": AutoXImage
}