From 0d45efb7d6809fae272a9ba68c3aac5f713347e7 Mon Sep 17 00:00:00 2001
From: shawnington <88048838+shawnington@users.noreply.github.com>
Date: Sat, 4 May 2024 02:32:41 -0500
Subject: [PATCH] Fixed Issue with LoadImage node when loading PNG files with
 embedded ICC profiles. (#3316)

* Fix issue with how PIL loads small PNG files nodes.py

Added flag to prevent ValueError: Decompressed Data Too Large
when loading PNG images with large meta data such as large embedded color profiles

* Update LoadImage node to fix error when loading PNG's in nodes.py

Fixed Value Error: Decompressed Data Too Large thrown by PIL when attempting to opening PNG files with large embedded ICC colorspaces by setting the follow flag to true when loading png images:  ImageFile.LOAD_TRUNCATED_IMAGES = True

* Update node_helpers.py to include open_image helper function

open_image includes try except to catch Pillow Value Errors that occur when large ICC profiles are embedded in images.

* Update LoadImage node to use open_image helper function inplace of Image.open

open_image helper function in node_helpers.py  fixes a Pillow error when attempting to open images with large embedded ICC profiles by adding an exception handler to load the image with truncated meta data if regular loading is not possible.
---
 node_helpers.py | 14 ++++++++++++++
 nodes.py        |  6 ++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/node_helpers.py b/node_helpers.py
index 8828a4ec..264bd4d5 100644
--- a/node_helpers.py
+++ b/node_helpers.py
@@ -1,3 +1,4 @@
+from PIL import Image, ImageFile
 
 def conditioning_set_values(conditioning, values={}):
     c = []
@@ -8,3 +9,16 @@ def conditioning_set_values(conditioning, values={}):
         c.append(n)
 
     return c
+
+def open_image(path):
+    try :
+        ImageFile.LOAD_TRUNCATED_IMAGES = False
+        img = Image.open(path)
+    
+    except:
+        ImageFile.LOAD_TRUNCATED_IMAGES = True
+        img = Image.open(path)
+        
+    finally:
+        ImageFile.LOAD_TRUNCATED_IMAGES = False
+        return img
diff --git a/nodes.py b/nodes.py
index acad256f..aa6d6fa9 100644
--- a/nodes.py
+++ b/nodes.py
@@ -12,12 +12,12 @@ import logging
 
 from PIL import Image, ImageOps, ImageSequence
 from PIL.PngImagePlugin import PngInfo
+
 import numpy as np
 import safetensors.torch
 
 sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
 
-
 import comfy.diffusers_load
 import comfy.samplers
 import comfy.sample
@@ -1456,7 +1456,9 @@ class LoadImage:
     FUNCTION = "load_image"
     def load_image(self, image):
         image_path = folder_paths.get_annotated_filepath(image)
-        img = Image.open(image_path)
+        
+        img = node_helpers.open_image(image_path)
+        
         output_images = []
         output_masks = []
         for i in ImageSequence.Iterator(img):