├── node.zip ├── img └── cocotools_rhino.exr ├── js ├── index.js ├── saver.js └── load_exr_layer_by_name.js ├── .gitignore ├── requirements.txt ├── pyproject.toml ├── LICENSE ├── __init__.py ├── README.md ├── utils ├── debug_utils.py ├── preview_utils.py ├── batch_utils.py ├── sequence_utils.py └── exr_utils.py ├── modules ├── load_exr.py ├── znormalize.py ├── image_loader.py ├── colorspace.py ├── saver.py ├── load_exr_layer_by_name.py └── load_exr_sequence.py ├── THIRD_PARTY_NOTICES.md └── workflows ├── coco_load_exr_layers.json └── cocotools_Load_EXR_Sequences.json /node.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Conor-Collins/ComfyUI-CoCoTools_IO/HEAD/node.zip -------------------------------------------------------------------------------- /img/cocotools_rhino.exr: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Conor-Collins/ComfyUI-CoCoTools_IO/HEAD/img/cocotools_rhino.exr -------------------------------------------------------------------------------- /js/index.js: -------------------------------------------------------------------------------- 1 | // CoCoTools_IO JavaScript Extensions Entry Point 2 | import "./saver.js"; 3 | import "./load_exr.js"; 4 | import "./load_exr_layer_by_name.js"; 5 | 6 | console.log("CoCoTools_IO extensions loaded successfully"); -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *.pyc 3 | *.DS_Store 4 | *.vscode/ 5 | *.code-workspace 6 | *.log 7 | *.bat 8 | 9 | CLAUDE.md 10 | .claude 11 | 12 | backup/ 13 | ref/ 14 | testing/ 15 | documentation/ 16 | deprecated/ 17 | image files/ 18 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # Core Dependencies 2 | torch>=1.10.0 3 | numpy>=1.21.0 4 | colour-science>=0.4.2 5 | 6 | # Image Processing 7 | Pillow>=9.0.0 8 | opencv-python>=4.5.0 9 | tifffile 10 | OpenImageIO>=2.4.13.0 11 | 12 | # Optional: Logging and Debugging 13 | rich>=10.0.0 14 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "cocotools_io" 3 | description = "Advanced image input and output: EXR, 32 bit support and more" 4 | version = "0.4.2" 5 | license = {file = "LICENSE"} 6 | dependencies = ["torch>=1.10.0", "numpy>=1.21.0", "colour-science>=0.4.2", "Pillow>=9.0.0", "opencv-python>=4.5.0", "tifffile", "OpenImageIO>=2.4.13.0", "rich>=10.0.0"] 7 | 8 | [project.urls] 9 | Repository = "https://github.com/Conor-Collins/ComfyUI-CoCoTools_IO" 10 | # Used by Comfy Registry https://comfyregistry.org 11 | 12 | [tool.comfy] 13 | PublisherId = "coco" 14 | DisplayName = "CoCoTools_IO" 15 | Icon = "" 16 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 Conor Collins 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | 2 | # Import all available modules 3 | from .modules.image_loader import ImageLoader 4 | from .modules.load_exr import LoadExr 5 | from .modules.load_exr_sequence import LoadExrSequence 6 | from .modules.saver import SaverNode 7 | from .modules.load_exr_layer_by_name import LoadExrLayerByName, CryptomatteLayer 8 | from .modules.colorspace import ColorspaceNode 9 | from .modules.znormalize import ZNormalizeNode 10 | 11 | # Initialize node mappings 12 | NODE_CLASS_MAPPINGS = {} 13 | NODE_DISPLAY_NAME_MAPPINGS = {} 14 | 15 | # Explicitly set the web directory path relative to this file 16 | import os 17 | NODE_DIR = os.path.dirname(os.path.realpath(__file__)) 18 | WEB_DIRECTORY = os.path.join(NODE_DIR, "js") 19 | 20 | # Add all available node classes 21 | NODE_CLASS_MAPPINGS.update({ 22 | "ImageLoader": ImageLoader, 23 | "LoadExr": LoadExr, 24 | "LoadExrSequence": LoadExrSequence, 25 | "SaverNode": SaverNode, 26 | "LoadExrLayerByName": LoadExrLayerByName, 27 | "CryptomatteLayer": CryptomatteLayer, 28 | "ColorspaceNode": ColorspaceNode, 29 | "ZNormalizeNode": ZNormalizeNode 30 | }) 31 | 32 | # Add display names for better UI presentation 33 | NODE_DISPLAY_NAME_MAPPINGS.update({ 34 | "ImageLoader": "CoCo Loader", 35 | "LoadExr": "CoCo Load EXR", 36 | "LoadExrSequence": "CoCo Load EXR Sequence", 37 | "SaverNode": "CoCo Saver", 38 | "LoadExrLayerByName": "CoCo Load EXR Layer by Name", 39 | "CryptomatteLayer": "CoCo Cryptomatte Layer", 40 | "ColorspaceNode": "CoCo Colorspace", 41 | "ZNormalizeNode": "CoCo Z Normalize" 42 | }) 43 | 44 | # Expose what ComfyUI needs 45 | __all__ = [ 46 | "NODE_CLASS_MAPPINGS", 47 | "NODE_DISPLAY_NAME_MAPPINGS", 48 | ] 49 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # CoCoTools_IO 2 | 3 | A set of nodes focused on advanced image I/O operations, particularly for EXR file handling. 4 | 5 | ## Features 6 | - Advanced EXR image input with multilayer support 7 | - EXR layer extraction and manipulation 8 | - High-quality image saving with format-specific options 9 | - Standard image format loading with bit depth awareness 10 | 11 | 12 | ## Installation for comfyui portable (tested on 0.3.44) 13 | 14 | 15 | 16 | ### Manual Installation 17 | 1. Clone the repository into your ComfyUI `custom_nodes` directory 18 | 2. Install dependencies 19 | from the python_embeded/ folder 20 | 21 | ```bash 22 | python.exe -m pip install -r ./ComfyUI/custom_nodes/ComfyUI-CoCoTools/requirements.txt 23 | ``` 24 | 3. Restart ComfyUI 25 | 26 | 27 | 28 | ## Current Nodes 29 | 30 | ### Image I/O 31 | - **Image Loader**: Load standard image formats (PNG, JPG, WebP, etc.) with proper bit depth handling 32 | - **Load EXR**: Comprehensive EXR file loading with support for multiple layers, channels, and cryptomatte data 33 | - **Load EXR Sequence**: Load EXR image sequences with #### frame patterns and batch processing 34 | - **Load EXR Layer by Name**: Extract specific layers from EXR files (similar to Nuke's Shuffle node) 35 | - **Cryptomatte Layer**: Specialized handling for cryptomatte layers in EXR files (WIP - not fully implemented) 36 | - **Image Saver**: Save images in various formats with format-specific options (bit depth, compression, etc.) 37 | 38 | ### Image Processing 39 | - **Colorspace Converter**: Convert between various colorspaces (sRGB, Linear, ACEScg, etc.) 40 | - **Z Normalize**: Normalize depth maps and other single-channel data 41 | 42 | 43 | ## To-Do 44 | #### IO 45 | - [x] Implement proper EXR loading 46 | - [x] Implement EXR sequence loader 47 | - [x] Implement EXR saver using OpenImageIO 48 | - [x] Implement multilayer EXR system (render passes, AOVs, embedded images, etc.) 49 | - [x] Add contextual menus based on selected file type in saver 50 | - [x] Add support for EXR sequences 51 | - [ ] Complete cryptomatte layer implementation 52 | - [ ] Adopt filename parsing for saving files into datestamped folders 53 | 54 | 55 | 56 | ## Third-Party Libraries and Licensing 57 | 58 | This project uses the following third-party libraries: 59 | 60 | - **Colour Science for Python**: Used for colorspace transformations in the Colorspace Converter node. Licensed under the New BSD License. 61 | - **OpenColorIO**: Used for color space transformations. Licensed under the BSD 3-Clause License. 62 | 63 | For detailed licensing information, please see the [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md) file. 64 | 65 | This project is licensed under the MIT License. The BSD 3-Clause License used by OpenColorIO and the New BSD License used by colour-science are compatible with the MIT License, allowing us to include and use these components within this MIT-licensed project. 66 | -------------------------------------------------------------------------------- /js/saver.js: -------------------------------------------------------------------------------- 1 | import { app } from "../../../scripts/app.js"; 2 | 3 | function chainCallback(object, property, callback) { 4 | if (object == undefined) { 5 | console.error("Tried to add callback to non-existant object"); 6 | return; 7 | } 8 | if (property in object && object[property]) { 9 | const callback_orig = object[property]; 10 | object[property] = function () { 11 | const r = callback_orig.apply(this, arguments); 12 | return callback.apply(this, arguments) ?? r; 13 | }; 14 | } else { 15 | object[property] = callback; 16 | } 17 | } 18 | 19 | function fitHeight(node) { 20 | requestAnimationFrame(() => { 21 | const size = node.computeSize(); 22 | node.setSize([node.size[0], size[1]]); 23 | app.canvas.setDirty(true); 24 | }); 25 | } 26 | 27 | function addFormatWidgets(nodeType, nodeData) { 28 | chainCallback(nodeType.prototype, "onNodeCreated", function() { 29 | var formatWidget = null; 30 | var formatWidgetIndex = -1; 31 | for(let i = 0; i < this.widgets.length; i++) { 32 | if (this.widgets[i].name === "file_type"){ 33 | formatWidget = this.widgets[i]; 34 | formatWidgetIndex = i+1; 35 | break; 36 | } 37 | } 38 | let formatWidgetsCount = 0; 39 | chainCallback(formatWidget, "callback", (value) => { 40 | const formats = (LiteGraph.registered_node_types[this.type] 41 | ?.nodeData?.input?.required?.file_type?.[1]?.formats); 42 | let newWidgets = []; 43 | if (formats?.[value]) { 44 | let formatWidgets = formats[value]; 45 | for (let wDef of formatWidgets) { 46 | let type = wDef[1]; 47 | if (Array.isArray(type)) { 48 | type = "COMBO"; 49 | } 50 | app.widgets[type](this, wDef[0], wDef.slice(1), app); 51 | let w = this.widgets.pop(); 52 | if (['INT', 'FLOAT'].includes(type)) { 53 | if (wDef.length > 2 && wDef[2]) { 54 | Object.assign(w.options, wDef[2]); 55 | } 56 | } 57 | w.config = wDef.slice(1); 58 | newWidgets.push(w); 59 | } 60 | } 61 | let removed = this.widgets.splice(formatWidgetIndex, 62 | formatWidgetsCount, ...newWidgets); 63 | for (let w of removed) { 64 | w?.onRemove?.(); 65 | } 66 | fitHeight(this); 67 | formatWidgetsCount = newWidgets.length; 68 | }); 69 | }); 70 | } 71 | 72 | app.registerExtension({ 73 | name: "CocoTools.Saver", 74 | async beforeRegisterNodeDef(nodeType, nodeData) { 75 | if (nodeData.name !== "SaverNode") { 76 | return; 77 | } 78 | 79 | addFormatWidgets(nodeType, nodeData); 80 | } 81 | }); -------------------------------------------------------------------------------- /utils/debug_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Debug utilities for COCO Tools with simple/verbose logging control. 3 | """ 4 | import logging 5 | from typing import List, Dict, Any, Optional 6 | 7 | # Global debug mode control - default to simple 8 | DEBUG_MODE = "simple" # "simple" or "verbose" 9 | 10 | # Flag to prevent multiple logging configurations 11 | _LOGGING_CONFIGURED = False 12 | 13 | def setup_logging(level: int = logging.INFO): 14 | """ 15 | Centralized logging configuration for all COCO Tools modules. 16 | Should be called once at module level to avoid conflicts. 17 | """ 18 | global _LOGGING_CONFIGURED 19 | if not _LOGGING_CONFIGURED: 20 | # Check if root logger already has handlers to prevent duplicate setup 21 | if not logging.getLogger().handlers: 22 | logging.basicConfig(level=level) 23 | _LOGGING_CONFIGURED = True 24 | 25 | def set_debug_mode(mode: str): 26 | """Set global debug mode""" 27 | global DEBUG_MODE 28 | DEBUG_MODE = mode.lower() if mode.lower() in ["simple", "verbose"] else "simple" 29 | 30 | def get_debug_mode() -> str: 31 | """Get current debug mode""" 32 | return DEBUG_MODE 33 | 34 | def debug_log(logger: logging.Logger, level: str, simple_msg: str, verbose_msg: Optional[str] = None, **kwargs): 35 | """ 36 | Log with debug verbosity control. 37 | 38 | Args: 39 | logger: Logger instance 40 | level: Log level ("info", "warning", "error", "debug") 41 | simple_msg: Short message for simple mode 42 | verbose_msg: Detailed message for verbose mode (if None, uses simple_msg) 43 | **kwargs: Additional context for verbose mode 44 | """ 45 | message = simple_msg if DEBUG_MODE == "simple" else (verbose_msg or simple_msg) 46 | 47 | # Add context info in verbose mode 48 | if DEBUG_MODE == "verbose" and kwargs: 49 | context_parts = [f"{k}={v}" for k, v in kwargs.items()] 50 | message += f" [{', '.join(context_parts)}]" 51 | 52 | # Log based on level 53 | if level.lower() == "info": 54 | logger.info(message) 55 | elif level.lower() == "warning": 56 | logger.warning(message) 57 | elif level.lower() == "error": 58 | logger.error(message) 59 | elif level.lower() == "debug": 60 | logger.debug(message) 61 | 62 | def format_layer_names(layer_names: List[str], max_simple: int = None) -> str: 63 | """Format layer names for logging - show all layer names as they're important for users""" 64 | # Always show all layer names as users need to know what's available 65 | return ', '.join(layer_names) 66 | 67 | def format_tensor_info(tensor_shape: tuple, tensor_dtype: Any, name: str = "") -> str: 68 | """Format tensor information for logging""" 69 | if DEBUG_MODE == "simple": 70 | return f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}" 71 | return f"{name} shape={tensor_shape}, dtype={tensor_dtype}" if name else f"shape={tensor_shape}, dtype={tensor_dtype}" 72 | 73 | def create_fallback_functions(): 74 | """ 75 | Create fallback functions for cases where debug_utils import fails. 76 | Returns a dictionary of fallback functions following single responsibility principle. 77 | """ 78 | return { 79 | 'debug_log': lambda logger, level, simple_msg, verbose_msg=None, **kwargs: getattr(logger, level.lower())(simple_msg), 80 | 'format_layer_names': lambda layer_names, max_simple=None: ', '.join(layer_names), 81 | 'format_tensor_info': lambda tensor_shape, tensor_dtype, name="": f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}", 82 | 'generate_preview_for_comfyui': lambda image_tensor, source_path="", is_sequence=False, frame_index=0: None 83 | } -------------------------------------------------------------------------------- /modules/load_exr.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | from typing import List 4 | 5 | # Import centralized logging setup 6 | try: 7 | from ..utils.debug_utils import setup_logging 8 | setup_logging() 9 | except ImportError: 10 | logging.basicConfig(level=logging.INFO) 11 | 12 | logger = logging.getLogger(__name__) 13 | 14 | # Import EXR utilities 15 | try: 16 | from ..utils.exr_utils import ExrProcessor 17 | except ImportError: 18 | raise ImportError("EXR utilities are required but not available. Please ensure utils are properly installed.") 19 | 20 | class LoadExr: 21 | @classmethod 22 | def INPUT_TYPES(cls): 23 | return { 24 | "required": { 25 | "image_path": ("STRING", { 26 | "default": "path/to/image.exr", 27 | "description": "Full path to the EXR file" 28 | }), 29 | "normalize": ("BOOLEAN", { 30 | "default": False, 31 | "description": "Normalize image values to the 0-1 range" 32 | }) 33 | }, 34 | "hidden": { 35 | "node_id": "UNIQUE_ID", 36 | "layer_data": "DICT" 37 | } 38 | } 39 | 40 | RETURN_TYPES = ("IMAGE", "MASK", "CRYPTOMATTE", "LAYERS", "STRING", "STRING", "STRING") 41 | RETURN_NAMES = ("image", "alpha", "cryptomatte", "layers", "layer names", "raw layer info", "metadata") 42 | 43 | FUNCTION = "load_image" 44 | CATEGORY = "Image/EXR" 45 | 46 | @classmethod 47 | def IS_CHANGED(cls, image_path, normalize=False, **kwargs): 48 | """ 49 | Smart caching based on file modification time and size. 50 | Only reload if file actually changed or parameters changed. 51 | """ 52 | try: 53 | if not os.path.isfile(image_path): 54 | return float("NaN") # File doesn't exist, always try to load 55 | 56 | stat = os.stat(image_path) 57 | # Create hash from file path, modification time, size, and normalize parameter 58 | return f"{image_path}_{stat.st_mtime}_{stat.st_size}_{normalize}" 59 | except Exception: 60 | # If we can't access file info, always try to load 61 | return float("NaN") 62 | 63 | def load_image(self, image_path: str, normalize: bool = False, 64 | node_id: str = None, layer_data: dict = None, **kwargs) -> List: 65 | """ 66 | Load a single EXR image with support for multiple layers/channel groups. 67 | Returns: 68 | - Base RGB image tensor (image) 69 | - Alpha channel tensor (alpha) 70 | - Dictionary of all cryptomatte layers as tensors (cryptomatte) 71 | - Dictionary of all non-cryptomatte layers as tensors (layers) 72 | - List of processed layer names matching keys in the returned dictionaries (layer names) 73 | - List of raw channel names from the file (raw layer info) 74 | - Metadata as JSON string (metadata) 75 | """ 76 | 77 | # Check for OIIO availability 78 | ExrProcessor.check_oiio_availability() 79 | 80 | try: 81 | # Validate single image path 82 | if not os.path.isfile(image_path): 83 | raise FileNotFoundError(f"Image not found: {image_path}") 84 | 85 | # Use shared EXR processing functionality 86 | return ExrProcessor.process_exr_data(image_path, normalize, node_id, layer_data) 87 | 88 | except Exception as e: 89 | logger.error(f"Error loading EXR file {image_path}: {str(e)}") 90 | raise 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | -------------------------------------------------------------------------------- /THIRD_PARTY_NOTICES.md: -------------------------------------------------------------------------------- 1 | # Third-Party Notices 2 | 3 | This project includes components from the following third-party libraries and resources: 4 | 5 | ## Colour Science for Python 6 | 7 | The colour-science library is used for colorspace transformations in the colorspace module. It is licensed under the New BSD License. 8 | 9 | ``` 10 | Copyright (c) 2013-2023, Colour Developers 11 | All rights reserved. 12 | 13 | Redistribution and use in source and binary forms, with or without 14 | modification, are permitted provided that the following conditions are met: 15 | * Redistributions of source code must retain the above copyright 16 | notice, this list of conditions and the following disclaimer. 17 | * Redistributions in binary form must reproduce the above copyright 18 | notice, this list of conditions and the following disclaimer in the 19 | documentation and/or other materials provided with the distribution. 20 | * Neither the name of the Colour Developers nor the 21 | names of its contributors may be used to endorse or promote products 22 | derived from this software without specific prior written permission. 23 | 24 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 25 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 26 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 27 | DISCLAIMED. IN NO EVENT SHALL COLOUR DEVELOPERS BE LIABLE FOR ANY 28 | DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 29 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 30 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND 31 | ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 32 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 33 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 34 | ``` 35 | 36 | For more information, visit: https://github.com/colour-science/colour 37 | 38 | ## OpenColorIO 39 | 40 | OpenColorIO is licensed under the BSD 3-Clause License. 41 | 42 | ``` 43 | BSD 3-Clause License 44 | 45 | Copyright (c) Contributors to the OpenColorIO Project. 46 | All rights reserved. 47 | 48 | Redistribution and use in source and binary forms, with or without 49 | modification, are permitted provided that the following conditions are met: 50 | 51 | 1. Redistributions of source code must retain the above copyright notice, this 52 | list of conditions and the following disclaimer. 53 | 54 | 2. Redistributions in binary form must reproduce the above copyright notice, 55 | this list of conditions and the following disclaimer in the documentation 56 | and/or other materials provided with the distribution. 57 | 58 | 3. Neither the name of the copyright holder nor the names of its 59 | contributors may be used to endorse or promote products derived from 60 | this software without specific prior written permission. 61 | 62 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 63 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 64 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 65 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 66 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 67 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 68 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 69 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 70 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 71 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 72 | ``` 73 | 74 | For more information, visit: https://github.com/AcademySoftwareFoundation/OpenColorIO 75 | 76 | ## License Compatibility 77 | 78 | This project is licensed under the MIT License. The BSD 3-Clause License used by OpenColorIO is compatible with the MIT License, allowing us to include and use OpenColorIO components within this MIT-licensed project. 79 | 80 | When using this project, please be aware of the license requirements for both the MIT License (covering our code) and the third-party licenses mentioned above. 81 | -------------------------------------------------------------------------------- /modules/znormalize.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import logging 3 | 4 | # Import centralized logging setup 5 | try: 6 | from ..utils.debug_utils import setup_logging 7 | setup_logging() 8 | except ImportError: 9 | logging.basicConfig(level=logging.INFO) 10 | 11 | logger = logging.getLogger(__name__) 12 | 13 | # Import debug utilities and batch processing 14 | try: 15 | from ..utils.debug_utils import debug_log, format_tensor_info 16 | from ..utils.batch_utils import validate_4d_batch, log_batch_processing 17 | except ImportError: 18 | # Fallback if utils not available 19 | def debug_log(logger, level, simple_msg, verbose_msg=None, **kwargs): 20 | getattr(logger, level.lower())(simple_msg) 21 | def format_tensor_info(tensor_shape, tensor_dtype, name=""): 22 | return f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}" 23 | def validate_4d_batch(tensor, name="input"): 24 | return tensor.shape 25 | def log_batch_processing(tensor, operation, name="tensor"): 26 | pass 27 | 28 | class ZNormalizeNode: 29 | @classmethod 30 | def INPUT_TYPES(cls): 31 | return { 32 | "required": { 33 | "image": ("IMAGE",), # Changed to accept IMAGE tensor 34 | "min_depth": ("FLOAT", { 35 | "default": 0.0, 36 | "min": -10000.0, 37 | "max": 10000.0, 38 | "step": 0.01, 39 | "description": "Minimum depth value for normalization" 40 | }), 41 | "max_depth": ("FLOAT", { 42 | "default": 1.0, 43 | "min": -10000.0, 44 | "max": 10000.0, 45 | "step": 0.01, 46 | "description": "Maximum depth value for normalization" 47 | }), 48 | } 49 | } 50 | 51 | RETURN_TYPES = ("IMAGE",) 52 | RETURN_NAMES = ("normalized_depth_image",) 53 | FUNCTION = "normalize_depth" 54 | CATEGORY = "COCO Tools/Processing" 55 | 56 | @classmethod 57 | def IS_CHANGED(cls, **kwargs): 58 | return float("NaN") # Always execute 59 | 60 | def normalize_depth(self, image, min_depth, max_depth): 61 | """ 62 | Normalize depth image tensor with full batch processing support. 63 | 64 | Args: 65 | image: Input tensor in [B,H,W,C] format 66 | min_depth: Minimum depth value for normalization 67 | max_depth: Maximum depth value for normalization 68 | 69 | Returns: 70 | Normalized tensor in [B,H,W,C] format 71 | """ 72 | try: 73 | # Log batch processing info using utility 74 | log_batch_processing(image, f"Normalizing depth range=[{min_depth}, {max_depth}]", "depth") 75 | 76 | # Validate input tensor using utility 77 | batch_size, height, width, channels = validate_4d_batch(image, "depth image") 78 | 79 | # Validate depth range 80 | if max_depth <= min_depth: 81 | raise ValueError(f"max_depth ({max_depth}) must be greater than min_depth ({min_depth})") 82 | 83 | # Create a copy to avoid modifying the input 84 | normalized = image.clone() 85 | 86 | # Log input value range for debugging 87 | input_min, input_max = normalized.min().item(), normalized.max().item() 88 | debug_log(logger, "info", f"Input range: [{input_min:.6f}, {input_max:.6f}]", 89 | f"Input depth values range from {input_min:.6f} to {input_max:.6f}") 90 | 91 | # Normalize depth values - this operation is automatically batch-aware 92 | depth_range = max_depth - min_depth 93 | normalized = (normalized - min_depth) / depth_range 94 | 95 | # Clip values to [0,1] range - also batch-aware 96 | normalized = torch.clamp(normalized, 0.0, 1.0) 97 | 98 | # Log normalized value range 99 | norm_min, norm_max = normalized.min().item(), normalized.max().item() 100 | debug_log(logger, "info", f"Normalized to: [{norm_min:.6f}, {norm_max:.6f}]", 101 | f"After normalization, values range from {norm_min:.6f} to {norm_max:.6f}") 102 | 103 | # Handle single channel depth maps by replicating to RGB 104 | if normalized.shape[-1] == 1: 105 | debug_log(logger, "info", "Converting single channel to RGB", 106 | "Single channel depth detected, replicating to RGB channels") 107 | normalized = normalized.repeat(1, 1, 1, 3) 108 | debug_log(logger, "info", f"RGB depth: {format_tensor_info(normalized.shape, normalized.dtype)}", 109 | f"Converted to RGB: {format_tensor_info(normalized.shape, normalized.dtype)}") 110 | 111 | debug_log(logger, "info", f"Depth normalization complete: {format_tensor_info(normalized.shape, normalized.dtype)}", 112 | f"Successfully normalized {batch_size} depth images with final shape {normalized.shape}") 113 | 114 | return (normalized,) 115 | 116 | except Exception as e: 117 | debug_log(logger, "error", "Depth normalization failed", f"Error normalizing depth image: {str(e)}") 118 | raise 119 | 120 | # Node registration 121 | # NODE_CLASS_MAPPINGS = { 122 | # "znormalize": znormalize 123 | # } 124 | 125 | # NODE_DISPLAY_NAME_MAPPINGS = { 126 | # "znormalize": "Z Normalize" 127 | # } -------------------------------------------------------------------------------- /modules/image_loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | import numpy as np 4 | import torch 5 | from PIL import Image, ImageOps 6 | from typing import Tuple 7 | 8 | # Import centralized logging setup 9 | try: 10 | from ..utils.debug_utils import setup_logging 11 | setup_logging() 12 | except ImportError: 13 | logging.basicConfig(level=logging.INFO) 14 | 15 | logger = logging.getLogger(__name__) 16 | 17 | 18 | class ImageLoader: 19 | @classmethod 20 | def INPUT_TYPES(cls): 21 | return { 22 | "required": { 23 | "image_path": ("STRING", { 24 | "default": "path/to/image.png", 25 | "description": "Full path to the image file" 26 | }), 27 | "normalize": ("BOOLEAN", { 28 | "default": True, 29 | "description": "Normalize image values to the 0-1 range" 30 | }) 31 | }, 32 | "hidden": {"node_id": "UNIQUE_ID"} 33 | } 34 | 35 | RETURN_TYPES = ("IMAGE", "MASK", "STRING") 36 | RETURN_NAMES = ("image", "mask", "metadata") 37 | FUNCTION = "load_regular_image" 38 | CATEGORY = "COCO Tools/Loaders" 39 | 40 | @classmethod 41 | def IS_CHANGED(cls, **kwargs): 42 | return float("NaN") # Always execute 43 | 44 | def load_regular_image( 45 | self, image_path: str, normalize: bool = True, node_id: str = None 46 | ) -> Tuple[torch.Tensor, torch.Tensor, str]: 47 | """ 48 | Main function to load and process a regular image. 49 | Supports formats like PNG, JPG, and WebP. 50 | """ 51 | if not os.path.exists(image_path): 52 | raise FileNotFoundError(f"Image path not found: {image_path}") 53 | 54 | try: 55 | with Image.open(image_path) as img: 56 | # Apply EXIF orientation and convert to RGB 57 | img = ImageOps.exif_transpose(img) 58 | rgb_image = img.convert("RGB") 59 | 60 | # Detect bit depth and convert to tensor 61 | info = self.detect_bit_depth(image_path, rgb_image) 62 | bit_depth = info["bit_depth"] 63 | rgb_tensor = self.pil2tensor(rgb_image, bit_depth) 64 | 65 | # Handle alpha channel if present 66 | alpha_tensor = ( 67 | self.pil2tensor(img.split()[-1], bit_depth).unsqueeze(-1) 68 | if img.mode == "RGBA" else torch.ones_like(rgb_tensor[:, :, :, :1]) 69 | ) 70 | 71 | # Normalize tensors if requested 72 | if normalize: 73 | rgb_tensor = self.normalize_image(rgb_tensor) 74 | alpha_tensor = self.normalize_image(alpha_tensor) 75 | 76 | # Prepare metadata 77 | metadata = { 78 | "file_path": image_path, 79 | "tensor_shape": tuple(rgb_tensor.shape), 80 | "format": os.path.splitext(image_path)[1].lower() 81 | } 82 | 83 | return rgb_tensor, alpha_tensor, str(metadata) 84 | 85 | except Exception as e: 86 | logger.error(f"Error loading image {image_path}: {e}") 87 | raise ValueError(f"Error loading image {image_path}: {e}") 88 | 89 | @staticmethod 90 | def normalize_image(image: torch.Tensor) -> torch.Tensor: 91 | """ 92 | Normalize a tensor to the 0-1 range. 93 | """ 94 | min_val, max_val = image.min(), image.max() 95 | return (image - min_val) / (max_val - min_val) if min_val != max_val else torch.zeros_like(image) 96 | 97 | @staticmethod 98 | def detect_bit_depth(image_path: str, image: Image.Image = None) -> dict: 99 | """ 100 | Detect the bit depth of an image. Supports a range of bit depths for accurate conversion. 101 | """ 102 | mode_to_bit_depth = { 103 | "1": 1, "L": 8, "P": 8, "RGB": 8, "RGBA": 8, 104 | "I;16": 16, "I": 32, "F": 32 105 | } 106 | 107 | if image is None: 108 | with Image.open(image_path) as img: 109 | mode = img.mode 110 | fmt = img.format 111 | else: 112 | mode = image.mode 113 | fmt = image.format 114 | 115 | bit_depth = mode_to_bit_depth.get(mode, 8) 116 | return {"bit_depth": bit_depth, "mode": mode, "format": fmt} 117 | 118 | @staticmethod 119 | def pil2tensor(image: Image.Image, bit_depth: int) -> torch.Tensor: 120 | """ 121 | Convert a PIL Image to a PyTorch tensor, scaled to the 0-1 range. 122 | """ 123 | image_np = np.array(image) 124 | if bit_depth == 8: 125 | image_tensor = torch.from_numpy(image_np.astype(np.float32) / 255.0) 126 | elif bit_depth == 16: 127 | image_tensor = torch.from_numpy(image_np.astype(np.float32) / 65535.0) 128 | elif bit_depth == 32: 129 | image_tensor = torch.from_numpy(image_np.astype(np.float32)) 130 | else: 131 | logger.warning(f"Unsupported bit depth: {bit_depth}. Defaulting to 8-bit normalization.") 132 | image_tensor = torch.from_numpy(image_np.astype(np.float32) / 255.0) 133 | 134 | # Add a batch dimension if not present 135 | if len(image_tensor.shape) == 3: 136 | image_tensor = image_tensor.unsqueeze(0) 137 | 138 | return image_tensor 139 | 140 | 141 | # NODE_CLASS_MAPPINGS = { 142 | # "coco_loader": coco_loader 143 | # } 144 | 145 | # NODE_DISPLAY_NAME_MAPPINGS = { 146 | # "coco_loader": "Load Image (supports jpg, png, tif, avif, webp)" 147 | # } 148 | -------------------------------------------------------------------------------- /workflows/coco_load_exr_layers.json: -------------------------------------------------------------------------------- 1 | {"last_node_id":63,"last_link_id":71,"nodes":[{"id":26,"type":"PreviewImage","pos":[2260,500],"size":[440.67724609375,452.8416442871094],"flags":{},"order":15,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":52}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.15","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":17,"type":"PreviewImage","pos":[2280,-730],"size":[627.10498046875,656.2786254882812],"flags":{},"order":9,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":58}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.15","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":57,"type":"PreviewImage","pos":[2740,-10],"size":[445.65460205078125,459.4781188964844],"flags":{},"order":17,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":64}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.15","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":19,"type":"ColorspaceNode","pos":[1433.7806396484375,-675.6625366210938],"size":[257.3529357910156,60.35293960571289],"flags":{},"order":2,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":37}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[58],"slot_index":0}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"ColorspaceNode"},"widgets_values":["Linear_to_sRGB"]},{"id":59,"type":"LoadExrLayerByName","pos":[1700,1060],"size":[239.62252807617188,102],"flags":{},"order":6,"mode":0,"inputs":[{"name":"layers","type":"LAYERS","link":66}],"outputs":[{"name":"image","type":"IMAGE","links":[65],"slot_index":0},{"name":"mask","type":"MASK","links":null}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"LoadExrLayerByName"},"widgets_values":["raw_direct_lighting","To RGB"]},{"id":39,"type":"LoadExrLayerByName","pos":[1710,1250],"size":[210,102],"flags":{},"order":3,"mode":0,"inputs":[{"name":"layers","type":"LAYERS","link":40}],"outputs":[{"name":"image","type":"IMAGE","links":[53],"slot_index":0},{"name":"mask","type":"MASK","links":null}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"LoadExrLayerByName"},"widgets_values":["normals","To RGB"]},{"id":54,"type":"ColorspaceNode","pos":[1670,-20],"size":[244.57708740234375,58],"flags":{},"order":12,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":60}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[69],"slot_index":0}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"ColorspaceNode"},"widgets_values":["Linear_to_sRGB"]},{"id":60,"type":"ZNormalizeNode","pos":[1970,-20],"size":[236.8000030517578,82],"flags":{},"order":16,"mode":0,"inputs":[{"name":"image","type":"IMAGE","link":69}],"outputs":[{"name":"normalized_depth_image","type":"IMAGE","links":[68],"slot_index":0}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"ZNormalizeNode"},"widgets_values":[0.04,0]},{"id":53,"type":"PreviewImage","pos":[2260,-20],"size":[440.67724609375,471.09197998046875],"flags":{},"order":18,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":68}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.15","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":40,"type":"LoadExrLayerByName","pos":[1710,880],"size":[221.73635864257812,102],"flags":{},"order":4,"mode":0,"inputs":[{"name":"layers","type":"LAYERS","link":41}],"outputs":[{"name":"image","type":"IMAGE","links":[46],"slot_index":0},{"name":"mask","type":"MASK","links":[],"slot_index":1}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"LoadExrLayerByName"},"widgets_values":["diffuse","Auto"]},{"id":58,"type":"ColorspaceNode","pos":[1980,1060],"size":[257.3529357910156,60.35293960571289],"flags":{},"order":13,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":65}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[64],"slot_index":0}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"ColorspaceNode"},"widgets_values":["Linear_to_sRGB"]},{"id":51,"type":"ColorspaceNode","pos":[1990,1250],"size":[257.3529357910156,60.35293960571289],"flags":{},"order":10,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":53}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[56],"slot_index":0}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"ColorspaceNode"},"widgets_values":["Linear_to_sRGB"]},{"id":49,"type":"ColorspaceNode","pos":[1970,880],"size":[245.85467529296875,59.07535171508789],"flags":{},"order":11,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":46}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[52],"slot_index":0}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"ColorspaceNode"},"widgets_values":["Linear_to_sRGB"]},{"id":55,"type":"LoadExrLayerByName","pos":[1390,-20],"size":[225.56910705566406,103.277587890625],"flags":{},"order":5,"mode":0,"inputs":[{"name":"layers","type":"LAYERS","link":61}],"outputs":[{"name":"image","type":"IMAGE","links":[60],"slot_index":0},{"name":"mask","type":"MASK","links":null}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"LoadExrLayerByName"},"widgets_values":["Z","To RGB"]},{"id":25,"type":"PreviewImage","pos":[2760,490],"size":[445.65460205078125,459.4781188964844],"flags":{},"order":14,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":56}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.15","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":62,"type":"WTFDebugNode","pos":[486.78900146484375,-353.044677734375],"size":[210,26],"flags":{},"order":7,"mode":0,"inputs":[{"name":"anything","localized_name":"anything","type":"*","link":71}],"outputs":[],"properties":{"cnr_id":"debugnode-comfyui","ver":"1.1.0","Node name for S&R":"WTFDebugNode"}},{"id":61,"type":"WTFDebugNode","pos":[822.1203002929688,-358.04486083984375],"size":[210,26],"flags":{},"order":8,"mode":0,"inputs":[{"name":"anything","localized_name":"anything","type":"*","link":70}],"outputs":[],"properties":{"cnr_id":"debugnode-comfyui","ver":"1.1.0","Node name for S&R":"WTFDebugNode"}},{"id":46,"type":"LoadExr","pos":[799.0757446289062,-760.0476684570312],"size":[341.5459899902344,202],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"image","type":"IMAGE","links":[37],"slot_index":0},{"name":"alpha","type":"MASK","links":null},{"name":"metadata","type":"STRING","links":[]},{"name":"layers","type":"LAYERS","links":[40,41,61,66],"slot_index":3},{"name":"cryptomatte","type":"CRYPTOMATTE","links":[],"slot_index":4},{"name":"layer names","type":"STRING","links":[71],"slot_index":5},{"name":"processed layer names","type":"STRING","links":[70],"slot_index":6}],"properties":{"aux_id":"Conor-Collins/coco_tools","ver":"e4cb6e618ec48d450f092cbe4ada8507e62d18fe","Node name for S&R":"LoadExr"},"widgets_values":["Full/Path/ToFile.exr",true],"layer_info":{"count":0,"types":{}},"cryptomatte_count":0,"current_exr_path":""},{"id":63,"type":"Note","pos":[1375.4871826171875,-330.90924072265625],"size":[329.2201843261719,223.99082946777344],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["The layer_name needs to be manually input\n\nthe Load_EXR node will output the string names of the channels found and processed, try them out as they are exported from the Load_EXR node.\n"],"color":"#432","bgcolor":"#653"}],"links":[[37,46,0,19,0,"IMAGE"],[40,46,3,39,0,"LAYERS"],[41,46,3,40,0,"LAYERS"],[46,40,0,49,0,"IMAGE"],[52,49,0,26,0,"IMAGE"],[53,39,0,51,0,"IMAGE"],[56,51,0,25,0,"IMAGE"],[58,19,0,17,0,"IMAGE"],[60,55,0,54,0,"IMAGE"],[61,46,3,55,0,"LAYERS"],[64,58,0,57,0,"IMAGE"],[65,59,0,58,0,"IMAGE"],[66,46,3,59,0,"LAYERS"],[68,60,0,53,0,"IMAGE"],[69,54,0,60,0,"IMAGE"],[70,46,6,61,0,"STRING"],[71,46,5,62,0,"STRING"]],"groups":[],"config":{},"extra":{"ds":{"scale":1.532727272727274,"offset":[-507.30755662430647,537.2035033873144]}},"version":0.4} -------------------------------------------------------------------------------- /workflows/cocotools_Load_EXR_Sequences.json: -------------------------------------------------------------------------------- 1 | {"id":"2c3b232d-a474-4d6a-8c95-5b38ac096283","revision":0,"last_node_id":151,"last_link_id":291,"nodes":[{"id":129,"type":"ZNormalizeNode","pos":[1180.0054931640625,-729.2772216796875],"size":[272.4375,82],"flags":{},"order":14,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":243},{"localized_name":"min_depth","name":"min_depth","type":"FLOAT","widget":{"name":"min_depth"},"link":null},{"localized_name":"max_depth","name":"max_depth","type":"FLOAT","widget":{"name":"max_depth"},"link":null}],"outputs":[{"localized_name":"normalized_depth_image","name":"normalized_depth_image","type":"IMAGE","links":[256]}],"properties":{"cnr_id":"cocotools_io","ver":"ff9dbf2f90159b8e08e12f4dc875aa8fbbcc4c00","Node name for S&R":"ZNormalizeNode"},"widgets_values":[0,0.10000000000000002]},{"id":142,"type":"PreviewImage","pos":[794.2015380859375,-1293.5029296875],"size":[323.0854797363281,295.1605529785156],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":268}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.44","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":137,"type":"WTFDebugNode","pos":[430.7371826171875,278.2013244628906],"size":[391.6265869140625,431.61383056640625],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"anything","name":"anything","type":"*","link":291}],"outputs":[],"properties":{"cnr_id":"debugnode-comfyui","ver":"1.1.1","Node name for S&R":"WTFDebugNode","aux_id":"webfiltered/DebugNode-ComfyUI"},"widgets_values":[]},{"id":69,"type":"PreviewImage","pos":[419.5577392578125,-1275.8236083984375],"size":[323.0854797363281,295.1605529785156],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":252}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.44","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":143,"type":"MaskToImage","pos":[828.5006103515625,-1403.9029541015625],"size":[184.62362670898438,26],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"mask","name":"mask","type":"MASK","link":287}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[268]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.44","Node name for S&R":"MaskToImage"},"widgets_values":[]},{"id":141,"type":"WTFDebugNode","pos":[423.2242126464844,-690.652099609375],"size":[288.96868896484375,199.78988647460938],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"anything","name":"anything","type":"*","link":289}],"outputs":[],"properties":{"cnr_id":"debugnode-comfyui","ver":"1.1.1","Node name for S&R":"WTFDebugNode","aux_id":"webfiltered/DebugNode-ComfyUI"},"widgets_values":[]},{"id":67,"type":"PreviewImage","pos":[773.77978515625,-687.7500610351562],"size":[261.5470275878906,253.4022979736328],"flags":{},"order":13,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":244}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.44","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":128,"type":"ColorspaceNode","pos":[1182.2362060546875,-863.476318359375],"size":[270,82],"flags":{},"order":12,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":242},{"localized_name":"from_colorspace","name":"from_colorspace","type":"COMBO","widget":{"name":"from_colorspace"},"link":null},{"localized_name":"to_colorspace","name":"to_colorspace","type":"COMBO","widget":{"name":"to_colorspace"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[243]}],"properties":{"cnr_id":"cocotools_io","ver":"ff9dbf2f90159b8e08e12f4dc875aa8fbbcc4c00","Node name for S&R":"ColorspaceNode"},"widgets_values":["Raw","Rec.709"]},{"id":134,"type":"PreviewImage","pos":[1189.1539306640625,-550.1947021484375],"size":[261.5470275878906,253.4022979736328],"flags":{},"order":15,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":256}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.44","Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":149,"type":"Note","pos":[434.67132568359375,-426.9302978515625],"size":[271.3953552246094,177.44186401367188],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["The layer names are what you should use to load the specific layers in a \"Load EXR Layer by Name\" node"],"color":"#432","bgcolor":"#653"},{"id":127,"type":"LoadExrLayerByName","pos":[414.4118957519531,-853.3021850585938],"size":[302.2232360839844,102],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"layers","name":"layers","type":"LAYERS","link":288},{"localized_name":"layer_name","name":"layer_name","type":"STRING","widget":{"name":"layer_name"},"link":null},{"localized_name":"conversion","name":"conversion","shape":7,"type":"COMBO","widget":{"name":"conversion"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[242,244]},{"localized_name":"mask","name":"mask","type":"MASK","links":null}],"title":"Load EXR Layer by Name","properties":{"cnr_id":"cocotools_io","ver":"ff9dbf2f90159b8e08e12f4dc875aa8fbbcc4c00","Node name for S&R":"LoadExrLayerByName"},"widgets_values":["ao","Auto"]},{"id":148,"type":"Note","pos":[861.25927734375,-125.77428436279297],"size":[330.5789489746094,209.6666717529297],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["EXR's save Layers and channels in a multitude of patterns and can be named any string given to them. \n\nThis raw layer info is here to elucidate what is happening in each file, show each layer and channel present.\n\n\nTo go even further there is a gigantic amount of info stored inside these files as well.\n\nThe metadata will export a json formatted dump of everything in the file/files"],"color":"#432","bgcolor":"#653"},{"id":126,"type":"ColorspaceNode","pos":[444.1439514160156,-1418.427978515625],"size":[270,82],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":286},{"localized_name":"from_colorspace","name":"from_colorspace","type":"COMBO","widget":{"name":"from_colorspace"},"link":null},{"localized_name":"to_colorspace","name":"to_colorspace","type":"COMBO","widget":{"name":"to_colorspace"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[252]}],"properties":{"cnr_id":"cocotools_io","ver":"ff9dbf2f90159b8e08e12f4dc875aa8fbbcc4c00","Node name for S&R":"ColorspaceNode"},"widgets_values":["Raw","sRGB"]},{"id":151,"type":"LoadExrSequence","pos":[-167.1421356201172,-883.9916381835938],"size":[270.3472595214844,274],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"sequence_path","name":"sequence_path","type":"STRING","widget":{"name":"sequence_path"},"link":null},{"localized_name":"start_frame","name":"start_frame","type":"INT","widget":{"name":"start_frame"},"link":null},{"localized_name":"end_frame","name":"end_frame","type":"INT","widget":{"name":"end_frame"},"link":null},{"localized_name":"frame_step","name":"frame_step","type":"INT","widget":{"name":"frame_step"},"link":null},{"localized_name":"normalize","name":"normalize","type":"BOOLEAN","widget":{"name":"normalize"},"link":null}],"outputs":[{"localized_name":"sequence","name":"sequence","type":"IMAGE","links":[286]},{"localized_name":"alpha","name":"alpha","type":"MASK","links":[287]},{"localized_name":"cryptomatte","name":"cryptomatte","type":"CRYPTOMATTE","links":null},{"localized_name":"layers","name":"layers","type":"LAYERS","links":[288]},{"localized_name":"layer names","name":"layer names","type":"STRING","links":[289]},{"localized_name":"raw layer info","name":"raw layer info","type":"STRING","links":[290]},{"localized_name":"metadata","name":"metadata","type":"STRING","links":[291]}],"properties":{"cnr_id":"cocotools_io","ver":"ff9dbf2f90159b8e08e12f4dc875aa8fbbcc4c00","Node name for S&R":"LoadExrSequence"},"widgets_values":["path/to/sequence_####.exr",1,100,1,false]},{"id":140,"type":"WTFDebugNode","pos":[426.3560485839844,-129.3087921142578],"size":[396.1719055175781,331.7623291015625],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"anything","name":"anything","type":"*","link":290}],"outputs":[],"properties":{"cnr_id":"debugnode-comfyui","ver":"1.1.1","Node name for S&R":"WTFDebugNode","aux_id":"webfiltered/DebugNode-ComfyUI"},"widgets_values":[]},{"id":150,"type":"Note","pos":[-172.98321533203125,-1106.8221435546875],"size":[283,155],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["path for sequences use #### as a wildcard format. \n\nThis will find your image sequence as long as they have some sort of sequence frame integers in the filenames\n\nstart frame 1 and ### = 001\nstart frame 10 and #### = 0010\nstart frame 1001 and #### = 1001"],"color":"#432","bgcolor":"#653"}],"links":[[242,127,0,128,0,"IMAGE"],[243,128,0,129,0,"IMAGE"],[244,127,0,67,0,"IMAGE"],[252,126,0,69,0,"IMAGE"],[256,129,0,134,0,"IMAGE"],[268,143,0,142,0,"IMAGE"],[286,151,0,126,0,"IMAGE"],[287,151,1,143,0,"MASK"],[288,151,3,127,0,"LAYERS"],[289,151,4,141,0,"STRING"],[290,151,5,140,0,"STRING"],[291,151,6,137,0,"STRING"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.56,"offset":[2070.2724014029964,1640.0318094917407]}},"version":0.4} -------------------------------------------------------------------------------- /utils/preview_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Preview utilities for COCO Tools nodes 3 | Provides standardized preview generation for ComfyUI nodes 4 | """ 5 | 6 | import os 7 | import uuid 8 | import logging 9 | import tempfile 10 | import numpy as np 11 | import torch 12 | from typing import List, Dict, Optional, Union, Tuple 13 | from PIL import Image 14 | 15 | logger = logging.getLogger(__name__) 16 | 17 | 18 | class PreviewGenerator: 19 | """Standardized preview generator for ComfyUI nodes""" 20 | 21 | def __init__(self, max_preview_size: int = 1024, enable_full_size: bool = False): 22 | self.max_preview_size = max_preview_size 23 | self.enable_full_size = enable_full_size 24 | self.temp_dir = None 25 | self._init_temp_dir() 26 | 27 | def _init_temp_dir(self): 28 | """Initialize ComfyUI temp directory""" 29 | try: 30 | import folder_paths 31 | self.temp_dir = folder_paths.get_temp_directory() 32 | except ImportError: 33 | # Fallback if not in ComfyUI environment 34 | self.temp_dir = tempfile.gettempdir() 35 | logger.warning("folder_paths not available, using system temp directory") 36 | 37 | def generate_preview_for_comfyui(self, image_tensor: torch.Tensor, 38 | source_path: str = "", 39 | is_sequence: bool = False, 40 | frame_index: int = 0, 41 | full_size: bool = False) -> Optional[List[Dict]]: 42 | """ 43 | Generate preview image for ComfyUI 44 | 45 | Args: 46 | image_tensor: Tensor with shape [batch, height, width, channels] 47 | source_path: Original file path (for generating unique names) 48 | is_sequence: Whether this is from a sequence 49 | frame_index: Frame index for sequences 50 | full_size: Whether to generate full resolution preview 51 | 52 | Returns: 53 | List of preview data dicts for ComfyUI UI 54 | """ 55 | try: 56 | if not self.temp_dir: 57 | return None 58 | 59 | # Override instance setting if full_size is explicitly requested 60 | original_full_size = self.enable_full_size 61 | if full_size: 62 | self.enable_full_size = True 63 | 64 | # Handle sequence vs single image 65 | if is_sequence and image_tensor.shape[0] > 1: 66 | # For sequences, use specified frame or first frame 67 | preview_tensor = image_tensor[min(frame_index, image_tensor.shape[0] - 1)] 68 | else: 69 | # For single images, use first (and only) frame 70 | preview_tensor = image_tensor[0] 71 | 72 | # Convert tensor to PIL Image 73 | pil_image = self._tensor_to_pil(preview_tensor) 74 | if pil_image is None: 75 | return None 76 | 77 | # Resize to preview size 78 | pil_image = self._resize_for_preview(pil_image) 79 | 80 | # Generate unique filename 81 | preview_filename = self._generate_preview_filename(source_path, is_sequence, frame_index) 82 | preview_path = os.path.join(self.temp_dir, preview_filename) 83 | 84 | # Save preview image with quality settings 85 | if self.enable_full_size: 86 | # Use highest quality for full resolution previews - no compression limit 87 | pil_image.save(preview_path, format='PNG', optimize=False, compress_level=0) 88 | else: 89 | # Standard quality for thumbnails 90 | pil_image.save(preview_path, format='PNG', optimize=True, compress_level=6) 91 | 92 | # Create preview data structure for ComfyUI 93 | preview_data = [{ 94 | "filename": preview_filename, 95 | "subfolder": "", 96 | "type": "temp" 97 | }] 98 | 99 | sequence_info = " (sequence)" if is_sequence else "" 100 | size_info = f"{pil_image.width}x{pil_image.height}" 101 | full_size_info = " (full resolution)" if self.enable_full_size else " (thumbnail)" 102 | logger.info(f"Generated preview{sequence_info}{full_size_info} {size_info}: {preview_filename}") 103 | 104 | return preview_data 105 | 106 | except Exception as e: 107 | logger.warning(f"Failed to generate preview for {source_path}: {e}") 108 | return None 109 | finally: 110 | # Restore original full_size setting 111 | if full_size: 112 | self.enable_full_size = original_full_size 113 | 114 | def _tensor_to_pil(self, tensor: torch.Tensor) -> Optional[Image.Image]: 115 | """Convert tensor to PIL Image""" 116 | try: 117 | # Convert tensor to numpy array 118 | if tensor.dim() == 3: # [H, W, C] 119 | image_array = tensor.cpu().numpy() 120 | else: 121 | logger.warning(f"Unexpected tensor shape: {tensor.shape}") 122 | return None 123 | 124 | # Ensure float32 and [0,1] range 125 | image_array = np.clip(image_array, 0, 1).astype(np.float32) 126 | 127 | # Convert to [0,255] uint8 128 | image_array = (image_array * 255).clip(0, 255).astype(np.uint8) 129 | 130 | # Handle different channel counts 131 | if image_array.shape[2] == 1: 132 | # Grayscale 133 | image_array = image_array.squeeze(2) 134 | return Image.fromarray(image_array, mode='L') 135 | elif image_array.shape[2] == 3: 136 | # RGB 137 | return Image.fromarray(image_array, mode='RGB') 138 | elif image_array.shape[2] == 4: 139 | # RGBA 140 | return Image.fromarray(image_array, mode='RGBA') 141 | else: 142 | # Use first 3 channels 143 | return Image.fromarray(image_array[:, :, :3], mode='RGB') 144 | 145 | except Exception as e: 146 | logger.warning(f"Failed to convert tensor to PIL: {e}") 147 | return None 148 | 149 | def _resize_for_preview(self, pil_image: Image.Image) -> Image.Image: 150 | """Resize image for preview while maintaining aspect ratio""" 151 | # If full size is enabled, never resize - preserve original resolution 152 | if self.enable_full_size: 153 | return pil_image 154 | 155 | # Standard resize logic for thumbnail mode 156 | if pil_image.width <= self.max_preview_size and pil_image.height <= self.max_preview_size: 157 | return pil_image 158 | 159 | # Calculate new size maintaining aspect ratio 160 | ratio = min(self.max_preview_size / pil_image.width, 161 | self.max_preview_size / pil_image.height) 162 | new_width = int(pil_image.width * ratio) 163 | new_height = int(pil_image.height * ratio) 164 | 165 | return pil_image.resize((new_width, new_height), Image.Resampling.LANCZOS) 166 | 167 | def _generate_preview_filename(self, source_path: str, is_sequence: bool, frame_index: int) -> str: 168 | """Generate unique preview filename""" 169 | # Create hash from source path 170 | file_hash = abs(hash(source_path)) % 1000000 171 | 172 | # Add sequence info 173 | seq_info = f"_seq_f{frame_index:04d}" if is_sequence else "" 174 | 175 | # Generate unique ID 176 | unique_id = uuid.uuid4().hex[:8] 177 | 178 | return f"coco_preview_{file_hash}{seq_info}_{unique_id}.png" 179 | 180 | def generate_saver_preview(self, saved_files: List[Dict], 181 | images: torch.Tensor) -> Optional[List[Dict]]: 182 | """ 183 | Generate preview for saver nodes showing saved files 184 | 185 | Args: 186 | saved_files: List of saved file info 187 | images: Original images tensor 188 | 189 | Returns: 190 | Preview data for ComfyUI UI 191 | """ 192 | try: 193 | if not saved_files or images.shape[0] == 0: 194 | return None 195 | 196 | # Use first image for preview 197 | preview_data = self.generate_preview_for_comfyui( 198 | images, 199 | source_path=f"saver_{len(saved_files)}_files", 200 | is_sequence=len(saved_files) > 1 201 | ) 202 | 203 | return preview_data 204 | 205 | except Exception as e: 206 | logger.warning(f"Failed to generate saver preview: {e}") 207 | return None 208 | 209 | 210 | # Global preview generator instance with full size enabled 211 | preview_generator = PreviewGenerator(max_preview_size=1024, enable_full_size=True) 212 | 213 | 214 | def generate_preview_for_comfyui(image_tensor: torch.Tensor, 215 | source_path: str = "", 216 | is_sequence: bool = False, 217 | frame_index: int = 0, 218 | full_size: bool = False) -> Optional[List[Dict]]: 219 | """Convenience function for generating ComfyUI previews""" 220 | # Create a temporary generator with full size if requested 221 | if full_size: 222 | temp_generator = PreviewGenerator(enable_full_size=True) 223 | return temp_generator.generate_preview_for_comfyui( 224 | image_tensor, source_path, is_sequence, frame_index 225 | ) 226 | else: 227 | return preview_generator.generate_preview_for_comfyui( 228 | image_tensor, source_path, is_sequence, frame_index 229 | ) 230 | 231 | 232 | def generate_saver_preview(saved_files: List[Dict], 233 | images: torch.Tensor) -> Optional[List[Dict]]: 234 | """Convenience function for generating saver previews""" 235 | return preview_generator.generate_saver_preview(saved_files, images) -------------------------------------------------------------------------------- /utils/batch_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Batch processing utilities for COCO Tools nodes 3 | Provides shared functionality for handling tensor batch operations 4 | """ 5 | 6 | import torch 7 | import numpy as np 8 | import logging 9 | from typing import Tuple, Union, Optional 10 | 11 | logger = logging.getLogger(__name__) 12 | 13 | try: 14 | from .debug_utils import debug_log, format_tensor_info 15 | except ImportError: 16 | # Use fallback functions 17 | debug_log = lambda logger, level, simple_msg, verbose_msg=None, **kwargs: getattr(logger, level.lower())(simple_msg) 18 | format_tensor_info = lambda tensor_shape, tensor_dtype, name="": f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}" 19 | 20 | 21 | class BatchProcessor: 22 | """Utility class for batch tensor operations""" 23 | 24 | @staticmethod 25 | def validate_batch_tensor(tensor: torch.Tensor, expected_dims: int = 4, 26 | tensor_name: str = "input") -> Tuple[int, int, int, int]: 27 | """ 28 | Validate and extract batch tensor dimensions 29 | 30 | Args: 31 | tensor: Input tensor to validate 32 | expected_dims: Expected number of dimensions (default 4 for [B,H,W,C]) 33 | tensor_name: Name for error messages 34 | 35 | Returns: 36 | Tuple of (batch_size, height, width, channels) 37 | 38 | Raises: 39 | ValueError: If tensor doesn't match expected format 40 | """ 41 | if len(tensor.shape) != expected_dims: 42 | raise ValueError(f"Expected {expected_dims}D {tensor_name} tensor, got {len(tensor.shape)}D with shape {tensor.shape}") 43 | 44 | if expected_dims == 4: 45 | batch_size, height, width, channels = tensor.shape 46 | debug_log(logger, "debug", f"Validated {tensor_name}: B={batch_size}, H={height}, W={width}, C={channels}", 47 | f"Validated {tensor_name} tensor: batch_size={batch_size}, height={height}, width={width}, channels={channels}") 48 | return batch_size, height, width, channels 49 | else: 50 | raise NotImplementedError(f"Validation for {expected_dims}D tensors not implemented") 51 | 52 | @staticmethod 53 | def reshape_for_processing(tensor: torch.Tensor, preserve_alpha: bool = True) -> Tuple[np.ndarray, Optional[np.ndarray], Tuple]: 54 | """ 55 | Reshape batch tensor for element-wise processing (like colorspace conversion) 56 | 57 | Args: 58 | tensor: Input tensor [B, H, W, C] 59 | preserve_alpha: Whether to separate alpha channel (for RGBA inputs) 60 | 61 | Returns: 62 | Tuple of (reshaped_rgb, alpha_channel_or_none, original_shape) 63 | """ 64 | # Convert to numpy and store original shape 65 | img_np = tensor.cpu().numpy() 66 | original_shape = img_np.shape 67 | batch_size, height, width, channels = original_shape 68 | 69 | debug_log(logger, "debug", f"Reshaping for processing: {format_tensor_info(original_shape, img_np.dtype)}", 70 | f"Reshaping tensor from {original_shape} for batch processing") 71 | 72 | # Reshape to [B*H*W, C] for element-wise processing 73 | reshaped = img_np.reshape(-1, channels) 74 | 75 | alpha_channel = None 76 | if preserve_alpha and channels == 4: 77 | # Separate alpha channel 78 | alpha_channel = reshaped[..., 3:4] 79 | rgb_data = reshaped[..., :3] 80 | debug_log(logger, "debug", "Separated alpha channel", 81 | f"Separated alpha channel: RGB shape={rgb_data.shape}, Alpha shape={alpha_channel.shape}") 82 | return rgb_data, alpha_channel, original_shape 83 | elif channels != 3 and not preserve_alpha: 84 | # Handle non-RGB channels 85 | if channels == 1: 86 | # Replicate single channel to RGB 87 | rgb_data = np.repeat(reshaped, 3, axis=-1) 88 | debug_log(logger, "debug", "Replicated single channel to RGB", 89 | f"Replicated single channel to RGB: {rgb_data.shape}") 90 | elif channels > 3: 91 | # Take first 3 channels 92 | rgb_data = reshaped[..., :3] 93 | debug_log(logger, "debug", f"Truncated to RGB from {channels} channels", 94 | f"Truncated from {channels} channels to RGB: {rgb_data.shape}") 95 | else: 96 | # Pad to 3 channels 97 | padding = np.zeros((reshaped.shape[0], 3 - channels)) 98 | rgb_data = np.concatenate([reshaped, padding], axis=-1) 99 | debug_log(logger, "debug", f"Padded {channels} channels to RGB", 100 | f"Padded from {channels} channels to RGB: {rgb_data.shape}") 101 | return rgb_data, None, original_shape 102 | else: 103 | return reshaped, None, original_shape 104 | 105 | @staticmethod 106 | def reshape_from_processing(processed_rgb: np.ndarray, alpha_channel: Optional[np.ndarray], 107 | original_shape: Tuple, target_device: torch.device) -> torch.Tensor: 108 | """ 109 | Reshape processed data back to original batch format 110 | 111 | Args: 112 | processed_rgb: Processed RGB data [B*H*W, 3] 113 | alpha_channel: Optional alpha channel [B*H*W, 1] 114 | original_shape: Original tensor shape (B, H, W, C) 115 | target_device: Device to place result tensor on 116 | 117 | Returns: 118 | Reshaped tensor [B, H, W, C] 119 | """ 120 | batch_size, height, width, original_channels = original_shape 121 | 122 | # Reshape RGB back to batch dimensions 123 | rgb_reshaped = processed_rgb.reshape(batch_size, height, width, 3) 124 | 125 | # Reattach alpha if present 126 | if alpha_channel is not None: 127 | alpha_reshaped = alpha_channel.reshape(batch_size, height, width, 1) 128 | final_array = np.concatenate([rgb_reshaped, alpha_reshaped], axis=-1) 129 | debug_log(logger, "debug", "Reattached alpha channel", 130 | f"Reattached alpha channel: final shape={final_array.shape}") 131 | else: 132 | final_array = rgb_reshaped 133 | 134 | # Convert back to torch tensor 135 | result_tensor = torch.from_numpy(final_array).to(target_device) 136 | 137 | debug_log(logger, "debug", f"Reshaped back to batch: {format_tensor_info(result_tensor.shape, result_tensor.dtype)}", 138 | f"Reshaped processed data back to batch format: {result_tensor.shape}") 139 | 140 | return result_tensor 141 | 142 | @staticmethod 143 | def normalize_batch_range(tensor: torch.Tensor, target_min: float = 0.0, target_max: float = 1.0, 144 | source_min: Optional[float] = None, source_max: Optional[float] = None) -> torch.Tensor: 145 | """ 146 | Normalize batch tensor values to target range 147 | 148 | Args: 149 | tensor: Input tensor to normalize 150 | target_min: Target minimum value 151 | target_max: Target maximum value 152 | source_min: Source minimum (if None, use tensor.min()) 153 | source_max: Source maximum (if None, use tensor.max()) 154 | 155 | Returns: 156 | Normalized tensor 157 | """ 158 | if source_min is None: 159 | source_min = tensor.min().item() 160 | if source_max is None: 161 | source_max = tensor.max().item() 162 | 163 | debug_log(logger, "debug", f"Normalizing range: [{source_min:.6f}, {source_max:.6f}] -> [{target_min:.6f}, {target_max:.6f}]", 164 | f"Batch normalization: source range [{source_min:.6f}, {source_max:.6f}] to target range [{target_min:.6f}, {target_max:.6f}]") 165 | 166 | # Avoid division by zero 167 | if source_max == source_min: 168 | debug_log(logger, "warning", "Source range is zero, returning target_min", 169 | f"Source min/max are equal ({source_min}), returning constant value {target_min}") 170 | return torch.full_like(tensor, target_min) 171 | 172 | # Normalize to [0,1] then scale to target range 173 | normalized = (tensor - source_min) / (source_max - source_min) 174 | scaled = normalized * (target_max - target_min) + target_min 175 | 176 | return scaled 177 | 178 | @staticmethod 179 | def log_batch_info(tensor: torch.Tensor, operation: str, tensor_name: str = "tensor"): 180 | """ 181 | Log information about a batch tensor 182 | 183 | Args: 184 | tensor: Tensor to log info about 185 | operation: Operation being performed 186 | tensor_name: Name of the tensor for logging 187 | """ 188 | if len(tensor.shape) == 4: 189 | batch_size = tensor.shape[0] 190 | debug_log(logger, "info", f"{operation}: {batch_size} images, {format_tensor_info(tensor.shape, tensor.dtype, tensor_name)}", 191 | f"{operation} - Batch processing {batch_size} images: {tensor_name} {format_tensor_info(tensor.shape, tensor.dtype)} " + 192 | f"range=[{tensor.min().item():.6f}, {tensor.max().item():.6f}]") 193 | else: 194 | debug_log(logger, "info", f"{operation}: {format_tensor_info(tensor.shape, tensor.dtype, tensor_name)}", 195 | f"{operation} - Processing {tensor_name}: {format_tensor_info(tensor.shape, tensor.dtype)} " + 196 | f"range=[{tensor.min().item():.6f}, {tensor.max().item():.6f}]") 197 | 198 | 199 | # Convenience functions for common operations 200 | def validate_4d_batch(tensor: torch.Tensor, name: str = "input") -> Tuple[int, int, int, int]: 201 | """Convenience function for 4D batch validation""" 202 | return BatchProcessor.validate_batch_tensor(tensor, 4, name) 203 | 204 | def log_batch_processing(tensor: torch.Tensor, operation: str, name: str = "tensor"): 205 | """Convenience function for batch logging""" 206 | BatchProcessor.log_batch_info(tensor, operation, name) -------------------------------------------------------------------------------- /utils/sequence_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Sequence utilities for COCO Tools nodes 3 | Provides shared functionality for handling image sequences with #### patterns 4 | """ 5 | 6 | import os 7 | import glob 8 | import re 9 | import logging 10 | from typing import List, Tuple, Dict, Optional 11 | 12 | logger = logging.getLogger(__name__) 13 | 14 | try: 15 | from .debug_utils import debug_log 16 | except ImportError: 17 | # Use fallback function 18 | debug_log = lambda logger, level, simple_msg, verbose_msg=None, **kwargs: getattr(logger, level.lower())(simple_msg) 19 | 20 | 21 | class SequenceHandler: 22 | """Shared sequence handling functionality for loader and saver nodes""" 23 | 24 | @staticmethod 25 | def detect_sequence_pattern(path: str) -> bool: 26 | """Detect if the path contains a sequence pattern (#### or ###)""" 27 | if not path: 28 | return False 29 | # Look for # patterns in the path 30 | return bool(re.search(r'#+', path)) 31 | 32 | @staticmethod 33 | def get_padding_from_template(template: str) -> int: 34 | """Extract padding length from #### pattern""" 35 | match = re.search(r'#+', template) 36 | if match: 37 | return len(match.group(0)) 38 | return 4 # Default padding 39 | 40 | @staticmethod 41 | def replace_frame_number(filename: str, frame_number: int, padding_length: int = None) -> str: 42 | """ 43 | Replace frame number in filename with proper padding 44 | 45 | Args: 46 | filename: Template filename with #### placeholder or existing frame number 47 | frame_number: The frame number to insert 48 | padding_length: Number of digits to pad to (auto-detected if None) 49 | 50 | Returns: 51 | Filename with properly padded frame number 52 | """ 53 | # Auto-detect padding if not specified 54 | if padding_length is None: 55 | padding_length = SequenceHandler.get_padding_from_template(filename) 56 | 57 | # Pattern to match either #### placeholder or existing frame numbers 58 | pattern = r'#+|\d+' 59 | 60 | # Create padded frame number 61 | padded_frame = str(frame_number).zfill(padding_length) 62 | 63 | # Replace with padded frame number 64 | def replacer(match): 65 | return padded_frame 66 | 67 | return re.sub(pattern, replacer, filename) 68 | 69 | @staticmethod 70 | def extract_frame_number_from_path(file_path: str) -> Optional[int]: 71 | """Extract frame number from a file path""" 72 | # Look for sequences of digits that could be frame numbers 73 | basename = os.path.basename(file_path) 74 | matches = re.findall(r'\d+', basename) 75 | 76 | # Take the last sequence of digits (usually the frame number) 77 | if matches: 78 | # Prioritize 3-4 digit numbers (common frame number lengths) 79 | for match in reversed(matches): 80 | if 3 <= len(match) <= 4: 81 | return int(match) 82 | # Fall back to last number found 83 | return int(matches[-1]) 84 | return None 85 | 86 | @staticmethod 87 | def find_sequence_files(pattern_path: str) -> List[str]: 88 | """Find all files matching the sequence pattern""" 89 | # Convert #### pattern to glob pattern 90 | glob_pattern = pattern_path.replace('####', '*') 91 | 92 | # Find all matching files 93 | matching_files = glob.glob(glob_pattern) 94 | 95 | # Create regex pattern - escape special regex characters but keep path separators 96 | # Convert pattern to regex: replace #### with exactly 4 digits 97 | escaped_pattern = re.escape(pattern_path) 98 | 99 | # Replace escaped #### with regex for 4 digits - handle both Windows and Unix escaping 100 | if '\\\\#\\\\#\\\\#\\\\#' in escaped_pattern: 101 | pattern_for_regex = escaped_pattern.replace('\\\\#\\\\#\\\\#\\\\#', r'\d{4}') 102 | elif '\\#\\#\\#\\#' in escaped_pattern: 103 | pattern_for_regex = escaped_pattern.replace('\\#\\#\\#\\#', r'\d{4}') 104 | else: 105 | # Direct replacement if no escaping occurred 106 | pattern_for_regex = escaped_pattern.replace('####', r'\d{4}') 107 | 108 | regex_pattern = re.compile(pattern_for_regex) 109 | 110 | # Debug logging 111 | debug_log(logger, "debug", f"Pattern matching debug", 112 | f"Original: {pattern_path}\\nEscaped: {escaped_pattern}\\nRegex: {pattern_for_regex}\\nMatching files: {len(matching_files)}") 113 | 114 | valid_files = [] 115 | for file_path in matching_files: 116 | if regex_pattern.match(file_path): 117 | valid_files.append(file_path) 118 | else: 119 | # Debug first few failures 120 | if len(valid_files) < 3: 121 | debug_log(logger, "debug", f"No match", f"File: {file_path}\\nPattern: {pattern_for_regex}") 122 | 123 | debug_log(logger, "info", f"Found {len(valid_files)} sequence files", 124 | f"Pattern: {pattern_path}, Found {len(valid_files)} files matching pattern") 125 | 126 | return sorted(valid_files) 127 | 128 | @staticmethod 129 | def extract_frame_numbers(file_paths: List[str]) -> List[Tuple[int, str]]: 130 | """Extract frame numbers from file paths and return sorted list of (frame_num, path) tuples""" 131 | frame_info = [] 132 | for file_path in file_paths: 133 | frame_num = SequenceHandler.extract_frame_number_from_path(file_path) 134 | if frame_num is not None: 135 | frame_info.append((frame_num, file_path)) 136 | 137 | frame_info.sort() # Sort by frame number 138 | return frame_info 139 | 140 | @staticmethod 141 | def generate_frame_paths(pattern_path: str, start_frame: int, end_frame: int, frame_step: int) -> List[str]: 142 | """Generate list of frame paths based on pattern and parameters using improved regex approach""" 143 | frame_paths = [] 144 | 145 | current_frame = start_frame 146 | while current_frame <= end_frame: 147 | frame_path = SequenceHandler.replace_frame_number(pattern_path, current_frame) 148 | frame_paths.append(frame_path) 149 | current_frame += frame_step 150 | 151 | return frame_paths 152 | 153 | @staticmethod 154 | def select_sequence_frames(available_frames: List[Tuple[int, str]], start_frame: int, 155 | end_frame: int, frame_step: int) -> List[str]: 156 | """ 157 | Select specific frames from available sequence based on parameters 158 | 159 | Args: 160 | available_frames: List of (frame_number, file_path) tuples 161 | start_frame: Starting frame number 162 | end_frame: Ending frame number 163 | frame_step: Step between frames 164 | 165 | Returns: 166 | List of selected file paths (strict selection - no fallbacks) 167 | """ 168 | selected_frames = [] 169 | 170 | # Generate expected frame numbers and find exact matches 171 | current_frame = start_frame 172 | while current_frame <= end_frame: 173 | # Find the exact matching frame 174 | found = False 175 | for frame_num, file_path in available_frames: 176 | if frame_num == current_frame: 177 | selected_frames.append(file_path) 178 | found = True 179 | break 180 | 181 | # If frame not found, append None to maintain sequence positions 182 | if not found: 183 | selected_frames.append(None) 184 | 185 | current_frame += frame_step 186 | 187 | expected_count = len(range(start_frame, end_frame + 1, frame_step)) 188 | debug_log(logger, "info", f"Selected {len([f for f in selected_frames if f is not None])} of {expected_count} frames", 189 | f"Selected {len([f for f in selected_frames if f is not None])} frames from {len(available_frames)} available " + 190 | f"(start={start_frame}, end={end_frame}, step={frame_step})") 191 | 192 | return selected_frames 193 | 194 | @staticmethod 195 | def validate_sequence_parameters(start_frame: int, end_frame: int, frame_step: int) -> Tuple[int, int, int]: 196 | """Validate and sanitize sequence parameters""" 197 | # Ensure valid values - no fallback defaults 198 | if start_frame is None: 199 | raise ValueError("start_frame parameter is required and cannot be None. " + 200 | "Make sure the node is in 'sequence' mode and the sequence widgets are visible.") 201 | if end_frame is None: 202 | raise ValueError("end_frame parameter is required and cannot be None. " + 203 | "Make sure the node is in 'sequence' mode and the sequence widgets are visible.") 204 | if frame_step is None: 205 | raise ValueError("frame_step parameter is required and cannot be None. " + 206 | "Make sure the node is in 'sequence' mode and the sequence widgets are visible.") 207 | 208 | start_frame = max(0, start_frame) 209 | end_frame = max(start_frame, end_frame) 210 | frame_step = max(1, frame_step) 211 | 212 | return start_frame, end_frame, frame_step 213 | 214 | @staticmethod 215 | def get_sequence_info(pattern_path: str) -> Dict: 216 | """Get information about an available sequence""" 217 | if not SequenceHandler.detect_sequence_pattern(pattern_path): 218 | return {"is_sequence": False} 219 | 220 | sequence_files = SequenceHandler.find_sequence_files(pattern_path) 221 | if not sequence_files: 222 | return {"is_sequence": True, "available_frames": 0} 223 | 224 | frame_info = SequenceHandler.extract_frame_numbers(sequence_files) 225 | frame_numbers = [fn for fn, fp in frame_info] 226 | 227 | return { 228 | "is_sequence": True, 229 | "available_frames": len(sequence_files), 230 | "first_frame": min(frame_numbers) if frame_numbers else 0, 231 | "last_frame": max(frame_numbers) if frame_numbers else 0, 232 | "frame_numbers": frame_numbers 233 | } 234 | 235 | 236 | class DynamicUIHelper: 237 | """Helper for creating dynamic UI widgets in ComfyUI nodes""" 238 | 239 | @staticmethod 240 | def create_sequence_widgets(default_start: int = 1, default_end: int = 100, default_step: int = 1) -> Dict: 241 | """Create standard sequence control widgets""" 242 | return { 243 | "sequence": [ 244 | ["start_frame", "INT", {"default": default_start, "min": 0, "max": 999999}], 245 | ["end_frame", "INT", {"default": default_end, "min": 1, "max": 999999}], 246 | ["frame_step", "INT", {"default": default_step, "min": 1, "max": 100}] 247 | ] 248 | } 249 | 250 | @staticmethod 251 | def create_versioning_widgets() -> Dict: 252 | """Create versioning control widgets""" 253 | return { 254 | "versioning": [ 255 | ["version", "INT", {"default": 1, "min": -1, "max": 999}] 256 | ] 257 | } 258 | 259 | @staticmethod 260 | def create_save_mode_widgets() -> Dict: 261 | """Create save mode control widgets for saver""" 262 | return { 263 | "sequence": [ 264 | ["start_frame", "INT", {"default": 1, "min": 0, "max": 999999}], 265 | ["frame_step", "INT", {"default": 1, "min": 1, "max": 100}] 266 | ] 267 | } -------------------------------------------------------------------------------- /js/load_exr_layer_by_name.js: -------------------------------------------------------------------------------- 1 | import { app } from "../../../scripts/app.js"; 2 | 3 | app.registerExtension({ 4 | name: "LoadExrLayerByName", 5 | async beforeRegisterNodeDef(nodeType, nodeData) { 6 | // Only handle the Load EXR Layer by Name nodes 7 | if (nodeData.name !== "LoadExrLayerByName" && nodeData.name !== "CryptomatteLayer") { 8 | return; 9 | } 10 | 11 | const isCryptomatte = nodeData.name === "CryptomatteLayer"; 12 | console.log(`Registering ${isCryptomatte ? "Cryptomatte " : ""}Load EXR Layer by Name node`); 13 | 14 | // Store original methods to call them later 15 | const onNodeCreated = nodeType.prototype.onNodeCreated; 16 | const onExecuted = nodeType.prototype.onExecuted; 17 | const onConnectionsChange = nodeType.prototype.onConnectionsChange; 18 | const onWidgetChange = nodeType.prototype.onWidgetChange; 19 | 20 | // Override onNodeCreated to set up node 21 | nodeType.prototype.onNodeCreated = function() { 22 | const result = onNodeCreated?.apply(this, arguments); 23 | 24 | // Initialize storage for layer information 25 | this.availableLayers = []; 26 | this.selectedLayer = ""; 27 | this.connectedNodes = {}; // Track connected nodes 28 | 29 | console.log(`${isCryptomatte ? "Cryptomatte " : ""}Load EXR Layer by Name node created`); 30 | 31 | return result; 32 | }; 33 | 34 | // Handle connections to monitor layer sources 35 | nodeType.prototype.onConnectionsChange = function(type, index, connected, link_info) { 36 | if (onConnectionsChange) { 37 | onConnectionsChange.apply(this, arguments); 38 | } 39 | 40 | // Only care about input connections 41 | if (type !== LiteGraph.INPUT || !link_info) { 42 | return; 43 | } 44 | 45 | // Check if it's connecting or disconnecting 46 | if (connected) { 47 | // Store the connected node info 48 | const inputName = this.inputs[index].name; 49 | const sourceNodeId = link_info.origin_id; 50 | const sourceNode = app.graph.getNodeById(sourceNodeId); 51 | 52 | if (!sourceNode) return; 53 | 54 | console.log(`Connection made to ${inputName} from node ${sourceNode.title || sourceNode.type}`); 55 | 56 | // Store connection in a structured way 57 | this.connectedNodes[inputName] = { 58 | nodeId: sourceNodeId, 59 | node: sourceNode, 60 | outputIndex: link_info.origin_slot 61 | }; 62 | 63 | // If this is the layers/cryptomatte input, try to get layer info immediately 64 | if (index === 0 && inputName.toLowerCase() === (isCryptomatte ? "cryptomatte" : "layers")) { 65 | this.updateLayerOptions(); 66 | } 67 | } else { 68 | // Remove connection info 69 | const inputName = this.inputs[index].name; 70 | delete this.connectedNodes[inputName]; 71 | 72 | // If the main input was disconnected, reset layers 73 | if (index === 0) { 74 | this.resetLayerOptions(); 75 | } 76 | } 77 | }; 78 | 79 | // Handler for when node is executed with fresh data 80 | nodeType.prototype.onExecuted = function(message) { 81 | // Call original method if it exists 82 | if (onExecuted) { 83 | onExecuted.apply(this, arguments); 84 | } 85 | 86 | try { 87 | // Check if the node execution was successful 88 | if (message && message.status === "executed") { 89 | console.log(`${isCryptomatte ? "Cryptomatte " : ""}Load EXR Layer by Name executed successfully`); 90 | 91 | // Update layer options based on connected nodes 92 | // This ensures the node has the most up-to-date layer information 93 | this.updateLayerOptions(); 94 | 95 | // Update node help after execution 96 | this.updateNodeHelp(); 97 | } 98 | 99 | } catch (error) { 100 | console.error("Error in Load EXR Layer by Name onExecuted:", error); 101 | } 102 | }; 103 | 104 | // Update layer options based on connected nodes 105 | nodeType.prototype.updateLayerOptions = function() { 106 | const inputName = isCryptomatte ? "cryptomatte" : "layers"; 107 | const connectionInfo = this.connectedNodes[inputName]; 108 | 109 | if (!connectionInfo || !connectionInfo.node) { 110 | return; 111 | } 112 | 113 | const sourceNode = connectionInfo.node; 114 | console.log(`Finding available layers from ${sourceNode.title || sourceNode.type}`); 115 | 116 | // Check if the source is a LoadExr node 117 | const isLoadExr = sourceNode.type.includes("LoadExr"); 118 | 119 | // Try multiple approaches to get layer names 120 | let layerNames = []; 121 | 122 | // First try: Get output data if available 123 | if (connectionInfo.outputIndex !== undefined) { 124 | const outputData = sourceNode.getOutputData(connectionInfo.outputIndex); 125 | if (outputData && typeof outputData === 'object') { 126 | layerNames = Object.keys(outputData); 127 | if (layerNames.length > 0) { 128 | console.log(`Found ${layerNames.length} layers in direct output data`); 129 | } 130 | } 131 | } 132 | 133 | // Second try: Get data from source node properties 134 | if (layerNames.length === 0 && sourceNode.layerInfo) { 135 | if (isCryptomatte && sourceNode.layerInfo.cryptomatte) { 136 | layerNames = Object.keys(sourceNode.layerInfo.cryptomatte); 137 | } else if (!isCryptomatte && sourceNode.layerInfo.layers) { 138 | layerNames = Object.keys(sourceNode.layerInfo.layers); 139 | } else if (sourceNode.layerInfo.types) { 140 | layerNames = Object.keys(sourceNode.layerInfo.types); 141 | } 142 | 143 | if (layerNames.length > 0) { 144 | console.log(`Found ${layerNames.length} layers in source node layerInfo`); 145 | } 146 | } 147 | 148 | // Third try: Extract from metadata 149 | if (layerNames.length === 0 && sourceNode.widgets) { 150 | const metadataWidget = sourceNode.widgets.find(w => w.name === "metadata"); 151 | if (metadataWidget && metadataWidget.value) { 152 | try { 153 | const metadata = JSON.parse(metadataWidget.value); 154 | if (metadata.layers) { 155 | layerNames = metadata.layers; 156 | } else if (metadata.layer_types) { 157 | layerNames = Object.keys(metadata.layer_types); 158 | } 159 | 160 | if (layerNames.length > 0) { 161 | console.log(`Found ${layerNames.length} layers in metadata`); 162 | } 163 | } catch (error) { 164 | console.error("Failed to parse metadata:", error); 165 | } 166 | } 167 | } 168 | 169 | // Filter and update layers 170 | if (layerNames && layerNames.length > 0) { 171 | const filteredLayers = this.filterLayerNames(layerNames); 172 | console.log(`After filtering: ${filteredLayers.length} layers available`); 173 | 174 | // Store the available layers for tooltip/help 175 | this.availableLayers = [...filteredLayers]; 176 | 177 | // Update the node tooltip and title 178 | this.updateNodeHelp(); 179 | } else { 180 | console.warn("No layers found by any method"); 181 | this.availableLayers = ["none"]; 182 | this.updateNodeHelp(); 183 | } 184 | }; 185 | 186 | // Filter layer names based on the node type 187 | nodeType.prototype.filterLayerNames = function(layerNames) { 188 | if (!layerNames || layerNames.length === 0) { 189 | return ["none"]; 190 | } 191 | 192 | // For regular Load EXR Layer by Name, exclude system and crypto layers 193 | if (!isCryptomatte) { 194 | return layerNames.filter(name => 195 | name !== "rgb" && 196 | name !== "alpha" && 197 | !name.toLowerCase().includes("cryptomatte") && 198 | !name.toLowerCase().startsWith("crypto")); 199 | } 200 | // For Cryptomatte Load EXR Layer by Name, include only crypto layers 201 | else { 202 | return layerNames.filter(name => 203 | name.toLowerCase().includes("cryptomatte") || 204 | name.toLowerCase().startsWith("crypto")); 205 | } 206 | }; 207 | 208 | // Reset to default layers 209 | nodeType.prototype.resetLayerOptions = function() { 210 | this.availableLayers = ["none"]; 211 | this.updateNodeHelp(); 212 | }; 213 | 214 | // Update the node tooltip and title 215 | nodeType.prototype.updateNodeHelp = function() { 216 | const layerList = this.availableLayers.join(", "); 217 | this.title = `${isCryptomatte ? "Cryptomatte " : ""}Load EXR Layer by Name`; 218 | this.help = `Available layers: ${layerList}`; 219 | }; 220 | 221 | // Track when widgets change 222 | nodeType.prototype.onWidgetChange = function(name, value) { 223 | if (onWidgetChange) { 224 | onWidgetChange.apply(this, arguments); 225 | } 226 | }; 227 | 228 | // Add a method to be called by other nodes (e.g., load_exr) 229 | // This allows direct communication between nodes 230 | nodeType.prototype.notifyLayersChanged = function(layerNames) { 231 | if (!layerNames || layerNames.length === 0) return; 232 | 233 | const filteredLayers = this.filterLayerNames(layerNames); 234 | this.availableLayers = [...filteredLayers]; 235 | this.updateNodeHelp(); 236 | }; 237 | }, 238 | 239 | // Setup global handler for all nodes 240 | async setup() { 241 | // Listen for graph execution 242 | app.addEventListener("graphExecuted", (e) => { 243 | try { 244 | // Get all matched nodes 245 | const matchedNodes = findNodes(isCryptomatte ? "CryptomatteLayer" : "LoadExrLayerByName"); 246 | 247 | if (matchedNodes.length === 0) { 248 | return; // No nodes to update 249 | } 250 | 251 | console.log(`Found ${matchedNodes.length} ${isCryptomatte ? "cryptomatte " : ""}layer nodes to update after execution`); 252 | 253 | // For each node, call updateLayerOptions 254 | for (const node of matchedNodes) { 255 | if (node.updateLayerOptions) { 256 | node.updateLayerOptions(); 257 | } 258 | } 259 | } catch (error) { 260 | console.error(`Error in ${isCryptomatte ? "cryptomatte " : ""}layer node graph execution handler:`, error); 261 | } 262 | }); 263 | 264 | function findNodes(type) { 265 | if (!app.graph || !app.graph._nodes) { 266 | return []; 267 | } 268 | 269 | return app.graph._nodes.filter(node => node.type === type); 270 | } 271 | } 272 | }); -------------------------------------------------------------------------------- /modules/colorspace.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import colour 4 | import logging 5 | from typing import Tuple 6 | 7 | # Import batch processing utilities 8 | try: 9 | from ..utils.batch_utils import BatchProcessor, log_batch_processing 10 | except ImportError: 11 | # Fallback if utils not available 12 | class BatchProcessor: 13 | @staticmethod 14 | def validate_batch_tensor(tensor, expected_dims=4, tensor_name="input"): 15 | return tensor.shape 16 | @staticmethod 17 | def reshape_for_processing(tensor, preserve_alpha=True): 18 | img_np = tensor.cpu().numpy() 19 | return img_np.reshape(-1, img_np.shape[-1]), None, img_np.shape 20 | @staticmethod 21 | def reshape_from_processing(processed_rgb, alpha_channel, original_shape, target_device): 22 | return torch.from_numpy(processed_rgb.reshape(original_shape)).to(target_device) 23 | def log_batch_processing(tensor, operation, name="tensor"): 24 | pass 25 | 26 | # Import centralized logging setup 27 | try: 28 | from ..utils.debug_utils import setup_logging 29 | setup_logging() 30 | except ImportError: 31 | logging.basicConfig(level=logging.INFO) 32 | 33 | logger = logging.getLogger(__name__) 34 | 35 | class ColorspaceNode: 36 | """Simplified colorspace converter using colour-science library.""" 37 | 38 | def __init__(self): 39 | self.type = "ColorspaceNode" 40 | 41 | # Map user-friendly names to colour-science names 42 | self.colorspace_mapping = { 43 | # ACES colorspaces 44 | "ACES2065-1": "ACES2065-1", # Linear scene-referred ACES 45 | "ACEScg": "ACEScg", # Linear scene-referred ACEScg 46 | "ACEScct": "ACEScct", # Log-encoded ACEScg with toe 47 | "ACEScc": "ACEScc", # Log-encoded ACES 48 | 49 | # sRGB and Rec.709 50 | "sRGB": "sRGB", # Standard display-referred sRGB (non-linear) 51 | "sRGB Linear": "sRGB", # Linear version of sRGB 52 | "Rec.709": "ITU-R BT.709", # Standard Rec.709 (scene-referred) 53 | "Rec.709 Linear": "ITU-R BT.709", # Linear version of Rec.709 54 | 55 | # Display P3 56 | "Display P3": "Display P3", # Apple's Display P3 (non-linear) 57 | "Display P3 Linear": "Display P3", # Linear version of Display P3 58 | 59 | # Rec.2020 60 | "Rec.2020": "ITU-R BT.2020", # Standard Rec.2020 (non-linear) 61 | "Rec.2020 Linear": "ITU-R BT.2020", # Linear version of Rec.2020 62 | 63 | # Adobe RGB 64 | "Adobe RGB": "Adobe RGB (1998)", # Standard Adobe RGB (non-linear) 65 | "Adobe RGB Linear": "Adobe RGB (1998)", # Linear version of Adobe RGB 66 | 67 | # Raw/passthrough 68 | "Raw": "Raw", # No colorspace conversion 69 | } 70 | 71 | # Track which colorspaces need encoding/decoding 72 | self.encoded_spaces = { 73 | "sRGB", # Standard sRGB is encoded (non-linear) 74 | "ACEScct", # ACEScct is an encoded version of ACEScg 75 | "ACEScc", # ACEScc is an encoded version of ACES2065-1 76 | } 77 | 78 | # Available colorspaces for the UI 79 | self.available_colorspaces = list(self.colorspace_mapping.keys()) 80 | 81 | logger.debug(f"Initialized with {len(self.available_colorspaces)} colorspaces") 82 | 83 | @classmethod 84 | def INPUT_TYPES(cls): 85 | """Define input types.""" 86 | instance = cls() 87 | return { 88 | "required": { 89 | "images": ("IMAGE",), 90 | "from_colorspace": (instance.available_colorspaces,), 91 | "to_colorspace": (instance.available_colorspaces,), 92 | }, 93 | } 94 | 95 | RETURN_TYPES = ("IMAGE",) 96 | FUNCTION = "convert_colorspace" 97 | CATEGORY = "COCO Tools/Processing" 98 | 99 | @classmethod 100 | def IS_CHANGED(cls, **kwargs): 101 | return float("NaN") # Always execute 102 | 103 | def _is_encoded_colorspace(self, colorspace_name: str) -> bool: 104 | """Check if a colorspace name indicates encoded (non-linear) data.""" 105 | # Explicitly encoded spaces 106 | if colorspace_name in self.encoded_spaces: 107 | return True 108 | 109 | # Linear spaces are not encoded 110 | if "Linear" in colorspace_name: 111 | return False 112 | 113 | # Default encoding status for common colorspaces 114 | default_encoded = { 115 | "sRGB": True, 116 | "Rec.709": True, 117 | "Display P3": True, 118 | "Rec.2020": True, 119 | "Adobe RGB": True, 120 | "ACES2065-1": False, # Linear by default 121 | "ACEScg": False, # Linear by default 122 | "Raw": False # Raw is always linear 123 | } 124 | 125 | # Check if the colorspace is in the default encoding status dictionary 126 | if colorspace_name in default_encoded: 127 | return default_encoded[colorspace_name] 128 | 129 | # For unknown colorspaces, assume they are linear 130 | return False 131 | 132 | def _apply_gamma_encoding(self, rgb: np.ndarray, colorspace: str) -> np.ndarray: 133 | """Apply appropriate gamma encoding based on colorspace.""" 134 | # Handle standard colorspaces 135 | if colorspace == "sRGB" or "sRGB" in colorspace and "Linear" not in colorspace: 136 | # Apply sRGB EOTF (gamma curve) 137 | return colour.models.eotf_sRGB(rgb) 138 | elif colorspace == "Rec.709" or "Rec.709" in colorspace and "Linear" not in colorspace: 139 | # Rec.709 uses the same EOTF as sRGB 140 | return colour.models.eotf_sRGB(rgb) 141 | elif colorspace == "Display P3" or "Display P3" in colorspace and "Linear" not in colorspace: 142 | # Display P3 uses the same EOTF as sRGB 143 | return colour.models.eotf_sRGB(rgb) 144 | elif colorspace == "Rec.2020" or "Rec.2020" in colorspace and "Linear" not in colorspace: 145 | # Rec.2020 uses a slightly different EOTF, but we'll use sRGB for simplicity 146 | return colour.models.eotf_sRGB(rgb) 147 | elif colorspace == "Adobe RGB" or "Adobe RGB" in colorspace and "Linear" not in colorspace: 148 | # Adobe RGB uses a gamma of 2.2 149 | return np.power(np.maximum(rgb, 0), 1/2.2) 150 | 151 | # Handle ACES colorspaces 152 | elif colorspace == "ACEScc": 153 | # Apply ACEScc encoding (log encoding for ACES) 154 | return colour.models.log_encoding_ACEScc(rgb) 155 | elif colorspace == "ACEScct": 156 | # Apply ACEScct encoding (log encoding with toe for ACEScg) 157 | return colour.models.log_encoding_ACEScct(rgb) 158 | 159 | # Handle other gamma values 160 | elif "Gamma 2.2" in colorspace: 161 | return np.power(np.maximum(rgb, 0), 1/2.2) 162 | elif "Gamma 2.4" in colorspace: 163 | return np.power(np.maximum(rgb, 0), 1/2.4) 164 | 165 | # Linear colorspaces don't need encoding 166 | else: 167 | return rgb 168 | 169 | def _apply_gamma_decoding(self, rgb: np.ndarray, colorspace: str) -> np.ndarray: 170 | """Apply appropriate gamma decoding based on colorspace.""" 171 | # Handle standard colorspaces 172 | if colorspace == "sRGB" or "sRGB" in colorspace and "Linear" not in colorspace: 173 | # Apply inverse sRGB EOTF 174 | return colour.models.eotf_inverse_sRGB(rgb) 175 | elif colorspace == "Rec.709" or "Rec.709" in colorspace and "Linear" not in colorspace: 176 | # Rec.709 uses the same EOTF as sRGB 177 | return colour.models.eotf_inverse_sRGB(rgb) 178 | elif colorspace == "Display P3" or "Display P3" in colorspace and "Linear" not in colorspace: 179 | # Display P3 uses the same EOTF as sRGB 180 | return colour.models.eotf_inverse_sRGB(rgb) 181 | elif colorspace == "Rec.2020" or "Rec.2020" in colorspace and "Linear" not in colorspace: 182 | # Rec.2020 uses a slightly different EOTF, but we'll use sRGB for simplicity 183 | return colour.models.eotf_inverse_sRGB(rgb) 184 | elif colorspace == "Adobe RGB" or "Adobe RGB" in colorspace and "Linear" not in colorspace: 185 | # Adobe RGB uses a gamma of 2.2 186 | return np.power(np.maximum(rgb, 0), 2.2) 187 | 188 | # Handle ACES colorspaces 189 | elif colorspace == "ACEScc": 190 | # Apply ACEScc decoding (inverse log encoding for ACES) 191 | return colour.models.log_decoding_ACEScc(rgb) 192 | elif colorspace == "ACEScct": 193 | # Apply ACEScct decoding (inverse log encoding with toe for ACEScg) 194 | return colour.models.log_decoding_ACEScct(rgb) 195 | 196 | # Handle other gamma values 197 | elif "Gamma 2.2" in colorspace: 198 | return np.power(np.maximum(rgb, 0), 2.2) 199 | elif "Gamma 2.4" in colorspace: 200 | return np.power(np.maximum(rgb, 0), 2.4) 201 | 202 | # Linear colorspaces don't need decoding 203 | else: 204 | return rgb 205 | 206 | def convert_colorspace(self, images: torch.Tensor, from_colorspace: str, to_colorspace: str) -> Tuple[torch.Tensor]: 207 | """ 208 | Convert images between colorspaces using colour-science library. 209 | 210 | Args: 211 | images: Input images as torch tensor [B, H, W, C] 212 | from_colorspace: Source colorspace name 213 | to_colorspace: Target colorspace name 214 | 215 | Returns: 216 | Tuple containing the converted images as torch tensor 217 | """ 218 | # Log batch processing info 219 | log_batch_processing(images, f"Converting from '{from_colorspace}' to '{to_colorspace}'", "input") 220 | 221 | # If source and target are the same, return original 222 | if from_colorspace == to_colorspace: 223 | logger.info("Source and target colorspaces are the same") 224 | return (images,) 225 | 226 | # Convert to numpy 227 | img_np = images.cpu().numpy() 228 | logger.info(f"Input range: min={img_np.min():.6f}, max={img_np.max():.6f}") 229 | 230 | # Handle problematic values 231 | if np.isnan(img_np).any() or np.isinf(img_np).any(): 232 | logger.warning("Input contains NaN/Inf values, cleaning...") 233 | img_np = np.nan_to_num(img_np, nan=0.0, posinf=1.0, neginf=0.0) 234 | images = torch.from_numpy(img_np).to(images.device) 235 | 236 | try: 237 | # Handle special cases first 238 | if from_colorspace == "Raw" or to_colorspace == "Raw": 239 | logger.info("Raw colorspace detected, returning input unchanged") 240 | return (images,) 241 | 242 | # Get the colour-science colorspace names 243 | from_cs = self.colorspace_mapping.get(from_colorspace) 244 | to_cs = self.colorspace_mapping.get(to_colorspace) 245 | 246 | if not from_cs or not to_cs: 247 | logger.error(f"Unsupported colorspaces: {from_colorspace} -> {to_colorspace}") 248 | return (images,) 249 | 250 | # Handle encoding/decoding 251 | working_img = img_np.copy() 252 | 253 | # Step 1: Decode input if it's encoded 254 | if self._is_encoded_colorspace(from_colorspace): 255 | logger.info(f"Decoding {from_colorspace}") 256 | working_img = self._apply_gamma_decoding(working_img, from_colorspace) 257 | 258 | # Step 2: Convert between colorspaces (linear to linear) 259 | if from_cs != to_cs and from_cs != "Raw" and to_cs != "Raw": 260 | logger.info(f"Converting colorspace: {from_cs} -> {to_cs}") 261 | 262 | # Handle special case where both map to same underlying space 263 | if from_cs == to_cs: 264 | logger.info("Same underlying colorspace, skipping conversion") 265 | else: 266 | # Use batch processing utilities for reshaping 267 | working_tensor = torch.from_numpy(working_img) 268 | rgb_data, alpha_channel, original_shape = BatchProcessor.reshape_for_processing( 269 | working_tensor, preserve_alpha=True 270 | ) 271 | 272 | # Apply the colorspace conversion 273 | try: 274 | converted_rgb = colour.RGB_to_RGB( 275 | rgb_data, 276 | input_colourspace=from_cs, 277 | output_colourspace=to_cs, 278 | apply_cctf_decoding=False, # We handle encoding separately 279 | apply_cctf_encoding=False 280 | ) 281 | logger.info("Colorspace conversion successful") 282 | except Exception as e: 283 | logger.error(f"Colour-science conversion failed: {e}") 284 | # Try with chromatic adaptation 285 | try: 286 | converted_rgb = colour.RGB_to_RGB( 287 | rgb_data, 288 | input_colourspace=from_cs, 289 | output_colourspace=to_cs, 290 | apply_cctf_decoding=False, 291 | apply_cctf_encoding=False, 292 | chromatic_adaptation_transform='CAT02' 293 | ) 294 | logger.info("Colorspace conversion with CAT02 successful") 295 | except Exception as e2: 296 | logger.error(f"All conversion attempts failed: {e2}") 297 | # Return original image 298 | return (images,) 299 | 300 | # Use batch processing utilities to reshape back 301 | working_tensor = BatchProcessor.reshape_from_processing( 302 | converted_rgb, alpha_channel, original_shape, images.device 303 | ) 304 | working_img = working_tensor.cpu().numpy() 305 | 306 | # Step 3: Encode output if needed 307 | if self._is_encoded_colorspace(to_colorspace): 308 | logger.info(f"Encoding to {to_colorspace}") 309 | working_img = self._apply_gamma_encoding(working_img, to_colorspace) 310 | 311 | # Handle clipping based on colorspace 312 | # For HDR colorspaces like ACES, we don't want to clip to 0-1 313 | is_hdr_colorspace = any(hdr_space in to_colorspace for hdr_space in ["ACES", "Raw", "Linear"]) 314 | 315 | if not is_hdr_colorspace: 316 | # For display-referred spaces, clip to 0-1 317 | working_img = np.clip(working_img, 0.0, 1.0) 318 | else: 319 | # For HDR/scene-referred spaces, just ensure no negative values 320 | working_img = np.maximum(working_img, 0.0) 321 | 322 | # Log if we have values > 1.0 (common in HDR) 323 | if np.any(working_img > 1.0): 324 | max_val = np.max(working_img) 325 | logger.info(f"HDR values detected: max={max_val:.6f}") 326 | 327 | # Convert back to torch tensor 328 | result_tensor = torch.from_numpy(working_img).to(images.device) 329 | 330 | logger.info(f"Output range: min={result_tensor.min().item():.6f}, max={result_tensor.max().item():.6f}") 331 | 332 | return (result_tensor,) 333 | 334 | except Exception as e: 335 | logger.error(f"Conversion failed: {e}") 336 | import traceback 337 | logger.error(traceback.format_exc()) 338 | return (images,) 339 | 340 | # Test function to verify colour-science setup 341 | def test_colour_science_setup(): 342 | """Test that colour-science is working correctly.""" 343 | try: 344 | # Test basic functionality 345 | test_rgb = np.array([[[0.18, 0.18, 0.18]]]) 346 | 347 | # Test sRGB encoding 348 | encoded = colour.models.eotf_sRGB(test_rgb) 349 | print(f"sRGB encoding test: {test_rgb.flatten()} -> {encoded.flatten()}") 350 | 351 | # Test colorspace conversion 352 | converted = colour.RGB_to_RGB( 353 | test_rgb, 354 | input_colourspace='ITU-R BT.709', 355 | output_colourspace='ACEScg', 356 | apply_cctf_decoding=False, 357 | apply_cctf_encoding=False 358 | ) 359 | print(f"Rec.709 to ACEScg: {test_rgb.flatten()} -> {converted.flatten()}") 360 | 361 | print("colour-science setup test passed!") 362 | return True 363 | 364 | except Exception as e: 365 | print(f"colour-science setup test failed: {e}") 366 | return False 367 | 368 | # Register the node 369 | # NODE_CLASS_MAPPINGS = { 370 | # "ColorspaceNode": colorspace, 371 | # } 372 | 373 | # NODE_DISPLAY_NAME_MAPPINGS = { 374 | # "ColorspaceNode": "Colorspace", 375 | # } 376 | 377 | if __name__ == "__main__": 378 | test_colour_science_setup() 379 | -------------------------------------------------------------------------------- /modules/saver.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import numpy as np 4 | import tifffile 5 | import folder_paths 6 | import logging 7 | from typing import Dict, Tuple, Optional, List 8 | import OpenImageIO as oiio 9 | from datetime import datetime 10 | 11 | os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" 12 | import cv2 as cv 13 | 14 | logger = logging.getLogger(__name__) 15 | 16 | # Import debug utilities, preview utilities, and sequence utilities 17 | try: 18 | from ..utils.debug_utils import debug_log, format_tensor_info 19 | from ..utils.preview_utils import generate_preview_for_comfyui 20 | from ..utils.sequence_utils import SequenceHandler, DynamicUIHelper 21 | except ImportError: 22 | # Fallback if utils not available 23 | def debug_log(logger, level, simple_msg, verbose_msg=None, **kwargs): 24 | getattr(logger, level.lower())(simple_msg) 25 | def format_tensor_info(tensor_shape, tensor_dtype, name=""): 26 | return f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}" 27 | def generate_preview_for_comfyui(image_tensor, source_path="", is_sequence=False, frame_index=0, full_size=False): 28 | return None 29 | 30 | # Fallback sequence handler 31 | class SequenceHandler: 32 | @staticmethod 33 | def detect_sequence_pattern(path): return '####' in path if path else False 34 | @staticmethod 35 | def generate_frame_paths(pattern, start, count, step): return [] 36 | 37 | class DynamicUIHelper: 38 | @staticmethod 39 | def create_save_mode_widgets(): return {"sequence": [["start_frame", "INT", {"default": 1}], ["frame_step", "INT", {"default": 1}]]} 40 | @staticmethod 41 | def create_versioning_widgets(): return {"versioning": [["version", "INT", {"default": 1}]]} 42 | 43 | class SaverNode: 44 | """Optimized image saver node with consistent bit depth handling""" 45 | 46 | # Format specifications 47 | FORMAT_SPECS = { 48 | "exr": {"depths": [16, 32], "opencv": False}, # EXR only supports half and full float 49 | "png": {"depths": [8, 16], "opencv": False}, # PNG only supports integer formats 50 | "jpg": {"depths": [8], "opencv": True}, 51 | "webp": {"depths": [8], "opencv": True}, 52 | "tiff": {"depths": [8, 16, 32], "opencv": False} 53 | } 54 | 55 | @classmethod 56 | def INPUT_TYPES(cls): 57 | # Define format-specific widgets 58 | format_widgets = { 59 | "exr": [ 60 | ["bit_depth", ["16", "32"], {"default": "32"}], 61 | ["exr_compression", ["none", "zip", "zips", "rle", "pxr24", "b44", "b44a", "dwaa", "dwab"], {"default": "zips"}], 62 | ["save_as_grayscale", "BOOLEAN", {"default": False}] 63 | ], 64 | "png": [ 65 | ["bit_depth", ["8", "16"], {"default": "16"}], 66 | ["save_as_grayscale", "BOOLEAN", {"default": False}] 67 | ], 68 | "tiff": [ 69 | ["bit_depth", ["8", "16", "32"], {"default": "16"}], 70 | ["save_as_grayscale", "BOOLEAN", {"default": False}] 71 | ], 72 | "jpg": [ 73 | ["quality", "INT", {"default": 95, "min": 1, "max": 100}] 74 | ], 75 | "webp": [ 76 | ["quality", "INT", {"default": 95, "min": 1, "max": 100}] 77 | ] 78 | } 79 | 80 | # Create sequence and versioning widgets using shared utilities 81 | save_mode_widgets = DynamicUIHelper.create_save_mode_widgets() 82 | versioning_widgets = DynamicUIHelper.create_versioning_widgets() 83 | 84 | return { 85 | "required": { 86 | "images": ("IMAGE",), 87 | "file_path": ("STRING", {"default": ""}), 88 | "filename": ("STRING", {"default": "ComfyUI"}), 89 | "save_mode": (["single", "sequence"], { 90 | "default": "single", 91 | "description": "Save mode: single files or sequence pattern", 92 | "formats": save_mode_widgets 93 | }), 94 | "use_versioning": ("BOOLEAN", { 95 | "default": False, 96 | "description": "Enable version numbering", 97 | "formats": versioning_widgets 98 | }), 99 | "file_type": (["exr", "png", "jpg", "webp", "tiff"], {"default": "png", "formats": format_widgets}), 100 | }, 101 | "hidden": { 102 | "prompt": "PROMPT", 103 | "extra_pnginfo": "EXTRA_PNGINFO", 104 | }, 105 | } 106 | 107 | RETURN_TYPES = () 108 | FUNCTION = "save_images" 109 | OUTPUT_NODE = True 110 | CATEGORY = "COCO Tools/Savers" 111 | 112 | @classmethod 113 | def IS_CHANGED(cls, **kwargs): 114 | return float("NaN") # Always execute 115 | 116 | def __init__(self): 117 | self.output_dir = folder_paths.get_output_directory() 118 | 119 | @staticmethod 120 | def is_grayscale_fast(image: np.ndarray, sample_rate: float = 0.1) -> bool: 121 | """Fast grayscale detection by sampling pixels""" 122 | if len(image.shape) == 2 or image.shape[-1] == 1: 123 | return True 124 | if image.shape[-1] == 3: 125 | # Sample 10% of pixels for large images 126 | total_pixels = image.shape[0] * image.shape[1] 127 | if total_pixels > 1000000: # 1MP 128 | indices = np.random.choice(total_pixels, int(total_pixels * sample_rate), replace=False) 129 | flat_img = image.reshape(-1, 3) 130 | sampled = flat_img[indices] 131 | return np.allclose(sampled[:, 0], sampled[:, 1], rtol=0.01) and \ 132 | np.allclose(sampled[:, 1], sampled[:, 2], rtol=0.01) 133 | else: 134 | return np.allclose(image[..., 0], image[..., 1], rtol=0.01) and \ 135 | np.allclose(image[..., 1], image[..., 2], rtol=0.01) 136 | return False 137 | 138 | def validate_bit_depth(self, file_type: str, bit_depth: int) -> int: 139 | """Validate and adjust bit depth for format""" 140 | valid_depths = self.FORMAT_SPECS[file_type]["depths"] 141 | if bit_depth not in valid_depths: 142 | return valid_depths[0] 143 | return bit_depth 144 | 145 | def convert_to_grayscale(self, img: np.ndarray) -> np.ndarray: 146 | """Proper grayscale conversion using luminance weights""" 147 | if len(img.shape) == 2 or img.shape[-1] == 1: 148 | return img if img.shape[-1] == 1 else img[..., np.newaxis] 149 | 150 | if img.shape[-1] >= 3: 151 | # ITU-R BT.709 luminance weights 152 | weights = np.array([0.2126, 0.7152, 0.0722]) 153 | gray = np.dot(img[..., :3], weights) 154 | return gray[..., np.newaxis] 155 | 156 | return img[..., 0:1] # Fallback for 2-channel images 157 | 158 | def prepare_image(self, img_tensor: torch.Tensor, save_as_grayscale: bool) -> np.ndarray: 159 | """Convert tensor to numpy and prepare channels""" 160 | # Convert from tensor 161 | if img_tensor.ndim == 4 and img_tensor.shape[0] == 1: 162 | img_np = img_tensor.squeeze(0).cpu().numpy() 163 | else: 164 | img_np = img_tensor.cpu().numpy() 165 | 166 | # Ensure float32 [0,1] range 167 | img_np = np.clip(img_np, 0, 1).astype(np.float32) 168 | 169 | # Handle channels 170 | if len(img_np.shape) == 2: 171 | img_np = img_np[..., np.newaxis] 172 | 173 | # Convert to grayscale if explicitly requested (don't auto-detect for solid colors) 174 | if save_as_grayscale: 175 | img_np = self.convert_to_grayscale(img_np) 176 | 177 | return img_np 178 | 179 | def save_exr(self, img: np.ndarray, path: str, bit_depth: int, compression: str) -> None: 180 | """Save EXR with proper float handling""" 181 | # Convert to appropriate type based on bit depth 182 | if bit_depth == 16: 183 | data = img.astype(np.float16) 184 | pixel_type = oiio.HALF 185 | else: # 32 186 | data = img.astype(np.float32) 187 | pixel_type = oiio.FLOAT 188 | 189 | data = np.ascontiguousarray(data) 190 | channels = 1 if data.ndim == 2 else data.shape[-1] 191 | 192 | spec = oiio.ImageSpec(data.shape[1], data.shape[0], channels, pixel_type) 193 | spec.attribute("compression", compression) 194 | spec.attribute("Software", "COCO Tools") 195 | 196 | buf = oiio.ImageBuf(spec) 197 | buf.set_pixels(oiio.ROI(), data) 198 | 199 | if not buf.write(path): 200 | raise RuntimeError(f"Failed to write EXR: {oiio.geterror()}") 201 | 202 | def save_png(self, img: np.ndarray, path: str, bit_depth: int) -> None: 203 | """Save PNG with proper bit depth handling""" 204 | # Convert to appropriate type based on bit depth 205 | if bit_depth == 8: 206 | data = (img * 255).astype(np.uint8) 207 | pixel_type = oiio.UINT8 208 | elif bit_depth == 16: 209 | data = (img * 65535).astype(np.uint16) 210 | pixel_type = oiio.UINT16 211 | else: 212 | # Fallback to 8-bit for unsupported depths 213 | data = (img * 255).astype(np.uint8) 214 | pixel_type = oiio.UINT8 215 | 216 | data = np.ascontiguousarray(data) 217 | channels = 1 if data.ndim == 2 else data.shape[-1] 218 | 219 | spec = oiio.ImageSpec(data.shape[1], data.shape[0], channels, pixel_type) 220 | spec.attribute("compression", "zip") 221 | spec.attribute("png:compressionLevel", 9) 222 | 223 | buf = oiio.ImageBuf(spec) 224 | buf.set_pixels(oiio.ROI(), data) 225 | 226 | if not buf.write(path): 227 | raise RuntimeError(f"Failed to write PNG: {oiio.geterror()}") 228 | 229 | def save_opencv_format(self, img: np.ndarray, path: str, quality: int = 95) -> None: 230 | """Save JPEG/WebP using OpenCV""" 231 | # Convert to 8-bit BGR 232 | data = (img * 255).astype(np.uint8) 233 | 234 | # Convert RGB to BGR only for 3+ channel images 235 | if data.shape[-1] >= 3: 236 | data = cv.cvtColor(data, cv.COLOR_RGB2BGR) 237 | 238 | # Save with quality setting 239 | if path.endswith(('.jpg', '.jpeg')): 240 | cv.imwrite(path, data, [cv.IMWRITE_JPEG_QUALITY, quality]) 241 | else: # webp 242 | cv.imwrite(path, data, [cv.IMWRITE_WEBP_QUALITY, quality]) 243 | 244 | def save_tiff(self, img: np.ndarray, path: str, bit_depth: int) -> None: 245 | """Save TIFF with proper bit depth""" 246 | if bit_depth == 8: 247 | data = (img * 255).astype(np.uint8) 248 | elif bit_depth == 16: 249 | data = (img * 65535).astype(np.uint16) 250 | else: # 32 251 | data = img.astype(np.float32) 252 | 253 | # Determine photometric interpretation 254 | photometric = 'minisblack' if data.shape[-1] == 1 else 'rgb' 255 | 256 | tifffile.imwrite(path, data, photometric=photometric) 257 | 258 | def get_unique_filepath(self, base_path: str) -> str: 259 | """Get unique filepath with incremental counter""" 260 | if not os.path.exists(base_path): 261 | return base_path 262 | 263 | dir_name = os.path.dirname(base_path) 264 | base_name = os.path.basename(base_path) 265 | name, ext = os.path.splitext(base_name) 266 | 267 | counter = 1 268 | while True: 269 | new_path = os.path.join(dir_name, f"{name}_{counter}{ext}") 270 | if not os.path.exists(new_path): 271 | return new_path 272 | counter += 1 273 | 274 | def save_images(self, images, file_path, filename, save_mode="single", file_type="png", 275 | bit_depth=None, quality=None, save_as_grayscale=None, use_versioning=False, 276 | version=1, start_frame=None, frame_step=None, prompt=None, extra_pnginfo=None, 277 | exr_compression=None, **kwargs): 278 | """Main save function with optimized pipeline - handles missing contextual inputs and sequence mode""" 279 | 280 | # Provide format-specific defaults for missing inputs 281 | if bit_depth is None: 282 | bit_depth = "16" if file_type in ["exr", "png", "tiff"] else "8" 283 | 284 | if quality is None: 285 | quality = 95 # Default for JPG/WebP 286 | 287 | if save_as_grayscale is None: 288 | save_as_grayscale = False 289 | 290 | if exr_compression is None: 291 | exr_compression = "zips" # Default for EXR 292 | 293 | # Handle sequence parameters with defaults 294 | if start_frame is None: 295 | start_frame = 1 296 | if frame_step is None: 297 | frame_step = 1 298 | 299 | try: 300 | # Validate inputs 301 | bit_depth = int(bit_depth) 302 | file_type = file_type.lower() 303 | 304 | # Determine if this is sequence mode and validate pattern 305 | is_sequence_mode = save_mode == "sequence" or SequenceHandler.detect_sequence_pattern(filename) 306 | 307 | debug_log(logger, "info", f"Saving {len(images)} images in {save_mode} mode", 308 | f"Save mode: {save_mode}, Is sequence: {is_sequence_mode}, File type: {file_type}") 309 | bit_depth = self.validate_bit_depth(file_type, bit_depth) 310 | 311 | # Log save operation 312 | debug_log(logger, "info", f"Saving {len(images)} image(s) as {file_type.upper()}", 313 | f"Saving {len(images)} image(s) as {file_type.upper()} {bit_depth}-bit to {filename}") 314 | 315 | # Build base path 316 | if file_path: 317 | full_path = os.path.join(self.output_dir, file_path) if not os.path.isabs(file_path) else file_path 318 | os.makedirs(full_path, exist_ok=True) 319 | base_path = os.path.join(full_path, filename) 320 | else: 321 | base_path = os.path.join(self.output_dir, filename) 322 | 323 | # Add version string 324 | version_str = f"_v{version:03d}" if use_versioning and version >= 0 else "" 325 | 326 | # Track saved files for preview 327 | saved_files = [] 328 | 329 | # Process each image - handle single vs sequence mode 330 | for i, img_tensor in enumerate(images): 331 | # Prepare image (all formats start from float32 [0,1]) 332 | img_np = self.prepare_image(img_tensor, save_as_grayscale) 333 | 334 | # Build output path based on mode 335 | if is_sequence_mode and SequenceHandler.detect_sequence_pattern(filename): 336 | # Sequence mode with #### pattern 337 | frame_number = start_frame + (i * frame_step) 338 | sequence_filename = filename.replace('####', f'{frame_number:04d}') 339 | out_path = f"{os.path.join(os.path.dirname(base_path), sequence_filename)}{version_str}.{file_type}" 340 | elif is_sequence_mode: 341 | # Sequence mode without pattern - use frame numbers 342 | frame_number = start_frame + (i * frame_step) 343 | out_path = f"{base_path}_{frame_number:04d}{version_str}.{file_type}" 344 | else: 345 | # Single mode - use index for multiple images 346 | frame_str = f"_{i}" if len(images) > 1 else "" 347 | out_path = f"{base_path}{version_str}{frame_str}.{file_type}" 348 | 349 | out_path = self.get_unique_filepath(out_path) 350 | 351 | # Save based on format 352 | if file_type == "exr": 353 | self.save_exr(img_np, out_path, bit_depth, exr_compression) 354 | elif file_type == "png": 355 | self.save_png(img_np, out_path, bit_depth) 356 | elif file_type in ["jpg", "jpeg", "webp"]: 357 | self.save_opencv_format(img_np, out_path, quality) 358 | elif file_type == "tiff": 359 | self.save_tiff(img_np, out_path, bit_depth) 360 | 361 | # Track saved file info 362 | saved_files.append({ 363 | "filename": os.path.basename(out_path), 364 | "fullPath": out_path, 365 | "format": file_type, 366 | "bitDepth": bit_depth, 367 | "index": i 368 | }) 369 | 370 | # Generate full resolution preview for saved images 371 | preview_data = generate_preview_for_comfyui( 372 | images, 373 | source_path=f"saver_{len(saved_files)}_files", 374 | is_sequence=is_sequence_mode or len(saved_files) > 1, 375 | frame_index=0 376 | ) 377 | 378 | # Log completion 379 | debug_log(logger, "info", f"Saved {len(saved_files)} files successfully", 380 | f"Successfully saved {len(saved_files)} files: {[f['filename'] for f in saved_files]}") 381 | 382 | # Return with preview and saved file info 383 | result = { 384 | "ui": { 385 | "images": preview_data or [], 386 | "saved_files": saved_files 387 | } 388 | } 389 | 390 | return result 391 | 392 | except Exception as e: 393 | debug_log(logger, "error", "Save operation failed", f"Saver error: {str(e)}") 394 | raise RuntimeError(f"Saver error: {str(e)}") from e 395 | 396 | # NODE_CLASS_MAPPINGS = { 397 | # "saver": saver, 398 | # } 399 | 400 | # NODE_DISPLAY_NAME_MAPPINGS = { 401 | # "saver": "Image Saver" 402 | # } 403 | -------------------------------------------------------------------------------- /modules/load_exr_layer_by_name.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import logging 3 | from typing import Dict, List, Union, Any 4 | 5 | # Import centralized logging setup 6 | try: 7 | from ..utils.debug_utils import setup_logging 8 | setup_logging() 9 | except ImportError: 10 | logging.basicConfig(level=logging.INFO) 11 | 12 | logger = logging.getLogger(__name__) 13 | 14 | # Import debug utilities 15 | try: 16 | from ..utils.debug_utils import debug_log, format_layer_names, format_tensor_info, DEBUG_MODE 17 | except ImportError: 18 | # Fallback if utils not available 19 | def debug_log(logger, level, simple_msg, verbose_msg=None, **kwargs): 20 | getattr(logger, level.lower())(simple_msg) 21 | def format_layer_names(layer_names, max_simple=None): 22 | return ', '.join(layer_names) 23 | def format_tensor_info(tensor_shape, tensor_dtype, name=""): 24 | return f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}" 25 | DEBUG_MODE = "simple" 26 | 27 | class LoadExrLayerByName: 28 | """ 29 | The Load EXR Layer by Name node allows selecting a specific layer from an EXR layer dictionary. 30 | It works like Nuke's Shuffle node, allowing users to pick a specific layer to output. 31 | """ 32 | 33 | # Class variables to store available layer names 34 | available_layers = ["none"] 35 | 36 | def __init__(self): 37 | debug_log(logger, "info", "Layer selector initialized", "load_exr_layer_by_name class initialized") 38 | 39 | @classmethod 40 | def INPUT_TYPES(cls): 41 | # debug_log(logger, "info", f"Available layers: {len(cls.available_layers)}", 42 | # f"INPUT_TYPES called - available layers: {cls.available_layers}") 43 | return { 44 | "required": { 45 | "layers": ("LAYERS",), 46 | "layer_name": ("STRING", { 47 | "default": "none", 48 | "multiline": False, 49 | "description": "Name of the layer to extract from the EXR. You can find layer names in the metadata output of the Load EXR node." 50 | }) 51 | }, 52 | "optional": { 53 | "conversion": (["Auto", "To RGB", "To Mask"], { 54 | "default": "Auto" 55 | }) 56 | } 57 | } 58 | 59 | RETURN_TYPES = ("IMAGE", "MASK") 60 | RETURN_NAMES = ("image", "mask") 61 | FUNCTION = "process_layer" 62 | CATEGORY = "Image/EXR" 63 | 64 | @classmethod 65 | def IS_CHANGED(cls, **kwargs): 66 | return float("NaN") # Always execute 67 | 68 | def process_layer(self, layers: Dict[str, torch.Tensor], layer_name: str, 69 | conversion: str = "Auto") -> List[Union[torch.Tensor, None]]: 70 | """ 71 | Extract a specific layer from the layers dictionary. 72 | 73 | Args: 74 | layers: Dictionary of layer names to tensors 75 | layer_name: Name of the layer to extract 76 | conversion: How to convert the layer (Auto, To RGB, To Mask) 77 | 78 | Returns: 79 | List containing [image, mask] tensors 80 | """ 81 | # Check if we have any layers at all 82 | if not layers or len(layers) == 0: 83 | debug_log(logger, "warning", "No layers available", "No layers available in the input") 84 | return [torch.zeros((1, 1, 1, 3)), torch.zeros((1, 1, 1))] 85 | 86 | # Log the available layers for debugging 87 | # debug_log(logger, "info", f"Found {len(layers)} layers: {format_layer_names(list(layers.keys()))}", 88 | # f"Available layers: {list(layers.keys())}") 89 | if DEBUG_MODE == "verbose": 90 | for layer_key, layer_tensor in layers.items(): 91 | debug_log(logger, "info", "", f"Layer '{layer_key}' has shape {layer_tensor.shape} and type {layer_tensor.dtype}") 92 | 93 | # Update the class variable with available layer names 94 | self.__class__.available_layers = ["none"] + sorted(list(layers.keys())) 95 | 96 | # If the layer doesn't exist, try to find a close match 97 | if layer_name not in layers and layer_name != "none": 98 | # Try to find an exact match ignoring case 99 | case_insensitive_matches = [l for l in layers.keys() if l.lower() == layer_name.lower()] 100 | if case_insensitive_matches: 101 | layer_name = case_insensitive_matches[0] 102 | debug_log(logger, "info", "Found layer with different case", 103 | f"Layer name '{layer_name}' found with different case: '{layer_name}'") 104 | else: 105 | # Try to find a partial match 106 | matches = [l for l in layers.keys() if layer_name.lower() in l.lower()] 107 | if matches: 108 | # Sort matches by length to find the closest match 109 | matches.sort(key=len) 110 | layer_name = matches[0] 111 | debug_log(logger, "info", "Using closest layer match", 112 | f"Layer name '{layer_name}' not found exactly, using closest match: '{layer_name}'") 113 | else: 114 | # Try to match hierarchical names (e.g., "CITY SCENE.AO" when user enters "AO") 115 | hierarchical_matches = [] 116 | for l in layers.keys(): 117 | if '.' in l: 118 | parts = l.split('.') 119 | # Check if any part matches the layer name 120 | if any(part.lower() == layer_name.lower() for part in parts): 121 | hierarchical_matches.append(l) 122 | 123 | if hierarchical_matches: 124 | layer_name = hierarchical_matches[0] 125 | debug_log(logger, "info", "Found hierarchical match", 126 | f"Found hierarchical layer match: '{layer_name}'") 127 | else: 128 | # Try to match subimage names (e.g., "AO" for a subimage) 129 | subimage_matches = [l for l in layers.keys() if l.split('.')[0].lower() == layer_name.lower()] 130 | if subimage_matches: 131 | layer_name = subimage_matches[0] 132 | debug_log(logger, "info", "Found subimage match", 133 | f"Found subimage match: '{layer_name}'") 134 | else: 135 | debug_log(logger, "warning", "Layer not found", 136 | f"Layer '{layer_name}' not found and no close matches") 137 | # Use the first available layer as fallback 138 | if len(layers) > 0: 139 | layer_name = list(layers.keys())[0] 140 | debug_log(logger, "info", "Using first available layer", 141 | f"Using first available layer: {layer_name}") 142 | else: 143 | return [torch.zeros((1, 1, 1, 3)), torch.zeros((1, 1, 1))] 144 | 145 | # If no layer is specified or "none" is selected, return empty tensors 146 | if not layer_name or layer_name == "none": 147 | debug_log(logger, "warning", "No layer specified", "No layer specified, returning empty tensors") 148 | return [torch.zeros((1, 1, 1, 3)), torch.zeros((1, 1, 1))] 149 | 150 | # Get the requested layer 151 | layer_tensor = layers[layer_name] 152 | 153 | # Log the layer processing 154 | debug_log(logger, "info", f"Processing layer '{layer_name}'", 155 | f"Processing layer '{layer_name}' with shape {layer_tensor.shape} and type {layer_tensor.dtype}") 156 | 157 | # Debug: Print the requested layer name 158 | debug_log(logger, "info", "", f"Requested layer: '{layer_name}'") 159 | 160 | # Special handling for alpha layers only (not depth or Z) 161 | is_alpha_layer = 'alpha' in layer_name.lower() 162 | 163 | # Check tensor shape to determine its type 164 | if len(layer_tensor.shape) == 4 and layer_tensor.shape[3] == 3: 165 | # It's an RGB tensor [1, H, W, 3] 166 | if conversion == "To Mask" or (conversion == "Auto" and is_alpha_layer): 167 | # Convert RGB to mask by taking the mean across channels 168 | mask_output = layer_tensor.mean(dim=3, keepdim=False) 169 | image_output = None 170 | debug_log(logger, "info", "Converted to mask", f"Converted RGB tensor to mask: shape={mask_output.shape}") 171 | else: 172 | # Keep as an image 173 | image_output = layer_tensor 174 | mask_output = None 175 | debug_log(logger, "info", "Using as RGB image", f"Using RGB tensor as image: shape={image_output.shape}") 176 | elif len(layer_tensor.shape) == 3: 177 | # It's a single-channel tensor [1, H, W] 178 | # Special handling for depth and Z channels 179 | is_depth_or_z = 'depth' in layer_name.lower() or layer_name.lower() == 'z' 180 | 181 | if conversion == "To RGB": 182 | # Convert to RGB by replicating to 3 channels 183 | image_output = torch.cat([layer_tensor.unsqueeze(3)] * 3, dim=3) 184 | mask_output = None 185 | debug_log(logger, "info", "Converted to RGB", f"Converted single-channel tensor to RGB: shape={image_output.shape}") 186 | elif is_depth_or_z: 187 | # For depth and Z channels, return as image by default 188 | image_output = torch.cat([layer_tensor.unsqueeze(3)] * 3, dim=3) 189 | mask_output = None 190 | debug_log(logger, "info", "Converted depth to RGB", f"Converted depth/Z tensor to RGB: shape={image_output.shape}") 191 | elif 'alpha' in layer_name.lower(): 192 | # For alpha channels, return as mask 193 | mask_output = layer_tensor 194 | image_output = None 195 | debug_log(logger, "info", "Using alpha as mask", f"Using alpha tensor as mask: shape={mask_output.shape}") 196 | else: 197 | # For other single-channel data, use the conversion setting 198 | if conversion == "To Mask": 199 | mask_output = layer_tensor 200 | image_output = None 201 | debug_log(logger, "info", "Using as mask", f"Using single-channel tensor as mask: shape={mask_output.shape}") 202 | else: 203 | # Default to RGB for Auto mode for non-alpha channels 204 | image_output = torch.cat([layer_tensor.unsqueeze(3)] * 3, dim=3) 205 | mask_output = None 206 | debug_log(logger, "info", "Using as RGB (Auto)", f"Using single-channel tensor as RGB (Auto): shape={image_output.shape}") 207 | # Special case for empty tensors or tensors with shape [1, 1, 1, 3] 208 | elif len(layer_tensor.shape) == 4 and layer_tensor.shape[1] == 1 and layer_tensor.shape[2] == 1: 209 | # This is likely an empty tensor or a placeholder 210 | # Check if it's an alpha or depth layer 211 | if 'alpha' in layer_name.lower() or 'depth' in layer_name.lower(): 212 | # Create a proper mask tensor 213 | # Get dimensions from another layer if possible 214 | height, width = 720, 1280 # Default size 215 | for other_name, other_tensor in layers.items(): 216 | if other_name != layer_name and len(other_tensor.shape) >= 3: 217 | if len(other_tensor.shape) == 4: # RGB tensor [1, H, W, 3] 218 | height, width = other_tensor.shape[1], other_tensor.shape[2] 219 | elif len(other_tensor.shape) == 3: # Mask tensor [1, H, W] 220 | height, width = other_tensor.shape[1], other_tensor.shape[2] 221 | break 222 | 223 | # Create a mask tensor with the correct dimensions 224 | mask_output = torch.zeros((1, height, width)) 225 | image_output = None 226 | else: 227 | # Keep as an image 228 | image_output = layer_tensor 229 | mask_output = None 230 | else: 231 | # Unknown format, log error 232 | debug_log(logger, "error", "Unsupported tensor shape", 233 | f"Layer '{layer_name}' has an unsupported tensor shape: {layer_tensor.shape}") 234 | return [torch.zeros((1, 1, 1, 3)), torch.zeros((1, 1, 1))] 235 | 236 | # Set placeholder for any None outputs 237 | if image_output is None: 238 | image_output = torch.zeros((1, 1, 1, 3)) 239 | if mask_output is None: 240 | mask_output = torch.zeros((1, 1, 1)) 241 | 242 | return [image_output, mask_output] 243 | 244 | # Define a copy of the main class for cryptomatte layers 245 | class CryptomatteLayer(LoadExrLayerByName): 246 | """ 247 | The Cryptomatte Shamble node allows selecting a specific cryptomatte layer from an EXR dictionary. 248 | It is identical to the Load EXR Layer by Name node but filters for cryptomatte layers only. 249 | """ 250 | 251 | @classmethod 252 | def INPUT_TYPES(cls): 253 | return { 254 | "required": { 255 | "cryptomatte": ("CRYPTOMATTE",), 256 | "layer_name": ("STRING", { 257 | "default": "none", 258 | "multiline": False, 259 | "description": "Name of the cryptomatte layer to extract. Look for names starting with 'crypto' in the metadata." 260 | }) 261 | } 262 | } 263 | 264 | RETURN_TYPES = ("IMAGE",) 265 | RETURN_NAMES = ("image",) 266 | FUNCTION = "process_cryptomatte" 267 | CATEGORY = "Image/EXR" 268 | 269 | @classmethod 270 | def IS_CHANGED(cls, **kwargs): 271 | return float("NaN") # Always execute 272 | 273 | def process_cryptomatte(self, cryptomatte: Dict[str, torch.Tensor], layer_name: str) -> List[torch.Tensor]: 274 | """ 275 | Extract a specific cryptomatte layer. 276 | 277 | Args: 278 | cryptomatte: Dictionary of cryptomatte layer names to tensors 279 | layer_name: Name of the cryptomatte layer to extract 280 | 281 | Returns: 282 | List containing the cryptomatte image tensor 283 | """ 284 | # Check if we have any layers at all 285 | if not cryptomatte or len(cryptomatte) == 0: 286 | debug_log(logger, "warning", "No cryptomatte layers available", "No cryptomatte layers available in the input") 287 | return [torch.zeros((1, 1, 1, 3))] 288 | 289 | # Update the class variable with available cryptomatte layer names 290 | self.__class__.available_layers = ["none"] + sorted(list(cryptomatte.keys())) 291 | 292 | # If the layer doesn't exist, try to find a close match 293 | if layer_name not in cryptomatte and layer_name != "none": 294 | # Try to find an exact match ignoring case 295 | case_insensitive_matches = [l for l in cryptomatte.keys() if l.lower() == layer_name.lower()] 296 | if case_insensitive_matches: 297 | layer_name = case_insensitive_matches[0] 298 | debug_log(logger, "info", "Found cryptomatte with different case", 299 | f"Cryptomatte layer name '{layer_name}' found with different case: '{layer_name}'") 300 | else: 301 | # Try to find a partial match 302 | matches = [l for l in cryptomatte.keys() if layer_name.lower() in l.lower()] 303 | if matches: 304 | # Sort matches by length to find the closest match 305 | matches.sort(key=len) 306 | layer_name = matches[0] 307 | debug_log(logger, "info", "Using closest cryptomatte match", 308 | f"Cryptomatte layer name '{layer_name}' not found exactly, using closest match: '{layer_name}'") 309 | else: 310 | # Try to match hierarchical names (e.g., "CITY SCENE.CryptoAsset00" when user enters "CryptoAsset") 311 | hierarchical_matches = [] 312 | for l in cryptomatte.keys(): 313 | if '.' in l: 314 | parts = l.split('.') 315 | # Check if any part matches the layer name 316 | if any(part.lower() == layer_name.lower() for part in parts): 317 | hierarchical_matches.append(l) 318 | 319 | if hierarchical_matches: 320 | layer_name = hierarchical_matches[0] 321 | debug_log(logger, "info", "Found hierarchical cryptomatte match", 322 | f"Found hierarchical cryptomatte layer match: '{layer_name}'") 323 | else: 324 | debug_log(logger, "warning", "Cryptomatte layer not found", 325 | f"Cryptomatte layer '{layer_name}' not found and no close matches") 326 | # Use the first available layer as fallback 327 | if len(cryptomatte) > 0: 328 | layer_name = list(cryptomatte.keys())[0] 329 | debug_log(logger, "info", "Using first available cryptomatte", 330 | f"Using first available cryptomatte layer: {layer_name}") 331 | else: 332 | return [torch.zeros((1, 1, 1, 3))] 333 | 334 | # If no layer is specified or "none" is selected, return an empty tensor 335 | if not layer_name or layer_name == "none": 336 | debug_log(logger, "warning", "No cryptomatte layer specified", "No cryptomatte layer specified, returning empty tensor") 337 | return [torch.zeros((1, 1, 1, 3))] 338 | 339 | # Return the requested cryptomatte layer 340 | return [cryptomatte[layer_name]] 341 | 342 | # NODE_CLASS_MAPPINGS = { 343 | # "load_exr_layer_by_name": load_exr_layer_by_name, 344 | # "shamble_cryptomatte": shamble_cryptomatte 345 | # } 346 | 347 | # NODE_DISPLAY_NAME_MAPPINGS = { 348 | # "load_exr_layer_by_name": "Load EXR Layer by Name", 349 | # "shamble_cryptomatte": "Cryptomatte Layer" 350 | # } 351 | -------------------------------------------------------------------------------- /modules/load_exr_sequence.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | import torch 4 | import json 5 | from typing import List 6 | 7 | # Import centralized logging setup 8 | try: 9 | from ..utils.debug_utils import setup_logging 10 | setup_logging() 11 | except ImportError: 12 | logging.basicConfig(level=logging.INFO) 13 | 14 | logger = logging.getLogger(__name__) 15 | 16 | # Import EXR and sequence utilities 17 | try: 18 | from ..utils.exr_utils import ExrProcessor 19 | from ..utils.sequence_utils import SequenceHandler 20 | from ..utils.debug_utils import debug_log, format_tensor_info 21 | from ..utils.preview_utils import generate_preview_for_comfyui 22 | except ImportError as e: 23 | raise ImportError(f"Required utilities are not available: {str(e)}. Please ensure utils are properly installed.") 24 | 25 | 26 | class LoadExrSequence: 27 | @classmethod 28 | def INPUT_TYPES(cls): 29 | return { 30 | "required": { 31 | "sequence_path": ("STRING", { 32 | "default": "path/to/sequence_####.exr", 33 | "description": "Path with #### pattern for frame numbers (e.g., render_####.exr)" 34 | }), 35 | "start_frame": ("INT", { 36 | "default": 1, 37 | "min": 0, 38 | "max": 999999, 39 | "description": "Starting frame number" 40 | }), 41 | "end_frame": ("INT", { 42 | "default": 100, 43 | "min": 1, 44 | "max": 999999, 45 | "description": "Ending frame number" 46 | }), 47 | "frame_step": ("INT", { 48 | "default": 1, 49 | "min": 1, 50 | "max": 100, 51 | "description": "Step between frames" 52 | }), 53 | "normalize": ("BOOLEAN", { 54 | "default": False, 55 | "description": "Normalize image values to the 0-1 range" 56 | }) 57 | }, 58 | "hidden": { 59 | "node_id": "UNIQUE_ID", 60 | "layer_data": "DICT" 61 | } 62 | } 63 | 64 | RETURN_TYPES = ("IMAGE", "MASK", "CRYPTOMATTE", "LAYERS", "STRING", "STRING", "STRING") 65 | RETURN_NAMES = ("sequence", "alpha", "cryptomatte", "layers", "layer names", "raw layer info", "metadata") 66 | 67 | FUNCTION = "load_sequence" 68 | CATEGORY = "Image/EXR" 69 | 70 | @classmethod 71 | def IS_CHANGED(cls, sequence_path, start_frame, end_frame, frame_step, normalize=False, **kwargs): 72 | """ 73 | Smart caching based on file modification times and sequence parameters. 74 | Only reload if files actually changed or parameters changed. 75 | """ 76 | try: 77 | # Validate sequence pattern first 78 | if not SequenceHandler.detect_sequence_pattern(sequence_path): 79 | return float("NaN") # Invalid pattern, always try 80 | 81 | # Validate and sanitize sequence parameters 82 | try: 83 | start_frame, end_frame, frame_step = SequenceHandler.validate_sequence_parameters( 84 | start_frame, end_frame, frame_step 85 | ) 86 | except Exception: 87 | return float("NaN") # Invalid parameters, always try 88 | 89 | # Find existing sequence files 90 | try: 91 | existing_sequence_files = SequenceHandler.find_sequence_files(sequence_path) 92 | if not existing_sequence_files: 93 | return float("NaN") # No files found, always try 94 | 95 | # Extract frame numbers and select frames that would be loaded 96 | frame_info = SequenceHandler.extract_frame_numbers(existing_sequence_files) 97 | selected_frames = SequenceHandler.select_sequence_frames( 98 | frame_info, start_frame, end_frame, frame_step 99 | ) 100 | 101 | if not selected_frames: 102 | return float("NaN") # No frames selected, always try 103 | except Exception: 104 | return float("NaN") # Error in sequence processing, always try 105 | 106 | # Create hash from all file stats and parameters 107 | file_stats = [] 108 | for frame_path in selected_frames: 109 | if frame_path and os.path.isfile(frame_path): 110 | try: 111 | stat = os.stat(frame_path) 112 | file_stats.append(f"{stat.st_mtime}_{stat.st_size}") 113 | except Exception: 114 | file_stats.append("error") 115 | else: 116 | file_stats.append("missing") 117 | 118 | # Include all parameters that affect the result 119 | param_hash = f"{sequence_path}_{start_frame}_{end_frame}_{frame_step}_{normalize}" 120 | files_hash = "_".join(file_stats) 121 | 122 | return f"{param_hash}_{files_hash}" 123 | 124 | except Exception: 125 | # If anything goes wrong, always try to load 126 | return float("NaN") 127 | 128 | def load_sequence(self, sequence_path: str, start_frame: int, end_frame: int, frame_step: int, 129 | normalize: bool = False, node_id: str = None, layer_data: dict = None, **kwargs) -> List: 130 | """ 131 | Load a sequence of EXR files and return batched tensors. 132 | Returns: 133 | - Batched RGB image tensors [B, H, W, 3] (sequence) 134 | - Batched Alpha channel tensors [B, H, W] (alpha) 135 | - Dictionary of all cryptomatte layers as batched tensors (cryptomatte) 136 | - Dictionary of all non-cryptomatte layers as batched tensors (layers) 137 | - List of processed layer names matching keys in the returned dictionaries (layer names) 138 | - List of raw channel names from the files (raw layer info) 139 | - Metadata as JSON string with sequence info (metadata) 140 | """ 141 | 142 | # Check for OIIO availability 143 | ExrProcessor.check_oiio_availability() 144 | 145 | try: 146 | # Validate sequence pattern 147 | if not SequenceHandler.detect_sequence_pattern(sequence_path): 148 | raise ValueError(f"Sequence path must contain #### pattern for frame numbers: {sequence_path}") 149 | 150 | # Validate and sanitize sequence parameters 151 | start_frame, end_frame, frame_step = SequenceHandler.validate_sequence_parameters( 152 | start_frame, end_frame, frame_step 153 | ) 154 | 155 | frame_count = len(range(start_frame, end_frame + 1, frame_step)) 156 | debug_log(logger, "info", f"Loading EXR sequence: {frame_count} frames", 157 | f"Loading EXR sequence: {sequence_path} (start={start_frame}, end={end_frame}, step={frame_step})") 158 | 159 | # Find existing sequence files 160 | existing_sequence_files = SequenceHandler.find_sequence_files(sequence_path) 161 | debug_log(logger, "info", f"Found {len(existing_sequence_files)} sequence files", 162 | f"SequenceHandler found {len(existing_sequence_files)} files for pattern: {sequence_path}") 163 | 164 | if not existing_sequence_files: 165 | logger.error(f"No sequence files found for pattern: {sequence_path}") 166 | logger.error(f"Pattern path: {sequence_path}") 167 | logger.error(f"Expected frame range: {start_frame} to {end_frame} step {frame_step}") 168 | raise FileNotFoundError(f"No sequence frames found for pattern: {sequence_path}") 169 | 170 | # Extract frame numbers and select frames 171 | frame_info = SequenceHandler.extract_frame_numbers(existing_sequence_files) 172 | selected_frames = SequenceHandler.select_sequence_frames( 173 | frame_info, start_frame, end_frame, frame_step 174 | ) 175 | 176 | if not selected_frames: 177 | logger.error(f"No frames selected from sequence") 178 | logger.error(f"Available frames: {len(existing_sequence_files)}") 179 | logger.error(f"Frame range requested: {start_frame} to {end_frame} step {frame_step}") 180 | if frame_info: 181 | available_frame_numbers = [frame_num for frame_num, _ in frame_info] 182 | logger.error(f"Available frame numbers: {sorted(available_frame_numbers)}") 183 | raise ValueError(f"No frames selected from sequence for range {start_frame}-{end_frame} step {frame_step}") 184 | 185 | debug_log(logger, "info", f"Selected {len(selected_frames)} frames for loading", 186 | f"Loading {len(selected_frames)} frames from sequence") 187 | 188 | # Find first valid frame to establish structure 189 | first_valid_frame = None 190 | first_frame_index = 0 191 | for i, frame_path in enumerate(selected_frames): 192 | if frame_path is not None: 193 | first_valid_frame = frame_path 194 | first_frame_index = i 195 | break 196 | 197 | if first_valid_frame is None: 198 | raise ValueError(f"No valid frames found in sequence range {start_frame}-{end_frame} step {frame_step}") 199 | 200 | # Load first valid frame to establish structure 201 | try: 202 | first_frame_result = ExrProcessor.process_exr_data(first_valid_frame, normalize, node_id, layer_data) 203 | except Exception as e: 204 | logger.error(f"Failed to load first frame: {first_valid_frame}") 205 | logger.error(f"Error: {str(e)}") 206 | raise 207 | 208 | # Handle result format (dict if preview generated, list if not) 209 | if isinstance(first_frame_result, dict): 210 | first_frame_data = first_frame_result["result"] 211 | else: 212 | first_frame_data = first_frame_result 213 | 214 | # Initialize batch tensors with first frame 215 | batch_rgb_list = [first_frame_data[0]] 216 | batch_alpha_list = [first_frame_data[1]] 217 | batch_layers_dict = {} 218 | batch_cryptomatte_dict = {} 219 | 220 | # Initialize layer dictionaries 221 | for layer_name, layer_tensor in first_frame_data[3].items(): 222 | batch_layers_dict[layer_name] = [layer_tensor] 223 | 224 | for crypto_name, crypto_tensor in first_frame_data[2].items(): 225 | batch_cryptomatte_dict[crypto_name] = [crypto_tensor] 226 | 227 | # Load remaining frames (skip the first valid frame we already loaded) 228 | for i, frame_path in enumerate(selected_frames): 229 | # Skip the first valid frame (already loaded) and None paths 230 | if i == first_frame_index or frame_path is None: 231 | if frame_path is None: 232 | logger.warning(f"Skipping missing frame {i+1}/{len(selected_frames)}: frame not found in sequence") 233 | # Create white placeholder frames for missing frames 234 | white_rgb = torch.ones_like(first_frame_data[0]) 235 | white_alpha = torch.ones_like(first_frame_data[1]) 236 | batch_rgb_list.append(white_rgb) 237 | batch_alpha_list.append(white_alpha) 238 | 239 | # Create white placeholders for layers 240 | for layer_name, layer_tensor in first_frame_data[3].items(): 241 | if layer_name in batch_layers_dict: 242 | white_layer = torch.ones_like(layer_tensor) 243 | batch_layers_dict[layer_name].append(white_layer) 244 | 245 | # Create white placeholders for cryptomatte 246 | for crypto_name, crypto_tensor in first_frame_data[2].items(): 247 | if crypto_name in batch_cryptomatte_dict: 248 | white_crypto = torch.ones_like(crypto_tensor) 249 | batch_cryptomatte_dict[crypto_name].append(white_crypto) 250 | continue 251 | 252 | try: 253 | frame_result = ExrProcessor.process_exr_data(frame_path, normalize, node_id, layer_data) 254 | 255 | # Handle result format (dict if preview generated, list if not) 256 | if isinstance(frame_result, dict): 257 | frame_data = frame_result["result"] 258 | else: 259 | frame_data = frame_result 260 | 261 | # Add to batch lists 262 | batch_rgb_list.append(frame_data[0]) 263 | batch_alpha_list.append(frame_data[1]) 264 | 265 | # Add layers to batch 266 | for layer_name, layer_tensor in frame_data[3].items(): 267 | if layer_name in batch_layers_dict: 268 | batch_layers_dict[layer_name].append(layer_tensor) 269 | else: 270 | logger.warning(f"Layer '{layer_name}' not found in first frame, skipping for this frame") 271 | 272 | # Add cryptomatte to batch 273 | for crypto_name, crypto_tensor in frame_data[2].items(): 274 | if crypto_name in batch_cryptomatte_dict: 275 | batch_cryptomatte_dict[crypto_name].append(crypto_tensor) 276 | else: 277 | logger.warning(f"Cryptomatte '{crypto_name}' not found in first frame, skipping for this frame") 278 | 279 | except Exception as e: 280 | logger.error(f"Failed to load frame {i+1}/{len(selected_frames)}: {frame_path}") 281 | logger.error(f"Error: {str(e)}") 282 | # Create white placeholder frame for failed loads 283 | white_rgb = torch.ones_like(first_frame_data[0]) 284 | white_alpha = torch.ones_like(first_frame_data[1]) 285 | batch_rgb_list.append(white_rgb) 286 | batch_alpha_list.append(white_alpha) 287 | 288 | # Create white placeholders for layers 289 | for layer_name, layer_tensor in first_frame_data[3].items(): 290 | if layer_name in batch_layers_dict: 291 | white_layer = torch.ones_like(layer_tensor) 292 | batch_layers_dict[layer_name].append(white_layer) 293 | 294 | # Create white placeholders for cryptomatte 295 | for crypto_name, crypto_tensor in first_frame_data[2].items(): 296 | if crypto_name in batch_cryptomatte_dict: 297 | white_crypto = torch.ones_like(crypto_tensor) 298 | batch_cryptomatte_dict[crypto_name].append(white_crypto) 299 | 300 | # Stack tensors into batches 301 | final_rgb = torch.cat(batch_rgb_list, dim=0) 302 | final_alpha = torch.cat(batch_alpha_list, dim=0) 303 | 304 | # Stack layer tensors 305 | final_layers = {} 306 | for layer_name, tensor_list in batch_layers_dict.items(): 307 | if tensor_list: 308 | final_layers[layer_name] = torch.cat(tensor_list, dim=0) 309 | 310 | # Stack cryptomatte tensors 311 | final_cryptomatte = {} 312 | for crypto_name, tensor_list in batch_cryptomatte_dict.items(): 313 | if tensor_list: 314 | final_cryptomatte[crypto_name] = torch.cat(tensor_list, dim=0) 315 | 316 | # Update metadata with sequence information 317 | metadata_str = first_frame_data[6] # metadata is at index 6 318 | metadata = json.loads(metadata_str) if metadata_str else {} 319 | metadata["sequence_info"] = { 320 | "is_sequence": True, 321 | "pattern": sequence_path, 322 | "start_frame": start_frame, 323 | "end_frame": end_frame, 324 | "frame_step": frame_step, 325 | "total_frames": len(selected_frames), 326 | "loaded_frames": [f for f in selected_frames if f is not None], 327 | "missing_frames": [i for i, f in enumerate(selected_frames) if f is None] 328 | } 329 | metadata_json = json.dumps(metadata) 330 | 331 | # Log final batch information 332 | debug_log(logger, "info", f"Sequence loaded: {len(selected_frames)} frames, RGB shape: {format_tensor_info(final_rgb.shape, final_rgb.dtype)}", 333 | f"Successfully loaded sequence: {len(selected_frames)} frames, RGB batch shape: {final_rgb.shape}, " 334 | f"Alpha batch shape: {final_alpha.shape}, {len(final_layers)} layer types, {len(final_cryptomatte)} cryptomatte types") 335 | 336 | # Generate preview for sequence (first frame) at full resolution 337 | preview_result = generate_preview_for_comfyui(final_rgb, sequence_path, is_sequence=True, frame_index=0) 338 | 339 | # Return same structure as single image but with batched tensors 340 | result = [ 341 | final_rgb, # sequence 342 | final_alpha, # alpha 343 | final_cryptomatte, # cryptomatte 344 | final_layers, # layers 345 | first_frame_data[4], # layer names (processed layer names) 346 | first_frame_data[5], # raw layer info (layer names) 347 | metadata_json # metadata 348 | ] 349 | 350 | # Return with preview if generated 351 | if preview_result: 352 | return {"ui": {"images": preview_result}, "result": result} 353 | else: 354 | return result 355 | 356 | except Exception as e: 357 | logger.error(f"Error loading EXR sequence {sequence_path}: {str(e)}") 358 | raise -------------------------------------------------------------------------------- /utils/exr_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | EXR utility functions for COCO Tools nodes 3 | Provides shared functionality for EXR file processing across loader nodes 4 | """ 5 | 6 | import os 7 | import logging 8 | import numpy as np 9 | import torch 10 | import json 11 | from typing import Tuple, Dict, List, Optional, Union, Any 12 | 13 | try: 14 | import OpenImageIO as oiio 15 | OIIO_AVAILABLE = True 16 | except ImportError: 17 | OIIO_AVAILABLE = False 18 | 19 | try: 20 | from .debug_utils import debug_log, format_layer_names, format_tensor_info, create_fallback_functions 21 | from .preview_utils import generate_preview_for_comfyui 22 | except ImportError: 23 | # Use fallback functions 24 | fallbacks = { 25 | 'debug_log': lambda logger, level, simple_msg, verbose_msg=None, **kwargs: getattr(logger, level.lower())(simple_msg), 26 | 'format_layer_names': lambda layer_names, max_simple=None: ', '.join(layer_names), 27 | 'format_tensor_info': lambda tensor_shape, tensor_dtype, name="": f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}", 28 | 'generate_preview_for_comfyui': lambda image_tensor, source_path="", is_sequence=False, frame_index=0: None 29 | } 30 | debug_log = fallbacks['debug_log'] 31 | format_layer_names = fallbacks['format_layer_names'] 32 | format_tensor_info = fallbacks['format_tensor_info'] 33 | generate_preview_for_comfyui = fallbacks['generate_preview_for_comfyui'] 34 | 35 | logger = logging.getLogger(__name__) 36 | 37 | 38 | class ExrProcessor: 39 | """Shared EXR processing functionality for loader nodes""" 40 | 41 | @staticmethod 42 | def check_oiio_availability(): 43 | """Check if OpenImageIO is available""" 44 | if not OIIO_AVAILABLE: 45 | raise ImportError("OpenImageIO is required for EXR loading but not available") 46 | 47 | @staticmethod 48 | def scan_exr_metadata(image_path: str) -> Dict[str, Any]: 49 | """ 50 | Scan the EXR file to extract metadata about available subimages without loading pixel data. 51 | Returns a dictionary of subimage information including names, channels, dimensions, etc. 52 | """ 53 | ExrProcessor.check_oiio_availability() 54 | 55 | if not os.path.exists(image_path): 56 | raise FileNotFoundError(f"EXR file not found: {image_path}") 57 | 58 | input_file = None 59 | try: 60 | input_file = oiio.ImageInput.open(image_path) 61 | if not input_file: 62 | raise IOError(f"Could not open {image_path}") 63 | 64 | metadata = {} 65 | subimages = [] 66 | 67 | current_subimage = 0 68 | more_subimages = True 69 | 70 | while more_subimages: 71 | spec = input_file.spec() 72 | 73 | width = spec.width 74 | height = spec.height 75 | channels = spec.nchannels 76 | channel_names = [spec.channel_name(i) for i in range(channels)] 77 | 78 | subimage_name = "default" 79 | if "name" in spec.extra_attribs: 80 | subimage_name = spec.getattribute("name") 81 | 82 | subimage_info = { 83 | "index": current_subimage, 84 | "name": subimage_name, 85 | "width": width, 86 | "height": height, 87 | "channels": channels, 88 | "channel_names": channel_names 89 | } 90 | 91 | extra_attribs = {} 92 | for i in range(len(spec.extra_attribs)): 93 | name = spec.extra_attribs[i].name 94 | value = spec.extra_attribs[i].value 95 | extra_attribs[name] = value 96 | 97 | subimage_info["extra_attributes"] = extra_attribs 98 | subimages.append(subimage_info) 99 | 100 | more_subimages = input_file.seek_subimage(current_subimage + 1, 0) 101 | current_subimage += 1 102 | 103 | metadata["subimages"] = subimages 104 | metadata["is_multipart"] = len(subimages) > 1 105 | metadata["subimage_count"] = len(subimages) 106 | metadata["file_path"] = image_path 107 | 108 | return metadata 109 | 110 | except Exception as e: 111 | debug_log(logger, "error", "Error scanning EXR metadata", f"Error scanning EXR metadata from {image_path}: {str(e)}") 112 | raise 113 | 114 | finally: 115 | if input_file: 116 | input_file.close() 117 | 118 | @staticmethod 119 | def load_all_data(image_path: str) -> Dict[int, np.ndarray]: 120 | """ 121 | Load all pixel data from all subimages in the EXR file. 122 | Returns a dictionary mapping subimage index to numpy array of shape (height, width, channels). 123 | """ 124 | input_file = None 125 | try: 126 | input_file = oiio.ImageInput.open(image_path) 127 | if not input_file: 128 | raise IOError(f"Could not open {image_path}") 129 | 130 | all_subimage_data = {} 131 | 132 | current_subimage = 0 133 | more_subimages = True 134 | 135 | while more_subimages: 136 | spec = input_file.spec() 137 | width = spec.width 138 | height = spec.height 139 | channels = spec.nchannels 140 | 141 | pixels = input_file.read_image() 142 | if pixels is None: 143 | debug_log(logger, "warning", "Failed to read subimage", 144 | f"Failed to read image data for subimage {current_subimage} from {image_path}") 145 | else: 146 | all_subimage_data[current_subimage] = np.array(pixels, dtype=np.float32).reshape(height, width, channels) 147 | 148 | more_subimages = input_file.seek_subimage(current_subimage + 1, 0) 149 | current_subimage += 1 150 | 151 | return all_subimage_data 152 | 153 | finally: 154 | if input_file: 155 | input_file.close() 156 | 157 | @staticmethod 158 | def is_cryptomatte_layer(group_name: str) -> bool: 159 | """Determine if a layer is a cryptomatte layer based on its name""" 160 | group_name_lower = group_name.lower() 161 | return ( 162 | "cryptomatte" in group_name_lower or 163 | group_name_lower.startswith("crypto") or 164 | any(crypto_key for crypto_key in ("cryptoasset", "cryptomaterial", "cryptoobject", "cryptoprimvar") 165 | if crypto_key in group_name_lower) or 166 | any(part.lower().startswith("crypto") for part in group_name.split('.')) 167 | ) 168 | 169 | @staticmethod 170 | def process_default_channels(all_data, channel_names, height, width, normalize): 171 | """Process default RGB and Alpha channels""" 172 | rgb_tensor = None 173 | if 'R' in channel_names and 'G' in channel_names and 'B' in channel_names: 174 | r_idx = channel_names.index('R') 175 | g_idx = channel_names.index('G') 176 | b_idx = channel_names.index('B') 177 | 178 | rgb_array = np.stack([ 179 | all_data[:, :, r_idx], 180 | all_data[:, :, g_idx], 181 | all_data[:, :, b_idx] 182 | ], axis=2) 183 | 184 | rgb_tensor = torch.from_numpy(rgb_array).float() 185 | rgb_tensor = rgb_tensor.unsqueeze(0) # [1, H, W, 3] 186 | 187 | if normalize: 188 | rgb_range = rgb_tensor.max() - rgb_tensor.min() 189 | if rgb_range > 0: 190 | rgb_tensor = (rgb_tensor - rgb_tensor.min()) / rgb_range 191 | else: 192 | if all_data.shape[2] >= 3: 193 | rgb_array = all_data[:, :, :3] 194 | else: 195 | rgb_array = np.stack([all_data[:, :, 0]] * 3, axis=2) 196 | 197 | rgb_tensor = torch.from_numpy(rgb_array).float() 198 | rgb_tensor = rgb_tensor.unsqueeze(0) # [1, H, W, 3] 199 | 200 | if normalize: 201 | rgb_range = rgb_tensor.max() - rgb_tensor.min() 202 | if rgb_range > 0: 203 | rgb_tensor = (rgb_tensor - rgb_tensor.min()) / rgb_range 204 | 205 | alpha_tensor = None 206 | if 'A' in channel_names: 207 | a_idx = channel_names.index('A') 208 | alpha_array = all_data[:, :, a_idx] 209 | 210 | alpha_tensor = torch.from_numpy(alpha_array).float() 211 | alpha_tensor = alpha_tensor.unsqueeze(0) # [1, H, W] 212 | 213 | if normalize: 214 | alpha_tensor = alpha_tensor.clamp(0, 1) 215 | else: 216 | alpha_tensor = torch.ones((1, height, width)) 217 | 218 | return rgb_tensor, alpha_tensor 219 | 220 | @staticmethod 221 | def process_rgb_type_layer(group_name, r_suffix, g_suffix, b_suffix, a_suffix, 222 | channel_names, all_data, normalize, is_cryptomatte, 223 | layers_dict, cryptomatte_dict): 224 | """Process RGB/RGBA type layers with various naming conventions""" 225 | try: 226 | r_channel = f"{group_name}.{r_suffix}" 227 | g_channel = f"{group_name}.{g_suffix}" 228 | b_channel = f"{group_name}.{b_suffix}" 229 | a_channel = f"{group_name}.{a_suffix}" 230 | 231 | try: 232 | r_idx = channel_names.index(r_channel) 233 | g_idx = channel_names.index(g_channel) 234 | b_idx = channel_names.index(b_channel) 235 | except ValueError: 236 | debug_log(logger, "warning", "Missing RGB channels", f"Could not find RGB channels for {group_name}") 237 | return 238 | 239 | has_alpha = False 240 | a_idx = -1 241 | try: 242 | a_idx = channel_names.index(a_channel) 243 | has_alpha = True 244 | except ValueError: 245 | pass 246 | 247 | rgb_array = np.stack([ 248 | all_data[:, :, r_idx], 249 | all_data[:, :, g_idx], 250 | all_data[:, :, b_idx] 251 | ], axis=2) 252 | 253 | rgb_tensor_layer = torch.from_numpy(rgb_array).float() 254 | rgb_tensor_layer = rgb_tensor_layer.unsqueeze(0) # [1, H, W, 3] 255 | 256 | if normalize: 257 | rgb_range = rgb_tensor_layer.max() - rgb_tensor_layer.min() 258 | if rgb_range > 0: 259 | rgb_tensor_layer = (rgb_tensor_layer - rgb_tensor_layer.min()) / rgb_range 260 | 261 | if is_cryptomatte: 262 | cryptomatte_dict[group_name] = rgb_tensor_layer 263 | else: 264 | layers_dict[group_name] = rgb_tensor_layer 265 | 266 | if has_alpha: 267 | alpha_array = all_data[:, :, a_idx] 268 | 269 | alpha_tensor_layer = torch.from_numpy(alpha_array).float() 270 | alpha_tensor_layer = alpha_tensor_layer.unsqueeze(0) # [1, H, W] 271 | 272 | if normalize: 273 | alpha_tensor_layer = alpha_tensor_layer.clamp(0, 1) 274 | 275 | alpha_layer_name = f"{group_name}_alpha" 276 | layers_dict[alpha_layer_name] = alpha_tensor_layer 277 | except ValueError as e: 278 | debug_log(logger, "warning", "Error processing RGB layer", f"Error processing RGB layer {group_name}: {str(e)}") 279 | 280 | @staticmethod 281 | def process_xyz_type_layer(group_name, x_suffix, y_suffix, z_suffix, 282 | channel_names, all_data, normalize, layers_dict): 283 | """Process XYZ type vector layers with various naming conventions""" 284 | try: 285 | x_channel = f"{group_name}.{x_suffix}" 286 | y_channel = f"{group_name}.{y_suffix}" 287 | z_channel = f"{group_name}.{z_suffix}" 288 | 289 | try: 290 | x_idx = channel_names.index(x_channel) 291 | y_idx = channel_names.index(y_channel) 292 | z_idx = channel_names.index(z_channel) 293 | except ValueError: 294 | debug_log(logger, "warning", "Missing XYZ channels", f"Could not find XYZ channels for {group_name}") 295 | return 296 | 297 | xyz_array = np.stack([ 298 | all_data[:, :, x_idx], 299 | all_data[:, :, y_idx], 300 | all_data[:, :, z_idx] 301 | ], axis=2) 302 | 303 | xyz_tensor = torch.from_numpy(xyz_array).float() 304 | xyz_tensor = xyz_tensor.unsqueeze(0) # [1, H, W, 3] 305 | 306 | if normalize: 307 | max_abs = xyz_tensor.abs().max() 308 | if max_abs > 0: 309 | xyz_tensor = xyz_tensor / max_abs 310 | 311 | layers_dict[group_name] = xyz_tensor 312 | except ValueError as e: 313 | debug_log(logger, "warning", "Error processing XYZ layer", f"Error processing XYZ layer {group_name}: {str(e)}") 314 | 315 | @staticmethod 316 | def process_single_channel(group_name, suffixes, group_indices, 317 | channel_names, all_data, normalize, layers_dict): 318 | """Process single channel data like depth maps or Z channels""" 319 | idx = -1 320 | if 'Z' in suffixes: 321 | z_channel = f"{group_name}.Z" 322 | z_channel_lower = f"{group_name}.z" 323 | try: 324 | idx = channel_names.index(z_channel) 325 | except ValueError: 326 | try: 327 | idx = channel_names.index(z_channel_lower) 328 | except ValueError: 329 | idx = group_indices[0] 330 | else: 331 | idx = group_indices[0] 332 | 333 | if idx >= 0: 334 | channel_array = all_data[:, :, idx] 335 | 336 | is_mask_type = any(keyword in group_name.lower() 337 | for keyword in ['depth', 'mask', 'matte', 'alpha', 'id', 'z']) 338 | 339 | if group_name == 'Z': 340 | is_mask_type = True 341 | debug_log(logger, "info", "Processing Z channel as mask", 342 | f"Processing Z channel as mask: shape={channel_array.shape}") 343 | 344 | if is_mask_type: 345 | mask_tensor = torch.from_numpy(channel_array).float().unsqueeze(0) # [1, H, W] 346 | 347 | if normalize: 348 | mask_range = mask_tensor.max() - mask_tensor.min() 349 | if mask_range > 0: 350 | mask_tensor = (mask_tensor - mask_tensor.min()) / mask_range 351 | 352 | debug_log(logger, "info", f"Created mask: {format_tensor_info(mask_tensor.shape, mask_tensor.dtype, group_name)}", 353 | f"Created mask tensor for {group_name}: shape={mask_tensor.shape}, " + 354 | f"min={mask_tensor.min().item():.6f}, max={mask_tensor.max().item():.6f}, " + 355 | f"mean={mask_tensor.mean().item():.6f}") 356 | 357 | layers_dict[group_name] = mask_tensor 358 | else: 359 | rgb_array = np.stack([channel_array] * 3, axis=2) 360 | 361 | channel_tensor = torch.from_numpy(rgb_array).float() 362 | channel_tensor = channel_tensor.unsqueeze(0) # [1, H, W, 3] 363 | 364 | if normalize: 365 | channel_range = channel_tensor.max() - channel_tensor.min() 366 | if channel_range > 0: 367 | channel_tensor = (channel_tensor - channel_tensor.min()) / channel_range 368 | 369 | debug_log(logger, "info", f"Created RGB: {format_tensor_info(channel_tensor.shape, channel_tensor.dtype, group_name)}", 370 | f"Created RGB tensor for {group_name}: shape={channel_tensor.shape}, " + 371 | f"min={channel_tensor.min().item():.6f}, max={channel_tensor.max().item():.6f}, " + 372 | f"mean={channel_tensor.mean().item():.6f}") 373 | 374 | layers_dict[group_name] = channel_tensor 375 | 376 | @staticmethod 377 | def process_multi_channel(group_name, group_indices, all_data, normalize, 378 | is_cryptomatte, layers_dict, cryptomatte_dict): 379 | """Process multi-channel data that doesn't fit standard patterns""" 380 | channels_to_use = min(3, len(group_indices)) 381 | array_channels = [] 382 | 383 | for i in range(channels_to_use): 384 | array_channels.append(all_data[:, :, group_indices[i]]) 385 | 386 | while len(array_channels) < 3: 387 | array_channels.append(array_channels[-1]) 388 | 389 | multi_array = np.stack(array_channels, axis=2) 390 | 391 | multi_tensor = torch.from_numpy(multi_array).float() 392 | multi_tensor = multi_tensor.unsqueeze(0) # [1, H, W, 3] 393 | 394 | if normalize: 395 | multi_range = multi_tensor.max() - multi_tensor.min() 396 | if multi_range > 0: 397 | multi_tensor = (multi_tensor - multi_tensor.min()) / multi_range 398 | 399 | if is_cryptomatte: 400 | cryptomatte_dict[group_name] = multi_tensor 401 | else: 402 | layers_dict[group_name] = multi_tensor 403 | 404 | @staticmethod 405 | def process_layer_groups(channel_groups, cryptomatte_dict, metadata): 406 | """Process groups of related layers (like cryptomatte layer groups)""" 407 | for group_name, suffixes in channel_groups.items(): 408 | if not group_name.endswith('_layer_group'): 409 | continue 410 | 411 | base_name = group_name[:-12] # Remove '_layer_group' 412 | 413 | if not suffixes: 414 | continue 415 | 416 | is_crypto_layer_group = ExrProcessor.is_cryptomatte_layer(base_name) 417 | 418 | in_crypto_dict = any(group_part in cryptomatte_dict for group_part in suffixes) 419 | 420 | if 'layer_groups' not in metadata: 421 | metadata['layer_groups'] = {} 422 | 423 | metadata['layer_groups'][base_name] = suffixes 424 | 425 | if is_crypto_layer_group or in_crypto_dict: 426 | cryptomatte_dict[group_name] = [cryptomatte_dict.get(part, None) for part in suffixes] 427 | 428 | @staticmethod 429 | def store_layer_type_metadata(layers_dict, metadata): 430 | """Store information about layer types in metadata""" 431 | layer_types = {} 432 | for layer_name, tensor in layers_dict.items(): 433 | if len(tensor.shape) >= 4 and tensor.shape[3] == 3: # It has 3 channels 434 | layer_types[layer_name] = "IMAGE" 435 | else: 436 | layer_types[layer_name] = "MASK" 437 | 438 | metadata["layer_types"] = layer_types 439 | 440 | @staticmethod 441 | def create_processed_layer_names(layers_dict, cryptomatte_dict): 442 | """Create a sorted list of processed layer names""" 443 | processed_layer_names = [] 444 | 445 | for layer_name in layers_dict.keys(): 446 | processed_layer_names.append(layer_name) 447 | 448 | for crypto_name in cryptomatte_dict.keys(): 449 | processed_layer_names.append(f"crypto:{crypto_name}") 450 | 451 | processed_layer_names.sort() 452 | 453 | return processed_layer_names 454 | 455 | @staticmethod 456 | def get_channel_groups(channel_names: List[str]) -> Dict[str, List[str]]: 457 | """ 458 | Group channel names by their prefix (before the dot). 459 | Returns a dictionary of groups with their respective channel suffixes. 460 | 461 | This method handles complex naming schemes including: 462 | - Standard RGB/XYZ channels 463 | - Cryptomatte layers and layer groups (like CryptoAsset00, CryptoMaterial00) 464 | - Depth channels with various naming conventions 465 | - Layer groups of related layers (e.g., segmentation, segmentation00, segmentation01) 466 | - Hierarchical naming with multiple dots (e.g., "CITY SCENE.AO.R") 467 | """ 468 | groups = {} 469 | layer_group_prefixes = set() 470 | 471 | for channel in channel_names: 472 | if '.' in channel: 473 | parts = channel.split('.') 474 | 475 | if len(parts) > 2: 476 | prefix = '.'.join(parts[:-1]) 477 | suffix = parts[-1] 478 | else: 479 | prefix, suffix = channel.split('.', 1) 480 | 481 | base_prefix = prefix 482 | if any(prefix.endswith(f"{i:02d}") for i in range(10)): 483 | for i in range(10): 484 | if prefix.endswith(f"{i:02d}"): 485 | base_prefix = prefix[:-2] 486 | layer_group_prefixes.add(base_prefix) 487 | break 488 | 489 | if prefix not in groups: 490 | groups[prefix] = [] 491 | groups[prefix].append(suffix) 492 | else: 493 | if channel not in groups: 494 | groups[channel] = [] 495 | groups[channel].append(None) 496 | 497 | if len(channel) == 1 and channel in 'RGBAXYZ': 498 | if all(c in channel_names for c in 'RGB'): 499 | if 'RGB' not in groups: 500 | groups['RGB'] = [] 501 | elif all(c in channel_names for c in 'XYZ'): 502 | if 'XYZ' not in groups: 503 | groups['XYZ'] = [] 504 | 505 | if all(c in channel_names for c in 'RGB'): 506 | groups['RGB'] = ['R', 'G', 'B'] 507 | 508 | if all(c in channel_names for c in 'XYZ'): 509 | groups['XYZ'] = ['X', 'Y', 'Z'] 510 | 511 | depth_channels = [c for c in channel_names if c in ('Z', 'zDepth', 'zDepth1') or 512 | (('depth' in c.lower() or 'z' in c.lower()) and not '.' in c)] 513 | if depth_channels: 514 | if 'Depth' not in groups: 515 | groups['Depth'] = [] 516 | for dc in depth_channels: 517 | groups['Depth'].append(dc) 518 | 519 | crypto_prefixes = set() 520 | for prefix in groups.keys(): 521 | if ('crypto' in prefix.lower() or prefix.startswith('Crypto')): 522 | base_name = prefix 523 | if any(prefix.endswith(f"{i:02d}") for i in range(10)): 524 | for i in range(10): 525 | if prefix.endswith(f"{i:02d}"): 526 | base_name = prefix[:-2] 527 | crypto_prefixes.add(base_name) 528 | break 529 | 530 | for crypto_base in crypto_prefixes: 531 | if f"{crypto_base}_layer_group" not in groups: 532 | groups[f"{crypto_base}_layer_group"] = [] 533 | 534 | for i in range(10): 535 | group_name = f"{crypto_base}{i:02d}" 536 | if group_name in groups: 537 | groups[f"{crypto_base}_layer_group"].append(group_name) 538 | 539 | for group_base in layer_group_prefixes: 540 | if group_base in crypto_prefixes: 541 | continue 542 | 543 | if f"{group_base}_layer_group" not in groups: 544 | groups[f"{group_base}_layer_group"] = [] 545 | 546 | for i in range(10): 547 | group_name = f"{group_base}{i:02d}" 548 | if group_name in groups: 549 | groups[f"{group_base}_layer_group"].append(group_name) 550 | 551 | return groups 552 | 553 | @staticmethod 554 | def process_exr_data(image_path: str, normalize: bool, node_id: str = None, layer_data: Dict = None) -> List: 555 | """ 556 | Main EXR processing function that handles all layer extraction and processing. 557 | This replaces the _load_single_image functionality from the original LoadExr class. 558 | """ 559 | if image_path is None: 560 | raise ValueError("image_path cannot be None. This may indicate a missing frame in the sequence.") 561 | 562 | try: 563 | metadata = layer_data if layer_data else ExrProcessor.scan_exr_metadata(image_path) 564 | 565 | all_subimage_data = ExrProcessor.load_all_data(image_path) 566 | 567 | layers_dict = {} 568 | cryptomatte_dict = {} 569 | all_channel_names = [] 570 | 571 | for subimage_idx, subimage_info in enumerate(metadata["subimages"]): 572 | if subimage_idx not in all_subimage_data: 573 | debug_log(logger, "warning", "No subimage data", f"No data found for subimage {subimage_idx}") 574 | continue 575 | 576 | subimage_data = all_subimage_data[subimage_idx] 577 | subimage_name = subimage_info["name"] 578 | channel_names = subimage_info["channel_names"] 579 | 580 | all_channel_names.extend(channel_names) 581 | 582 | height, width, channels = subimage_data.shape 583 | 584 | if subimage_idx == 0: 585 | channel_groups = ExrProcessor.get_channel_groups(channel_names) 586 | metadata["channel_groups"] = channel_groups 587 | 588 | rgb_tensor, alpha_tensor = ExrProcessor.process_default_channels( 589 | subimage_data, channel_names, height, width, normalize 590 | ) 591 | 592 | if subimage_name != "default": 593 | if channels >= 3: 594 | rgb_array = subimage_data[:, :, :3] 595 | 596 | rgb_tensor_layer = torch.from_numpy(rgb_array).float() 597 | rgb_tensor_layer = rgb_tensor_layer.unsqueeze(0) # [1, H, W, 3] 598 | 599 | if normalize: 600 | rgb_range = rgb_tensor_layer.max() - rgb_tensor_layer.min() 601 | if rgb_range > 0: 602 | rgb_tensor_layer = (rgb_tensor_layer - rgb_tensor_layer.min()) / rgb_range 603 | 604 | layers_dict[subimage_name] = rgb_tensor_layer 605 | 606 | if channels >= 4: 607 | alpha_array = subimage_data[:, :, 3] 608 | 609 | alpha_tensor_layer = torch.from_numpy(alpha_array).float() 610 | alpha_tensor_layer = alpha_tensor_layer.unsqueeze(0) # [1, H, W] 611 | 612 | if normalize: 613 | alpha_tensor_layer = alpha_tensor_layer.clamp(0, 1) 614 | 615 | layers_dict[f"{subimage_name}_alpha"] = alpha_tensor_layer 616 | 617 | elif channels == 1: 618 | channel_array = subimage_data[:, :, 0] 619 | 620 | is_mask_type = any(keyword in subimage_name.lower() 621 | for keyword in ['depth', 'mask', 'matte', 'alpha', 'id', 'z']) 622 | 623 | if is_mask_type or subimage_name == 'depth': 624 | mask_tensor = torch.from_numpy(channel_array).float().unsqueeze(0) # [1, H, W] 625 | 626 | if normalize: 627 | mask_range = mask_tensor.max() - mask_tensor.min() 628 | if mask_range > 0: 629 | mask_tensor = (mask_tensor - mask_tensor.min()) / mask_range 630 | 631 | if mask_tensor.numel() > 1: 632 | layers_dict[subimage_name] = mask_tensor 633 | else: 634 | layers_dict[subimage_name] = torch.zeros((1, height, width)) 635 | else: 636 | rgb_array = np.stack([channel_array] * 3, axis=2) 637 | 638 | channel_tensor = torch.from_numpy(rgb_array).float() 639 | channel_tensor = channel_tensor.unsqueeze(0) # [1, H, W, 3] 640 | 641 | if normalize: 642 | channel_range = channel_tensor.max() - channel_tensor.min() 643 | if channel_range > 0: 644 | channel_tensor = (channel_tensor - channel_tensor.min()) / channel_range 645 | 646 | if channel_tensor.numel() > 3: 647 | layers_dict[subimage_name] = channel_tensor 648 | else: 649 | layers_dict[subimage_name] = torch.zeros((1, height, width, 3)) 650 | 651 | if subimage_idx == 0: 652 | channel_groups = ExrProcessor.get_channel_groups(channel_names) 653 | 654 | for group_name, suffixes in channel_groups.items(): 655 | if group_name in ('R', 'G', 'B', 'A', 'RGB', 'XYZ'): 656 | continue 657 | 658 | if group_name.endswith('_layer_group'): 659 | continue 660 | 661 | is_cryptomatte = ExrProcessor.is_cryptomatte_layer(group_name) 662 | 663 | group_indices = [] 664 | for i, channel in enumerate(channel_names): 665 | if (channel == group_name) or (channel.startswith(f"{group_name}.")): 666 | group_indices.append(i) 667 | 668 | if not group_indices: 669 | continue 670 | 671 | if all(suffix in suffixes for suffix in ['R', 'G', 'B']): 672 | ExrProcessor.process_rgb_type_layer( 673 | group_name, 'R', 'G', 'B', 'A', channel_names, subimage_data, 674 | normalize, is_cryptomatte, layers_dict, cryptomatte_dict 675 | ) 676 | 677 | elif all(suffix in suffixes for suffix in ['r', 'g', 'b']): 678 | ExrProcessor.process_rgb_type_layer( 679 | group_name, 'r', 'g', 'b', 'a', channel_names, subimage_data, 680 | normalize, is_cryptomatte, layers_dict, cryptomatte_dict 681 | ) 682 | 683 | elif all(suffix in suffixes for suffix in ['X', 'Y', 'Z']): 684 | ExrProcessor.process_xyz_type_layer( 685 | group_name, 'X', 'Y', 'Z', channel_names, subimage_data, 686 | normalize, layers_dict 687 | ) 688 | 689 | elif all(suffix in suffixes for suffix in ['x', 'y', 'z']): 690 | ExrProcessor.process_xyz_type_layer( 691 | group_name, 'x', 'y', 'z', channel_names, subimage_data, 692 | normalize, layers_dict 693 | ) 694 | 695 | elif len(group_indices) == 1 or 'Z' in suffixes: 696 | ExrProcessor.process_single_channel( 697 | group_name, suffixes, group_indices, channel_names, 698 | subimage_data, normalize, layers_dict 699 | ) 700 | 701 | else: 702 | ExrProcessor.process_multi_channel( 703 | group_name, group_indices, subimage_data, normalize, 704 | is_cryptomatte, layers_dict, cryptomatte_dict 705 | ) 706 | 707 | ExrProcessor.process_layer_groups( 708 | channel_groups, cryptomatte_dict, metadata 709 | ) 710 | 711 | ExrProcessor.store_layer_type_metadata(layers_dict, metadata) 712 | 713 | metadata_json = json.dumps(metadata) 714 | 715 | debug_log(logger, "info", f"Loaded {len(layers_dict)} layers: {format_layer_names(list(layers_dict.keys()))}", 716 | f"Available EXR layers: {list(layers_dict.keys())}") 717 | if cryptomatte_dict: 718 | debug_log(logger, "info", f"Loaded {len(cryptomatte_dict)} cryptomatte layers", 719 | f"Available cryptomatte layers: {list(cryptomatte_dict.keys())}") 720 | 721 | layer_names = all_channel_names 722 | 723 | processed_layer_names = ExrProcessor.create_processed_layer_names(layers_dict, cryptomatte_dict) 724 | 725 | preview_result = generate_preview_for_comfyui(rgb_tensor, image_path, is_sequence=False, frame_index=0) 726 | 727 | result = [rgb_tensor, alpha_tensor, cryptomatte_dict, layers_dict, processed_layer_names, layer_names, metadata_json] 728 | 729 | if preview_result: 730 | return {"ui": {"images": preview_result}, "result": result} 731 | else: 732 | return result 733 | 734 | except Exception as e: 735 | debug_log(logger, "error", "Error loading EXR", f"Error loading EXR file {image_path}: {str(e)}") 736 | raise --------------------------------------------------------------------------------