Fixed Issue with LoadImage node when loading PNG files with embedded ICC profiles.

[shawnington]

Fixed Value Error: Decompressed Data Too Large thrown by PIL when attempting to opening PNG files with large embedded ICC colorspaces.

Error occurs because PIL misinterprets large embedded ICC profiles in PNG images as malicious decompression bombs, and throws

By setting the follow flag to true when loading PNG images PIL discards meta data, preventing the error:

ImageFile.LOAD_TRUNCATED_IMAGES = True

[shawnington]

Additional Context:

Large ICC profiles are embedded when systems have had their display color profiled with a colorimeter creating an ICC profile that adjusts for accurate color display for use in a color managed workflow. The ICC profiles can be several MB in size depending on the profiling method (large LUTs for example).

The issue is with the Pillow library and how it treats embedded ICC profiles in PNG files. As some ICC profiles can actually be larger in size than the image data, PIL assumes it is a potential decompression bomb DDOS attack, and throws an error preventing the image from being loaded.

Setting the ImageFile.LOAD_TRUNCATED_IMAGES flag to true when the image is of type PNG allow these images to be used in the Load Image Node.

The flag instructs PIL to discard meta data before loading the image. Setting this flag does not expose any vulnerabilities to decompression bomb attacks, but does allow the image to be loaded.

ICC profiles are handled correctly for JPG images in PIL.

Metadata is destroyed anyways in the process of VAE encoding, so there are no negative effects to the proposed fix, as neither ComfyUI nor Stable Diffusion include any kind of color management system, and the ICC profile will not be used by the VAE when encoding, nor will it be passed through to the output.

[asagi4]

Should the node restore the old value after the image has been loaded? It seems to be a module global and that looks like setting it might affect nodes that do want to read metadata.

[comfyanonymous]

Yeah the best way to do this might be with a custom context manager that sets ImageFile.LOAD_TRUNCATED_IMAGES = True and restores it to the previous value afterwards.

This is a good place to put that context manager so it can be reused by other nodes if needed: https://github.com/comfyanonymous/ComfyUI/blob/master/node_helpers.py

[shawnington]

Should the node restore the old value after the image has been loaded? It seems to be a module global and that looks like setting it might affect nodes that do want to read metadata.

Fair, I did not consider resetting the flag after loading. The current issue is the nodes reliance on Image.open without any try except error handling when Image.open can throw in a blocking way.

perhaps a helper function is the way to go.

Presumably downstream behavior of nodes wanting to use the meta data can be ignored in cases where the nodes are unreachable due to image loading failure, so loading the image with the flag set in such case would be the preferred behavior.

[shawnington]

Yeah the best way to do this might be with a custom context manager that sets ImageFile.LOAD_TRUNCATED_IMAGES = True and restores it to the previous value afterwards.

This is a good place to put that context manager so it can be reused by other nodes if needed: https://github.com/comfyanonymous/ComfyUI/blob/master/node_helpers.py

As the issue seems to be with Pillows treatment of ICC profiles in PNG images, the use of Image.open directly without a try except in the LoadImage node itself hard errors on failure. Perhaps a helper node_helpers.open_image function that uses try except to catch errors when opening images such as follows is a palatable workaround to the behavior.

Something like this in node_helpers.py

def open_image(path):
    try : 
        img = Image.open(path)
    
    except:
        ImageFile.LOAD_TRUNCATED_IMAGES = True
        img = Image.open(path)
        ImageFile.LOAD_TRUNCATED_IMAGES = False
        
    finally: 
        return img

and replace in LoadImage in nodes.py

img = Image.open(image_path)

with

img = node_helpers.open_image(image_path)

Just tested this implementation method locally, preservation of meta-data is achieved except in instances where Value Errors are thrown by Pillow and truncation of meta data is the only way to open the image.

It behaves the same as Image.open but would attempt to open the image using the flag if Image.open throws to due to pillows incorrect handling of ICC profiles for png files.

Presumably truncation of metadata is acceptable downstream behavior in cases where Pillow throws and halts the workflow, as nodes wanting to use the meta data are unreachable in cases where an image is unable to be loaded without truncated meta data.

If this appears to be a palatable solution, I will commit it to my fork.

[comfyanonymous]

Sure that works.

[shawnington]

Sure that works.

The suggested changed have been committed to nodes.py and node_helpers.py