298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-07-18 05:27:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

Only add PNG chunks if metadata is enabled

Browse Source 
Prevents a hypothetical scenario of a VAE being crafted to inject
undesired data into PNG chunks.
This commit is contained in:
catboxanon 2025-01-31 11:40:28 -05:00
parent a4aba18d29
commit 49aaad7c21
1 changed files with 8 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
nodes.py
View File

@ -1615,23 +1615,23 @@ class SaveImage:
for (batch_number, image) in enumerate(images): for (batch_number, image) in enumerate(images):
i = 255. * image.cpu().numpy() i = 255. * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
metadata = PngInfo() metadata = None
if not args.disable_metadata: if not args.disable_metadata:
metadata = PngInfo()
if prompt is not None: if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt)) metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None: if extra_pnginfo is not None:
for x in extra_pnginfo: for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x])) metadata.add_text(x, json.dumps(extra_pnginfo[x]))
if hasattr(images, "png_chunks"): if hasattr(images, "png_chunks"):
for name, data in images.png_chunks.items(): for name, data in images.png_chunks.items():
if name in self.extra_chunks: if name in self.extra_chunks:
metadata.add(name.lower(), data) metadata.add(name.lower(), data)
else: else:
metadata.add(name, data) metadata.add(name, data)
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
file = f"{filename_with_batch_num}_{counter:05}_.png" file = f"{filename_with_batch_num}_{counter:05}_.png"
#TODO: revert to using img.save once Pillow supports cICP chunk #TODO: revert to using img.save once Pillow supports cICP chunk
img.encoderinfo = {"pnginfo": metadata, "compress_level": self.compress_level} img.encoderinfo = {"pnginfo": metadata, "compress_level": self.compress_level}
with open(os.path.join(full_output_folder, file), 'wb') as fp: with open(os.path.join(full_output_folder, file), 'wb') as fp: