Patch WebP Workflow Support

WAS_Node_Suite.py

@@ -7166,14 +7166,30 @@ def was_save_images(self, images, output_path='', filename_prefix="ComfyUI", fil
             i = 255. * image.cpu().numpy()
             img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
 
-            metadata = PngInfo()
-            if embed_workflow == 'true':
+            # Delegate metadata/pnginfo
+            if extension == 'webp':
+                img_exif = img.getexif()
+                workflow_metadata = ''
+                prompt_str = ''
                 if prompt is not None:
-                    metadata.add_text("prompt", json.dumps(prompt))
+                    prompt_str = json.dumps(prompt)
+                    img_exif[0x010f] = "Prompt:" + prompt_str
                 if extra_pnginfo is not None:
                     for x in extra_pnginfo:
-                        metadata.add_text(x, json.dumps(extra_pnginfo[x]))
-
+                        workflow_metadata += json.dumps(extra_pnginfo[x])
+                img_exif[0x010e] = "Workflow:" + workflow_metadata
+                exif_data = img_exif.tobytes()
+            else:
+                metadata = PngInfo()
+                if embed_workflow == 'true':
+                    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]))
+                exif_data = metadata
+
+            # Delegate the filename stuffs
             if overwrite_mode == 'prefix_as_filename':
                 file = f"{filename_prefix}{file_extension}"
             else:
@@ -7183,22 +7199,27 @@ def was_save_images(self, images, output_path='', filename_prefix="ComfyUI", fil
                     file = f"{filename_prefix}{delimiter}{counter:0{number_padding}}{file_extension}"
                 if os.path.exists(os.path.join(output_path, file)):
                     counter += 1
+
+            # Save the images
             try:
                 output_file = os.path.abspath(os.path.join(output_path, file))
-                if extension == 'png':
-                    img.save(output_file,
-                             pnginfo=metadata, optimize=True)
-                elif extension in ["jpg", "jpeg"]:
-                    img.save(output_file,
-                             quality=quality, optimize=True)
-                elif extension == 'tiff':
+                if extension in ["jpg", "jpeg"]:
                     img.save(output_file,
                              quality=quality, optimize=True)
                 elif extension == 'webp':
-                    img.save(output_file, quality=quality, 
-                                lossless=lossless_webp, exif=metadata)
+                    img.save(output_file, 
+                             quality=quality, lossless=lossless_webp, exif=exif_data)
+                elif extension == 'png':
+                    img.save(output_file, 
+                             pnginfo=exif_data, optimize=True)
                 elif extension == 'bmp':
                     img.save(output_file)
+                elif extension == 'tiff':
+                    img.save(output_file,
+                             quality=quality, optimize=True)
+                else:
+                    img.save(output_file, 
+                             pnginfo=exif_data, optimize=True)
 
                 cstr(f"Image file saved to: {output_file}").msg.print()
 
@@ -7264,7 +7285,6 @@ def was_save_images(self, images, output_path='', filename_prefix="ComfyUI", fil
             return {"ui": {"images": results}}
         else:
             return {"ui": {"images": []}}
-
 
     def get_subfolder_path(self, image_path, output_path):
         output_parts = output_path.strip(os.sep).split(os.sep)
@@ -7319,8 +7339,9 @@ def load_image(self, image_path, RGBA='false', filename_text_extension="true"):
         # Update history
         update_history_images(image_path)
 
+        image = i
         if not RGBA:
-            image = i.convert('RGB')
+            image = image.convert('RGB')
         image = np.array(image).astype(np.float32) / 255.0
         image = torch.from_numpy(image)[None,]