From 0d140aa2b68e95e409e2e9a372cf62ac229048a4 Mon Sep 17 00:00:00 2001 From: Regi Ellis Date: Wed, 25 Jun 2025 17:53:36 -0400 Subject: [PATCH] Backen code --- __init__.py | 6 +- src/__init__.py | 1219 ++++++++++++++++++++++++++++++++--------------- 2 files changed, 834 insertions(+), 391 deletions(-) diff --git a/__init__.py b/__init__.py index 837e3bc..820b485 100644 --- a/__init__.py +++ b/__init__.py @@ -13,9 +13,9 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { "EasyColorCorrection": "Color Corrector", - "BatchColorCorrection": "Batch Color Corrector", - "RawImageProcessor": "RAW Image Processor", - "ColorCorrectionViewer": "Color Correction Viewer", + "BatchColorCorrection": "Batch Color Corrector (beta)", + "RawImageProcessor": "RAW Image Processor (beta)", + "ColorCorrectionViewer": "Color Correction Viewer (beta)", } __all__ = [ diff --git a/src/__init__.py b/src/__init__.py index 5ed401f..d9b7093 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -23,6 +23,7 @@ except ImportError as e: # RAW IMAGE PROCESSING try: import rawpy + RAW_PROCESSING_AVAILABLE = True except ImportError as e: print(f"EasyColorCorrection: RAW processing disabled - missing rawpy library: {e}") @@ -32,16 +33,22 @@ except ImportError as e: try: import OpenEXR import Imath + EXR_PROCESSING_AVAILABLE = True except ImportError as e: - print(f"EasyColorCorrection: EXR processing disabled - missing OpenEXR library: {e}") + print( + f"EasyColorCorrection: EXR processing disabled - missing OpenEXR library: {e}" + ) EXR_PROCESSING_AVAILABLE = False try: import imageio + IMAGEIO_AVAILABLE = True except ImportError as e: - print(f"EasyColorCorrection: Extended format support disabled - missing imageio library: {e}") + print( + f"EasyColorCorrection: Extended format support disabled - missing imageio library: {e}" + ) IMAGEIO_AVAILABLE = False @@ -211,18 +218,12 @@ def match_to_reference_colors( img_mean = np.mean(lab_image[:, :, i]) ref_mean = np.mean(lab_reference[:, :, i]) - color_shift = ( - (ref_mean - img_mean) * strength * 0.3 - ) + color_shift = (ref_mean - img_mean) * strength * 0.3 lab_image[:, :, i] = lab_image[:, :, i] + color_shift lab_image[:, :, 0] = np.clip(lab_image[:, :, 0], 0, 100) - lab_image[:, :, 1] = np.clip( - lab_image[:, :, 1], -127, 127 - ) - lab_image[:, :, 2] = np.clip( - lab_image[:, :, 2], -127, 127 - ) + lab_image[:, :, 1] = np.clip(lab_image[:, :, 1], -127, 127) + lab_image[:, :, 2] = np.clip(lab_image[:, :, 2], -127, 127) lab_image_uint8 = lab_image.astype(np.uint8) matched_image = cv2.cvtColor(lab_image_uint8, cv2.COLOR_LAB2RGB) @@ -244,6 +245,7 @@ def get_preferred_device(use_gpu: bool = True): return torch.device("cpu") return torch.device("cuda" if torch.cuda.is_available() else "cpu") + def analyze_image_content( image_np: np.ndarray, device: torch.device = None ) -> typing.Dict[str, typing.Any]: @@ -428,7 +430,7 @@ def intelligent_white_balance( image_np: np.ndarray, strength: float = 0.6 ) -> np.ndarray: """Advanced white balance using professional color science. - + Args: strength: -1.0 to 1.0, where negative values enhance cool/blue tones, positive values enhance warm/orange tones, 0.0 = neutral @@ -447,16 +449,20 @@ def intelligent_white_balance( a_shift = np.median(a_channel) b_shift = np.median(b_channel) # Reduce tint influence since we have separate tint control - lab[:, :, 1] = np.clip(lab[:, :, 1] - a_shift * strength * 0.3, 0, 255) # Reduced tint correction - lab[:, :, 2] = np.clip(lab[:, :, 2] - b_shift * strength, 0, 255) # Full temperature correction + lab[:, :, 1] = np.clip( + lab[:, :, 1] - a_shift * strength * 0.3, 0, 255 + ) # Reduced tint correction + lab[:, :, 2] = np.clip( + lab[:, :, 2] - b_shift * strength, 0, 255 + ) # Full temperature correction else: # Negative strength: enhance cool tones (blue/cyan) abs_strength = abs(strength) # Focus primarily on b channel (temperature) for cool adjustment cool_b_shift = -25 * abs_strength # Toward blue (cooler temperature) # Minimal a channel adjustment since we have separate tint control - cool_a_shift = -5 * abs_strength # Slight toward green (cooler) - + cool_a_shift = -5 * abs_strength # Slight toward green (cooler) + lab[:, :, 1] = np.clip(lab[:, :, 1] + cool_a_shift, 0, 255) lab[:, :, 2] = np.clip(lab[:, :, 2] + cool_b_shift, 0, 255) @@ -475,9 +481,7 @@ def intelligent_white_balance( "D65" ] - if ( - np.linalg.norm(illuminant_xy - d65_xy) > 0.01 - ): + if np.linalg.norm(illuminant_xy - d65_xy) > 0.01: adapted_xyz = colour.chromatic_adaptation_VonKries( xyz, illuminant_xy[0], d65_xy, transform="Bradford" ) @@ -607,15 +611,9 @@ def enhance_faces( a_channel = face_lab[:, :, 1].astype(np.float32) b_channel = face_lab[:, :, 2].astype(np.float32) - very_dark_skin = (l_channel < 60) & ( - a_channel > 132 - ) - dark_skin = ( - (l_channel < 90) & (l_channel >= 60) & (a_channel > 128) - ) - very_light_skin = ( - l_channel > 180 - ) + very_dark_skin = (l_channel < 60) & (a_channel > 132) + dark_skin = (l_channel < 90) & (l_channel >= 60) & (a_channel > 128) + very_light_skin = l_channel > 180 brightness_adjustment = np.where( very_dark_skin, @@ -1286,7 +1284,10 @@ class EasyColorCorrection: if realtime_preview: # Check if the input image has changed (cache invalidation) image_changed = False - if not hasattr(self, '_cached_original_image') or self._cached_original_image is None: + if ( + not hasattr(self, "_cached_original_image") + or self._cached_original_image is None + ): image_changed = True else: # Compare image tensors to detect changes from upstream @@ -1296,34 +1297,40 @@ class EasyColorCorrection: except: # Different shapes or other comparison issues = definitely changed image_changed = True - + if image_changed: self._cached_original_image = original_image.clone() self._cached_analysis = None print("๐Ÿ”„ Realtime Preview: New image detected, updating cache") - + processed_image = self._cached_original_image.clone() else: processed_image = image.clone() # Clear cache when not in realtime mode - if hasattr(self, '_cached_original_image'): + if hasattr(self, "_cached_original_image"): self._cached_original_image = None self._cached_analysis = None # --- SHARED AI ANALYSIS (available to all modes) --- # Always create image_np for potential use in processing functions image_np = (processed_image[0].cpu().numpy() * 255).astype(np.uint8) - + analysis = None if ai_analysis and ADVANCED_LIBS_AVAILABLE: # Use cached analysis in realtime preview mode - if realtime_preview and hasattr(self, '_cached_analysis') and self._cached_analysis is not None: + if ( + realtime_preview + and hasattr(self, "_cached_analysis") + and self._cached_analysis is not None + ): analysis = self._cached_analysis else: - analysis = analyze_image_content(image_np, get_preferred_device(use_gpu)) + analysis = analyze_image_content( + image_np, get_preferred_device(use_gpu) + ) if realtime_preview: self._cached_analysis = analysis - + print( f"๐Ÿค– AI Analysis{'(cached)' if realtime_preview and hasattr(self, '_cached_analysis') else ''} for {mode} Mode: {analysis['scene_type']} scene, {analysis['lighting']} lighting, {len(analysis['faces'])} faces detected" ) @@ -1380,7 +1387,7 @@ class EasyColorCorrection: print("๐Ÿ”ง Basic Auto Mode: AI analysis disabled") if white_balance_strength != 0.0: B, C = processed_image.shape[0], processed_image.shape[3] - + if white_balance_strength > 0.0: # Positive: traditional white balance (neutralize) flat_image = processed_image.view(B, -1, C) @@ -1404,11 +1411,17 @@ class EasyColorCorrection: # Negative: enhance cool tones (shift toward blue) abs_strength = abs(white_balance_strength) # Cool shift: reduce red, enhance blue - cool_scale = torch.tensor([0.95, 1.0, 1.08], device=processed_image.device, dtype=processed_image.dtype) - cool_scale = torch.lerp(torch.ones_like(cool_scale), cool_scale, abs_strength) + cool_scale = torch.tensor( + [0.95, 1.0, 1.08], + device=processed_image.device, + dtype=processed_image.dtype, + ) + cool_scale = torch.lerp( + torch.ones_like(cool_scale), cool_scale, abs_strength + ) cool_scale = cool_scale.view(1, 1, 1, 3) processed_image = processed_image * cool_scale - + processed_image = torch.clamp(processed_image, 0.0, 1.0) hsv_enhanced = rgb_to_hsv(processed_image) h_enh, s_enh, v_enh = ( @@ -1497,9 +1510,7 @@ class EasyColorCorrection: # --- PRESET MODE --- if mode == "Preset": - p_vals = self.PRESETS.get( - preset, {} - ).copy() + p_vals = self.PRESETS.get(preset, {}).copy() if analysis: scene_type = analysis["scene_type"] lighting = analysis["lighting"] @@ -1507,9 +1518,7 @@ class EasyColorCorrection: if scene_type == "concept_art": p_vals["vibrancy"] = p_vals.get("vibrancy", 0.0) * 1.4 p_vals["contrast"] = p_vals.get("contrast", 0.0) * 1.3 - p_vals["brightness"] = ( - p_vals.get("brightness", 0.0) + 0.05 - ) + p_vals["brightness"] = p_vals.get("brightness", 0.0) + 0.05 elif scene_type == "anime": p_vals["vibrancy"] = p_vals.get("vibrancy", 0.0) * 1.3 p_vals["contrast"] = p_vals.get("contrast", 0.0) * 1.2 @@ -1521,9 +1530,7 @@ class EasyColorCorrection: p_vals["contrast"] = p_vals.get("contrast", 0.0) * 1.1 elif scene_type == "portrait" and analysis["faces"]: p_vals["vibrancy"] = p_vals.get("vibrancy", 0.0) * 0.8 - p_vals["warmth"] = ( - p_vals.get("warmth", 0.0) + 0.05 - ) + p_vals["warmth"] = p_vals.get("warmth", 0.0) + 0.05 elif scene_type == "realistic_photo": p_vals["contrast"] = p_vals.get("contrast", 0.0) * 1.1 @@ -1554,8 +1561,12 @@ class EasyColorCorrection: # Apply preset values with intelligent variation (only for Preset mode) if mode == "Preset": warmth = p_vals.get("warmth", 0.0) + (torch.randn(1).item() * v_factor) - vibrancy = p_vals.get("vibrancy", 0.0) + (torch.randn(1).item() * v_factor) - contrast = p_vals.get("contrast", 0.0) + (torch.randn(1).item() * v_factor) + vibrancy = p_vals.get("vibrancy", 0.0) + ( + torch.randn(1).item() * v_factor + ) + contrast = p_vals.get("contrast", 0.0) + ( + torch.randn(1).item() * v_factor + ) brightness = p_vals.get("brightness", 0.0) + ( torch.randn(1).item() * v_factor ) @@ -1564,7 +1575,11 @@ class EasyColorCorrection: # --- ADVANCED COLOR PROCESSING (Preset and Manual modes) --- if mode != "Auto": # === INTELLIGENT WHITE BALANCE (if enabled) === - if white_balance_strength != 0.0 and ai_analysis and ADVANCED_LIBS_AVAILABLE: + if ( + white_balance_strength != 0.0 + and ai_analysis + and ADVANCED_LIBS_AVAILABLE + ): wb_corrected = intelligent_white_balance( image_np, white_balance_strength ) @@ -1612,45 +1627,65 @@ class EasyColorCorrection: # === TEMPERATURE & TINT PROCESSING (LAB color space) === if (warmth != 0.0 or tint != 0.0) and mode == "Manual": # Convert to numpy for LAB processing - image_np_for_color = (processed_image[0].cpu().numpy() * 255).astype(np.uint8) - + image_np_for_color = (processed_image[0].cpu().numpy() * 255).astype( + np.uint8 + ) + if ADVANCED_LIBS_AVAILABLE: try: # Convert to LAB color space lab = cv2.cvtColor(image_np_for_color, cv2.COLOR_RGB2LAB) - + # Apply temperature adjustment to 'b' channel (blue-yellow axis) if warmth != 0.0: # 'b' channel: values around 128 are neutral, <128 is blue, >128 is yellow - temperature_shift = warmth * 35 # Scale factor for visible temperature effect - lab[:, :, 2] = np.clip(lab[:, :, 2] + temperature_shift, 0, 255) - + temperature_shift = ( + warmth * 35 + ) # Scale factor for visible temperature effect + lab[:, :, 2] = np.clip( + lab[:, :, 2] + temperature_shift, 0, 255 + ) + # Apply tint adjustment to 'a' channel (green-magenta axis) if tint != 0.0: # 'a' channel: values around 128 are neutral, <128 is green, >128 is magenta - tint_shift = tint * 30 # Scale factor for visible tint effect + tint_shift = ( + tint * 30 + ) # Scale factor for visible tint effect lab[:, :, 1] = np.clip(lab[:, :, 1] + tint_shift, 0, 255) - + # Convert back to RGB color_corrected_rgb = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB) processed_image = ( - torch.from_numpy(color_corrected_rgb.astype(np.float32) / 255.0) + torch.from_numpy( + color_corrected_rgb.astype(np.float32) / 255.0 + ) .unsqueeze(0) .to(processed_image.device) ) processed_image = torch.clamp(processed_image, 0.0, 1.0) - + adjustments = [] if warmth != 0.0: - adjustments.append(f"temperature: {warmth:.2f} ({'warmer' if warmth > 0 else 'cooler'})") + adjustments.append( + f"temperature: {warmth:.2f} ({'warmer' if warmth > 0 else 'cooler'})" + ) if tint != 0.0: - adjustments.append(f"tint: {tint:.2f} ({'magenta' if tint > 0 else 'green'})") - print(f"๐ŸŒก๏ธ Applied professional color adjustments: {', '.join(adjustments)}") - + adjustments.append( + f"tint: {tint:.2f} ({'magenta' if tint > 0 else 'green'})" + ) + print( + f"๐ŸŒก๏ธ Applied professional color adjustments: {', '.join(adjustments)}" + ) + except Exception as e: - print(f"โš ๏ธ Temperature/tint processing failed, using original: {e}") + print( + f"โš ๏ธ Temperature/tint processing failed, using original: {e}" + ) else: - print("โš ๏ธ Professional temperature/tint requires advanced libraries (OpenCV), skipping") + print( + "โš ๏ธ Professional temperature/tint requires advanced libraries (OpenCV), skipping" + ) # --- MANUAL MODE --- if mode == "Manual": @@ -1674,13 +1709,9 @@ class EasyColorCorrection: "stylized_art", "detailed_illustration", ]: - shadows_mask = 1.0 - torch.clamp( - v * 2.5, 0.0, 1.0 - ) + shadows_mask = 1.0 - torch.clamp(v * 2.5, 0.0, 1.0) midtones_mask = 1.0 - torch.abs(v - 0.5) * 1.8 - highlights_mask = torch.clamp( - (v - 0.6) * 2.5, 0.0, 1.0 - ) + highlights_mask = torch.clamp((v - 0.6) * 2.5, 0.0, 1.0) else: shadows_mask = 1.0 - torch.clamp(v * 3.0, 0.0, 1.0) midtones_mask = 1.0 - torch.abs(v - 0.5) * 2.0 @@ -1694,7 +1725,9 @@ class EasyColorCorrection: if gamma != 0.0: # Midtones (gamma) - reduced strength for better control - gamma_exp = 1.0 / (1.0 + gamma * 0.6) # Reduced from 0.8-1.0 for more subtle control + gamma_exp = 1.0 / ( + 1.0 + gamma * 0.6 + ) # Reduced from 0.8-1.0 for more subtle control v_gamma = torch.pow(torch.clamp(v, 0.001, 1.0), gamma_exp) v = torch.lerp(v, v_gamma, midtones_mask) @@ -1722,23 +1755,46 @@ class EasyColorCorrection: # Smart detection to avoid colorizing art/anime if analysis: scene_type = analysis["scene_type"] - if scene_type in ["anime", "concept_art", "stylized_art", "detailed_illustration"]: - print(f"๐Ÿšซ Colorize Mode: Skipping colorization for {scene_type} content") + if scene_type in [ + "anime", + "concept_art", + "stylized_art", + "detailed_illustration", + ]: + print( + f"๐Ÿšซ Colorize Mode: Skipping colorization for {scene_type} content" + ) # Return original image for art content processed_image = original_image.clone() else: - print(f"๐ŸŽจ Colorize Mode: Processing {scene_type} for intelligent colorization") + print( + f"๐ŸŽจ Colorize Mode: Processing {scene_type} for intelligent colorization" + ) processed_image = self._apply_colorization( - original_image, processed_image, analysis, - colorize_strength, skin_warmth, sky_saturation, - vegetation_green, sepia_tone, colorize_mode + original_image, + processed_image, + analysis, + colorize_strength, + skin_warmth, + sky_saturation, + vegetation_green, + sepia_tone, + colorize_mode, ) else: - print("๐ŸŽจ Colorize Mode: Applying general colorization (no AI analysis)") + print( + "๐ŸŽจ Colorize Mode: Applying general colorization (no AI analysis)" + ) processed_image = self._apply_colorization( - original_image, processed_image, None, - colorize_strength, skin_warmth, sky_saturation, - vegetation_green, sepia_tone, colorize_mode + original_image, + processed_image, + None, + colorize_strength, + skin_warmth, + sky_saturation, + vegetation_green, + sepia_tone, + colorize_mode, ) # --- SKIN TONE PROTECTION --- @@ -1759,9 +1815,7 @@ class EasyColorCorrection: h = torch.lerp(h, h_orig, skin_mask * 0.8) s = torch.lerp(s, s_orig, skin_mask * 0.6) - v = torch.lerp( - v, v_orig, skin_mask * 0.3 - ) + v = torch.lerp(v, v_orig, skin_mask * 0.3) s = torch.clamp(s, 0.0, 1.0) v = torch.clamp(v, 0.0, 1.0) @@ -1856,116 +1910,133 @@ class EasyColorCorrection: return (processed_image, palette_data, histogram_tensor, palette_image_tensor) - def _apply_colorization(self, original_image, processed_image, analysis, - colorize_strength, skin_warmth, sky_saturation, - vegetation_green, sepia_tone, colorize_mode): + def _apply_colorization( + self, + original_image, + processed_image, + analysis, + colorize_strength, + skin_warmth, + sky_saturation, + vegetation_green, + sepia_tone, + colorize_mode, + ): """ Apply intelligent colorization to grayscale or desaturated photos. Uses GPU-optimized tensor operations for efficient processing. """ device = original_image.device - + # Convert to HSV for color manipulation hsv_image = rgb_to_hsv(processed_image) h, s, v = hsv_image[..., 0], hsv_image[..., 1], hsv_image[..., 2] - + # Check if image is grayscale or very desaturated avg_saturation = torch.mean(s).item() if avg_saturation > 0.3: - print(f"โš ๏ธ Image already has color (avg saturation: {avg_saturation:.2f}), applying gentle enhancement") + print( + f"โš ๏ธ Image already has color (avg saturation: {avg_saturation:.2f}), applying gentle enhancement" + ) colorize_strength *= 0.3 # Reduce strength for already colored images - + # Create region masks based on luminance and edge detection luminance = torch.mean(processed_image, dim=-1, keepdim=True) - + # Sky detection (upper regions with high luminance) height = processed_image.shape[1] sky_region = torch.zeros_like(luminance, device=device) sky_upper_third = height // 3 sky_region[:, :sky_upper_third, :, :] = 1.0 - + # Enhance sky detection with luminance sky_luminance_mask = (luminance > 0.7).float() sky_mask = sky_region * sky_luminance_mask - + # Vegetation detection (mid-luminance areas, typically green regions) vegetation_mask = ((luminance > 0.2) & (luminance < 0.8)).float() vegetation_mask = vegetation_mask * (1.0 - sky_mask) # Exclude sky areas - + # Skin tone detection (mid-luminance warm areas) skin_mask = torch.zeros_like(luminance, device=device) if analysis and analysis.get("faces"): # Use luminance-based approximation for skin areas skin_mask = ((luminance > 0.25) & (luminance < 0.85)).float() skin_mask = skin_mask * (1.0 - sky_mask) * (1.0 - vegetation_mask) - + # Apply colorization based on mode if colorize_mode == "vintage": # Vintage sepia-toned colorization base_hue = 0.08 # Warm sepia hue h = torch.full_like(h, base_hue) s = s + sepia_tone * 0.4 * colorize_strength - + elif colorize_mode == "portrait": # Portrait-focused colorization # Warm skin tones skin_hue = 0.08 # Warm skin hue h = torch.where(skin_mask > 0.3, skin_hue, h) s = torch.where(skin_mask > 0.3, s + skin_warmth * colorize_strength, s) - + # Subtle sky blues sky_hue = 0.58 # Blue hue h = torch.where(sky_mask > 0.5, sky_hue, h) - s = torch.where(sky_mask > 0.5, s + sky_saturation * 0.3 * colorize_strength, s) - + s = torch.where( + sky_mask > 0.5, s + sky_saturation * 0.3 * colorize_strength, s + ) + elif colorize_mode == "landscape": # Landscape-focused colorization # Green vegetation vegetation_hue = 0.25 # Green hue h = torch.where(vegetation_mask > 0.4, vegetation_hue, h) - s = torch.where(vegetation_mask > 0.4, s + vegetation_green * colorize_strength, s) - + s = torch.where( + vegetation_mask > 0.4, s + vegetation_green * colorize_strength, s + ) + # Blue skies sky_hue = 0.58 # Blue hue h = torch.where(sky_mask > 0.5, sky_hue, h) s = torch.where(sky_mask > 0.5, s + sky_saturation * colorize_strength, s) - + else: # auto mode # Intelligent auto colorization # Sky areas -> blue sky_hue = 0.58 h = torch.where(sky_mask > 0.5, sky_hue, h) s = torch.where(sky_mask > 0.5, s + sky_saturation * colorize_strength, s) - + # Vegetation areas -> green vegetation_hue = 0.25 h = torch.where(vegetation_mask > 0.4, vegetation_hue, h) - s = torch.where(vegetation_mask > 0.4, s + vegetation_green * colorize_strength, s) - + s = torch.where( + vegetation_mask > 0.4, s + vegetation_green * colorize_strength, s + ) + # Skin areas -> warm tones if torch.sum(skin_mask) > 0: skin_hue = 0.08 h = torch.where(skin_mask > 0.3, skin_hue, h) s = torch.where(skin_mask > 0.3, s + skin_warmth * colorize_strength, s) - + # Apply sepia tone if specified if sepia_tone > 0: sepia_hue = 0.08 # Warm sepia h = torch.lerp(h, torch.full_like(h, sepia_hue), sepia_tone) s = s + sepia_tone * 0.3 - + # Ensure values stay in valid range h = h % 1.0 # Wrap hue s = torch.clamp(s, 0.0, 1.0) v = torch.clamp(v, 0.0, 1.0) - + # Convert back to RGB colorized_hsv = torch.stack([h, s, v], dim=-1) colorized_rgb = hsv_to_rgb(colorized_hsv) - + # Blend with original based on colorize_strength final_image = torch.lerp(processed_image, colorized_rgb, colorize_strength) - + return torch.clamp(final_image, 0.0, 1.0) @@ -1974,7 +2045,7 @@ class BatchColorCorrection: Batch Color Correction node for processing video frame sequences from VHS upload nodes. Processes multiple frames efficiently while maintaining consistency across the sequence. """ - + # Share presets with main EasyColorCorrection class PRESETS = EasyColorCorrection.PRESETS @@ -1984,50 +2055,126 @@ class BatchColorCorrection: "required": { "images": ("IMAGE",), "mode": (["Auto", "Preset", "Manual"], {"default": "Auto"}), - "frames_per_batch": ("INT", {"default": 16, "min": 1, "max": 64, "step": 1}), - "use_gpu": ("BOOLEAN", {"default": False, "tooltip": "โš ๏ธ GPU: Faster processing but uses significant VRAM (2-8GB+ for large batches). CPU: Slower but uses system RAM instead of VRAM."}), + "frames_per_batch": ( + "INT", + {"default": 16, "min": 1, "max": 64, "step": 1}, + ), + "use_gpu": ( + "BOOLEAN", + { + "default": False, + "tooltip": "โš ๏ธ GPU: Faster processing but uses significant VRAM (2-8GB+ for large batches). CPU: Slower but uses system RAM instead of VRAM.", + }, + ), }, "optional": { "ai_analysis": ("BOOLEAN", {"default": True}), "preset": ( [ - "Natural", "Warm", "Cool", "High Key", "Dramatic", - "Epic Fantasy", "Sci-Fi Chrome", "Dark Fantasy", "Vibrant Concept", "Matte Painting", "Digital Art", - "Anime Bright", "Anime Moody", "Cyberpunk", "Pastel Dreams", "Neon Nights", "Comic Book", - "Cinematic", "Teal & Orange", "Film Noir", "Vintage Film", "Bleach Bypass", - "Golden Hour", "Blue Hour", "Sunny Day", "Overcast", - "Sepia", "Black & White", "Faded", "Moody" + "Natural", + "Warm", + "Cool", + "High Key", + "Dramatic", + "Epic Fantasy", + "Sci-Fi Chrome", + "Dark Fantasy", + "Vibrant Concept", + "Matte Painting", + "Digital Art", + "Anime Bright", + "Anime Moody", + "Cyberpunk", + "Pastel Dreams", + "Neon Nights", + "Comic Book", + "Cinematic", + "Teal & Orange", + "Film Noir", + "Vintage Film", + "Bleach Bypass", + "Golden Hour", + "Blue Hour", + "Sunny Day", + "Overcast", + "Sepia", + "Black & White", + "Faded", + "Moody", ], - {"default": "Natural"} + {"default": "Natural"}, + ), + "effect_strength": ( + "FLOAT", + {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.1}, + ), + "enhancement_strength": ( + "FLOAT", + {"default": 0.8, "min": 0.0, "max": 1.5, "step": 0.1}, ), - "effect_strength": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.1}), - "enhancement_strength": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.5, "step": 0.1}), "adjust_for_skin_tone": ("BOOLEAN", {"default": True}), - "white_balance_strength": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "warmth": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "vibrancy": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "brightness": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "contrast": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "lift": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "gamma": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "gain": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "noise": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.1}), + "white_balance_strength": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "warmth": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "vibrancy": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "brightness": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "contrast": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "lift": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "gamma": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "gain": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "noise": ( + "FLOAT", + {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.1}, + ), "extract_palette": ("BOOLEAN", {"default": False}), "reference_image": ("IMAGE",), - "reference_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}), + "reference_strength": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}, + ), "mask": ("MASK",), - } + }, } RETURN_TYPES = ("IMAGE", "STRING", "IMAGE", "IMAGE", "INT") - RETURN_NAMES = ("images", "palette_data", "histogram", "palette_image", "frame_count") + RETURN_NAMES = ( + "images", + "palette_data", + "histogram", + "palette_image", + "frame_count", + ) FUNCTION = "batch_color_correct" CATEGORY = "itsjustregi / Easy Color Correction" OUTPUT_NODE = True - + @classmethod def IS_CHANGED(cls, **kwargs): return float("NaN") # Always update for video previews + DISPLAY_NAME = "Batch Color Correction" def batch_color_correct( @@ -2059,25 +2206,25 @@ class BatchColorCorrection: GPU-optimized batch processing for video frame sequences. Processes multiple frames efficiently while keeping tensors on GPU. """ - + # Get batch dimensions and device total_frames = images.shape[0] frame_height = images.shape[1] frame_width = images.shape[2] device = images.device - + # Debug device and user preferences print(f"๐Ÿ”ง CUDA Available: {torch.cuda.is_available()}") print(f"๐Ÿ”ง User GPU preference: {use_gpu}") print(f"๐Ÿ”ง Initial device: {device}") print(f"๐Ÿ”ง Input tensor device: {images.device} | dtype: {images.dtype}") - + # Handle GPU processing based on user choice if use_gpu and torch.cuda.is_available(): gpu_memory_before = torch.cuda.memory_allocated() / 1024**3 # GB print(f"๐Ÿš€ GPU Memory Before: {gpu_memory_before:.2f} GB") - - if not str(device).startswith('cuda'): + + if not str(device).startswith("cuda"): print("๐Ÿš€ User enabled GPU - moving tensors to GPU...") images = images.cuda() device = images.device @@ -2087,48 +2234,69 @@ class BatchColorCorrection: print(f"โœ… Mask moved to: {device}") else: print(f"โœ… Tensors already on GPU: {device}") - + elif use_gpu and not torch.cuda.is_available(): print("โŒ User requested GPU but CUDA not available - falling back to CPU") gpu_memory_before = 0 - + else: print("๐Ÿ’ป User selected CPU processing - keeping tensors on CPU") gpu_memory_before = 0 - - print(f"๐ŸŽฌ GPU Batch Color Correction: Processing {total_frames} frames ({frame_width}x{frame_height}) on {device}") + + print( + f"๐ŸŽฌ GPU Batch Color Correction: Processing {total_frames} frames ({frame_width}x{frame_height}) on {device}" + ) print(f"๐Ÿ“Š Processing in batches of {frames_per_batch} frames") - + # Prepare output containers on GPU processed_frames = [] all_palette_data = [] all_histograms = [] all_palette_images = [] - + # Process frames in GPU-optimized batches with interruption support try: for batch_start in range(0, total_frames, frames_per_batch): # Check for interruption requests import comfy.model_management as model_management + if model_management.interrupt_processing: print("๐Ÿ›‘ Batch processing interrupted by user") # Return partially processed results if processed_frames: partial_images = torch.cat(processed_frames, dim=0) partial_count = partial_images.shape[0] - print(f"โš ๏ธ Partial result: {partial_count}/{total_frames} frames processed") - return (partial_images, "", torch.zeros((1, 512, 768, 3), device=device), torch.zeros((1, 120, 600, 3), device=device), partial_count) + print( + f"โš ๏ธ Partial result: {partial_count}/{total_frames} frames processed" + ) + return ( + partial_images, + "", + torch.zeros((1, 512, 768, 3), device=device), + torch.zeros((1, 120, 600, 3), device=device), + partial_count, + ) else: - return (images, "", torch.zeros((1, 512, 768, 3), device=device), torch.zeros((1, 120, 600, 3), device=device), 0) - + return ( + images, + "", + torch.zeros((1, 512, 768, 3), device=device), + torch.zeros((1, 120, 600, 3), device=device), + 0, + ) + batch_end = min(batch_start + frames_per_batch, total_frames) - - print(f"๐Ÿ”„ GPU Processing batch {batch_start//frames_per_batch + 1}/{(total_frames + frames_per_batch - 1)//frames_per_batch}: frames {batch_start}-{batch_end-1}") - + + print( + f"๐Ÿ”„ GPU Processing batch {batch_start//frames_per_batch + 1}/{(total_frames + frames_per_batch - 1)//frames_per_batch}: frames {batch_start}-{batch_end-1}" + ) + # Extract current batch - keep on GPU - batch_frames = images[batch_start:batch_end] # Shape: (batch_size, H, W, C) + batch_frames = images[ + batch_start:batch_end + ] # Shape: (batch_size, H, W, C) batch_masks = mask[batch_start:batch_end] if mask is not None else None - + # Process entire batch on GPU using vectorized operations batch_processed = self._process_batch_gpu( batch_frames=batch_frames, @@ -2145,74 +2313,129 @@ class BatchColorCorrection: gamma=gamma, gain=gain, noise=noise, - device=device + device=device, ) - + processed_frames.append(batch_processed) - + # Only extract palette from middle frame to avoid CPU overhead if extract_palette and batch_start <= total_frames // 2 < batch_end: middle_frame_idx = (total_frames // 2) - batch_start - middle_frame = batch_processed[middle_frame_idx:middle_frame_idx+1] - + middle_frame = batch_processed[ + middle_frame_idx : middle_frame_idx + 1 + ] + # Generate GPU-based simplified palette without CPU transfer palette_data = "GPU_BATCH_MODE" # Simplified for performance histogram_tensor = torch.zeros((1, 512, 768, 3), device=device) palette_img_tensor = torch.zeros((1, 120, 600, 3), device=device) - + all_palette_data.append(palette_data) all_histograms.append(histogram_tensor) all_palette_images.append(palette_img_tensor) - + except KeyboardInterrupt: print("๐Ÿ›‘ Batch processing interrupted by KeyboardInterrupt") if processed_frames: partial_images = torch.cat(processed_frames, dim=0) partial_count = partial_images.shape[0] - print(f"โš ๏ธ Partial result: {partial_count}/{total_frames} frames processed") - return (partial_images, "", torch.zeros((1, 512, 768, 3), device=device), torch.zeros((1, 120, 600, 3), device=device), partial_count) + print( + f"โš ๏ธ Partial result: {partial_count}/{total_frames} frames processed" + ) + return ( + partial_images, + "", + torch.zeros((1, 512, 768, 3), device=device), + torch.zeros((1, 120, 600, 3), device=device), + partial_count, + ) else: - return (images, "", torch.zeros((1, 512, 768, 3), device=device), torch.zeros((1, 120, 600, 3), device=device), 0) + return ( + images, + "", + torch.zeros((1, 512, 768, 3), device=device), + torch.zeros((1, 120, 600, 3), device=device), + 0, + ) except Exception as e: print(f"โŒ Error during batch processing: {e}") if processed_frames: partial_images = torch.cat(processed_frames, dim=0) partial_count = partial_images.shape[0] - print(f"โš ๏ธ Partial result after error: {partial_count}/{total_frames} frames processed") - return (partial_images, "", torch.zeros((1, 512, 768, 3), device=device), torch.zeros((1, 120, 600, 3), device=device), partial_count) + print( + f"โš ๏ธ Partial result after error: {partial_count}/{total_frames} frames processed" + ) + return ( + partial_images, + "", + torch.zeros((1, 512, 768, 3), device=device), + torch.zeros((1, 120, 600, 3), device=device), + partial_count, + ) else: - return (images, "", torch.zeros((1, 512, 768, 3), device=device), torch.zeros((1, 120, 600, 3), device=device), 0) - + return ( + images, + "", + torch.zeros((1, 512, 768, 3), device=device), + torch.zeros((1, 120, 600, 3), device=device), + 0, + ) + # Combine all processed frames on GPU if processed_frames: final_images = torch.cat(processed_frames, dim=0) - + # Use middle frame data for representation representative_palette = all_palette_data[0] if all_palette_data else "" - representative_histogram = all_histograms[0] if all_histograms else torch.zeros((1, 512, 768, 3), device=device) - representative_palette_img = all_palette_images[0] if all_palette_images else torch.zeros((1, 120, 600, 3), device=device) - + representative_histogram = ( + all_histograms[0] + if all_histograms + else torch.zeros((1, 512, 768, 3), device=device) + ) + representative_palette_img = ( + all_palette_images[0] + if all_palette_images + else torch.zeros((1, 120, 600, 3), device=device) + ) + # Memory cleanup and final GPU status if use_gpu and torch.cuda.is_available(): gpu_memory_after = torch.cuda.memory_allocated(device) / 1024**3 # GB - print(f"๐Ÿš€ GPU Memory After: {gpu_memory_after:.2f} GB (Delta: {gpu_memory_after - gpu_memory_before:.2f} GB)") - + print( + f"๐Ÿš€ GPU Memory After: {gpu_memory_after:.2f} GB (Delta: {gpu_memory_after - gpu_memory_before:.2f} GB)" + ) + # Force garbage collection to free memory import gc + gc.collect() torch.cuda.empty_cache() - + gpu_memory_final = torch.cuda.memory_allocated(device) / 1024**3 # GB print(f"๐Ÿงน GPU Memory After Cleanup: {gpu_memory_final:.2f} GB") - - print(f"โœ… GPU Batch processing complete: {total_frames} frames processed on {device}") + + print( + f"โœ… GPU Batch processing complete: {total_frames} frames processed on {device}" + ) if representative_palette: print(f"๐ŸŽจ Representative palette: {representative_palette}") - - return (final_images, representative_palette, representative_histogram, representative_palette_img, total_frames) + + return ( + final_images, + representative_palette, + representative_histogram, + representative_palette_img, + total_frames, + ) else: # Fallback if no frames processed - return (images, "", torch.zeros((1, 512, 768, 3), device=device), torch.zeros((1, 120, 600, 3), device=device), 0) + return ( + images, + "", + torch.zeros((1, 512, 768, 3), device=device), + torch.zeros((1, 120, 600, 3), device=device), + 0, + ) def _process_batch_gpu( self, @@ -2230,47 +2453,49 @@ class BatchColorCorrection: gamma, gain, noise, - device + device, ): """ GPU-optimized batch processing that processes multiple frames simultaneously. """ batch_size = batch_frames.shape[0] - + # Debug GPU utilization in core processing print(f"๐Ÿ”ฅ Processing batch of {batch_size} frames on {device}") print(f"๐ŸŽฏ Batch tensor device: {batch_frames.device}") - print(f"๐ŸŽ›๏ธ Parameters: warmth={warmth}, vibrancy={vibrancy}, brightness={brightness}, contrast={contrast}") - + print( + f"๐ŸŽ›๏ธ Parameters: warmth={warmth}, vibrancy={vibrancy}, brightness={brightness}, contrast={contrast}" + ) + # Ensure batch is on correct GPU device if str(batch_frames.device) != str(device): print(f"โš ๏ธ Moving batch from {batch_frames.device} to {device}") batch_frames = batch_frames.to(device) if batch_masks is not None: batch_masks = batch_masks.to(device) - + original_batch = batch_frames.clone() - + # Process all frames in batch simultaneously using vectorized operations processed_batch = batch_frames.clone() - + # AI Analysis (only on first frame to save computation) analysis = None if ai_analysis: # Use GPU-based analysis instead of CPU-heavy OpenCV operations analysis = self._analyze_image_gpu(batch_frames[0], device) - + # Apply preset modifications if in Preset mode if mode == "Preset": # Map simplified batch preset names to full preset names preset_mapping = { "Natural": "Natural Portrait", - "Warm": "Warm Portrait", + "Warm": "Warm Portrait", "Cool": "Cool Portrait", "High Key": "High Key Portrait", - "Dramatic": "Dramatic Portrait" + "Dramatic": "Dramatic Portrait", } - + full_preset_name = preset_mapping.get(preset, preset) if full_preset_name in self.PRESETS: preset_values = self.PRESETS[full_preset_name] @@ -2281,62 +2506,91 @@ class BatchColorCorrection: print(f"๐ŸŽจ Applied batch preset: {preset} -> {full_preset_name}") else: print(f"โš ๏ธ Preset '{preset}' not found in PRESETS dictionary") - - if analysis and analysis["scene_type"] in ["concept_art", "anime", "stylized_art"]: + + if analysis and analysis["scene_type"] in [ + "concept_art", + "anime", + "stylized_art", + ]: vibrancy *= 1.4 contrast *= 1.25 - + # Auto mode specific processing if mode == "Auto": - print(f"๐Ÿค– Batch Auto Mode: Applying intelligent enhancements to {batch_size} frames") - + print( + f"๐Ÿค– Batch Auto Mode: Applying intelligent enhancements to {batch_size} frames" + ) + # Apply white balance to entire batch if enabled if white_balance_strength > 0.0: if ai_analysis and ADVANCED_LIBS_AVAILABLE: # AI-based white balance on first frame, apply to all - first_frame_np = (batch_frames[0].cpu().numpy() * 255).astype(np.uint8) - wb_corrected_first = intelligent_white_balance(first_frame_np, white_balance_strength) - + first_frame_np = (batch_frames[0].cpu().numpy() * 255).astype( + np.uint8 + ) + wb_corrected_first = intelligent_white_balance( + first_frame_np, white_balance_strength + ) + # Calculate the correction factors from first frame original_mean = torch.mean(batch_frames[0], dim=(0, 1)) - corrected_mean = torch.mean(torch.from_numpy(wb_corrected_first.astype(np.float32) / 255.0).to(device), dim=(0, 1)) + corrected_mean = torch.mean( + torch.from_numpy( + wb_corrected_first.astype(np.float32) / 255.0 + ).to(device), + dim=(0, 1), + ) wb_factors = corrected_mean / (original_mean + 1e-6) - + # Apply same factors to entire batch processed_batch = processed_batch * wb_factors.view(1, 1, 1, 3) processed_batch = torch.clamp(processed_batch, 0.0, 1.0) - print(f"๐Ÿ”ง Applied AI white balance to batch (factors: {wb_factors})") + print( + f"๐Ÿ”ง Applied AI white balance to batch (factors: {wb_factors})" + ) else: # Tensor-based white balance for entire batch B, H, W, C = processed_batch.shape flat_batch = processed_batch.view(B, -1, C) - percentile_40 = torch.quantile(flat_batch, 0.40, dim=1, keepdim=True) - percentile_60 = torch.quantile(flat_batch, 0.60, dim=1, keepdim=True) + percentile_40 = torch.quantile( + flat_batch, 0.40, dim=1, keepdim=True + ) + percentile_60 = torch.quantile( + flat_batch, 0.60, dim=1, keepdim=True + ) midtone_mean = (percentile_40 + percentile_60) / 2.0 avg_gray = torch.mean(midtone_mean, dim=-1, keepdim=True) scale = avg_gray / (midtone_mean + 1e-6) - scale = torch.lerp(torch.ones_like(scale), scale, white_balance_strength) + scale = torch.lerp( + torch.ones_like(scale), scale, white_balance_strength + ) scale = scale.view(B, 1, 1, C) processed_batch = processed_batch * scale processed_batch = torch.clamp(processed_batch, 0.0, 1.0) print(f"๐Ÿ”ง Applied tensor white balance to batch") - + # Apply enhancement based on scene analysis if enhancement_strength > 0.2: hsv_temp = rgb_to_hsv(processed_batch) - h_temp, s_temp, v_temp = hsv_temp[..., 0], hsv_temp[..., 1], hsv_temp[..., 2] - + h_temp, s_temp, v_temp = ( + hsv_temp[..., 0], + hsv_temp[..., 1], + hsv_temp[..., 2], + ) + if analysis: scene_type = analysis["scene_type"] lighting = analysis["lighting"] - + # Scene-specific enhancements if scene_type == "anime": contrast_boost = 0.18 * enhancement_strength saturation_boost = 0.55 * enhancement_strength v_temp = 0.5 + (v_temp - 0.5) * (1.0 + contrast_boost) s_temp = s_temp * (1.0 + saturation_boost) - print(f"๐ŸŽจ Applied anime enhancement (contrast: {contrast_boost:.3f}, saturation: {saturation_boost:.3f})") + print( + f"๐ŸŽจ Applied anime enhancement (contrast: {contrast_boost:.3f}, saturation: {saturation_boost:.3f})" + ) elif scene_type == "concept_art": contrast_boost = 0.25 * enhancement_strength saturation_boost = 0.40 * enhancement_strength @@ -2355,7 +2609,7 @@ class BatchColorCorrection: v_temp = 0.5 + (v_temp - 0.5) * (1.0 + contrast_boost) s_temp = s_temp * (1.0 + saturation_boost) print(f"๐ŸŽจ Applied general enhancement") - + # Lighting adjustments if lighting == "low_light": brightness += 0.1 * enhancement_strength @@ -2367,94 +2621,112 @@ class BatchColorCorrection: elif lighting == "flat": contrast += 0.5 * enhancement_strength print(f"๐Ÿ’ก Applied flat lighting enhancement") - + s_temp = torch.clamp(s_temp, 0.0, 1.0) v_temp = torch.clamp(v_temp, 0.0, 1.0) - processed_batch = hsv_to_rgb(torch.stack([h_temp, s_temp, v_temp], dim=-1)) - + processed_batch = hsv_to_rgb( + torch.stack([h_temp, s_temp, v_temp], dim=-1) + ) + # Face enhancement for Auto mode (if faces detected and adjust_for_skin_tone enabled) - if analysis and analysis.get("faces") and adjust_for_skin_tone and ADVANCED_LIBS_AVAILABLE: - print(f"๐Ÿ‘ค Applying face enhancement to batch with {len(analysis['faces'])} faces detected") + if ( + analysis + and analysis.get("faces") + and adjust_for_skin_tone + and ADVANCED_LIBS_AVAILABLE + ): + print( + f"๐Ÿ‘ค Applying face enhancement to batch with {len(analysis['faces'])} faces detected" + ) face_enhanced_frames = [] - + for i in range(batch_size): frame_np = (processed_batch[i].cpu().numpy() * 255).astype(np.uint8) - enhanced_frame_np = enhance_faces(frame_np, analysis["faces"], enhancement_strength * 0.5) - enhanced_frame_tensor = torch.from_numpy(enhanced_frame_np.astype(np.float32) / 255.0).to(device) + enhanced_frame_np = enhance_faces( + frame_np, analysis["faces"], enhancement_strength * 0.5 + ) + enhanced_frame_tensor = torch.from_numpy( + enhanced_frame_np.astype(np.float32) / 255.0 + ).to(device) face_enhanced_frames.append(enhanced_frame_tensor) - + processed_batch = torch.stack(face_enhanced_frames, dim=0) print(f"โœ… Applied face enhancement to {batch_size} frames") - + # Convert to HSV for batch processing hsv_batch = rgb_to_hsv(processed_batch) h, s, v = hsv_batch[..., 0], hsv_batch[..., 1], hsv_batch[..., 2] - + # Track changes to confirm processing is happening original_mean = torch.mean(processed_batch).item() print(f"๐Ÿ“Š Original batch mean: {original_mean:.4f}") - + # Apply color corrections to entire batch if warmth != 0.0: h = (h + warmth * 0.1) % 1.0 - + if vibrancy != 0.0: saturation_mask = 1.0 - s s = s * (1.0 + vibrancy) + (vibrancy * 0.3 * saturation_mask * s) - + if brightness != 0.0: v = v + brightness * (1.0 - v * 0.5) - + if contrast != 0.0: v = 0.5 + (v - 0.5) * (1.0 + contrast) - + # Manual mode 3-way color correction (applied to entire batch) if mode == "Manual": # Create masks for entire batch shadows_mask = 1.0 - torch.clamp(v * 3.0, 0.0, 1.0) midtones_mask = 1.0 - torch.abs(v - 0.5) * 2.0 highlights_mask = torch.clamp((v - 0.66) * 3.0, 0.0, 1.0) - + if lift != 0.0: v = v + (lift * 0.4 * shadows_mask) - + if gamma != 0.0: gamma_exp = 1.0 / (1.0 + gamma * 0.8) v_gamma = torch.pow(torch.clamp(v, 0.001, 1.0), gamma_exp) v = torch.lerp(v, v_gamma, midtones_mask) - + if gain != 0.0: v = v + (gain * 0.4 * highlights_mask) - + # Add noise to entire batch if specified if noise > 0.0: mono_noise = torch.randn( - (batch_size, processed_batch.shape[1], processed_batch.shape[2], 1), - device=device + (batch_size, processed_batch.shape[1], processed_batch.shape[2], 1), + device=device, ) luminance_mask = 1.0 - torch.abs(v - 0.5) * 2.0 luminance_mask = torch.clamp(luminance_mask, 0.0, 1.0).unsqueeze(-1) - + rgb_temp = hsv_to_rgb(torch.stack([h, s, v], dim=-1)) rgb_temp += mono_noise * noise * 0.15 * luminance_mask rgb_temp = torch.clamp(rgb_temp, 0.0, 1.0) - + hsv_temp = rgb_to_hsv(rgb_temp) h, s, v = hsv_temp[..., 0], hsv_temp[..., 1], hsv_temp[..., 2] - + # Clamp and convert back to RGB s = torch.clamp(s, 0.0, 1.0) v = torch.clamp(v, 0.0, 1.0) processed_hsv = torch.stack([h, s, v], dim=-1) processed_batch = hsv_to_rgb(processed_hsv) - + # Apply effect strength to entire batch if mode in ["Auto", "Preset"]: - processed_batch = torch.lerp(original_batch, processed_batch, effect_strength) - + processed_batch = torch.lerp( + original_batch, processed_batch, effect_strength + ) + # Apply masks to entire batch if provided if batch_masks is not None: - if batch_masks.shape[1:] != (processed_batch.shape[1], processed_batch.shape[2]): + if batch_masks.shape[1:] != ( + processed_batch.shape[1], + processed_batch.shape[2], + ): batch_masks = F.interpolate( batch_masks.unsqueeze(1), size=(processed_batch.shape[1], processed_batch.shape[2]), @@ -2463,9 +2735,9 @@ class BatchColorCorrection: ).squeeze(1) batch_masks = batch_masks.unsqueeze(-1) processed_batch = torch.lerp(original_batch, processed_batch, batch_masks) - + processed_batch = torch.clamp(processed_batch, 0.0, 1.0) - + # Confirm processing happened final_mean = torch.mean(processed_batch).item() change = abs(final_mean - original_mean) @@ -2474,7 +2746,7 @@ class BatchColorCorrection: print("โœ… Color correction applied successfully") else: print("โš ๏ธ Minimal/no changes detected") - + return processed_batch def _analyze_image_gpu(self, image_tensor, device): @@ -2485,11 +2757,11 @@ class BatchColorCorrection: # Convert to HSV for analysis hsv = rgb_to_hsv(image_tensor.unsqueeze(0))[0] # Remove batch dim h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2] - + # Analyze brightness distribution (replaces histogram analysis) brightness_mean = torch.mean(v).item() brightness_std = torch.std(v).item() - + # Determine lighting condition based on brightness stats if brightness_mean < 0.3: lighting_condition = "low_light" @@ -2499,11 +2771,11 @@ class BatchColorCorrection: lighting_condition = "flat" else: lighting_condition = "optimal" - + # Analyze saturation for scene type detection saturation_mean = torch.mean(s).item() saturation_std = torch.std(s).item() - + # Simple scene classification based on color statistics if saturation_mean > 0.6 and saturation_std > 0.25: scene_type = "concept_art" @@ -2513,22 +2785,28 @@ class BatchColorCorrection: scene_type = "portrait" else: scene_type = "realistic_photo" - + # Edge detection using Sobel filters on GPU gray = torch.mean(image_tensor, dim=-1, keepdim=True) - + # Sobel kernels - sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32, device=device).view(1, 1, 3, 3) - sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32, device=device).view(1, 1, 3, 3) - + sobel_x = torch.tensor( + [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32, device=device + ).view(1, 1, 3, 3) + sobel_y = torch.tensor( + [[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32, device=device + ).view(1, 1, 3, 3) + # Apply convolution for edge detection - gray_padded = F.pad(gray.permute(2, 0, 1).unsqueeze(0), (1, 1, 1, 1), mode='reflect') + gray_padded = F.pad( + gray.permute(2, 0, 1).unsqueeze(0), (1, 1, 1, 1), mode="reflect" + ) edges_x = F.conv2d(gray_padded, sobel_x) edges_y = F.conv2d(gray_padded, sobel_y) edges = torch.sqrt(edges_x**2 + edges_y**2) - + edge_density = torch.mean(edges).item() - + return { "scene_type": scene_type, "lighting_condition": lighting_condition, @@ -2536,7 +2814,7 @@ class BatchColorCorrection: "saturation_mean": saturation_mean, "edge_density": edge_density, "has_faces": False, # Simplified - no CPU face detection - "skin_tone_areas": [] # Simplified - avoid CPU processing + "skin_tone_areas": [], # Simplified - avoid CPU processing } @@ -2552,29 +2830,78 @@ class RawImageProcessor: return { "required": { "file_path": ("STRING", {"default": "", "multiline": False}), - "format_type": (["auto", "raw", "exr", "hdr", "tiff16"], {"default": "auto"}), + "format_type": ( + ["auto", "raw", "exr", "hdr", "tiff16"], + {"default": "auto"}, + ), }, "optional": { # RAW-specific controls - "white_balance": (["auto", "camera", "daylight", "cloudy", "shade", "tungsten", "fluorescent", "flash"], {"default": "auto"}), - "demosaic_algorithm": (["AHD", "VNG", "PPG", "AAHD"], {"default": "AHD"}), - + "white_balance": ( + [ + "auto", + "camera", + "daylight", + "cloudy", + "shade", + "tungsten", + "fluorescent", + "flash", + ], + {"default": "auto"}, + ), + "demosaic_algorithm": ( + ["AHD", "VNG", "PPG", "AAHD"], + {"default": "AHD"}, + ), # HDR/EXR tone mapping controls - "tone_mapping": (["none", "reinhard", "drago", "aces"], {"default": "aces"}), - "hdr_exposure": ("FLOAT", {"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.1}), - "hdr_gamma": ("FLOAT", {"default": 2.2, "min": 0.5, "max": 4.0, "step": 0.1}), - + "tone_mapping": ( + ["none", "reinhard", "drago", "aces"], + {"default": "aces"}, + ), + "hdr_exposure": ( + "FLOAT", + {"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.1}, + ), + "hdr_gamma": ( + "FLOAT", + {"default": 2.2, "min": 0.5, "max": 4.0, "step": 0.1}, + ), # Universal controls - "exposure": ("FLOAT", {"default": 0.0, "min": -3.0, "max": 3.0, "step": 0.1}), - "highlights": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "shadows": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "brightness": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "contrast": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "saturation": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}), - "noise_reduction": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.1}), - "output_colorspace": (["sRGB", "Adobe RGB", "ProPhoto RGB"], {"default": "sRGB"}), + "exposure": ( + "FLOAT", + {"default": 0.0, "min": -3.0, "max": 3.0, "step": 0.1}, + ), + "highlights": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "shadows": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "brightness": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "contrast": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "saturation": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.1}, + ), + "noise_reduction": ( + "FLOAT", + {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.1}, + ), + "output_colorspace": ( + ["sRGB", "Adobe RGB", "ProPhoto RGB"], + {"default": "sRGB"}, + ), "output_gamma": (["sRGB", "linear", "1.8", "2.2"], {"default": "sRGB"}), - } + }, } RETURN_TYPES = ("IMAGE", "STRING") @@ -2606,98 +2933,142 @@ class RawImageProcessor: """ if not file_path or not file_path.strip(): raise ValueError("File path is required") - + # Auto-detect format based on file extension - file_extension = file_path.lower().split('.')[-1] + file_extension = file_path.lower().split(".")[-1] if format_type == "auto": - if file_extension in ['dng', 'arw', 'cr2', 'nef', 'orf', 'rw2', 'raf']: + if file_extension in ["dng", "arw", "cr2", "nef", "orf", "rw2", "raf"]: format_type = "raw" - elif file_extension == 'exr': + elif file_extension == "exr": format_type = "exr" - elif file_extension == 'hdr': + elif file_extension == "hdr": format_type = "hdr" - elif file_extension in ['tiff', 'tif']: + elif file_extension in ["tiff", "tif"]: format_type = "tiff16" else: format_type = "raw" # Default fallback - + try: if format_type == "raw": rgb_array, metadata_info = self._process_raw_format( - file_path, white_balance, demosaic_algorithm, exposure, - highlights, shadows, brightness, noise_reduction, - output_colorspace, output_gamma + file_path, + white_balance, + demosaic_algorithm, + exposure, + highlights, + shadows, + brightness, + noise_reduction, + output_colorspace, + output_gamma, ) - + elif format_type == "exr": rgb_array, metadata_info = self._process_exr_format( - file_path, tone_mapping, hdr_exposure, hdr_gamma, - exposure, highlights, shadows, brightness, contrast, saturation + file_path, + tone_mapping, + hdr_exposure, + hdr_gamma, + exposure, + highlights, + shadows, + brightness, + contrast, + saturation, ) - + elif format_type == "hdr": rgb_array, metadata_info = self._process_hdr_format( - file_path, tone_mapping, hdr_exposure, hdr_gamma, - exposure, highlights, shadows, brightness, contrast, saturation + file_path, + tone_mapping, + hdr_exposure, + hdr_gamma, + exposure, + highlights, + shadows, + brightness, + contrast, + saturation, ) - + elif format_type == "tiff16": rgb_array, metadata_info = self._process_tiff16_format( - file_path, exposure, highlights, shadows, - brightness, contrast, saturation + file_path, + exposure, + highlights, + shadows, + brightness, + contrast, + saturation, ) - + else: raise ValueError(f"Unsupported format type: {format_type}") - + except Exception as e: error_msg = f"Error processing {format_type.upper()} file: {str(e)}" print(error_msg) # Return a black image as fallback rgb_array = np.zeros((512, 512, 3), dtype=np.uint8) metadata_info = {"error": error_msg} - + # Convert to PyTorch tensor format expected by ComfyUI if rgb_array.dtype != np.uint8: # Normalize HDR data to 0-1 range before converting to uint8 if rgb_array.max() > 1.0: rgb_array = rgb_array / rgb_array.max() rgb_array = (rgb_array * 255).astype(np.uint8) - + # Convert to float and normalize to 0-1 range image_tensor = torch.from_numpy(rgb_array.astype(np.float32) / 255.0) - + # Add batch dimension: (H, W, C) -> (1, H, W, C) image_tensor = image_tensor.unsqueeze(0) - + # Format metadata as string metadata_str = ", ".join([f"{k}: {v}" for k, v in metadata_info.items()]) - - print(f"โœ… {format_type.upper()} Image Processed: {rgb_array.shape[1]}x{rgb_array.shape[0]} from {file_path}") + + print( + f"โœ… {format_type.upper()} Image Processed: {rgb_array.shape[1]}x{rgb_array.shape[0]} from {file_path}" + ) if metadata_str: print(f"๐Ÿ“ท Metadata: {metadata_str}") - + return (image_tensor, metadata_str) - def _process_raw_format(self, file_path, white_balance, demosaic_algorithm, exposure, - highlights, shadows, brightness, noise_reduction, - output_colorspace, output_gamma): + def _process_raw_format( + self, + file_path, + white_balance, + demosaic_algorithm, + exposure, + highlights, + shadows, + brightness, + noise_reduction, + output_colorspace, + output_gamma, + ): """Process RAW camera formats using rawpy.""" if not RAW_PROCESSING_AVAILABLE: - raise ValueError("RAW processing not available. Install rawpy: pip install rawpy") - + raise ValueError( + "RAW processing not available. Install rawpy: pip install rawpy" + ) + with rawpy.imread(file_path) as raw: params = rawpy.Params() - + # Demosaic algorithm demosaic_map = { "AHD": rawpy.DemosaicAlgorithm.AHD, - "VNG": rawpy.DemosaicAlgorithm.VNG, + "VNG": rawpy.DemosaicAlgorithm.VNG, "PPG": rawpy.DemosaicAlgorithm.PPG, - "AAHD": rawpy.DemosaicAlgorithm.AAHD + "AAHD": rawpy.DemosaicAlgorithm.AAHD, } - params.demosaic_algorithm = demosaic_map.get(demosaic_algorithm, rawpy.DemosaicAlgorithm.AHD) - + params.demosaic_algorithm = demosaic_map.get( + demosaic_algorithm, rawpy.DemosaicAlgorithm.AHD + ) + # White balance if white_balance == "auto": params.use_auto_wb = True @@ -2710,182 +3081,241 @@ class RawImageProcessor: "shade": [1.4, 1.0, 0.7, 1.0], "tungsten": [0.6, 1.0, 1.8, 1.0], "fluorescent": [0.8, 1.0, 1.4, 1.0], - "flash": [1.1, 1.0, 0.9, 1.0] + "flash": [1.1, 1.0, 0.9, 1.0], } if white_balance in wb_presets: params.user_wb = wb_presets[white_balance] - + # Exposure and processing params.exp_correc = True params.exp_shift = exposure - params.highlight_mode = rawpy.HighlightMode.Clip if highlights == 0 else rawpy.HighlightMode.Reconstruct + params.highlight_mode = ( + rawpy.HighlightMode.Clip + if highlights == 0 + else rawpy.HighlightMode.Reconstruct + ) params.bright = 1.0 + brightness - params.output_color = rawpy.ColorSpace.sRGB if output_colorspace == "sRGB" else rawpy.ColorSpace.Adobe + params.output_color = ( + rawpy.ColorSpace.sRGB + if output_colorspace == "sRGB" + else rawpy.ColorSpace.Adobe + ) params.gamma = (1.0, 1.0) if output_gamma == "linear" else (2.2, 4.5) - + if noise_reduction > 0: params.median_filter_passes = int(noise_reduction * 5) - + rgb_array = raw.postprocess(params) - + metadata_info = { "format": "RAW", "width": rgb_array.shape[1], "height": rgb_array.shape[0], "demosaic": demosaic_algorithm, - "white_balance": white_balance + "white_balance": white_balance, } - + return rgb_array, metadata_info - def _process_exr_format(self, file_path, tone_mapping, hdr_exposure, hdr_gamma, - exposure, highlights, shadows, brightness, contrast, saturation): + def _process_exr_format( + self, + file_path, + tone_mapping, + hdr_exposure, + hdr_gamma, + exposure, + highlights, + shadows, + brightness, + contrast, + saturation, + ): """Process EXR HDR format using OpenEXR.""" if not EXR_PROCESSING_AVAILABLE: # Fallback to imageio if OpenEXR not available if IMAGEIO_AVAILABLE: - return self._process_with_imageio(file_path, "EXR", tone_mapping, hdr_exposure, hdr_gamma, - exposure, highlights, shadows, brightness, contrast, saturation) + return self._process_with_imageio( + file_path, + "EXR", + tone_mapping, + hdr_exposure, + hdr_gamma, + exposure, + highlights, + shadows, + brightness, + contrast, + saturation, + ) else: raise ValueError("EXR processing requires OpenEXR or imageio library") - + exr_file = OpenEXR.InputFile(file_path) header = exr_file.header() - + # Get image dimensions - dw = header['displayWindow'] + dw = header["displayWindow"] width = dw.max.x - dw.min.x + 1 height = dw.max.y - dw.min.y + 1 - + # Read RGB channels - channels = exr_file.channels(['R', 'G', 'B'], Imath.PixelType(Imath.PixelType.FLOAT)) - + channels = exr_file.channels( + ["R", "G", "B"], Imath.PixelType(Imath.PixelType.FLOAT) + ) + # Convert to numpy arrays r_channel = np.frombuffer(channels[0], dtype=np.float32).reshape(height, width) g_channel = np.frombuffer(channels[1], dtype=np.float32).reshape(height, width) b_channel = np.frombuffer(channels[2], dtype=np.float32).reshape(height, width) - + # Combine channels rgb_array = np.stack([r_channel, g_channel, b_channel], axis=2) - + # Apply HDR exposure if hdr_exposure != 0: - rgb_array = rgb_array * (2 ** hdr_exposure) - + rgb_array = rgb_array * (2**hdr_exposure) + # Apply tone mapping rgb_array = self._apply_tone_mapping(rgb_array, tone_mapping, hdr_gamma) - + metadata_info = { "format": "EXR", "width": width, "height": height, "tone_mapping": tone_mapping, - "hdr_exposure": hdr_exposure + "hdr_exposure": hdr_exposure, } - + return rgb_array, metadata_info - def _process_hdr_format(self, file_path, tone_mapping, hdr_exposure, hdr_gamma, - exposure, highlights, shadows, brightness, contrast, saturation): + def _process_hdr_format( + self, + file_path, + tone_mapping, + hdr_exposure, + hdr_gamma, + exposure, + highlights, + shadows, + brightness, + contrast, + saturation, + ): """Process HDR format using imageio.""" if not IMAGEIO_AVAILABLE: raise ValueError("HDR processing requires imageio library") - + # Read HDR image - rgb_array = imageio.imread(file_path, format='HDR-FI') - + rgb_array = imageio.imread(file_path, format="HDR-FI") + # Apply HDR exposure if hdr_exposure != 0: - rgb_array = rgb_array * (2 ** hdr_exposure) - + rgb_array = rgb_array * (2**hdr_exposure) + # Apply tone mapping rgb_array = self._apply_tone_mapping(rgb_array, tone_mapping, hdr_gamma) - + metadata_info = { "format": "HDR", "width": rgb_array.shape[1], "height": rgb_array.shape[0], "tone_mapping": tone_mapping, - "hdr_exposure": hdr_exposure + "hdr_exposure": hdr_exposure, } - + return rgb_array, metadata_info - def _process_tiff16_format(self, file_path, exposure, highlights, shadows, - brightness, contrast, saturation): + def _process_tiff16_format( + self, file_path, exposure, highlights, shadows, brightness, contrast, saturation + ): """Process 16-bit TIFF format.""" if not IMAGEIO_AVAILABLE: raise ValueError("TIFF 16-bit processing requires imageio library") - + # Read 16-bit TIFF rgb_array = imageio.imread(file_path) - + # Convert to float and normalize if rgb_array.dtype == np.uint16: rgb_array = rgb_array.astype(np.float32) / 65535.0 elif rgb_array.dtype == np.uint8: rgb_array = rgb_array.astype(np.float32) / 255.0 - + # Apply basic adjustments if exposure != 0: - rgb_array = rgb_array * (2 ** exposure) - + rgb_array = rgb_array * (2**exposure) + rgb_array = np.clip(rgb_array, 0, 1) - + metadata_info = { "format": "TIFF 16-bit", "width": rgb_array.shape[1], "height": rgb_array.shape[0], - "bit_depth": "16-bit" if rgb_array.dtype == np.uint16 else "8-bit" + "bit_depth": "16-bit" if rgb_array.dtype == np.uint16 else "8-bit", } - + return rgb_array, metadata_info - def _process_with_imageio(self, file_path, format_name, tone_mapping, hdr_exposure, hdr_gamma, - exposure, highlights, shadows, brightness, contrast, saturation): + def _process_with_imageio( + self, + file_path, + format_name, + tone_mapping, + hdr_exposure, + hdr_gamma, + exposure, + highlights, + shadows, + brightness, + contrast, + saturation, + ): """Fallback processing using imageio.""" rgb_array = imageio.imread(file_path) - + # Convert to float if needed if rgb_array.dtype == np.uint8: rgb_array = rgb_array.astype(np.float32) / 255.0 elif rgb_array.dtype == np.uint16: rgb_array = rgb_array.astype(np.float32) / 65535.0 - + # Apply HDR exposure if it's an HDR format if hdr_exposure != 0: - rgb_array = rgb_array * (2 ** hdr_exposure) - + rgb_array = rgb_array * (2**hdr_exposure) + # Apply tone mapping for HDR formats if format_name in ["EXR", "HDR"]: rgb_array = self._apply_tone_mapping(rgb_array, tone_mapping, hdr_gamma) - + metadata_info = { "format": f"{format_name} (imageio fallback)", "width": rgb_array.shape[1], - "height": rgb_array.shape[0] + "height": rgb_array.shape[0], } - + return rgb_array, metadata_info def _apply_tone_mapping(self, rgb_array, tone_mapping, hdr_gamma): """Apply tone mapping to HDR image data.""" if tone_mapping == "none": return np.clip(rgb_array, 0, 1) - + elif tone_mapping == "reinhard": # Simple Reinhard tone mapping return rgb_array / (1.0 + rgb_array) - + elif tone_mapping == "drago": # Drago tone mapping approximation luminance = np.dot(rgb_array, [0.299, 0.587, 0.114]) max_lum = np.max(luminance) if max_lum > 0: - scale = np.log10(max_lum + 1) / np.log10(2.0 + 8.0 * ((luminance / max_lum) ** (np.log10(0.5) / np.log10(0.85)))) + scale = np.log10(max_lum + 1) / np.log10( + 2.0 + + 8.0 * ((luminance / max_lum) ** (np.log10(0.5) / np.log10(0.85))) + ) return rgb_array * scale[..., np.newaxis] return rgb_array - + elif tone_mapping == "aces": # ACES tone mapping curve approximation a = 2.51 @@ -2893,12 +3323,17 @@ class RawImageProcessor: c = 2.43 d = 0.59 e = 0.14 - return np.clip((rgb_array * (a * rgb_array + b)) / (rgb_array * (c * rgb_array + d) + e), 0, 1) - + return np.clip( + (rgb_array * (a * rgb_array + b)) + / (rgb_array * (c * rgb_array + d) + e), + 0, + 1, + ) + # Apply gamma correction if hdr_gamma != 2.2: rgb_array = np.power(np.clip(rgb_array, 0, 1), 1.0 / hdr_gamma) - + return np.clip(rgb_array, 0, 1) @@ -2913,13 +3348,16 @@ class ColorCorrectionViewer: return { "required": { "images": ("IMAGE",), - "fps": ("FLOAT", {"default": 24.0, "min": 1.0, "max": 60.0, "step": 0.1}), + "fps": ( + "FLOAT", + {"default": 24.0, "min": 1.0, "max": 60.0, "step": 0.1}, + ), }, "optional": { "auto_play": ("BOOLEAN", {"default": True}), "loop": ("BOOLEAN", {"default": True}), "frame_skip": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}), - } + }, } RETURN_TYPES = ("IMAGE", "INT", "FLOAT") @@ -2931,7 +3369,7 @@ class ColorCorrectionViewer: @classmethod def IS_CHANGED(cls, **kwargs): return float("NaN") # Always update for video display - + @classmethod def VALIDATE_INPUTS(cls, **kwargs): return True @@ -2943,27 +3381,32 @@ class ColorCorrectionViewer: if images is None or images.shape[0] == 0: print("โš ๏ธ No images provided to viewer") return (images, 0, 0.0) - + total_frames = images.shape[0] - + # Apply frame skipping if specified if frame_skip > 1: selected_indices = torch.arange(0, total_frames, frame_skip) images = images[selected_indices] total_frames = images.shape[0] - + duration = total_frames / fps - + print(f"๐ŸŽฌ Color Correction Viewer: {total_frames} frames at {fps} FPS") print(f"โฑ๏ธ Duration: {duration:.2f} seconds") print(f"๐Ÿ”„ Auto-play: {auto_play}, Loop: {loop}") if frame_skip > 1: print(f"โญ๏ธ Frame skip: every {frame_skip} frames") - + # For ComfyUI preview, we need to return the images # The frontend will handle the video player widget return (images, total_frames, duration) # Export all classes -__all__ = ["EasyColorCorrection", "BatchColorCorrection", "RawImageProcessor", "ColorCorrectionViewer"] +__all__ = [ + "EasyColorCorrection", + "BatchColorCorrection", + "RawImageProcessor", + "ColorCorrectionViewer", +]