From a4aba18d290d9e29ab1584fb1876513cee427743 Mon Sep 17 00:00:00 2001 From: catboxanon <122327233+catboxanon@users.noreply.github.com> Date: Thu, 30 Jan 2025 15:44:41 -0500 Subject: [PATCH 1/7] PNG cICP chunk support --- comfy/sd.py | 11 ++++++++++- nodes.py | 36 +++++++++++++++++++++++++++++------- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index d7e89f726..a3713c260 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -246,7 +246,7 @@ class CLIP: return self.patcher.get_key_patches() class VAE: - def __init__(self, sd=None, device=None, config=None, dtype=None): + def __init__(self, sd=None, device=None, config=None, dtype=None, meta=None): if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format sd = diffusers_convert.convert_vae_state_dict(sd) @@ -416,6 +416,15 @@ class VAE: self.first_stage_model.to(self.vae_dtype) self.output_device = model_management.intermediate_device() + self.png_chunks = {} + + if meta is not None: + meta_color_space = meta.get("modelspec.color_space") + if str(meta_color_space).lower().startswith("cicp:"): + cicp_chunk = meta_color_space.split("cicp:")[-1].split(",") + cicp_chunk = bytes([1 if b.lower() == 'true' else 0 if b.lower() == 'false' else int(b) for b in cicp_chunk]) + self.png_chunks[b"cICP"] = cicp_chunk + self.patcher = comfy.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device) logging.info("VAE load device: {}, offload device: {}, dtype: {}".format(self.device, offload_device, self.vae_dtype)) diff --git a/nodes.py b/nodes.py index 968f0f9ad..0c2ef1eb4 100644 --- a/nodes.py +++ b/nodes.py @@ -11,7 +11,7 @@ import time import random import logging -from PIL import Image, ImageOps, ImageSequence +from PIL import Image, ImageOps, ImageSequence, PngImagePlugin from PIL.PngImagePlugin import PngInfo import numpy as np @@ -283,10 +283,12 @@ class VAEDecode: CATEGORY = "latent" DESCRIPTION = "Decodes latent images back into pixel space images." - def decode(self, vae, samples): + def decode(self, vae: comfy.sd.VAE, samples): images = vae.decode(samples["samples"]) if len(images.shape) == 5: #Combine batches images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) + if vae.png_chunks is not None: + images.png_chunks = vae.png_chunks return (images, ) class VAEDecodeTiled: @@ -769,7 +771,8 @@ class VAELoader: else: vae_path = folder_paths.get_full_path_or_raise("vae", vae_name) sd = comfy.utils.load_torch_file(vae_path) - vae = comfy.sd.VAE(sd=sd) + meta = json.loads(comfy.utils.safetensors_header(vae_path, max_size=1024*1024) or "{}").get("__metadata__") + vae = comfy.sd.VAE(sd=sd, meta=meta) return (vae,) class ControlNetLoader: @@ -1576,6 +1579,7 @@ class SaveImage: self.type = "output" self.prefix_append = "" self.compress_level = 4 + self.extra_chunks = [b"cICP"] @classmethod def INPUT_TYPES(s): @@ -1597,6 +1601,13 @@ class SaveImage: CATEGORY = "image" DESCRIPTION = "Saves the input images to your ComfyUI output directory." + def putchunk_patched(self, fp, cid, *data): + for chunk in self.extra_chunks: + if cid == chunk.lower(): + cid = chunk + break + return PngImagePlugin.putchunk(fp, cid, *data) + def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) @@ -1604,18 +1615,28 @@ class SaveImage: for (batch_number, image) in enumerate(images): i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) - metadata = None + metadata = PngInfo() if not args.disable_metadata: - metadata = PngInfo() if prompt is not None: metadata.add_text("prompt", json.dumps(prompt)) if extra_pnginfo is not None: for x in extra_pnginfo: metadata.add_text(x, json.dumps(extra_pnginfo[x])) - + if hasattr(images, "png_chunks"): + for name, data in images.png_chunks.items(): + if name in self.extra_chunks: + metadata.add(name.lower(), data) + else: + metadata.add(name, data) filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) file = f"{filename_with_batch_num}_{counter:05}_.png" - img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) + + + #TODO: revert to using img.save once Pillow supports cICP chunk + img.encoderinfo = {"pnginfo": metadata, "compress_level": self.compress_level} + with open(os.path.join(full_output_folder, file), 'wb') as fp: + PngImagePlugin._save(img, fp, None, chunk=self.putchunk_patched) + results.append({ "filename": file, "subfolder": subfolder, @@ -1627,6 +1648,7 @@ class SaveImage: class PreviewImage(SaveImage): def __init__(self): + super().__init__() self.output_dir = folder_paths.get_temp_directory() self.type = "temp" self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) From 49aaad7c21017838cf549c6c50f9853ce1a4c1a9 Mon Sep 17 00:00:00 2001 From: catboxanon <122327233+catboxanon@users.noreply.github.com> Date: Fri, 31 Jan 2025 11:40:28 -0500 Subject: [PATCH 2/7] Only add PNG chunks if metadata is enabled Prevents a hypothetical scenario of a VAE being crafted to inject undesired data into PNG chunks. --- nodes.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/nodes.py b/nodes.py index 0c2ef1eb4..641869496 100644 --- a/nodes.py +++ b/nodes.py @@ -1615,23 +1615,23 @@ class SaveImage: for (batch_number, image) in enumerate(images): i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) - metadata = PngInfo() + metadata = None if not args.disable_metadata: + metadata = PngInfo() if prompt is not None: metadata.add_text("prompt", json.dumps(prompt)) if extra_pnginfo is not None: for x in extra_pnginfo: metadata.add_text(x, json.dumps(extra_pnginfo[x])) - if hasattr(images, "png_chunks"): - for name, data in images.png_chunks.items(): - if name in self.extra_chunks: - metadata.add(name.lower(), data) - else: - metadata.add(name, data) + if hasattr(images, "png_chunks"): + for name, data in images.png_chunks.items(): + if name in self.extra_chunks: + metadata.add(name.lower(), data) + else: + metadata.add(name, data) filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) file = f"{filename_with_batch_num}_{counter:05}_.png" - #TODO: revert to using img.save once Pillow supports cICP chunk img.encoderinfo = {"pnginfo": metadata, "compress_level": self.compress_level} with open(os.path.join(full_output_folder, file), 'wb') as fp: From 3a18b3136c25ed970751ab56139fbaf4e7b74b86 Mon Sep 17 00:00:00 2001 From: catboxanon <122327233+catboxanon@users.noreply.github.com> Date: Mon, 14 Apr 2025 12:47:12 -0400 Subject: [PATCH 3/7] Remove cICP chunk hack Now supported in Pillow as of 11.2.1 https://github.com/python-pillow/Pillow/pull/8704 --- nodes.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/nodes.py b/nodes.py index 641869496..520a50889 100644 --- a/nodes.py +++ b/nodes.py @@ -1601,13 +1601,6 @@ class SaveImage: CATEGORY = "image" DESCRIPTION = "Saves the input images to your ComfyUI output directory." - def putchunk_patched(self, fp, cid, *data): - for chunk in self.extra_chunks: - if cid == chunk.lower(): - cid = chunk - break - return PngImagePlugin.putchunk(fp, cid, *data) - def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) @@ -1632,10 +1625,7 @@ class SaveImage: filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) file = f"{filename_with_batch_num}_{counter:05}_.png" - #TODO: revert to using img.save once Pillow supports cICP chunk - img.encoderinfo = {"pnginfo": metadata, "compress_level": self.compress_level} - with open(os.path.join(full_output_folder, file), 'wb') as fp: - PngImagePlugin._save(img, fp, None, chunk=self.putchunk_patched) + img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) results.append({ "filename": file, From 5ad716d8c86e33558e04e9780cd276e2b4a6e774 Mon Sep 17 00:00:00 2001 From: catboxanon <122327233+catboxanon@users.noreply.github.com> Date: Mon, 14 Apr 2025 12:51:15 -0400 Subject: [PATCH 4/7] Remove extra newlines --- nodes.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/nodes.py b/nodes.py index 33178621a..a17002b44 100644 --- a/nodes.py +++ b/nodes.py @@ -1634,9 +1634,7 @@ class SaveImage: metadata.add(name, data) filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) file = f"{filename_with_batch_num}_{counter:05}_.png" - img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) - results.append({ "filename": file, "subfolder": subfolder, From 5d6bca8e78207b29c7a108c0b0c0b096ae610034 Mon Sep 17 00:00:00 2001 From: catboxanon <122327233+catboxanon@users.noreply.github.com> Date: Mon, 14 Apr 2025 12:51:50 -0400 Subject: [PATCH 5/7] Remove extra chunks variable --- nodes.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/nodes.py b/nodes.py index a17002b44..74c92b507 100644 --- a/nodes.py +++ b/nodes.py @@ -1589,7 +1589,6 @@ class SaveImage: self.type = "output" self.prefix_append = "" self.compress_level = 4 - self.extra_chunks = [b"cICP"] @classmethod def INPUT_TYPES(s): @@ -1646,7 +1645,6 @@ class SaveImage: class PreviewImage(SaveImage): def __init__(self): - super().__init__() self.output_dir = folder_paths.get_temp_directory() self.type = "temp" self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) From 8060017fe0d9c1510f8f3f62d6f826987e510fe4 Mon Sep 17 00:00:00 2001 From: catboxanon <122327233+catboxanon@users.noreply.github.com> Date: Mon, 14 Apr 2025 12:52:52 -0400 Subject: [PATCH 6/7] Fix extra chunk metadata add --- nodes.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/nodes.py b/nodes.py index 74c92b507..06a2d1fb6 100644 --- a/nodes.py +++ b/nodes.py @@ -1627,10 +1627,7 @@ class SaveImage: metadata.add_text(x, json.dumps(extra_pnginfo[x])) if hasattr(images, "png_chunks"): for name, data in images.png_chunks.items(): - if name in self.extra_chunks: - metadata.add(name.lower(), data) - else: - metadata.add(name, data) + metadata.add(name, data) filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) file = f"{filename_with_batch_num}_{counter:05}_.png" img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) From 1ffed0f41cd733237f409afaba0a9ca40c8b50b5 Mon Sep 17 00:00:00 2001 From: catboxanon <122327233+catboxanon@users.noreply.github.com> Date: Mon, 14 Apr 2025 12:53:47 -0400 Subject: [PATCH 7/7] Remove unused PngImagePlugin import --- nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index 06a2d1fb6..48eb356fa 100644 --- a/nodes.py +++ b/nodes.py @@ -11,7 +11,7 @@ import time import random import logging -from PIL import Image, ImageOps, ImageSequence, PngImagePlugin +from PIL import Image, ImageOps, ImageSequence from PIL.PngImagePlugin import PngInfo import numpy as np