v0.3

save_api_image.py

@@ -1,6 +1,8 @@
 import os
 import numpy as np
 from PIL import Image
+import json
+from PIL.PngImagePlugin import PngInfo
 
 class SaveApiImage:
     @classmethod
@@ -8,7 +10,8 @@ class SaveApiImage:
         return {
             "required": {
                 "image": ("IMAGE", {"forceInput": True}),
-            }
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
         }
 
     FUNCTION = "save_api_image"
@@ -16,7 +19,7 @@ class SaveApiImage:
     OUTPUT_NODE = True
     CATEGORY = "Bjornulf"
 
-    def save_api_image(self, image):
+    def save_api_image(self, image, prompt=None, extra_pnginfo=None):
         # Ensure the output directory exists
         os.makedirs("./output/", exist_ok=True)
 
@@ -37,11 +40,21 @@ class SaveApiImage:
                 break
             counter += 1
 
-        # Save the image with the determined filename
-        img.save(filename, format="PNG")
+        # Prepare metadata
+        metadata = PngInfo()
+        if prompt is not None:
+            metadata.add_text("prompt", json.dumps(prompt))
+        if extra_pnginfo is not None:
+            for k, v in extra_pnginfo.items():
+                metadata.add_text(k, json.dumps(v))
+
+        # Save the image with the determined filename and metadata
+        img.save(filename, format="PNG", pnginfo=metadata)
 
         # Write the number of the last image to a text file with leading zeroes
         with open("./output/api_next_image.txt", "w") as f:
             f.write(f"api_{counter+1:05d}.png")
 
-        return ()
+        print(f"Image saved as: {filename}")
+
+        return {"ui": {"images": [{"filename": filename, "type": "output"}]}}