ComfyUIのカスタムノードをChatGPT 4o先生と作る

ChatGpt4o先生に相談して、ComfyUIのカスタムノードが簡単にできたので、ご報告。

(生成画像の直前のノードがカスタムノード）

AIで画像生成し、それをモノクロ化して、さらに黄色っぽくするという作業を
効率的にしたかったのでカスタムノードを作る必要が生じた。
このくらいのことは後でPhotoshopでもできる作業だが、AI画像を大量に生成し
それを評価するためには、なるべく工程を減らすのが吉。

そこで、既にChatGpt4o先生に相談して作った画像をモノクロ化して、それを
黄色に寄せる画像処理コードを作成した後、

ComfyUIで上のコードを実行して画像処理をするノードのコードを生成せよ

と指示。すると、以下のようなコードが生成された。

import numpy as np
from PIL import Image

class GrayscaleAddColorNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "add_color": ("RGB", {"default": [255, 255, 0]}),  # カラー加算するRGB
            }
        }

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

    CATEGORY = "Image Processing"

    def process(self, image, add_color):
        # 画像をNumPy配列に変換
        img_array = np.array(image)

        # グレースケールに変換
        gray_img = Image.fromarray(img_array).convert('L')
        gray_img_rgb = gray_img.convert('RGB')

        # NumPy配列に変換
        img_array = np.array(gray_img_rgb)

        # 明るさを基準にカラーを加算（明るさに応じてスケーリング）
        brightness = np.mean(img_array, axis=-1, keepdims=True) / 255.0
        color_to_add = np.array(add_color)
        img_array = np.clip(img_array + (color_to_add * brightness), 0, 255).astype(np.uint8)

        # 結果を画像に変換
        output_image = Image.fromarray(img_array)

        return (output_image,)

これをComfyUI内のどのフォルダに入れるかも教えてくれ。
起動して実行し、検索するが、このノードがみつからない！

そこで、既存のノードと比べると、NODE_CLASS_MAPPINGS がないので
これを以下のように追加する。

NODE_CLASS_MAPPINGS = {
    "GrayscaleAddColor": GrayscaleAddColor,
}

これでめでたく検索してノードが配置できるが、いくつかエラーや問題が起こるが
とくに、

GrayscaleAddColor
Cannot handle this data type: (1, 1, 512, 3), <f4

というエラーが発生。そこで、これも先生に相談すると、

エラーメッセージ「Cannot handle this data type: (1, 1, 512, 3), <f4」は、画像データの型や形状が期待される形式と異なる場合に発生します。ComfyUI内で使用される画像データの形式が、PILやNumPyが期待している形式とは異なるため、このエラーが発生している可能性があります。

ComfyUIの画像データは通常、torch.Tensor形式で保持されていることがあります。この場合、PILで扱える形式（通常はNumPy配列）に変換してから処理する必要があります。

以下の解決策を提案します。ComfyUIで使用されている画像データがtorch.Tensor形式で渡されていると仮定して、torch.TensorをNumPy配列に変換する処理を追加します。

と教えてくれて、最終的に以下のコードができた。

import numpy as np
from PIL import Image
import torch

class GrayscaleAddColor:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "red": ("INT", {"default": 255, "min": 0, "max": 255}),
                "green": ("INT", {"default": 255, "min": 0, "max": 255}),
                "blue": ("INT", {"default": 255, "min": 0, "max": 255}),
            }
        }

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

    CATEGORY = "Image"

    def process(self, image, red, green, blue):
        # Torch TensorをNumPy配列に変換
        if isinstance(image, torch.Tensor):
            # (C, H, W)の形状にしてからNumPy配列に変換
            image = image.squeeze(0).cpu().numpy()
            # 0〜1の範囲から0〜255にスケール変換
            image = (image * 255).astype(np.uint8)

        # グレースケールに変換
        gray_img = Image.fromarray(image).convert('L')
        gray_img_rgb = gray_img.convert('RGB')

        # NumPy配列に変換
        img_array = np.array(gray_img_rgb)

        # 明るさを基準にカラーを加算（明るさに応じてスケーリング）
        brightness = np.mean(img_array, axis=-1, keepdims=True) / 255.0
        color_to_add = np.array([red, green, blue])

        # 明るさを考慮してカラーを加算（色のバランスを保つ）
        img_array = img_array.astype(np.float32)
        img_array += color_to_add * brightness
        img_array = np.clip(img_array, 0, 255).astype(np.uint8)

        # 結果を画像に変換
        output_image = Image.fromarray(img_array)

        # NumPy配列をTorch Tensorに戻して返す
        output_image = torch.from_numpy(np.array(output_image)).float() / 255.0
        output_image = output_image.unsqueeze(0)  # バッチサイズの次元を戻す

        return (output_image,)

    
NODE_CLASS_MAPPINGS = {
    "GrayscaleAddColor": GrayscaleAddColor,
}