diff --git a/comfy-nodes/external_exr.py b/comfy-nodes/external_exr.py new file mode 100644 index 0000000..183e66a --- /dev/null +++ b/comfy-nodes/external_exr.py @@ -0,0 +1,109 @@ +import os +import io +import cv2 as cv +import numpy as np +import torch +import requests +from folder_paths import get_annotated_filepath + +class ComfyUIDeployExternalEXR: + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + FUNCTION = "load_exr" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_id": ( + "STRING", + {"multiline": False, "default": "input_exr"}, + ), + "exr_file": ("STRING", {"default": ""}), + "tonemap": (["linear", "sRGB", "Reinhard"], {"default": "sRGB"}), + }, + "optional": { + "default_image": ("IMAGE",), + "default_mask": ("MASK",), + "display_name": ( + "STRING", + {"multiline": False, "default": ""}, + ), + "description": ( + "STRING", + {"multiline": False, "default": ""}, + ), + } + } + + @classmethod + 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): + try: + if exr_file and exr_file != "": + if exr_file.startswith(('http://', 'https://')): + # Handle URL input + response = requests.get(exr_file) + # Write to temp buffer + buffer = io.BytesIO(response.content) + nparr = np.frombuffer(buffer.getvalue(), np.uint8) + image = cv.imdecode(nparr, cv.IMREAD_UNCHANGED).astype(np.float32) + else: + # Handle local file + exr_path = get_annotated_filepath(exr_file) + 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) + + # Extract RGB and flip channels + rgb = np.flip(image[:,:,:3], 2).copy() + + # Apply tonemapping + if tonemap == "sRGB": + self.linearToSRGB(rgb) + rgb = np.clip(rgb, 0, 1) + elif tonemap == "Reinhard": + rgb = np.clip(rgb, 0, None) + rgb = rgb / (rgb + 1) + self.linearToSRGB(rgb) + rgb = np.clip(rgb, 0, 1) + + rgb = torch.unsqueeze(torch.from_numpy(rgb), 0) + + # Handle alpha/mask + mask = torch.zeros((1, image.shape[0], image.shape[1]), dtype=torch.float32) + if image.shape[2] > 3: + mask[0] = torch.from_numpy(np.clip(image[:,:,3], 0, 1)) + + return (rgb, mask) + else: + # Return defaults if no file provided + return (default_image, default_mask) + + except Exception as e: + print(f"Error loading EXR: {str(e)}") + # Return defaults on error + return (default_image, default_mask) + +NODE_CLASS_MAPPINGS = { + "ComfyUIDeployExternalEXR": ComfyUIDeployExternalEXR +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ComfyUIDeployExternalEXR": "External EXR (ComfyUI Deploy)" +} diff --git a/custom_routes.py b/custom_routes.py index 25f2715..7f7069a 100644 --- a/custom_routes.py +++ b/custom_routes.py @@ -389,6 +389,9 @@ def apply_inputs_to_workflow(workflow_api: Any, inputs: Any, sid: str = None): if value["class_type"] == "ComfyUIDeployExternalAudio": value["inputs"]["audio_file"] = new_value + if value["class_type"] == "ComfyUIDeployExternalEXR": + value["inputs"]["exr_file"] = new_value + def send_prompt(sid: str, inputs: StreamingPrompt): # workflow_api = inputs.workflow_api