diff --git a/comfy-nodes/external_exr.py b/comfy-nodes/external_exr.py index d0ff4e5..f136993 100644 --- a/comfy-nodes/external_exr.py +++ b/comfy-nodes/external_exr.py @@ -1,4 +1,5 @@ import os +os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" import io import cv2 as cv import numpy as np @@ -6,6 +7,16 @@ import torch import requests from folder_paths import get_annotated_filepath +def sRGBtoLinear(npArray): + less = npArray <= 0.0404482362771082 + npArray[less] = npArray[less] / 12.92 + npArray[~less] = np.power((npArray[~less] + 0.055) / 1.055, 2.4) + +def linearToSRGB(npArray): + less = npArray <= 0.0031308 + npArray[less] = npArray[less] * 12.92 + npArray[~less] = np.power(npArray[~less], 1/2.4) * 1.055 - 0.055 + class ComfyUIDeployExternalEXR: RETURN_TYPES = ("IMAGE", "MASK") RETURN_NAMES = ("image", "mask") @@ -41,16 +52,6 @@ class ComfyUIDeployExternalEXR: def VALIDATE_INPUTS(s, exr_file, **kwargs): return True - def sRGBtoLinear(self, npArray): - less = npArray <= 0.0404482362771082 - npArray[less] = npArray[less] / 12.92 - npArray[~less] = np.power((npArray[~less] + 0.055) / 1.055, 2.4) - - def linearToSRGB(self, npArray): - less = npArray <= 0.0031308 - npArray[less] = npArray[less] * 12.92 - npArray[~less] = np.power(npArray[~less], 1/2.4) * 1.055 - 0.055 - def load_exr(self, input_id, exr_file, tonemap="sRGB", default_image=None, default_mask=None, display_name=None, description=None): @@ -76,12 +77,12 @@ class ComfyUIDeployExternalEXR: # Apply tonemapping if tonemap == "sRGB": - self.linearToSRGB(rgb) + linearToSRGB(rgb) rgb = np.clip(rgb, 0, 1) elif tonemap == "Reinhard": rgb = np.clip(rgb, 0, None) rgb = rgb / (rgb + 1) - self.linearToSRGB(rgb) + linearToSRGB(rgb) rgb = np.clip(rgb, 0, 1) rgb = torch.unsqueeze(torch.from_numpy(rgb), 0) @@ -108,3 +109,72 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "ComfyUIDeployExternalEXR": "External EXR (ComfyUI Deploy)" } + +class ComfyUIDeployExternalEXRFrames: + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + FUNCTION = "load_exr_frames" + CATEGORY = "🔗ComfyDeploy" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "exr_file_pattern": ("STRING", {"default": "path/to/frame%04d.exr"}), + "tonemap": (["linear", "sRGB", "Reinhard"], {"default": "sRGB"}), + "start_frame": ("INT", {"default": 1, "min": 0}), + "end_frame": ("INT", {"default": 1, "min": 0}), + }, + "optional": { + "default_image": ("IMAGE",), + "default_mask": ("MASK",), + } + } + + def load_exr_frames(self, exr_file_pattern, tonemap, start_frame, end_frame, + default_image=None, default_mask=None): + if "%04d" not in exr_file_pattern: + raise Exception("Filepath needs to contain a frame pattern like %04d") + + rgb_list = [] + mask_list = [] + + for frame_num in range(start_frame, end_frame + 1): + frame_path = exr_file_pattern.replace("%04d", f"{frame_num:04}") + exr_path = get_annotated_filepath(frame_path) + + if not os.path.exists(exr_path): + print(f"Frame not found, skipping: {exr_path}") + continue + + image = cv.imread(exr_path, cv.IMREAD_UNCHANGED).astype(np.float32) + + if len(image.shape) == 2: + image = np.repeat(image[..., np.newaxis], 3, axis=2) + + rgb = np.flip(image[:,:,:3], 2).copy() + + if tonemap == "sRGB": + linearToSRGB(rgb) + rgb = np.clip(rgb, 0, 1) + elif tonemap == "Reinhard": + rgb = np.clip(rgb, 0, None) + rgb = rgb / (rgb + 1) + linearToSRGB(rgb) + rgb = np.clip(rgb, 0, 1) + + rgb_list.append(torch.from_numpy(rgb)) + + single_mask = torch.zeros((image.shape[0], image.shape[1]), dtype=torch.float32) + if image.shape[2] > 3: + single_mask = torch.from_numpy(np.clip(image[:,:,3], 0, 1)) + mask_list.append(single_mask) + + if not rgb_list: + print("No frames loaded, returning default values.") + return (default_image, default_mask) + + return (torch.stack(rgb_list, 0), torch.stack(mask_list, 0)) + +NODE_CLASS_MAPPINGS["ComfyUIDeployExternalEXRFrames"] = ComfyUIDeployExternalEXRFrames +NODE_DISPLAY_NAME_MAPPINGS["ComfyUIDeployExternalEXRFrames"] = "External EXR Frames (ComfyUI Deploy)" diff --git a/comfy-nodes/input_websocket_exr.py b/comfy-nodes/input_websocket_exr.py new file mode 100644 index 0000000..61dc671 --- /dev/null +++ b/comfy-nodes/input_websocket_exr.py @@ -0,0 +1,106 @@ +import os +os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" +import cv2 as cv +import numpy as np +import torch +from server import PromptServer +from globals import streaming_prompt_metadata, max_output_id_length + +def linearToSRGB(npArray): + less = npArray <= 0.0031308 + npArray[less] = npArray[less] * 12.92 + npArray[~less] = np.power(npArray[~less], 1/2.4) * 1.055 - 0.055 + +class ComfyDeployWebscoketEXRInput: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "input_id": ( + "STRING", + {"multiline": False, "default": "input_exr"}, + ), + "tonemap": (["linear", "sRGB", "Reinhard"], {"default": "sRGB"}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }, + "optional": { + "default_image": ("IMAGE", ), + "default_mask": ("MASK", ), + "client_id": ( + "STRING", + {"multiline": False, "default": ""}, + ), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image","mask",) + + FUNCTION = "run" + CATEGORY = "🔗ComfyDeploy" + + @classmethod + def VALIDATE_INPUTS(s, input_id): + try: + if len(input_id.encode('ascii')) > max_output_id_length: + raise ValueError(f"input_id size is greater than {max_output_id_length} bytes") + except UnicodeEncodeError: + raise ValueError("input_id is not ASCII encodable") + + return True + + def run(self, input_id, tonemap, seed, default_image=None, default_mask=None, client_id=None): + if client_id in streaming_prompt_metadata and input_id in streaming_prompt_metadata[client_id].inputs: + exr_input = streaming_prompt_metadata[client_id].inputs[input_id] + + exr_byte_list = [] + if isinstance(exr_input, list): + print("Received EXR sequence from websocket input") + exr_byte_list = exr_input + elif isinstance(exr_input, bytes): + print("Received single EXR from websocket input") + exr_byte_list = [exr_input] + + if exr_byte_list: + rgb_batch = [] + mask_batch = [] + + for exr_bytes in exr_byte_list: + if not isinstance(exr_bytes, bytes): + print(f"Skipping non-bytes item in input list for {input_id}") + continue + + nparr = np.frombuffer(exr_bytes, np.uint8) + image = cv.imdecode(nparr, cv.IMREAD_UNCHANGED).astype(np.float32) + + if len(image.shape) == 2: + image = np.repeat(image[..., np.newaxis], 3, axis=2) + + rgb = np.flip(image[:,:,:3], 2).copy() + + if tonemap == "sRGB": + linearToSRGB(rgb) + rgb = np.clip(rgb, 0, 1) + elif tonemap == "Reinhard": + rgb = np.clip(rgb, 0, None) + rgb = rgb / (rgb + 1) + linearToSRGB(rgb) + rgb = np.clip(rgb, 0, 1) + + rgb_tensor = torch.from_numpy(rgb).unsqueeze(0) + rgb_batch.append(rgb_tensor) + + mask_tensor = torch.zeros((1, image.shape[0], image.shape[1]), dtype=torch.float32) + if image.shape[2] > 3: + mask_tensor[0] = torch.from_numpy(np.clip(image[:,:,3], 0, 1)) + mask_batch.append(mask_tensor) + + if rgb_batch: + print(f"Loaded {len(rgb_batch)} frames from websocket.") + return (torch.cat(rgb_batch, 0), torch.cat(mask_batch, 0)) + + print("Returning default EXR value") + return (default_image, default_mask) + +NODE_CLASS_MAPPINGS = {"ComfyDeployWebscoketEXRInput": ComfyDeployWebscoketEXRInput} +NODE_DISPLAY_NAME_MAPPINGS = {"ComfyDeployWebscoketEXRInput": "EXR Websocket Input (ComfyDeploy)"} \ No newline at end of file diff --git a/comfy-nodes/output_exr.py b/comfy-nodes/output_exr.py new file mode 100644 index 0000000..2992e53 --- /dev/null +++ b/comfy-nodes/output_exr.py @@ -0,0 +1,159 @@ +import os +os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" +import cv2 as cv +import torch +import numpy as np +import folder_paths + +def sRGBtoLinear(npArray): + less = npArray <= 0.0404482362771082 + npArray[less] = npArray[less] / 12.92 + npArray[~less] = np.power((npArray[~less] + 0.055) / 1.055, 2.4) + +class ComfyDeployOutputEXR: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", {"tooltip": "The images to save."}), + "filename_prefix": ( + "STRING", + { + "default": "ComfyUI", + "tooltip": "The prefix for the file to save.", + }, + ), + "tonemap": (["linear", "sRGB", "Reinhard"], {"default": "sRGB"}), + }, + "optional": { + "output_id": ( + "STRING", + {"multiline": False, "default": "output_exr"}, + ), + }, + } + + RETURN_TYPES = () + FUNCTION = "run" + + OUTPUT_NODE = True + CATEGORY = "🔗ComfyDeploy" + DESCRIPTION = "Saves the input images as EXR files to your ComfyUI output directory." + + def run( + self, + images, + filename_prefix="ComfyUI", + tonemap="sRGB", + output_id="output_exr", + ): + filename_prefix += self.prefix_append + full_output_folder, filename, counter, subfolder, filename_prefix = ( + folder_paths.get_save_image_path( + filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0] + ) + ) + results = list() + + linear = images.cpu().numpy().astype(np.float32) + if tonemap != "linear": + sRGBtoLinear(linear[...,:3]) + if tonemap == "Reinhard": + linear[...,:3] = np.clip(linear[...,:3], 0, 0.999999) + linear[...,:3] = -linear[...,:3] / (linear[...,:3] - 1) + + bgr = linear.copy() + bgr[:,:,:,0] = linear[:,:,:,2] + bgr[:,:,:,2] = linear[:,:,:,0] + if bgr.shape[-1] > 3: + bgr[:,:,:,3] = np.clip(1 - linear[:,:,:,3], 0, 1) + + for i, image in enumerate(bgr): + filename_with_batch_num = filename.replace("%batch_num%", str(i)) + file = f"{filename_with_batch_num}_{counter:05}_.exr" + file_path = os.path.join(full_output_folder, file) + + cv.imwrite(file_path, image) + + results.append( + { + "filename": file, + "subfolder": subfolder, + "type": self.type, + "output_id": output_id, + } + ) + counter += 1 + + return {"ui": {"images": results}} + +class ComfyDeployOutputEXRFrames: + def __init__(self): + self.type = "output" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "filepath_pattern": ("STRING", {"default": "path/to/frame%04d.exr"}), + "tonemap": (["linear", "sRGB", "Reinhard"], {"default": "sRGB"}), + "start_frame": ("INT", {"default": 1, "min": 0}), + }, + "optional": { + "output_id": ("STRING", {"multiline": False, "default": "output_exr_frames"}), + }, + } + + RETURN_TYPES = () + FUNCTION = "run" + OUTPUT_NODE = True + CATEGORY = "🔗ComfyDeploy" + DESCRIPTION = "Saves the input image sequence as EXR files to a specified path." + + def run(self, images, filepath_pattern, tonemap, start_frame, output_id="output_exr_frames"): + if "%04d" not in filepath_pattern: + raise Exception("Filepath pattern must contain '%04d'") + + if not os.path.isabs(filepath_pattern): + raise Exception("Filepath pattern must be an absolute path.") + + os.makedirs(os.path.dirname(filepath_pattern), exist_ok=True) + + linear = images.cpu().numpy().astype(np.float32) + if tonemap != "linear": + sRGBtoLinear(linear[...,:3]) + if tonemap == "Reinhard": + linear[...,:3] = np.clip(linear[...,:3], 0, 0.999999) + linear[...,:3] = -linear[...,:3] / (linear[...,:3] - 1) + + bgr = linear.copy() + bgr[:,:,:,0] = linear[:,:,:,2] + bgr[:,:,:,2] = linear[:,:,:,0] + if bgr.shape[-1] > 3: + bgr[:,:,:,3] = np.clip(1 - linear[:,:,:,3], 0, 1) + + results = list() + + for i, image in enumerate(bgr): + frame_num = start_frame + i + file_path = filepath_pattern.replace("%04d", f"{frame_num:04}") + + cv.imwrite(file_path, image) + + results.append({ + "filename": os.path.basename(file_path), + "subfolder": os.path.dirname(file_path), + "type": self.type, + "output_id": output_id, + }) + + return {"ui": {"images": results}} + +NODE_CLASS_MAPPINGS = {"ComfyDeployOutputEXR": ComfyDeployOutputEXR, "ComfyDeployOutputEXRFrames": ComfyDeployOutputEXRFrames} +NODE_DISPLAY_NAME_MAPPINGS = {"ComfyDeployOutputEXR": "EXR Output (ComfyDeploy)", "ComfyDeployOutputEXRFrames": "EXR Frames Output (ComfyDeploy)"} \ No newline at end of file diff --git a/comfy-nodes/output_websocket_exr.py b/comfy-nodes/output_websocket_exr.py new file mode 100644 index 0000000..1d04f86 --- /dev/null +++ b/comfy-nodes/output_websocket_exr.py @@ -0,0 +1,87 @@ +import os +os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" +import cv2 as cv +import torch +import numpy as np +import folder_paths +from server import PromptServer +import asyncio +from .globals import send_exr, max_output_id_length + +def sRGBtoLinear(npArray): + less = npArray <= 0.0404482362771082 + npArray[less] = npArray[less] / 12.92 + npArray[~less] = np.power((npArray[~less] + 0.055) / 1.055, 2.4) + +class ComfyDeployWebscoketEXROutput: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "output_id": ( + "STRING", + {"multiline": False, "default": "output_id"}, + ), + "images": ("IMAGE", ), + "tonemap": (["linear", "sRGB", "Reinhard"], {"default": "sRGB"}), + }, + "optional": { + "start_frame": ("INT", {"default": 1, "min": 0}), + "client_id": ( + "STRING", + {"multiline": False, "default": ""}, + ), + } + } + + OUTPUT_NODE = True + RETURN_TYPES = () + FUNCTION = "run" + CATEGORY = "🔗ComfyDeploy" + + @classmethod + def VALIDATE_INPUTS(s, output_id): + try: + if len(output_id.encode('ascii')) > max_output_id_length - 5: # 5 for frame number + raise ValueError(f"output_id size is too large for frame sequences") + except UnicodeEncodeError: + raise ValueError("output_id is not ASCII encodable") + + return True + + def run(self, output_id, images, tonemap, start_frame=1, client_id=None): + prompt_server = PromptServer.instance + loop = prompt_server.loop + + def schedule_coroutine_blocking(target, *args): + future = asyncio.run_coroutine_threadsafe(target(*args), loop) + return future.result() + + linear = images.cpu().numpy().astype(np.float32) + if tonemap != "linear": + sRGBtoLinear(linear[...,:3]) + if tonemap == "Reinhard": + linear[...,:3] = np.clip(linear[...,:3], 0, 0.999999) + linear[...,:3] = -linear[...,:3] / (linear[...,:3] - 1) + + bgr = linear.copy() + bgr[:,:,:,0] = linear[:,:,:,2] + bgr[:,:,:,2] = linear[:,:,:,0] + if bgr.shape[-1] > 3: + bgr[:,:,:,3] = np.clip(1 - linear[:,:,:,3], 0, 1) + + for i, image in enumerate(bgr): + success, buffer = cv.imencode(".exr", image) + if not success: + raise Exception("Failed to encode EXR") + + frame_num = start_frame + i + frame_output_id = f"{output_id}_{frame_num:04d}" + + schedule_coroutine_blocking(send_exr, buffer, client_id, frame_output_id) + print(f"EXR sent for frame {frame_num}") + + return {"ui": {}} + +NODE_CLASS_MAPPINGS = {"ComfyDeployWebscoketEXROutput": ComfyDeployWebscoketEXROutput} +NODE_DISPLAY_NAME_MAPPINGS = {"ComfyDeployWebscoketEXROutput": "EXR Websocket Output (ComfyDeploy)"} \ No newline at end of file diff --git a/globals.py b/globals.py index 8588c42..be08a1f 100644 --- a/globals.py +++ b/globals.py @@ -56,11 +56,27 @@ streaming_prompt_metadata: dict[str, StreamingPrompt] = {} class BinaryEventTypes: PREVIEW_IMAGE = 1 UNENCODED_PREVIEW_IMAGE = 2 + EXR_IMAGE = 4 max_output_id_length = 24 +async def send_exr(image_data, sid=None, output_id: str = None): + max_length = max_output_id_length + output_id = output_id[:max_length] + padded_output_id = output_id.ljust(max_length, "\x00") + encoded_output_id = padded_output_id.encode("ascii", "replace") + + bytesIO = BytesIO() + # 10 bytes for the output_id + bytesIO.write(encoded_output_id) + bytesIO.write(image_data) + + preview_bytes = bytesIO.getvalue() + await send_bytes(BinaryEventTypes.EXR_IMAGE, preview_bytes, sid=sid) + + async def send_image(image_data, sid=None, output_id: str = None): max_length = max_output_id_length output_id = output_id[:max_length]