126 lines
4.7 KiB
Python
126 lines
4.7 KiB
Python
import os
|
|
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
|
|
import cv2 as cv
|
|
import numpy as np
|
|
import torch
|
|
import requests
|
|
import json
|
|
|
|
def linear_to_srgb(np_array):
|
|
"""Converts a linear RGB numpy array to sRGB."""
|
|
less = np_array <= 0.0031308
|
|
np_array[less] = np_array[less] * 12.92
|
|
np_array[~less] = np.power(np_array[~less], 1/2.4) * 1.055 - 0.055
|
|
return np_array
|
|
|
|
class HttpExrSequenceInput:
|
|
"""
|
|
Node to load a sequence of EXR images from a list of URLs provided as a JSON string.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"urls_json": ("STRING", {"multiline": True, "default": "[]"}),
|
|
"tonemap": (["linear", "sRGB", "Reinhard"], {"default": "sRGB"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"default_image": ("IMAGE",),
|
|
"default_mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("image", "mask",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "🔗ComfyDeploy/EXR"
|
|
|
|
def load_exr_from_data(self, exr_data):
|
|
try:
|
|
nparr = np.frombuffer(exr_data, np.uint8)
|
|
image = cv.imdecode(nparr, cv.IMREAD_UNCHANGED)
|
|
if image is None:
|
|
raise ValueError("Failed to decode EXR data.")
|
|
return image.astype(np.float32)
|
|
except Exception as e:
|
|
print(f"Error decoding EXR data: {e}")
|
|
return None
|
|
|
|
def run(self, urls_json, tonemap, seed, default_image=None, default_mask=None):
|
|
try:
|
|
urls = json.loads(urls_json)
|
|
if not isinstance(urls, list) or not all(isinstance(u, str) for u in urls):
|
|
raise ValueError("urls_json must be a JSON array of URL strings.")
|
|
except (json.JSONDecodeError, ValueError) as e:
|
|
print(f"Error parsing urls_json: {e}. Using default image if available.")
|
|
urls = []
|
|
|
|
if not urls:
|
|
if default_image is not None and default_mask is not None:
|
|
return (default_image, default_mask)
|
|
|
|
print("Warning: No valid URLs and no default image. Returning a black image.")
|
|
blank_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
blank_mask = torch.zeros((1, 64, 64), dtype=torch.float32)
|
|
return (blank_image, blank_mask)
|
|
|
|
rgb_frames = []
|
|
mask_frames = []
|
|
|
|
for url in urls:
|
|
image = None
|
|
try:
|
|
print(f"Fetching EXR from URL: {url}")
|
|
response = requests.get(url)
|
|
response.raise_for_status()
|
|
image = self.load_exr_from_data(response.content)
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error fetching EXR from URL {url}: {e}")
|
|
|
|
if image is None:
|
|
print(f"Warning: Could not decode EXR from {url}. Skipping frame.")
|
|
continue
|
|
|
|
if len(image.shape) == 2: # Grayscale
|
|
image = np.repeat(image[..., np.newaxis], 3, axis=2)
|
|
|
|
rgb = np.flip(image[:, :, :3], 2).copy() # BGR to RGB
|
|
|
|
if tonemap == "sRGB":
|
|
rgb = linear_to_srgb(rgb)
|
|
rgb = np.clip(rgb, 0, 1)
|
|
elif tonemap == "Reinhard":
|
|
rgb = np.clip(rgb, 0, None)
|
|
rgb = rgb / (rgb + 1)
|
|
rgb = linear_to_srgb(rgb)
|
|
rgb = np.clip(rgb, 0, 1)
|
|
|
|
rgb_frames.append(torch.from_numpy(rgb))
|
|
|
|
if image.shape[2] > 3:
|
|
mask = np.clip(image[:, :, 3], 0, 1)
|
|
else:
|
|
mask = np.ones_like(rgb[:, :, 0])
|
|
mask_frames.append(torch.from_numpy(mask))
|
|
|
|
if not rgb_frames:
|
|
print("Could not load any frames. Returning default image if available.")
|
|
if default_image is not None and default_mask is not None:
|
|
return (default_image, default_mask)
|
|
|
|
print("Warning: Failed to load any frames and no default image. Returning a black image.")
|
|
blank_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
blank_mask = torch.zeros((1, 64, 64), dtype=torch.float32)
|
|
return (blank_image, blank_mask)
|
|
|
|
print(f"Loaded {len(rgb_frames)} frames successfully.")
|
|
return (torch.stack(rgb_frames, 0), torch.stack(mask_frames, 0))
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"HttpExrSequenceInput": HttpExrSequenceInput
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"HttpExrSequenceInput": "HTTP EXR Sequence Input (ComfyDeploy)"
|
|
} |