Feature request: batch images

[agawrylak]

Is it possible to add an option to return a batch of images?

[aleph65]

Was also looking for this

I guess for now workaround is to do the batch image logic locally... change seed as required

[TimPietrusky]

@agawrylak @aleph65 so you mean that your workflow is creating multiple images and you want all of them to be returned?

[aleph65]

yes exactly!

[TimPietrusky]

@aleph65 would you mind providing me with an example workflow, so that I can use it for testing & development?

[vesper8]

I'm also interested in this. Just using the default test input from https://raw.githubusercontent.com/blib-la/runpod-worker-comfy/main/test_input.json

If I increase the batch_size from 1 to say 5, how can I get the 5 images separately from the API response?

[TimPietrusky]

Just using the default test input from https://raw.githubusercontent.com/blib-la/runpod-worker-comfy/main/test_input.json

If I increase the batch_size from 1 to say 5, how can I get the 5 images separately from the API response?

Ok thank you! I imagine we would deliver an array of images then. But we have to make sure that this is not breaking the current behavior OR release a breaking change.

[aleph65]

sorry for not sending this yet, my examples require additional packages and I want to send a out-of-box one

Basically, there are 2 scenarios:

1 - there is a batch of images in a Save Image or Display Image node

2 - There are multiple Save Image or Display Image nodes

If you query comfyui api today you will receive all of those outputs (all batches of all output image nodes)

Personally if backcompat is important you can just return the original first image in the same place and return the full array as well. And add a flag for people who have already migrated to this new signature

[RaresKeY]

I'm also looking for this feature. Here's an example of the default flux workflow with changed batch latents to 2: workflow.json

The way comfyui handles this is by saving 2 images sequentially, you guys send the last image in ["output"]["message"] so the easiest way is to read how many files where generated in folder before and after comfyui gen and send those as ["output"]["message1"], ["output"]["message2"] ... ["output"]["messageN"]

The reason I suggest reading the files instead of the workflow is in case the workflow contains multiple empty latents or custom nodes but for a default workflow you can just read "batch_size" to get the batch number.

this is some relevant part of a log from my runpod testing:

...
2024-09-13T10:07:20.027113249Z runpod-worker-comfy - image generation is done
2024-09-13T10:07:20.027156841Z runpod-worker-comfy - /comfyui/output/ComfyUI_00040_.png
...
2024-09-13T10:17:43.164717148Z runpod-worker-comfy - image generation is done
2024-09-13T10:17:43.164772808Z runpod-worker-comfy - /comfyui/output/ComfyUI_00042_.png
...

let me know if you need testing, I would be happy to help

PS: A workaround for users is to make a custom workflow that combines all images and saves them as one and then split them locally

PPS: Modifying src/rp_handler.py, process_output_images function like this should do the trick (only tested with strings):

def process_output_images(outputs, job_id):
    """
    This function processes multiple output images generated by ComfyUI.
    It returns either the URLs (if AWS S3 is configured) or base64 strings for each image.

    Args:
        outputs (dict): A dictionary containing the outputs from image generation,
                        typically includes node IDs and their respective output data.
        job_id (str): The unique identifier for the job.

    Returns:
        dict: A dictionary containing a list of images (URLs or base64) and the status.
    """

    # The path where ComfyUI stores the generated images
    COMFY_OUTPUT_PATH = os.environ.get("COMFY_OUTPUT_PATH", "/comfyui/output")

    # To store all the images (either URLs or base64 encoded strings)
    output_images = []

    # Iterate through each output node and its respective images
    for node_id, node_output in outputs.items():
        if "images" in node_output:
            for image in node_output["images"]:
                # Construct the local image path
                local_image_path = os.path.join(COMFY_OUTPUT_PATH, image["subfolder"], image["filename"])

                # Check if the image exists
                if os.path.exists(local_image_path):
                    if os.environ.get("BUCKET_ENDPOINT_URL", False):
                        # Upload image to AWS S3 and get the URL
                        image_url = rp_upload.upload_image(job_id, local_image_path)
                        output_images.append(image_url)
                        print(f"runpod-worker-comfy - Uploaded image {image['filename']} to S3")
                    else:
                        # Convert image to base64
                        base64_image = base64_encode(local_image_path)
                        output_images.append(base64_image)
                        print(f"runpod-worker-comfy - Converted image {image['filename']} to base64")
                else:
                    print(f"runpod-worker-comfy - Image {image['filename']} does not exist in output folder")

    if output_images:
        return {
            "status": "success",
            "message": output_images
        }
    else:
        print("runpod-worker-comfy - the image does not exist in the output folder")
        return {
            "status": "error",
            "message": "No images were found or generated."
        }