exr handling nodes
This commit is contained in:
impactframes
+82
-12
@@ -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)"
|
||||
|
||||
@@ -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)"}
|
||||
@@ -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)"}
|
||||
@@ -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)"}
|
||||
+16
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user