44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
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+1: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, )
|
|
|
|
|