Custom Save-to-Image Node: Expert bug finder required

[George713]

Hi everyone,

I am trying to build a custom Save-to-Image node. The aim is for the node to additionally take as input:

1. a custom output dir
2. the seed used for generating the picture
3. an iteration-id (I loop over a list of prompts and want to match the pictures with the prompt id)

Here is my code:

class SaveImageToDirectory:
    CATEGORY = "GeorgeCustomNodes"

    def __init__(self):
        self.compress_level = 4
    
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE", ),
                "directory": ("STRING", {
                    "multiline": False,
                    "default": r"C:\Users\george\Desktop\output",
                }),
                "seed": ("INT", {"default": -1, "forceInput": True}),
                "iteration": ("INT", {"default": -1, "forceInput": True}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "save_image"
    OUTPUT_NODE = True

    def save_image(self, image, directory, seed, iteration):
        i = 255. * image.cpu().numpy().reshape(-1, image.shape[-1]) # HERE LIES AT LEAST ONE ERROR
        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))

        dir_with_seed = Path(directory) / f"{seed}"
        dir_with_seed.mkdir(parents=True, exist_ok=True)

        img.save(dir_with_seed / f"{iteration}.png", pnginfo=None, compress_level=self.compress_level)
        return { "ui": { "images": [{
                "filename": f"{iteration}.png",
                "subfolder": str(dir_with_seed),
                "type": "output"
            }] } }

I copied most from ComfyUI's own SaveImage node.

You will easily spot the line i = 255. * ima.... The .reshape... part was added, as I was getting this error otherwise:

Error occurred when executing SaveImageToDirectory:

Cannot handle this data type: (1, 1, 768, 3), |u1

  ...

  File "C:\Users\froeb\Documents\Projects\ComfyUI_windows_portable\ComfyUI\custom_nodes\prompts_from_file.py", line 83, in save_image
    img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\froeb\Documents\Projects\ComfyUI_windows_portable\python_embeded\Lib\site-packages\PIL\Image.py", line 3102, in fromarray
    raise TypeError(msg) from e

However, now I produce an output png, that cannot be displayed due to format errors using Windows' Photo app. Using Firefox, the file appears to be 1px wide. ComfyUI's own SaveImage and PreviewImage work fine, leaving me wondering what's going on.

[blepping]

@George713

The .reshape... part was added, as I was getting this error

not an expert but i'm going to guess that is because the IMAGE type actually supports batches of images. so you're taking the batch of images (even if it's just one) and reshaping it. you'd probably be better off doing something like:

    def save_images(self, images, directory, seed, iteration):
        i = 255. * images[0].cpu().numpy()
        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
        # logic to save it here

code untested, but i believe that should get you a lot closer to the existing SaveImage node behavior. be aware that it doesn't handle batches: it'll crash on an empty batch of images and it'll only deal with the first one if you pass it a batch of multiple images.

[alexisrolland]

Thank you for this thread that puts me on the right track... I have a custom node and facing a similar issue. I'm wondering how could I tweak @blepping code above to be able to support batches of images. Currently when I input multiple images I get the error:

Error occurred when executing DummyTestNode:
Cannot handle this data type: (1, 1, 1383, 3), |u1

Here is the node I have today, very basic for troubleshooting:

import numpy as np
import torch
from PIL import Image

class DummyTestNode:
    # Node setup for ComfyUI
    CATEGORY = "Test"
    FUNCTION = "execute"
    OUTPUT_NODE = False
    RETURN_TYPES = ("IMAGE",)

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
                "IMAGE": ("IMAGE",),
            }
        }

    def execute(self, IMAGE,):
        # Convert tensor to pil image
        pil_img = Image.fromarray(np.clip(255. * IMAGE.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

        # Convert pil to tensor image
        tensor = torch.from_numpy(np.array(pil_img).astype(np.float32) / 255.0).unsqueeze(0)

        return (tensor,)

NODE_CLASS_MAPPINGS = {
    "DummyTestNode": DummyTestNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DummyTestNode": "DummyTestNode"
}

[blepping]

i don't really remember the context here, but assuming images[0] was the first image in the batch, the simple adaption to handle batches would be:

    def execute(self, IMAGE):
        output = []
        for idx in range(images.shape[0]):
          i = 255. * IMAGE[idx].cpu().numpy()
          output.append(Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)))
        # at this point output is a list of the processed images

i'm not sure your types makes sense though, your input type is IMAGE but i don't think the result after that conversion is still IMAGE so you'd probably have to return a different type

[alexisrolland]

Thank you @blepping I was able to find and implement a solution. Example code in this repo: https://github.com/alexisrolland/ComfyUI-AuraSR/blob/main/__init__.py#L34

I reused two functions from https://github.com/melMass/comfy_mtb/blob/4eebdd8b8bff73c3db4f0248da8dac7d67cb310b/utils.py#L504

• tensor2pil
• pil2tensor

class RunAuraSR:
    """Node to run AuraSR model."""

    # Node setup for ComfyUI
    CATEGORY = "AuraSR"
    FUNCTION = "execute"
    OUTPUT_NODE = False
    RETURN_TYPES = ("IMAGE",)

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
                "AURASR_MODEL": ("AURASR_MODEL",),
                "IMAGE": ("IMAGE",),
            }
        }

    def tensor2pil(self, image):
        batch_count = image.size(0) if len(image.shape) > 3 else 1
        if batch_count > 1:
            out = []
            for i in range(batch_count):
                out.extend(self.tensor2pil(image[i]))
            return out
        return [Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))]

    def pil2tensor(self, image):
        if isinstance(image, list):
            return torch.cat([self.pil2tensor(img) for img in image], dim=0)
        return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

    def execute(self, AURASR_MODEL, IMAGE,):
        # Convert tensor to PIL image
        images = self.tensor2pil(IMAGE)

        # Upscale images
        upscaled_images = list()
        for image in images:
            upscaled_images.append(AURASR_MODEL.upscale_4x(image))

        # Convert PIL images to tensor
        upscaled_images = self.pil2tensor(upscaled_images)

        return (upscaled_images,)