update: save node

Exposure2HDR.py

@@ -3,7 +3,7 @@ import torch
 
 class Exposure2HDR:
     """
-    DiffusionLight Ball2Envmap class
+    DiffusionLight Exposure2HDR class
 
     """
     def __init__(self):

README.md

@@ -4,5 +4,5 @@ DiffusionLight node implmented to ComfyUI
 ## PLAN
 - [X] Ball2Envmap
 - [X] Exposure2HDR
-- [ ] save HDR image
+- [X] save HDR image
 - [ ] view HDR image using somekind of slider or generate list of differnet EVs

SaveHDR.py

@@ -0,0 +1,43 @@
+import numpy as np
+import torch
+import cv2
+import imageio
+import folder_paths
+
+class SaveHDR:
+    """
+    DiffusionLight SaveHDR class
+
+    """
+    def __init__(self):
+        pass 
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "hdr_image": ("IMAGE",),
+                "filename_prefix": ("STRING", {"default": "DiffusionLight"}), 
+                "file_extension": (["hdr","npy", "exr"], {"default": "hdr"}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+
+    FUNCTION = "save_hdr"
+
+    def save_hdr(self, hdr_image, filename_prefix, file_extension):
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory())
+        filename = f"{filename_prefix}_{counter:04d}.{file_extension}"
+        full_path = f"{full_output_folder}/{filename}"
+        print(f"Saving HDR image to {full_path}")
+        out_image = hdr_image[0].cpu().numpy().astype(np.float32)
+        if file_extension == "npy":
+            np.save(full_path, out_image)
+        elif file_extension == "hdr":
+            imageio.imwrite(full_path, out_image)
+        else:
+            cv2.imwrite(full_path, out_image)  # save in HDR format
+        return (hdr_image, )
+    
+

__init__.py

@@ -1,16 +1,19 @@
 from .Ball2Envmap import Ball2Envmap
 from .Exposure2HDR import Exposure2HDR
+from .SaveHDR import SaveHDR
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
     "DiffusionLightBall2Envmap": Ball2Envmap,
-    "DiffusionLightExposure2HDR": Exposure2HDR
+    "DiffusionLightExposure2HDR": Exposure2HDR,
+    "DiffusionLightSaveHDR": SaveHDR
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     "DiffusionLightBall2Envmap": "Ball2Envmap",
-    "DiffusionLightExposure2HDR": "Exposure2HDR"
+    "DiffusionLightExposure2HDR": "Exposure2HDR",
+    "DiffusionLightSaveHDR": "SaveHDR"
 }