""" Base node class for Eric's Image Processing Nodes Handles common ComfyUI integration, tensor conversions, and device management Author: Eric Hiss (GitHub: EricRollei) License: See LICENSE file in repository root This module provides the base functionality for all image processing nodes: - ComfyUI tensor format conversion (NHWC [0-1] ↔ numpy [0-255]) - GPU device detection and management - Memory cleanup utilities - OpenCV 4.11+ compatibility (ensures contiguous arrays) Dependencies: - PyTorch (BSD 3-Clause License) - NumPy (BSD 3-Clause License) """ import torch import numpy as np from typing import Tuple, Optional, Union import gc class BaseImageProcessingNode: """Base class for all image processing nodes with ComfyUI integration""" CATEGORY = "Eric's Image Processing" @staticmethod def tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray: """Convert ComfyUI image tensor (NHWC) to numpy array (NHWC or HWC) Args: tensor: ComfyUI image tensor in format [N, H, W, C] or [H, W, C] with values 0-1 Returns: numpy array in format [N, H, W, C] or [H, W, C] with values 0-255 (uint8) GUARANTEED to be contiguous for OpenCV 4.11+ compatibility """ # CRITICAL: Ensure tensor is contiguous before numpy conversion # PyTorch operations like permute() can create non-contiguous tensors # OpenCV 4.11+ strictly requires contiguous arrays tensor = tensor.contiguous() # Keep batch dimension if present img_np = tensor.cpu().numpy() # Convert to numpy and scale to 0-255 img_np = (img_np * 255).astype(np.uint8) # CRITICAL: Ensure final array is contiguous # Arithmetic operations and astype() can create non-contiguous arrays img_np = np.ascontiguousarray(img_np) return img_np @staticmethod def numpy_to_tensor(img_np: np.ndarray) -> torch.Tensor: """Convert numpy array back to ComfyUI tensor format Args: img_np: numpy array in format [N, H, W, C] or [H, W, C] with values 0-255 (uint8) or 0-1 (float32) Returns: ComfyUI tensor in format [N, H, W, C] or [1, H, W, C] with values 0-1 """ # CRITICAL: Ensure input array is contiguous before ANY operations # Processing functions may return non-contiguous arrays img_np = np.ascontiguousarray(img_np) # Handle different input formats if img_np.dtype == np.uint8: # Convert uint8 [0,255] to float32 [0,1] img_tensor = torch.from_numpy(img_np.astype(np.float32) / 255.0) elif img_np.dtype in [np.float32, np.float64]: # Handle float input - check if it's in [0,1] or [0,255] range if img_np.max() <= 1.0: # Already in [0,1] range img_tensor = torch.from_numpy(img_np.astype(np.float32)) else: # Assume it's in [0,255] range img_tensor = torch.from_numpy(img_np.astype(np.float32) / 255.0) else: # Convert to uint8 first, then to float32 img_np = np.clip(img_np, 0, 255).astype(np.uint8) img_np = np.ascontiguousarray(img_np) # Ensure contiguous after clip img_tensor = torch.from_numpy(img_np.astype(np.float32) / 255.0) # Add batch dimension if not present (single image) if len(img_tensor.shape) == 3: img_tensor = img_tensor.unsqueeze(0) return img_tensor @staticmethod def get_device_info() -> Tuple[str, bool]: """Get optimal device for processing Returns: Tuple of (device_name, gpu_available) """ if torch.cuda.is_available(): device = f"cuda:{torch.cuda.current_device()}" gpu_available = True else: device = "cpu" gpu_available = False return device, gpu_available @staticmethod def cleanup_memory(): """Clean up GPU memory after processing""" if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() @staticmethod def validate_image_tensor(tensor: torch.Tensor) -> bool: """Validate that tensor is a proper ComfyUI image tensor Args: tensor: Input tensor to validate Returns: True if valid, raises ValueError if not """ if not isinstance(tensor, torch.Tensor): raise ValueError("Input must be a torch.Tensor") if len(tensor.shape) not in [3, 4]: raise ValueError(f"Image tensor must have 3 or 4 dimensions, got {len(tensor.shape)}") if len(tensor.shape) == 4: n, h, w, c = tensor.shape if c not in [1, 3, 4]: raise ValueError(f"Image must have 1, 3, or 4 channels, got {c}") else: h, w, c = tensor.shape if c not in [1, 3, 4]: raise ValueError(f"Image must have 1, 3, or 4 channels, got {c}") if tensor.min() < 0 or tensor.max() > 1: raise ValueError("Image tensor values must be in range [0, 1]") return True def process_image_safe(self, image: torch.Tensor, processing_func, **kwargs) -> torch.Tensor: """Safely process an image with error handling and memory cleanup Args: image: Input image tensor [N, H, W, C] or [H, W, C] processing_func: Function to apply to the image (expects [H, W, C]) **kwargs: Additional arguments for processing_func Returns: Processed image tensor [N, H, W, C] or [H, W, C] """ try: # Validate input self.validate_image_tensor(image) # Handle batch dimension properly is_batch = len(image.shape) == 4 if is_batch: # Process each image in the batch batch_size = image.shape[0] results = [] for i in range(batch_size): # Get single image [H, W, C] single_img = image[i] # Convert to numpy img_np = self.tensor_to_numpy(single_img) # Apply processing (expects 2D or 3D array) processed_np = processing_func(img_np, **kwargs) # Convert back to tensor result_tensor = self.numpy_to_tensor(processed_np) # Remove batch dimension if added by numpy_to_tensor if len(result_tensor.shape) == 4: result_tensor = result_tensor[0] results.append(result_tensor) # Stack results back into batch result = torch.stack(results, dim=0) else: # Single image [H, W, C] # Convert to numpy img_np = self.tensor_to_numpy(image) # Apply processing processed_np = processing_func(img_np, **kwargs) # Convert back to tensor result = self.numpy_to_tensor(processed_np) # Remove batch dimension if added by numpy_to_tensor if len(result.shape) == 4 and result.shape[0] == 1: result = result[0] return result except Exception as e: print(f"Error in image processing: {str(e)}") raise finally: # Clean up memory self.cleanup_memory()