v0.3

save_tmp_image.py

@@ -1,6 +1,8 @@
 import os
 import numpy as np
 from PIL import Image
+import json
+from PIL.PngImagePlugin import PngInfo
 
 class SaveTmpImage:
     @classmethod
@@ -8,7 +10,8 @@ class SaveTmpImage:
         return {
             "required": {
                 "image": ("IMAGE", {"forceInput": True}),
-            }
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
         }
 
     FUNCTION = "save_image"
@@ -16,12 +19,11 @@ class SaveTmpImage:
     OUTPUT_NODE = True
     CATEGORY = "Bjornulf"
 
-    def save_image(self, image):
+    def save_image(self, image, prompt=None, extra_pnginfo=None):
         # Ensure the output directory exists
-        # os.makedirs("./output", exist_ok=True)
+        os.makedirs("./output", exist_ok=True)
 
         # Convert the image from ComfyUI format to PIL Image
-        # Assuming the first two dimensions are extra, and we need to keep the last two
         i = 255. * image.cpu().numpy()
         # Reshape the image if it's not in the expected format, remove any leading dimensions of size 1
         if i.ndim > 3:
@@ -32,7 +34,18 @@ class SaveTmpImage:
 
         img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
 
-        # Save the image, overwriting if it exists
-        img.save("./output/tmp_api.png", 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))
 
-        return ()
+        # Save the image with metadata, overwriting if it exists
+        filename = "./output/tmp_api.png"
+        img.save(filename, format="PNG", pnginfo=metadata)
+
+        print(f"Temporary image saved as: {filename}")
+
+        return {"ui": {"images": [{"filename": filename, "type": "output"}]}}