diff --git a/__init__.py b/__init__.py index e5166fb..3364d2f 100644 --- a/__init__.py +++ b/__init__.py @@ -217,6 +217,88 @@ try: except Exception as e: print(f"⚠️ Could not load modular UniversalVisionCloner: {e}") +# APNext FX Nodes +try: + from .nodes.image_fx.bloom import APNextBloom + NEW_MAPPINGS["APNextBloom"] = APNextBloom + NEW_DISPLAY_MAPPINGS["APNextBloom"] = "APNext Bloom FX" + # print("✅ Loaded modular APNextBloom") +except Exception as e: + print(f"⚠️ Could not load modular APNextBloom: {e}") + +try: + from .nodes.image_fx.sharpen import APNextSharpen + NEW_MAPPINGS["APNextSharpen"] = APNextSharpen + NEW_DISPLAY_MAPPINGS["APNextSharpen"] = "APNext Sharpen FX" + # print("✅ Loaded modular APNextSharpen") +except Exception as e: + print(f"⚠️ Could not load modular APNextSharpen: {e}") + +try: + from .nodes.image_fx.noise import APNextNoise + NEW_MAPPINGS["APNextNoise"] = APNextNoise + NEW_DISPLAY_MAPPINGS["APNextNoise"] = "APNext Noise FX" + # print("✅ Loaded modular APNextNoise") +except Exception as e: + print(f"⚠️ Could not load modular APNextNoise: {e}") + +try: + from .nodes.image_fx.rough import APNextRough + NEW_MAPPINGS["APNextRough"] = APNextRough + NEW_DISPLAY_MAPPINGS["APNextRough"] = "APNext Rough FX" + # print("✅ Loaded modular APNextRough") +except Exception as e: + print(f"⚠️ Could not load modular APNextRough: {e}") + +# APNext Advanced FX Nodes +try: + from .nodes.image_fx.color_grading import APNextColorGrading + NEW_MAPPINGS["APNextColorGrading"] = APNextColorGrading + NEW_DISPLAY_MAPPINGS["APNextColorGrading"] = "APNext Color Grading FX" + # print("✅ Loaded modular APNextColorGrading") +except Exception as e: + print(f"⚠️ Could not load modular APNextColorGrading: {e}") + +try: + from .nodes.image_fx.cross_processing import APNextCrossProcessing + NEW_MAPPINGS["APNextCrossProcessing"] = APNextCrossProcessing + NEW_DISPLAY_MAPPINGS["APNextCrossProcessing"] = "APNext Cross Processing FX" + # print("✅ Loaded modular APNextCrossProcessing") +except Exception as e: + print(f"⚠️ Could not load modular APNextCrossProcessing: {e}") + +try: + from .nodes.image_fx.split_toning import APNextSplitToning + NEW_MAPPINGS["APNextSplitToning"] = APNextSplitToning + NEW_DISPLAY_MAPPINGS["APNextSplitToning"] = "APNext Split Toning FX" + # print("✅ Loaded modular APNextSplitToning") +except Exception as e: + print(f"⚠️ Could not load modular APNextSplitToning: {e}") + +try: + from .nodes.image_fx.hdr_tone_mapping import APNextHDRToneMapping + NEW_MAPPINGS["APNextHDRToneMapping"] = APNextHDRToneMapping + NEW_DISPLAY_MAPPINGS["APNextHDRToneMapping"] = "APNext HDR Tone Mapping FX" + # print("✅ Loaded modular APNextHDRToneMapping") +except Exception as e: + print(f"⚠️ Could not load modular APNextHDRToneMapping: {e}") + +try: + from .nodes.image_fx.glitch_art import APNextGlitchArt + NEW_MAPPINGS["APNextGlitchArt"] = APNextGlitchArt + NEW_DISPLAY_MAPPINGS["APNextGlitchArt"] = "APNext Glitch Art FX" + # print("✅ Loaded modular APNextGlitchArt") +except Exception as e: + print(f"⚠️ Could not load modular APNextGlitchArt: {e}") + +try: + from .nodes.image_fx.film_halation import APNextFilmHalation + NEW_MAPPINGS["APNextFilmHalation"] = APNextFilmHalation + NEW_DISPLAY_MAPPINGS["APNextFilmHalation"] = "APNext Film Halation FX" + # print("✅ Loaded modular APNextFilmHalation") +except Exception as e: + print(f"⚠️ Could not load modular APNextFilmHalation: {e}") + # Combine mappings (modular nodes + dynamic nodes) NODE_CLASS_MAPPINGS = {**NEW_MAPPINGS, **DYNAMIC_MAPPINGS} NODE_DISPLAY_NAME_MAPPINGS = {**NEW_DISPLAY_MAPPINGS, **DYNAMIC_DISPLAY_MAPPINGS} diff --git a/nodes/__init__.py b/nodes/__init__.py index 063dcbf..1f72470 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -1 +1,11 @@ # Node modules for comfyui_dagthomas + +# Import FX nodes +from .image_fx import NODE_CLASS_MAPPINGS as FX_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as FX_NODE_DISPLAY_NAME_MAPPINGS + +# Add FX nodes to the global mappings +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +NODE_CLASS_MAPPINGS.update(FX_NODE_CLASS_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(FX_NODE_DISPLAY_NAME_MAPPINGS) \ No newline at end of file diff --git a/nodes/image_fx/__init__.py b/nodes/image_fx/__init__.py new file mode 100644 index 0000000..ce39825 --- /dev/null +++ b/nodes/image_fx/__init__.py @@ -0,0 +1,38 @@ +# APNext FX Image Effects Nodes + +from .bloom import APNextBloom +from .sharpen import APNextSharpen +from .noise import APNextNoise +from .rough import APNextRough +from .color_grading import APNextColorGrading +from .cross_processing import APNextCrossProcessing +from .split_toning import APNextSplitToning +from .hdr_tone_mapping import APNextHDRToneMapping +from .glitch_art import APNextGlitchArt +from .film_halation import APNextFilmHalation + +NODE_CLASS_MAPPINGS = { + "APNextBloom": APNextBloom, + "APNextSharpen": APNextSharpen, + "APNextNoise": APNextNoise, + "APNextRough": APNextRough, + "APNextColorGrading": APNextColorGrading, + "APNextCrossProcessing": APNextCrossProcessing, + "APNextSplitToning": APNextSplitToning, + "APNextHDRToneMapping": APNextHDRToneMapping, + "APNextGlitchArt": APNextGlitchArt, + "APNextFilmHalation": APNextFilmHalation, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "APNextBloom": "APNext Bloom FX", + "APNextSharpen": "APNext Sharpen FX", + "APNextNoise": "APNext Noise FX", + "APNextRough": "APNext Rough FX", + "APNextColorGrading": "APNext Color Grading FX", + "APNextCrossProcessing": "APNext Cross Processing FX", + "APNextSplitToning": "APNext Split Toning FX", + "APNextHDRToneMapping": "APNext HDR Tone Mapping FX", + "APNextGlitchArt": "APNext Glitch Art FX", + "APNextFilmHalation": "APNext Film Halation FX", +} diff --git a/nodes/image_fx/bloom.py b/nodes/image_fx/bloom.py new file mode 100644 index 0000000..b7fbf2b --- /dev/null +++ b/nodes/image_fx/bloom.py @@ -0,0 +1,123 @@ +# APNext Bloom Effect Node + +import torch +import torch.nn.functional as F +import numpy as np +from ...utils.constants import CUSTOM_CATEGORY + + +class APNextBloom: + """ + APNext Bloom Effect Node + Creates a bloom effect by making bright areas glow using gaussian blur and blend modes + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}), + "threshold": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}), + "blur_radius": ("FLOAT", {"default": 15.0, "min": 1.0, "max": 50.0, "step": 0.5}), + "blend_mode": (["additive", "screen", "overlay"], {"default": "additive"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_bloom" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX" + + def apply_bloom(self, images, intensity, threshold, blur_radius, blend_mode): + """Apply bloom effect to images using optimized tensor operations""" + # Work directly with tensors for better performance + batch_size = images.shape[0] + device = images.device + dtype = images.dtype + + # Process all images in batch + result = self._create_bloom_effect_tensor( + images, intensity, threshold, blur_radius, blend_mode, device, dtype + ) + + return (result,) + + def _create_bloom_effect_tensor(self, images, intensity, threshold, blur_radius, blend_mode, device, dtype): + """Create bloom effect using tensor operations for better performance""" + # Apply threshold to isolate bright areas + bloom_layer = self._apply_threshold_tensor(images, threshold) + + # Apply gaussian blur using torch operations + bloom_layer = self._gaussian_blur_tensor(bloom_layer, blur_radius) + + # Blend the bloom layer with the original image + result = self._blend_images_tensor(images, bloom_layer, blend_mode, intensity) + + return result + + def _apply_threshold_tensor(self, images, threshold): + """Apply threshold to isolate bright areas using tensor operations""" + # Calculate luminance using tensor operations + luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2] + + # Create mask for bright areas + bright_mask = luminance > threshold + + # Apply threshold - keep only pixels above threshold + result = images.clone() + bright_mask = bright_mask.unsqueeze(-1).expand_as(images) + result = torch.where(bright_mask, images, torch.zeros_like(images)) + + return result + + def _blend_images_tensor(self, base, overlay, blend_mode, intensity): + """Blend two images using tensor operations""" + overlay = overlay * intensity + + if blend_mode == "additive": + result = base + overlay + elif blend_mode == "screen": + result = 1 - (1 - base) * (1 - overlay) + elif blend_mode == "overlay": + mask = base < 0.5 + result = torch.where( + mask, + 2 * base * overlay, + 1 - 2 * (1 - base) * (1 - overlay) + ) + + # Clamp values to [0, 1] + result = torch.clamp(result, 0, 1) + return result + + def _gaussian_blur_tensor(self, images, radius): + """Apply gaussian blur using torch operations""" + # Convert radius to sigma (approximation) + sigma = radius / 3.0 + + # Create gaussian kernel + kernel_size = int(2 * radius + 1) + if kernel_size % 2 == 0: + kernel_size += 1 + + # Create 1D gaussian kernel + x = torch.arange(kernel_size, dtype=images.dtype, device=images.device) + x = x - kernel_size // 2 + kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2) + kernel_1d = kernel_1d / kernel_1d.sum() + + # Apply separable gaussian blur + # Reshape for conv operations: [batch, channels, height, width] + images_conv = images.permute(0, 3, 1, 2) + + # Horizontal blur + kernel_h = kernel_1d.view(1, 1, 1, -1).expand(3, 1, 1, -1) + blurred_h = F.conv2d(images_conv, kernel_h, padding=(0, kernel_size//2), groups=3) + + # Vertical blur + kernel_v = kernel_1d.view(1, 1, -1, 1).expand(3, 1, -1, 1) + blurred = F.conv2d(blurred_h, kernel_v, padding=(kernel_size//2, 0), groups=3) + + # Reshape back to original format: [batch, height, width, channels] + return blurred.permute(0, 2, 3, 1) diff --git a/nodes/image_fx/color_grading.py b/nodes/image_fx/color_grading.py new file mode 100644 index 0000000..df95f85 --- /dev/null +++ b/nodes/image_fx/color_grading.py @@ -0,0 +1,533 @@ +# APNext Color Grading Effect Node + +import torch +import numpy as np +from PIL import Image +import os +import glob +import folder_paths +from ...utils.constants import CUSTOM_CATEGORY +from ...utils.image_utils import tensor2pil, pil2tensor + + +class APNextColorGrading: + """ + APNext Color Grading Effect Node + Applies color grading using LUT files or manual controls + + Supported LUT formats: + - .cube files (Adobe/Blackmagic) + - .3dl files (Autodesk/Flame) + - Image LUTs (.png, .jpg, .tiff, .exr) + + Example paths: + - C:/LUTs/cinematic.cube + - ./assets/vintage_film.3dl + - /path/to/lut_image.png + """ + + @classmethod + def INPUT_TYPES(cls): + # Register LUT file types with ComfyUI's folder_paths system + cls._register_lut_paths() + + # Get LUT files using ComfyUI's standard system + lut_files = folder_paths.get_filename_list("luts") + if not lut_files: + lut_files = ["None"] + + return { + "required": { + "images": ("IMAGE",), + "method": (["manual", "lut_file"], {"default": "manual"}), + "lut_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + # Manual controls + "exposure": ("FLOAT", {"default": 0.0, "min": -3.0, "max": 3.0, "step": 0.01}), + "contrast": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), + "highlights": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), + "shadows": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), + "saturation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), + "temperature": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), + "tint": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), + }, + "optional": { + "lut_file": (lut_files, {"default": "None"}), + "custom_lut_path": ("STRING", { + "default": "", + "multiline": False, + "placeholder": "Custom LUT path (overrides picker)" + }), + } + } + + @classmethod + def _register_lut_paths(cls): + """Register LUT file paths with ComfyUI's folder system""" + try: + # Add LUT directory to ComfyUI's folder paths + if "luts" not in folder_paths.folder_names_and_paths: + lut_dirs = [ + os.path.join(folder_paths.models_dir, "luts"), + os.path.join(folder_paths.input_directory, "luts"), + "./luts", + "./LUTs" + ] + + # Create luts directory in models if it doesn't exist + main_lut_dir = os.path.join(folder_paths.models_dir, "luts") + if not os.path.exists(main_lut_dir): + os.makedirs(main_lut_dir, exist_ok=True) + print(f"📁 Created LUT directory: {main_lut_dir}") + + # Register with ComfyUI + folder_paths.folder_names_and_paths["luts"] = (lut_dirs, {".cube", ".3dl", ".png", ".jpg", ".jpeg", ".tiff", ".exr"}) + print("✅ Registered LUT file paths with ComfyUI") + except Exception as e: + print(f"⚠️ Could not register LUT paths: {e}") + # Fallback to manual scanning + pass + + @classmethod + def _scan_for_lut_files(cls): + """Scan common directories for LUT files""" + lut_files = ["None"] # Default option + + # Common LUT directories to scan + search_dirs = [ + "./luts", + "./LUTs", + "./assets/luts", + "./custom_nodes/comfyui_dagthomas/luts", + os.path.expanduser("~/LUTs"), + os.path.expanduser("~/Documents/LUTs"), + "C:/LUTs" if os.name == 'nt' else "/usr/local/share/luts", + ] + + # LUT file extensions + extensions = ['*.cube', '*.3dl', '*.png', '*.jpg', '*.jpeg', '*.tiff', '*.exr'] + + for directory in search_dirs: + if os.path.exists(directory): + for ext in extensions: + pattern = os.path.join(directory, '**', ext) + files = glob.glob(pattern, recursive=True) + for file_path in files: + # Create a display name (relative path or just filename) + if len(file_path) > 60: + display_name = f"...{file_path[-57:]}" + else: + display_name = file_path + lut_files.append(f"{display_name}|{file_path}") + + # Also scan current working directory + for ext in extensions: + files = glob.glob(ext) + for file_path in files: + lut_files.append(f"{file_path}|{os.path.abspath(file_path)}") + + return lut_files + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_color_grading" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced" + + def apply_color_grading(self, images, method, lut_strength, exposure, contrast, + highlights, shadows, saturation, temperature, tint, + lut_file="None", custom_lut_path=""): + """Apply color grading to images""" + device = images.device + dtype = images.dtype + + if method == "lut_file": + # Determine which LUT path to use + lut_path = None + + # Custom path takes priority + if custom_lut_path and custom_lut_path.strip(): + lut_path = custom_lut_path.strip() + print(f"📁 Using custom LUT path: {lut_path}") + + # Otherwise use ComfyUI file selection + elif lut_file and lut_file != "None": + lut_path = folder_paths.get_full_path("luts", lut_file) + print(f"📁 Using selected LUT: {lut_path}") + + # Apply LUT if we have a valid path + if lut_path and os.path.exists(lut_path): + result = self._apply_lut_grading(images, lut_path, lut_strength) + elif lut_path: + print(f"⚠️ LUT file not found: {lut_path}") + print("Using manual grading instead.") + result = self._apply_manual_grading( + images, exposure, contrast, highlights, shadows, + saturation, temperature, tint + ) + else: + print("⚠️ No LUT file selected. Using manual grading instead.") + result = self._apply_manual_grading( + images, exposure, contrast, highlights, shadows, + saturation, temperature, tint + ) + else: + result = self._apply_manual_grading( + images, exposure, contrast, highlights, shadows, + saturation, temperature, tint + ) + + return (result,) + + def _apply_lut_grading(self, images, lut_path, strength): + """Apply LUT-based color grading""" + try: + # Load LUT file + lut = self._load_lut_file(lut_path) + if lut is None: + return images + + # Apply LUT to images + result = self._apply_lut_tensor(images, lut, strength) + return result + + except Exception as e: + print(f"Error applying LUT: {e}") + return images + + def _apply_manual_grading(self, images, exposure, contrast, highlights, shadows, + saturation, temperature, tint): + """Apply manual color grading controls""" + result = images.clone() + + # Apply exposure + if exposure != 0.0: + exposure_factor = 2.0 ** exposure + result = result * exposure_factor + + # Apply contrast + if contrast != 1.0: + # Contrast around middle gray (0.18 in linear, ~0.5 in sRGB) + mid_gray = 0.5 + result = (result - mid_gray) * contrast + mid_gray + + # Apply highlights and shadows + if highlights != 0.0 or shadows != 0.0: + result = self._apply_highlight_shadow_tensor(result, highlights, shadows) + + # Apply saturation + if saturation != 1.0: + result = self._apply_saturation_tensor(result, saturation) + + # Apply temperature and tint + if temperature != 0.0 or tint != 0.0: + result = self._apply_white_balance_tensor(result, temperature, tint) + + # Clamp values + result = torch.clamp(result, 0, 1) + return result + + def _apply_highlight_shadow_tensor(self, images, highlights, shadows): + """Apply highlight and shadow adjustments using tensor operations""" + # Calculate luminance + luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2] + luminance = luminance.unsqueeze(-1) + + # Create masks for highlights and shadows + highlight_mask = torch.sigmoid((luminance - 0.7) * 10) # Soft transition around 0.7 + shadow_mask = torch.sigmoid((0.3 - luminance) * 10) # Soft transition around 0.3 + + # Apply adjustments + highlight_adj = 1.0 + highlights * highlight_mask + shadow_adj = 1.0 + shadows * shadow_mask + + result = images * highlight_adj * shadow_adj + return result + + def _apply_saturation_tensor(self, images, saturation): + """Apply saturation adjustment using tensor operations""" + # Calculate luminance + luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2] + luminance = luminance.unsqueeze(-1) + + # Blend between grayscale and original based on saturation + result = luminance + (images - luminance) * saturation + return result + + def _apply_white_balance_tensor(self, images, temperature, tint): + """Apply white balance (temperature and tint) using tensor operations""" + # Temperature affects blue-yellow balance + if temperature != 0.0: + # Warm up (positive) or cool down (negative) + temp_matrix = torch.tensor([ + [1.0 + temperature * 0.2, 0.0, -temperature * 0.1], + [0.0, 1.0, 0.0], + [-temperature * 0.1, 0.0, 1.0 - temperature * 0.2] + ], device=images.device, dtype=images.dtype) + + # Apply color matrix + images_flat = images.view(-1, 3) + images_adjusted = torch.matmul(images_flat, temp_matrix.t()) + images = images_adjusted.view(images.shape) + + # Tint affects green-magenta balance + if tint != 0.0: + tint_matrix = torch.tensor([ + [1.0 - tint * 0.1, tint * 0.2, 0.0], + [-tint * 0.1, 1.0 + tint * 0.1, 0.0], + [0.0, 0.0, 1.0] + ], device=images.device, dtype=images.dtype) + + # Apply color matrix + images_flat = images.view(-1, 3) + images_adjusted = torch.matmul(images_flat, tint_matrix.t()) + images = images_adjusted.view(images.shape) + + return images + + def _load_lut_file(self, lut_path): + """Load LUT file (supports .cube, .3dl, and image formats)""" + try: + file_ext = lut_path.lower() + if file_ext.endswith('.cube'): + return self._load_cube_lut(lut_path) + elif file_ext.endswith('.3dl'): + return self._load_3dl_lut(lut_path) + elif file_ext.endswith(('.png', '.jpg', '.jpeg', '.tiff', '.exr')): + return self._load_image_lut(lut_path) + else: + print(f"Unsupported LUT format: {lut_path}") + print("Supported formats: .cube, .3dl, .png, .jpg, .jpeg, .tiff, .exr") + return None + except Exception as e: + print(f"Error loading LUT file {lut_path}: {e}") + return None + + def _load_cube_lut(self, cube_path): + """Load .cube format LUT file with enhanced compatibility""" + try: + # Try different encodings for better compatibility + encodings = ['utf-8', 'latin-1', 'cp1252'] + lines = None + + for encoding in encodings: + try: + with open(cube_path, 'r', encoding=encoding) as f: + lines = f.readlines() + break + except UnicodeDecodeError: + continue + + if lines is None: + print(f"Could not read .cube file with any supported encoding: {cube_path}") + return None + + lut_size = 33 # Default size + lut_data = [] + domain_min = [0.0, 0.0, 0.0] + domain_max = [1.0, 1.0, 1.0] + + for line in lines: + line = line.strip() + + # Skip empty lines and comments + if not line or line.startswith('#'): + continue + + # Parse header information + if line.startswith('LUT_3D_SIZE'): + lut_size = int(line.split()[-1]) + print(f"Loading .cube LUT with size: {lut_size}x{lut_size}x{lut_size}") + + elif line.startswith('DOMAIN_MIN'): + domain_min = [float(x) for x in line.split()[1:4]] + + elif line.startswith('DOMAIN_MAX'): + domain_max = [float(x) for x in line.split()[1:4]] + + elif line.startswith('TITLE'): + title = line.split('"')[1] if '"' in line else line.split()[1] + print(f"Loading LUT: {title}") + + # Parse LUT data + elif not any(line.startswith(keyword) for keyword in ['LUT_1D_SIZE', 'LUT_1D_INPUT_RANGE']): + try: + values = [float(x) for x in line.split()] + if len(values) == 3: + # Normalize values if they're outside 0-1 range + normalized_values = [] + for i, val in enumerate(values): + if domain_max[i] != domain_min[i]: + normalized = (val - domain_min[i]) / (domain_max[i] - domain_min[i]) + else: + normalized = val + normalized_values.append(max(0.0, min(1.0, normalized))) + + lut_data.append(normalized_values) + except (ValueError, IndexError): + continue + + # Validate LUT data + expected_size = lut_size ** 3 + if len(lut_data) == expected_size: + lut_array = np.array(lut_data, dtype=np.float32).reshape(lut_size, lut_size, lut_size, 3) + print(f"Successfully loaded .cube LUT: {len(lut_data)} entries") + return torch.from_numpy(lut_array) + else: + print(f"Invalid LUT data size: expected {expected_size}, got {len(lut_data)}") + + # Try to handle common size mismatches + if len(lut_data) > 0: + # Find the closest cube root + cube_root = round(len(lut_data) ** (1/3)) + if cube_root ** 3 == len(lut_data): + print(f"Adjusting LUT size to {cube_root}x{cube_root}x{cube_root}") + lut_array = np.array(lut_data, dtype=np.float32).reshape(cube_root, cube_root, cube_root, 3) + return torch.from_numpy(lut_array) + + return None + + except Exception as e: + print(f"Error loading .cube file {cube_path}: {e}") + return None + + def _load_3dl_lut(self, path_3dl): + """Load .3dl format LUT file""" + try: + with open(path_3dl, 'r') as f: + lines = f.readlines() + + lut_data = [] + for line in lines: + line = line.strip() + if line and not line.startswith('#'): + try: + values = [float(x) for x in line.split()] + if len(values) == 3: + lut_data.append(values) + except: + continue + + # 3dl files are typically 32x32x32 + if len(lut_data) == 32768: # 32^3 + lut_array = np.array(lut_data, dtype=np.float32).reshape(32, 32, 32, 3) + print(f"Successfully loaded .3dl LUT: 32x32x32") + return torch.from_numpy(lut_array) + else: + # Try to determine size + cube_root = round(len(lut_data) ** (1/3)) + if cube_root ** 3 == len(lut_data): + lut_array = np.array(lut_data, dtype=np.float32).reshape(cube_root, cube_root, cube_root, 3) + print(f"Successfully loaded .3dl LUT: {cube_root}x{cube_root}x{cube_root}") + return torch.from_numpy(lut_array) + else: + print(f"Invalid .3dl LUT size: {len(lut_data)} entries") + return None + + except Exception as e: + print(f"Error loading .3dl file {path_3dl}: {e}") + return None + + def _load_image_lut(self, image_path): + """Load LUT from image (supports various image formats including Photoshop LUT images)""" + try: + lut_image = Image.open(image_path).convert('RGB') + lut_array = np.array(lut_image, dtype=np.float32) / 255.0 + height, width = lut_array.shape[:2] + + print(f"Loading image LUT: {width}x{height}") + + # Common LUT image formats: + # 512x512 for 64^3 LUT (8x8 grid) + # 1024x32 for 32^3 LUT (32x1 strip) + # 256x16 for 16^3 LUT (16x1 strip) + + # Try different standard formats + formats = [ + (512, 512, 64), # 64^3 in 8x8 grid + (1024, 32, 32), # 32^3 in horizontal strip + (256, 16, 16), # 16^3 in horizontal strip + (2048, 64, 64), # 64^3 in horizontal strip + ] + + for expected_w, expected_h, lut_size in formats: + if width == expected_w and height == expected_h: + if expected_w == expected_h: + # Square format (grid layout) + grid_size = int(lut_size ** (1/3)) + lut_3d = np.zeros((lut_size, lut_size, lut_size, 3), dtype=np.float32) + + for r in range(lut_size): + for g in range(lut_size): + for b in range(lut_size): + # Calculate position in grid + grid_x = (b % grid_size) * lut_size + g + grid_y = (b // grid_size) * lut_size + r + + if grid_x < width and grid_y < height: + lut_3d[r, g, b] = lut_array[grid_y, grid_x] + else: + # Strip format + lut_3d = lut_array.reshape(lut_size, lut_size, lut_size, 3) + + print(f"Successfully loaded image LUT: {lut_size}x{lut_size}x{lut_size}") + return torch.from_numpy(lut_3d) + + # If no standard format matches, try to auto-detect + total_pixels = width * height + cube_root = round(total_pixels ** (1/3)) + + if cube_root ** 3 == total_pixels: + print(f"Auto-detected LUT size: {cube_root}x{cube_root}x{cube_root}") + lut_3d = lut_array.reshape(cube_root, cube_root, cube_root, 3) + return torch.from_numpy(lut_3d) + + print(f"Could not determine LUT format for {width}x{height} image") + return None + + except Exception as e: + print(f"Error loading image LUT: {e}") + return None + + def _apply_lut_tensor(self, images, lut, strength): + """Apply 3D LUT using tensor operations with trilinear interpolation""" + batch_size, height, width, channels = images.shape + lut_size = lut.shape[0] + + # Scale input to LUT coordinate space + coords = images * (lut_size - 1) + + # Get integer and fractional parts + coords_floor = torch.floor(coords).long() + coords_frac = coords - coords_floor.float() + + # Clamp coordinates + coords_floor = torch.clamp(coords_floor, 0, lut_size - 2) + coords_ceil = torch.clamp(coords_floor + 1, 0, lut_size - 1) + + # Trilinear interpolation + # Get 8 corner values + c000 = lut[coords_floor[:, :, :, 0], coords_floor[:, :, :, 1], coords_floor[:, :, :, 2]] + c001 = lut[coords_floor[:, :, :, 0], coords_floor[:, :, :, 1], coords_ceil[:, :, :, 2]] + c010 = lut[coords_floor[:, :, :, 0], coords_ceil[:, :, :, 1], coords_floor[:, :, :, 2]] + c011 = lut[coords_floor[:, :, :, 0], coords_ceil[:, :, :, 1], coords_ceil[:, :, :, 2]] + c100 = lut[coords_ceil[:, :, :, 0], coords_floor[:, :, :, 1], coords_floor[:, :, :, 2]] + c101 = lut[coords_ceil[:, :, :, 0], coords_floor[:, :, :, 1], coords_ceil[:, :, :, 2]] + c110 = lut[coords_ceil[:, :, :, 0], coords_ceil[:, :, :, 1], coords_floor[:, :, :, 2]] + c111 = lut[coords_ceil[:, :, :, 0], coords_ceil[:, :, :, 1], coords_ceil[:, :, :, 2]] + + # Interpolate + xf, yf, zf = coords_frac[:, :, :, 0:1], coords_frac[:, :, :, 1:2], coords_frac[:, :, :, 2:3] + + c00 = c000 * (1 - xf) + c100 * xf + c01 = c001 * (1 - xf) + c101 * xf + c10 = c010 * (1 - xf) + c110 * xf + c11 = c011 * (1 - xf) + c111 * xf + + c0 = c00 * (1 - yf) + c10 * yf + c1 = c01 * (1 - yf) + c11 * yf + + result_lut = c0 * (1 - zf) + c1 * zf + + # Blend with original based on strength + result = images * (1 - strength) + result_lut * strength + + return torch.clamp(result, 0, 1) diff --git a/nodes/image_fx/cross_processing.py b/nodes/image_fx/cross_processing.py new file mode 100644 index 0000000..899e3e7 --- /dev/null +++ b/nodes/image_fx/cross_processing.py @@ -0,0 +1,174 @@ +# APNext Cross Processing Effect Node + +import torch +import numpy as np +from ...utils.constants import CUSTOM_CATEGORY + + +class APNextCrossProcessing: + """ + APNext Cross Processing Effect Node + Simulates cross processing film techniques with channel manipulation and curves + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "process_type": ([ + "E6_in_C41", "C41_in_E6", "Custom", + "Vintage_Warm", "Cold_Blue", "Green_Magenta" + ], {"default": "E6_in_C41"}), + "intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + "contrast": ("FLOAT", {"default": 1.2, "min": 0.5, "max": 3.0, "step": 0.01}), + "saturation": ("FLOAT", {"default": 1.3, "min": 0.0, "max": 3.0, "step": 0.01}), + # Custom channel mixing + "red_in_red": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01}), + "green_in_red": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}), + "blue_in_red": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}), + "red_in_green": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}), + "green_in_green": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01}), + "blue_in_green": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}), + "red_in_blue": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}), + "green_in_blue": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}), + "blue_in_blue": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_cross_processing" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced" + + def apply_cross_processing(self, images, process_type, intensity, contrast, saturation, + red_in_red, green_in_red, blue_in_red, + red_in_green, green_in_green, blue_in_green, + red_in_blue, green_in_blue, blue_in_blue): + """Apply cross processing effect to images""" + device = images.device + dtype = images.dtype + + # Get the appropriate color matrix + if process_type == "Custom": + color_matrix = torch.tensor([ + [red_in_red, green_in_red, blue_in_red], + [red_in_green, green_in_green, blue_in_green], + [red_in_blue, green_in_blue, blue_in_blue] + ], device=device, dtype=dtype) + else: + color_matrix = self._get_preset_matrix(process_type, device, dtype) + + # Apply cross processing + result = self._apply_cross_process_tensor( + images, color_matrix, intensity, contrast, saturation + ) + + return (result,) + + def _get_preset_matrix(self, process_type, device, dtype): + """Get predefined color matrices for different cross processing types""" + matrices = { + "E6_in_C41": torch.tensor([ # Slide film in negative chemistry + [1.2, -0.1, 0.1], + [0.0, 1.1, -0.1], + [-0.1, 0.2, 1.3] + ], device=device, dtype=dtype), + + "C41_in_E6": torch.tensor([ # Negative film in slide chemistry + [0.9, 0.1, -0.1], + [0.1, 1.0, 0.1], + [0.2, -0.2, 0.8] + ], device=device, dtype=dtype), + + "Vintage_Warm": torch.tensor([ + [1.3, 0.1, -0.1], + [-0.1, 1.1, 0.2], + [-0.2, -0.1, 0.9] + ], device=device, dtype=dtype), + + "Cold_Blue": torch.tensor([ + [0.8, -0.1, 0.2], + [0.1, 1.0, -0.1], + [0.3, 0.2, 1.4] + ], device=device, dtype=dtype), + + "Green_Magenta": torch.tensor([ + [1.1, -0.2, 0.1], + [-0.3, 1.4, -0.1], + [0.1, -0.2, 1.0] + ], device=device, dtype=dtype), + } + + return matrices.get(process_type, torch.eye(3, device=device, dtype=dtype)) + + def _apply_cross_process_tensor(self, images, color_matrix, intensity, contrast, saturation): + """Apply cross processing using tensor operations""" + batch_size, height, width, channels = images.shape + + # Apply color matrix transformation + images_flat = images.view(-1, 3) + processed_flat = torch.matmul(images_flat, color_matrix.t()) + processed = processed_flat.view(batch_size, height, width, channels) + + # Apply contrast adjustment + if contrast != 1.0: + mid_point = 0.5 + processed = (processed - mid_point) * contrast + mid_point + + # Apply saturation adjustment + if saturation != 1.0: + processed = self._apply_saturation_tensor(processed, saturation) + + # Apply cross processing curves (S-curves for each channel) + processed = self._apply_film_curves_tensor(processed) + + # Blend with original based on intensity + result = images * (1 - intensity) + processed * intensity + + # Clamp values + result = torch.clamp(result, 0, 1) + return result + + def _apply_saturation_tensor(self, images, saturation): + """Apply saturation adjustment using tensor operations""" + # Calculate luminance + luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2] + luminance = luminance.unsqueeze(-1) + + # Blend between grayscale and original based on saturation + result = luminance + (images - luminance) * saturation + return result + + def _apply_film_curves_tensor(self, images): + """Apply characteristic film curves that create the cross-processing look""" + # Create S-curve for each channel with different characteristics + # Red channel: lifted shadows, compressed highlights + red = images[:, :, :, 0:1] + red_curved = self._s_curve_tensor(red, shadows=0.1, highlights=-0.1, contrast=1.1) + + # Green channel: more aggressive S-curve + green = images[:, :, :, 1:2] + green_curved = self._s_curve_tensor(green, shadows=0.05, highlights=-0.05, contrast=1.2) + + # Blue channel: inverted curve characteristics + blue = images[:, :, :, 2:3] + blue_curved = self._s_curve_tensor(blue, shadows=-0.05, highlights=0.15, contrast=1.1) + + return torch.cat([red_curved, green_curved, blue_curved], dim=3) + + def _s_curve_tensor(self, channel, shadows=0.0, highlights=0.0, contrast=1.0): + """Apply S-curve to a single channel""" + # Apply shadow/highlight adjustments + adjusted = channel + shadows * (1 - channel) + highlights * channel + + # Apply contrast with S-curve + # Use a sigmoid-based S-curve + mid_point = 0.5 + normalized = (adjusted - mid_point) * contrast + s_curved = torch.sigmoid(normalized * 2) + + # Scale back to 0-1 range + result = s_curved * (1 - mid_point) + mid_point + + return result diff --git a/nodes/image_fx/film_halation.py b/nodes/image_fx/film_halation.py new file mode 100644 index 0000000..b781af0 --- /dev/null +++ b/nodes/image_fx/film_halation.py @@ -0,0 +1,205 @@ +# APNext Film Halation Effect Node + +import torch +import torch.nn.functional as F +import numpy as np +from ...utils.constants import CUSTOM_CATEGORY + + +class APNextFilmHalation: + """ + APNext Film Halation Effect Node + Simulates film halation - light bleeding around bright objects + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "intensity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}), + "threshold": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}), + "radius": ("FLOAT", {"default": 20.0, "min": 1.0, "max": 100.0, "step": 0.5}), + "color_tint": (["none", "warm", "cool", "red", "blue", "custom"], {"default": "warm"}), + "custom_color_r": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + "custom_color_g": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 2.0, "step": 0.01}), + "custom_color_b": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01}), + "falloff": ("FLOAT", {"default": 2.0, "min": 0.5, "max": 4.0, "step": 0.1}), + "film_type": (["color_negative", "slide_film", "vintage", "modern"], {"default": "color_negative"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_film_halation" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced" + + def apply_film_halation(self, images, intensity, threshold, radius, color_tint, + custom_color_r, custom_color_g, custom_color_b, falloff, film_type): + """Apply film halation effect to images""" + device = images.device + dtype = images.dtype + + # Apply halation effect + result = self._create_halation_effect( + images, intensity, threshold, radius, color_tint, + custom_color_r, custom_color_g, custom_color_b, falloff, film_type + ) + + return (result,) + + def _create_halation_effect(self, images, intensity, threshold, radius, color_tint, + custom_color_r, custom_color_g, custom_color_b, falloff, film_type): + """Create film halation effect using tensor operations""" + batch_size, height, width, channels = images.shape + device = images.device + dtype = images.dtype + + # Calculate luminance for threshold detection + luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2] + + # Create mask for bright areas + bright_mask = (luminance > threshold).float() + + # Extract bright areas + bright_areas = images * bright_mask.unsqueeze(-1) + + # Apply film-specific characteristics + bright_areas = self._apply_film_characteristics(bright_areas, film_type) + + # Create halation glow + halation_glow = self._create_glow_effect(bright_areas, radius, falloff) + + # Apply color tinting + tinted_glow = self._apply_color_tint( + halation_glow, color_tint, custom_color_r, custom_color_g, custom_color_b + ) + + # Blend halation with original image + result = self._blend_halation(images, tinted_glow, intensity) + + return torch.clamp(result, 0, 1) + + def _apply_film_characteristics(self, bright_areas, film_type): + """Apply film-specific characteristics to bright areas""" + if film_type == "color_negative": + # Color negative film has strong red/orange halation + bright_areas[:, :, :, 0] *= 1.3 # Enhance red + bright_areas[:, :, :, 1] *= 1.1 # Slightly enhance green + bright_areas[:, :, :, 2] *= 0.8 # Reduce blue + + elif film_type == "slide_film": + # Slide film has more neutral halation with slight magenta cast + bright_areas[:, :, :, 0] *= 1.1 # Slight red enhancement + bright_areas[:, :, :, 1] *= 0.95 # Slight green reduction + bright_areas[:, :, :, 2] *= 1.05 # Slight blue enhancement + + elif film_type == "vintage": + # Vintage film has warm, yellowish halation + bright_areas[:, :, :, 0] *= 1.4 # Strong red + bright_areas[:, :, :, 1] *= 1.2 # Enhanced green (yellow) + bright_areas[:, :, :, 2] *= 0.6 # Reduced blue + + elif film_type == "modern": + # Modern film has more controlled, neutral halation + bright_areas *= 1.1 # Slight overall enhancement + + return torch.clamp(bright_areas, 0, 2) # Allow overexposure + + def _create_glow_effect(self, bright_areas, radius, falloff): + """Create glow effect using multiple gaussian blurs""" + # Create multiple blur layers for realistic halation + glow_layers = [] + + # Primary glow (tight) + primary_glow = self._gaussian_blur_tensor(bright_areas, radius * 0.3) + glow_layers.append(primary_glow * 0.6) + + # Secondary glow (medium) + secondary_glow = self._gaussian_blur_tensor(bright_areas, radius * 0.7) + glow_layers.append(secondary_glow * 0.3) + + # Tertiary glow (wide) + tertiary_glow = self._gaussian_blur_tensor(bright_areas, radius) + glow_layers.append(tertiary_glow * 0.1) + + # Combine glow layers + combined_glow = sum(glow_layers) + + # Apply falloff + if falloff != 1.0: + # Create distance-based falloff + luminance = 0.299 * combined_glow[:, :, :, 0] + 0.587 * combined_glow[:, :, :, 1] + 0.114 * combined_glow[:, :, :, 2] + falloff_factor = torch.pow(luminance.unsqueeze(-1), 1.0 / falloff) + combined_glow = combined_glow * falloff_factor + + return combined_glow + + def _apply_color_tint(self, glow, color_tint, custom_r, custom_g, custom_b): + """Apply color tinting to the halation glow""" + if color_tint == "none": + return glow + + tint_colors = { + "warm": torch.tensor([1.2, 1.0, 0.8], device=glow.device, dtype=glow.dtype), + "cool": torch.tensor([0.8, 1.0, 1.2], device=glow.device, dtype=glow.dtype), + "red": torch.tensor([1.5, 0.8, 0.8], device=glow.device, dtype=glow.dtype), + "blue": torch.tensor([0.8, 0.9, 1.4], device=glow.device, dtype=glow.dtype), + "custom": torch.tensor([custom_r, custom_g, custom_b], device=glow.device, dtype=glow.dtype) + } + + if color_tint in tint_colors: + tint = tint_colors[color_tint].view(1, 1, 1, 3) + tinted_glow = glow * tint + else: + tinted_glow = glow + + return tinted_glow + + def _blend_halation(self, original, halation, intensity): + """Blend halation effect with original image""" + # Use screen blend mode for realistic halation + # Screen formula: 1 - (1 - base) * (1 - overlay) + halation_scaled = halation * intensity + + # Apply screen blending + result = 1 - (1 - original) * (1 - halation_scaled) + + return result + + def _gaussian_blur_tensor(self, images, radius): + """Apply gaussian blur using torch operations""" + if radius < 0.5: + return images + + # Convert radius to sigma + sigma = radius / 3.0 + + # Create gaussian kernel + kernel_size = int(2 * radius + 1) + if kernel_size % 2 == 0: + kernel_size += 1 + + # Limit kernel size for performance + kernel_size = min(kernel_size, 99) + + # Create 1D gaussian kernel + x = torch.arange(kernel_size, dtype=images.dtype, device=images.device) + x = x - kernel_size // 2 + kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2) + kernel_1d = kernel_1d / kernel_1d.sum() + + # Apply separable gaussian blur + # Reshape for conv operations: [batch, channels, height, width] + images_conv = images.permute(0, 3, 1, 2) + + # Horizontal blur + kernel_h = kernel_1d.view(1, 1, 1, -1).expand(3, 1, 1, -1) + blurred_h = F.conv2d(images_conv, kernel_h, padding=(0, kernel_size//2), groups=3) + + # Vertical blur + kernel_v = kernel_1d.view(1, 1, -1, 1).expand(3, 1, -1, 1) + blurred = F.conv2d(blurred_h, kernel_v, padding=(kernel_size//2, 0), groups=3) + + # Reshape back to original format: [batch, height, width, channels] + return blurred.permute(0, 2, 3, 1) diff --git a/nodes/image_fx/glitch_art.py b/nodes/image_fx/glitch_art.py new file mode 100644 index 0000000..3dbd140 --- /dev/null +++ b/nodes/image_fx/glitch_art.py @@ -0,0 +1,315 @@ +# APNext Glitch Art Effect Node + +import torch +import numpy as np +import random +from ...utils.constants import CUSTOM_CATEGORY + + +class APNextGlitchArt: + """ + APNext Glitch Art Effect Node + Creates digital corruption and glitch effects + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "glitch_type": ([ + "RGB_Shift", "Data_Moshing", "Pixel_Sort", "Scanlines", + "Digital_Noise", "Color_Channel_Shift", "Compression_Artifacts" + ], {"default": "RGB_Shift"}), + "intensity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}), + "randomness": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_glitch_art" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced" + + def apply_glitch_art(self, images, glitch_type, intensity, randomness, seed): + """Apply glitch art effects to images""" + # Set random seed + random.seed(seed) + torch.manual_seed(seed % (2**32)) + + device = images.device + dtype = images.dtype + + if glitch_type == "RGB_Shift": + result = self._rgb_channel_shift(images, intensity, randomness) + elif glitch_type == "Data_Moshing": + result = self._data_moshing_effect(images, intensity, randomness) + elif glitch_type == "Pixel_Sort": + result = self._pixel_sorting_effect(images, intensity, randomness) + elif glitch_type == "Scanlines": + result = self._scanlines_effect(images, intensity, randomness) + elif glitch_type == "Digital_Noise": + result = self._digital_noise_effect(images, intensity, randomness) + elif glitch_type == "Color_Channel_Shift": + result = self._color_channel_shift(images, intensity, randomness) + elif glitch_type == "Compression_Artifacts": + result = self._compression_artifacts(images, intensity, randomness) + else: + result = images + + return (result,) + + def _rgb_channel_shift(self, images, intensity, randomness): + """Create RGB channel shift glitch effect""" + batch_size, height, width, channels = images.shape + device = images.device + + # Calculate shift amounts + max_shift = int(intensity * 20 + randomness * 10) + + # Create random shifts for each channel + r_shift_x = random.randint(-max_shift, max_shift) + r_shift_y = random.randint(-max_shift, max_shift) + g_shift_x = random.randint(-max_shift, max_shift) + g_shift_y = random.randint(-max_shift, max_shift) + b_shift_x = random.randint(-max_shift, max_shift) + b_shift_y = random.randint(-max_shift, max_shift) + + result = images.clone() + + # Apply shifts to each channel + for batch_idx in range(batch_size): + # Red channel shift + if abs(r_shift_x) > 0 or abs(r_shift_y) > 0: + shifted_red = torch.roll(images[batch_idx, :, :, 0], shifts=(r_shift_y, r_shift_x), dims=(0, 1)) + result[batch_idx, :, :, 0] = shifted_red + + # Green channel shift + if abs(g_shift_x) > 0 or abs(g_shift_y) > 0: + shifted_green = torch.roll(images[batch_idx, :, :, 1], shifts=(g_shift_y, g_shift_x), dims=(0, 1)) + result[batch_idx, :, :, 1] = shifted_green + + # Blue channel shift + if abs(b_shift_x) > 0 or abs(b_shift_y) > 0: + shifted_blue = torch.roll(images[batch_idx, :, :, 2], shifts=(b_shift_y, b_shift_x), dims=(0, 1)) + result[batch_idx, :, :, 2] = shifted_blue + + return torch.clamp(result, 0, 1) + + def _data_moshing_effect(self, images, intensity, randomness): + """Create data moshing effect by corrupting image data""" + result = images.clone() + batch_size, height, width, channels = images.shape + + # Number of corruption blocks + num_blocks = int(intensity * 50 + randomness * 20) + + for batch_idx in range(batch_size): + for _ in range(num_blocks): + # Random block position and size + block_size = random.randint(5, int(intensity * 50 + 10)) + x = random.randint(0, width - block_size) + y = random.randint(0, height - block_size) + + # Corruption type + corruption_type = random.choice(['duplicate', 'scramble', 'invert', 'noise']) + + if corruption_type == 'duplicate': + # Duplicate nearby block + source_x = max(0, min(width - block_size, x + random.randint(-30, 30))) + source_y = max(0, min(height - block_size, y + random.randint(-30, 30))) + result[batch_idx, y:y+block_size, x:x+block_size] = \ + images[batch_idx, source_y:source_y+block_size, source_x:source_x+block_size] + + elif corruption_type == 'scramble': + # Scramble pixels in block + block = result[batch_idx, y:y+block_size, x:x+block_size].clone() + flat_block = block.view(-1, channels) + indices = torch.randperm(flat_block.size(0)) + scrambled = flat_block[indices].view(block_size, block_size, channels) + result[batch_idx, y:y+block_size, x:x+block_size] = scrambled + + elif corruption_type == 'invert': + # Invert colors in block + result[batch_idx, y:y+block_size, x:x+block_size] = \ + 1.0 - result[batch_idx, y:y+block_size, x:x+block_size] + + elif corruption_type == 'noise': + # Add noise to block + noise = torch.randn_like(result[batch_idx, y:y+block_size, x:x+block_size]) * 0.1 + result[batch_idx, y:y+block_size, x:x+block_size] += noise + + return torch.clamp(result, 0, 1) + + def _pixel_sorting_effect(self, images, intensity, randomness): + """Create pixel sorting glitch effect""" + result = images.clone() + batch_size, height, width, channels = images.shape + + # Number of sorting operations + num_sorts = int(intensity * 20 + randomness * 10) + + for batch_idx in range(batch_size): + for _ in range(num_sorts): + # Random sorting direction and area + if random.random() < 0.5: + # Horizontal sorting + row = random.randint(0, height - 1) + start_col = random.randint(0, width - 20) + end_col = min(width, start_col + random.randint(10, int(intensity * 100 + 20))) + + # Sort pixels by brightness + row_data = result[batch_idx, row, start_col:end_col] + brightness = 0.299 * row_data[:, 0] + 0.587 * row_data[:, 1] + 0.114 * row_data[:, 2] + sorted_indices = torch.argsort(brightness) + result[batch_idx, row, start_col:end_col] = row_data[sorted_indices] + else: + # Vertical sorting + col = random.randint(0, width - 1) + start_row = random.randint(0, height - 20) + end_row = min(height, start_row + random.randint(10, int(intensity * 100 + 20))) + + # Sort pixels by brightness + col_data = result[batch_idx, start_row:end_row, col] + brightness = 0.299 * col_data[:, 0] + 0.587 * col_data[:, 1] + 0.114 * col_data[:, 2] + sorted_indices = torch.argsort(brightness) + result[batch_idx, start_row:end_row, col] = col_data[sorted_indices] + + return result + + def _scanlines_effect(self, images, intensity, randomness): + """Create scanlines and CRT-like glitch effects""" + result = images.clone() + batch_size, height, width, channels = images.shape + + # Scanline parameters + line_spacing = max(2, int(10 - intensity * 8)) + line_intensity = intensity * 0.5 + + for batch_idx in range(batch_size): + # Add horizontal scanlines + for y in range(0, height, line_spacing): + if random.random() < 0.7: # Not every line + # Darken scanline + result[batch_idx, y] *= (1.0 - line_intensity) + + # Add some color shift + if random.random() < randomness: + shift = random.randint(-2, 2) + if shift != 0: + result[batch_idx, y] = torch.roll(result[batch_idx, y], shift, dims=0) + + # Add random glitch lines + num_glitch_lines = int(intensity * 5 + randomness * 3) + for _ in range(num_glitch_lines): + y = random.randint(0, height - 1) + # Duplicate or shift line + if random.random() < 0.5: + # Duplicate from nearby line + source_y = max(0, min(height - 1, y + random.randint(-5, 5))) + result[batch_idx, y] = result[batch_idx, source_y] + else: + # Color shift + result[batch_idx, y, :, 0] = torch.roll(result[batch_idx, y, :, 0], random.randint(-10, 10)) + + return torch.clamp(result, 0, 1) + + def _digital_noise_effect(self, images, intensity, randomness): + """Add digital noise and artifacts""" + result = images.clone() + batch_size, height, width, channels = images.shape + + # Add random digital noise + noise_amount = intensity * randomness * 0.1 + noise = torch.randn_like(result) * noise_amount + result += noise + + # Add salt and pepper noise + salt_pepper_prob = intensity * randomness * 0.05 + mask = torch.rand_like(result[:, :, :, 0:1]) < salt_pepper_prob + salt_pepper = torch.rand_like(result) > 0.5 + result = torch.where(mask, salt_pepper.float(), result) + + # Add compression-like artifacts + if intensity > 0.3: + # Quantize colors + quantization_levels = max(8, int(256 - intensity * 200)) + result = torch.round(result * quantization_levels) / quantization_levels + + return torch.clamp(result, 0, 1) + + def _color_channel_shift(self, images, intensity, randomness): + """Shift color channels in different directions""" + result = images.clone() + batch_size, height, width, channels = images.shape + + # Create different shift patterns for each channel + for batch_idx in range(batch_size): + # Red channel - horizontal shift + h_shift = int(intensity * randomness * 20) + if h_shift > 0: + result[batch_idx, :, :, 0] = torch.roll(result[batch_idx, :, :, 0], h_shift, dims=1) + + # Green channel - vertical shift + v_shift = int(intensity * randomness * 15) + if v_shift > 0: + result[batch_idx, :, :, 1] = torch.roll(result[batch_idx, :, :, 1], v_shift, dims=0) + + # Blue channel - diagonal effect + if intensity > 0.5: + # Create diagonal shift effect + for y in range(height): + shift_amount = int((y / height) * intensity * 10) + if shift_amount > 0: + result[batch_idx, y, :, 2] = torch.roll(result[batch_idx, y, :, 2], shift_amount) + + return torch.clamp(result, 0, 1) + + def _compression_artifacts(self, images, intensity, randomness): + """Simulate compression artifacts and blocking""" + result = images.clone() + batch_size, height, width, channels = images.shape + + # Block size for compression artifacts + block_size = max(4, int(16 - intensity * 10)) + + for batch_idx in range(batch_size): + # Create blocking artifacts + for y in range(0, height, block_size): + for x in range(0, width, block_size): + end_y = min(height, y + block_size) + end_x = min(width, x + block_size) + + if random.random() < intensity * randomness: + # Average the block (compression-like effect) + block = result[batch_idx, y:end_y, x:end_x] + avg_color = torch.mean(block, dim=(0, 1), keepdim=True) + result[batch_idx, y:end_y, x:end_x] = avg_color + + elif random.random() < intensity * 0.3: + # Add ringing artifacts around edges + block = result[batch_idx, y:end_y, x:end_x] + # Simple edge detection + if block.size(0) > 1 and block.size(1) > 1: + # Calculate edges safely with proper size handling + vertical_edges = torch.abs(block[1:, :, :] - block[:-1, :, :]) + horizontal_edges = torch.abs(block[:, 1:, :] - block[:, :-1, :]) + + # Pad edges to match original block size + vertical_edges_padded = torch.zeros_like(block) + horizontal_edges_padded = torch.zeros_like(block) + + vertical_edges_padded[:-1, :, :] = vertical_edges + horizontal_edges_padded[:, :-1, :] = horizontal_edges + + edges = vertical_edges_padded + horizontal_edges_padded + edge_strength = torch.mean(edges) + + if edge_strength > 0.1: + # Add ringing + ringing = torch.sin(torch.arange(block.size(0), device=images.device).float() * 3.14159) * 0.05 * intensity + ringing = ringing.view(-1, 1, 1).expand_as(block) + result[batch_idx, y:end_y, x:end_x] += ringing + + return torch.clamp(result, 0, 1) diff --git a/nodes/image_fx/hdr_tone_mapping.py b/nodes/image_fx/hdr_tone_mapping.py new file mode 100644 index 0000000..3eb3e4b --- /dev/null +++ b/nodes/image_fx/hdr_tone_mapping.py @@ -0,0 +1,237 @@ +# APNext HDR Tone Mapping Effect Node + +import torch +import numpy as np +from ...utils.constants import CUSTOM_CATEGORY + + +class APNextHDRToneMapping: + """ + APNext HDR Tone Mapping Effect Node + Simulates HDR tone mapping effects for dramatic and surreal looks + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "method": ([ + "Reinhard", "Drago", "Mantiuk", "Photographic", + "Adaptive_Log", "Filmic", "ACES" + ], {"default": "Reinhard"}), + "exposure": ("FLOAT", {"default": 0.0, "min": -3.0, "max": 3.0, "step": 0.01}), + "gamma": ("FLOAT", {"default": 2.2, "min": 0.5, "max": 4.0, "step": 0.01}), + "intensity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01}), + # Method-specific parameters + "white_point": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.01}), + "key_value": ("FLOAT", {"default": 0.18, "min": 0.01, "max": 1.0, "step": 0.01}), + "adaptation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "color_correction": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}), + "local_adaptation": ("BOOLEAN", {"default": False}), + "radius": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_hdr_tone_mapping" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced" + + def apply_hdr_tone_mapping(self, images, method, exposure, gamma, intensity, + white_point, key_value, adaptation, color_correction, + local_adaptation, radius): + """Apply HDR tone mapping effect to images""" + device = images.device + dtype = images.dtype + + # Apply exposure adjustment first + if exposure != 0.0: + exposure_factor = 2.0 ** exposure + images = images * exposure_factor + + # Apply tone mapping based on selected method + if method == "Reinhard": + result = self._reinhard_tone_mapping(images, white_point, local_adaptation, radius) + elif method == "Drago": + result = self._drago_tone_mapping(images, adaptation, white_point) + elif method == "Mantiuk": + result = self._mantiuk_tone_mapping(images, color_correction, adaptation) + elif method == "Photographic": + result = self._photographic_tone_mapping(images, key_value, white_point) + elif method == "Adaptive_Log": + result = self._adaptive_log_tone_mapping(images, adaptation) + elif method == "Filmic": + result = self._filmic_tone_mapping(images, white_point) + elif method == "ACES": + result = self._aces_tone_mapping(images) + else: + result = images + + # Apply gamma correction + if gamma != 1.0: + result = torch.pow(torch.clamp(result, 0.001, 1.0), 1.0 / gamma) + + # Blend with original based on intensity + result = images * (1 - intensity) + result * intensity + + # Clamp final result + result = torch.clamp(result, 0, 1) + return (result,) + + def _reinhard_tone_mapping(self, images, white_point, local_adaptation, radius): + """Reinhard tone mapping operator""" + # Calculate luminance + luminance = self._calculate_luminance(images) + + if local_adaptation: + # Local adaptation version (simplified) + local_lum = self._gaussian_blur_luminance(luminance, radius) + adaptation_lum = local_lum / (local_lum + 1.0) + tone_mapped_lum = luminance / (luminance + adaptation_lum) + else: + # Global version + tone_mapped_lum = luminance / (luminance + white_point ** 2) + + # Apply to color channels + result = self._apply_luminance_change(images, luminance, tone_mapped_lum) + return result + + def _drago_tone_mapping(self, images, adaptation, white_point): + """Drago tone mapping operator""" + luminance = self._calculate_luminance(images) + + # Drago's adaptive logarithmic mapping + log_lum = torch.log10(torch.clamp(luminance, 1e-6, float('inf'))) + log_white = torch.log10(white_point) + + # Bias function + bias = torch.pow(adaptation, log_lum / log_white) + + # Apply tone mapping + tone_mapped_lum = (log_lum / log_white) / torch.log10(2 + 8 * bias) + tone_mapped_lum = torch.clamp(tone_mapped_lum, 0, 1) + + result = self._apply_luminance_change(images, luminance, tone_mapped_lum) + return result + + def _mantiuk_tone_mapping(self, images, color_correction, adaptation): + """Mantiuk tone mapping operator (simplified version)""" + luminance = self._calculate_luminance(images) + + # Logarithmic compression + log_lum = torch.log(torch.clamp(luminance, 1e-6, float('inf'))) + + # Adaptive factor + adaptive_factor = adaptation * torch.tanh(log_lum) + + # Tone mapping + tone_mapped_lum = torch.exp(log_lum * adaptive_factor) + tone_mapped_lum = tone_mapped_lum / (tone_mapped_lum + 1.0) + + # Color correction + if color_correction != 1.0: + saturation_factor = 1.0 + (color_correction - 1.0) * (1.0 - tone_mapped_lum) + result = self._apply_luminance_change(images, luminance, tone_mapped_lum) + result = self._apply_saturation_tensor(result, saturation_factor.mean().item()) + else: + result = self._apply_luminance_change(images, luminance, tone_mapped_lum) + + return result + + def _photographic_tone_mapping(self, images, key_value, white_point): + """Photographic tone reproduction operator""" + luminance = self._calculate_luminance(images) + + # Calculate log-average luminance + log_avg = torch.exp(torch.mean(torch.log(torch.clamp(luminance, 1e-6, float('inf'))))) + + # Scale luminance + scaled_lum = (key_value / log_avg) * luminance + + # Apply tone mapping with white point + tone_mapped_lum = scaled_lum * (1.0 + scaled_lum / (white_point ** 2)) / (1.0 + scaled_lum) + + result = self._apply_luminance_change(images, luminance, tone_mapped_lum) + return result + + def _adaptive_log_tone_mapping(self, images, adaptation): + """Adaptive logarithmic tone mapping""" + luminance = self._calculate_luminance(images) + + # Logarithmic mapping with adaptation + log_lum = torch.log(torch.clamp(luminance, 1e-6, float('inf'))) + max_log = torch.max(log_lum) + min_log = torch.min(log_lum) + + # Normalize and apply adaptation + normalized = (log_lum - min_log) / (max_log - min_log + 1e-6) + adapted = torch.pow(normalized, adaptation) + + result = self._apply_luminance_change(images, luminance, adapted) + return result + + def _filmic_tone_mapping(self, images, white_point): + """Filmic tone mapping (Uncharted 2 style)""" + def filmic_curve(x): + A, B, C, D, E, F = 0.15, 0.50, 0.10, 0.20, 0.02, 0.30 + return ((x * (A * x + C * B) + D * E) / (x * (A * x + B) + D * F)) - E / F + + # Apply filmic curve + tone_mapped = filmic_curve(images * 2.0) / filmic_curve(torch.tensor(white_point)) + + return torch.clamp(tone_mapped, 0, 1) + + def _aces_tone_mapping(self, images): + """ACES filmic tone mapping""" + # ACES RRT/ODT approximation + a, b, c, d, e = 2.51, 0.03, 2.43, 0.59, 0.14 + + tone_mapped = (images * (a * images + b)) / (images * (c * images + d) + e) + + return torch.clamp(tone_mapped, 0, 1) + + def _calculate_luminance(self, images): + """Calculate luminance from RGB""" + return 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2] + + def _apply_luminance_change(self, images, old_luminance, new_luminance): + """Apply luminance change while preserving color ratios""" + old_luminance = old_luminance.unsqueeze(-1) + new_luminance = new_luminance.unsqueeze(-1) + + # Avoid division by zero + ratio = torch.where(old_luminance > 1e-6, new_luminance / old_luminance, torch.ones_like(old_luminance)) + + result = images * ratio + return torch.clamp(result, 0, 1) + + def _gaussian_blur_luminance(self, luminance, radius): + """Apply gaussian blur to luminance for local adaptation""" + # Simple box blur approximation for performance + kernel_size = max(3, int(radius * 20)) + if kernel_size % 2 == 0: + kernel_size += 1 + + # Create simple averaging kernel + kernel = torch.ones(1, 1, kernel_size, kernel_size, device=luminance.device) / (kernel_size ** 2) + + # Add batch and channel dimensions for conv2d + lum_for_conv = luminance.unsqueeze(1) # [batch, 1, height, width] + + # Apply convolution + padding = kernel_size // 2 + blurred = torch.nn.functional.conv2d(lum_for_conv, kernel, padding=padding) + + # Remove extra dimension + return blurred.squeeze(1) + + def _apply_saturation_tensor(self, images, saturation): + """Apply saturation adjustment using tensor operations""" + # Calculate luminance + luminance = self._calculate_luminance(images) + luminance = luminance.unsqueeze(-1) + + # Blend between grayscale and original based on saturation + result = luminance + (images - luminance) * saturation + return torch.clamp(result, 0, 1) diff --git a/nodes/image_fx/noise.py b/nodes/image_fx/noise.py new file mode 100644 index 0000000..adba052 --- /dev/null +++ b/nodes/image_fx/noise.py @@ -0,0 +1,257 @@ +# APNext Noise Effect Node + +import torch +from PIL import Image, ImageDraw +import numpy as np +import random +from ...utils.constants import CUSTOM_CATEGORY +from ...utils.image_utils import tensor2pil, pil2tensor + +try: + from sklearn.cluster import KMeans + SKLEARN_AVAILABLE = True +except ImportError: + SKLEARN_AVAILABLE = False + + +class APNextNoise: + """ + APNext Noise Effect Node + Adds noise to images - can be B&W with transparency or colored based on dominant image colors + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "intensity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}), + "noise_type": (["monochrome", "colored", "film_grain"], {"default": "monochrome"}), + "blend_mode": (["overlay", "multiply", "screen", "soft_light"], {"default": "overlay"}), + "grain_size": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 5.0, "step": 0.1}), + "color_count": ("INT", {"default": 3, "min": 2, "max": 8, "step": 1}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_noise" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX" + + def apply_noise(self, images, intensity, noise_type, blend_mode, grain_size, color_count, seed): + """Apply noise effect to images""" + result_images = [] + + # Set random seed for reproducible results + random.seed(seed) + # For numpy, we need to clamp to 32-bit range + np.random.seed(seed % (2**32)) + + for image in images: + # Convert tensor to PIL + pil_image = tensor2pil(image) + + # Apply noise effect + noisy_image = self._create_noise_effect( + pil_image, intensity, noise_type, blend_mode, grain_size, color_count + ) + + # Convert back to tensor + result_tensor = pil2tensor(noisy_image) + result_images.append(result_tensor) + + # Stack all processed images + return (torch.cat(result_images, dim=0),) + + def _create_noise_effect(self, image, intensity, noise_type, blend_mode, grain_size, color_count): + """Create noise effect on a PIL image""" + # Convert to RGB if not already + if image.mode != 'RGB': + image = image.convert('RGB') + + width, height = image.size + + if noise_type == "monochrome": + noise_layer = self._create_monochrome_noise(width, height, grain_size) + elif noise_type == "colored": + dominant_colors = self._extract_dominant_colors(image, color_count) + noise_layer = self._create_colored_noise(width, height, dominant_colors, grain_size) + elif noise_type == "film_grain": + noise_layer = self._create_film_grain(width, height, grain_size) + + # Blend the noise with the original image + result = self._blend_noise(image, noise_layer, blend_mode, intensity) + + return result + + def _create_monochrome_noise(self, width, height, grain_size): + """Create black and white noise with transparency""" + # Create noise array + noise_size = int(max(width, height) / grain_size) + noise = np.random.randint(0, 256, (noise_size, noise_size), dtype=np.uint8) + + # Create PIL image from noise + noise_img = Image.fromarray(noise, mode='L') + + # Resize to match target dimensions + noise_img = noise_img.resize((width, height), Image.NEAREST if grain_size > 1 else Image.LANCZOS) + + # Convert to RGBA with transparency based on noise intensity + noise_rgba = Image.new('RGBA', (width, height)) + noise_array = np.array(noise_img) + + # Create RGBA array where alpha is based on the noise value + rgba_array = np.zeros((height, width, 4), dtype=np.uint8) + rgba_array[:, :, 0] = noise_array # R + rgba_array[:, :, 1] = noise_array # G + rgba_array[:, :, 2] = noise_array # B + rgba_array[:, :, 3] = noise_array # A (transparency based on noise) + + return Image.fromarray(rgba_array, 'RGBA') + + def _create_colored_noise(self, width, height, colors, grain_size): + """Create colored noise using dominant colors from the image""" + # Create noise array + noise_size = int(max(width, height) / grain_size) + + # Create RGB noise using the dominant colors + noise_r = np.random.randint(0, 256, (noise_size, noise_size), dtype=np.uint8) + noise_g = np.random.randint(0, 256, (noise_size, noise_size), dtype=np.uint8) + noise_b = np.random.randint(0, 256, (noise_size, noise_size), dtype=np.uint8) + + # Map noise values to dominant colors + color_indices = np.random.choice(len(colors), (noise_size, noise_size)) + + # Create colored noise + colored_noise = np.zeros((noise_size, noise_size, 3), dtype=np.uint8) + for i, color in enumerate(colors): + mask = color_indices == i + colored_noise[mask] = color + + # Add some randomness to the colors + variation = np.random.randint(-30, 31, (noise_size, noise_size, 3), dtype=np.int16) + colored_noise = np.clip(colored_noise.astype(np.int16) + variation, 0, 255).astype(np.uint8) + + # Create PIL image + noise_img = Image.fromarray(colored_noise) + + # Resize to match target dimensions + noise_img = noise_img.resize((width, height), Image.NEAREST if grain_size > 1 else Image.LANCZOS) + + return noise_img + + def _create_film_grain(self, width, height, grain_size): + """Create film grain effect""" + # Create multiple layers of noise for more realistic grain + noise_size = int(max(width, height) / grain_size) + + # Fine grain + fine_grain = np.random.normal(0, 0.3, (noise_size, noise_size)) + # Coarse grain + coarse_grain = np.random.normal(0, 0.1, (noise_size // 2, noise_size // 2)) + coarse_grain = np.repeat(np.repeat(coarse_grain, 2, axis=0), 2, axis=1)[:noise_size, :noise_size] + + # Combine grains + combined_grain = fine_grain + coarse_grain + + # Convert to 0-255 range + grain_normalized = ((combined_grain + 1) * 127.5).clip(0, 255).astype(np.uint8) + + # Create RGB grain + grain_rgb = np.stack([grain_normalized] * 3, axis=2) + + # Create PIL image + grain_img = Image.fromarray(grain_rgb) + + # Resize to match target dimensions + grain_img = grain_img.resize((width, height), Image.NEAREST if grain_size > 1 else Image.LANCZOS) + + return grain_img + + def _extract_dominant_colors(self, image, n_colors): + """Extract dominant colors from the image using optimized method""" + # Convert image to numpy array and reshape for clustering + img_array = np.array(image) + pixels = img_array.reshape(-1, 3) + + # Sample pixels for performance (use every 20th pixel for speed) + sampled_pixels = pixels[::20] + + if SKLEARN_AVAILABLE: + # Use K-means clustering if available + try: + kmeans = KMeans(n_clusters=n_colors, random_state=42, n_init=10) + kmeans.fit(sampled_pixels) + colors = kmeans.cluster_centers_.astype(int) + return colors.tolist() + except: + pass + + # Fallback: use histogram-based color extraction (faster than K-means) + return self._extract_colors_histogram(sampled_pixels, n_colors) + + def _extract_colors_histogram(self, pixels, n_colors): + """Extract colors using histogram-based method (faster fallback)""" + # Quantize colors to reduce search space + quantized = (pixels // 32) * 32 # Reduce to ~8 levels per channel + + # Find unique colors and their counts + unique_colors, counts = np.unique(quantized.reshape(-1, quantized.shape[-1]), + axis=0, return_counts=True) + + # Get the most frequent colors + most_frequent_indices = np.argsort(counts)[-n_colors:] + dominant_colors = unique_colors[most_frequent_indices] + + # Add some variation to avoid too uniform colors + variation = np.random.randint(-16, 17, dominant_colors.shape) + dominant_colors = np.clip(dominant_colors + variation, 0, 255) + + return dominant_colors.tolist() + + def _blend_noise(self, base_image, noise_layer, blend_mode, intensity): + """Blend noise with the base image using specified blend mode""" + # Convert to numpy arrays + base_array = np.array(base_image, dtype=np.float32) / 255.0 + + # Handle noise layer (might be RGBA or RGB) + if noise_layer.mode == 'RGBA': + noise_array = np.array(noise_layer, dtype=np.float32) + noise_rgb = noise_array[:, :, :3] / 255.0 + alpha = noise_array[:, :, 3] / 255.0 + else: + noise_rgb = np.array(noise_layer, dtype=np.float32) / 255.0 + alpha = np.ones((noise_rgb.shape[0], noise_rgb.shape[1])) * 0.5 + + # Apply intensity to alpha + alpha = alpha * intensity + + # Apply blend mode + if blend_mode == "overlay": + # Overlay blend mode + mask = base_array < 0.5 + blended = np.where( + mask, + 2 * base_array * noise_rgb, + 1 - 2 * (1 - base_array) * (1 - noise_rgb) + ) + elif blend_mode == "multiply": + blended = base_array * noise_rgb + elif blend_mode == "screen": + blended = 1 - (1 - base_array) * (1 - noise_rgb) + elif blend_mode == "soft_light": + blended = np.where( + noise_rgb <= 0.5, + base_array - (1 - 2 * noise_rgb) * base_array * (1 - base_array), + base_array + (2 * noise_rgb - 1) * (np.sqrt(base_array) - base_array) + ) + + # Apply alpha blending + alpha_expanded = np.expand_dims(alpha, axis=2) + result = base_array * (1 - alpha_expanded) + blended * alpha_expanded + + # Clamp and convert back to 0-255 + result = np.clip(result, 0, 1) * 255 + + return Image.fromarray(result.astype(np.uint8)) diff --git a/nodes/image_fx/rough.py b/nodes/image_fx/rough.py new file mode 100644 index 0000000..4821932 --- /dev/null +++ b/nodes/image_fx/rough.py @@ -0,0 +1,153 @@ +# APNext Rough Effect Node + +import torch +import numpy as np +from PIL import Image +from ...utils.constants import CUSTOM_CATEGORY +from ...utils.image_utils import tensor2pil, pil2tensor + + +class APNextRough: + """ + APNext Rough Effect Node + Creates a rough, posterized effect by reducing the image to a limited color palette + Based on numpy implementation for better performance + """ + + def __init__(self): + # Simple cache for color palettes + self.palette_cache = {} + self.max_cache_size = 10 + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "color_count": ("INT", {"default": 24, "min": 4, "max": 64, "step": 1}), + "aa_factor": ("INT", {"default": 4, "min": 1, "max": 8, "step": 1}), + "resize_method": (["BICUBIC", "LANCZOS", "BILINEAR", "NEAREST"], {"default": "BICUBIC"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_rough" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX" + + def apply_rough(self, images, color_count, aa_factor, resize_method): + """Apply rough/posterized effect to images""" + result_images = [] + + for image in images: + # Convert tensor to PIL + pil_image = tensor2pil(image) + + # Apply rough effect + rough_image = self._create_rough_effect( + pil_image, color_count, aa_factor, resize_method + ) + + # Convert back to tensor + result_tensor = pil2tensor(rough_image) + result_images.append(result_tensor) + + # Stack all processed images + return (torch.cat(result_images, dim=0),) + + def _create_rough_effect(self, image, color_count, aa_factor, resize_method): + """Create rough effect on a PIL image using numpy for efficiency""" + # Convert to RGB if not already + if image.mode != 'RGB': + image = image.convert('RGB') + + # Convert PIL to numpy array + img = np.asarray(image, dtype='int64') + height, width = img.shape[:2] + + # Try to get cached palette + cache_key = self._get_palette_cache_key(img, color_count) + if cache_key in self.palette_cache: + colors = self.palette_cache[cache_key] + print("Using cached palette.") + else: + # Generate new palette + colors = self._generate_palette(img, color_count) + + # Cache the palette + if len(self.palette_cache) >= self.max_cache_size: + # Remove oldest entry + oldest_key = next(iter(self.palette_cache)) + del self.palette_cache[oldest_key] + self.palette_cache[cache_key] = colors + + print("Created palette.") + print(colors) + + # Compute color distance of each pixel to each palette color (vectorized) + colors_array = np.array(colors) + img_flat = img.reshape(-1, 3) + + # Vectorized distance computation + distances = np.abs(img_flat[:, None, :] - colors_array[None, :, :]).sum(axis=2) + closest_indices = np.argmin(distances, axis=1) + closest_indices = closest_indices.reshape(height, width) + + # Create new image with anti-aliasing + new_img = np.zeros((height * aa_factor, width * aa_factor, 3), dtype=np.uint8) + + # Vectorized assignment of colors (much faster than nested loops) + for y in range(height): + for x in range(width): + col = colors[closest_indices[y, x]] + # Draw rectangles using numpy slicing + new_img[aa_factor * y : aa_factor * y + aa_factor, + aa_factor * x : aa_factor * x + aa_factor, :] = col + + print("Finished rough image.") + + # Convert back to PIL and resize + nim = Image.fromarray(new_img) + resize_filter = getattr(Image, resize_method) + aim = nim.resize((width, height), resize_filter) + + return aim + + def _get_palette_cache_key(self, img, color_count): + """Generate cache key for palette based on image statistics and color count""" + # Use image statistics for cache key (much faster than hashing entire image) + img_sample = img[::10, ::10] # Sample every 10th pixel + stats = ( + img_sample.mean(), + img_sample.std(), + img_sample.min(), + img_sample.max(), + color_count + ) + return hash(stats) + + def _generate_palette(self, img, color_count): + """Generate color palette using optimized algorithm""" + height, width = img.shape[:2] + + # Initialize color palette with black and white + colors = [(255, 255, 255), (0, 0, 0)] + colors_array = np.array(colors) + + for i in range(color_count - 2): + # Vectorized computation of distances to all existing colors + img_expanded = img.reshape(-1, 3) # Flatten image + + # Compute Manhattan distance to all existing colors at once + distances = np.abs(img_expanded[:, None, :] - colors_array[None, :, :]).sum(axis=2) + min_distances = distances.min(axis=1) + + # Find pixel with maximum minimum distance (furthest from all existing colors) + max_idx = np.argmax(min_distances) + new_color = img_expanded[max_idx] + + print(f"Added color: {new_color}") + colors.append(tuple(new_color)) + colors_array = np.array(colors) + + return colors diff --git a/nodes/image_fx/sharpen.py b/nodes/image_fx/sharpen.py new file mode 100644 index 0000000..5ebb28e --- /dev/null +++ b/nodes/image_fx/sharpen.py @@ -0,0 +1,147 @@ +# APNext Sharpen Effect Node + +import torch +import torch.nn.functional as F +import numpy as np +from ...utils.constants import CUSTOM_CATEGORY + + +class APNextSharpen: + """ + APNext Sharpen Effect Node + Sharpens images using unsharp mask technique for professional results + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}), + "radius": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}), + "threshold": ("INT", {"default": 0, "min": 0, "max": 255, "step": 1}), + "method": (["unsharp_mask", "high_pass", "edge_enhance"], {"default": "unsharp_mask"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_sharpen" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX" + + def apply_sharpen(self, images, strength, radius, threshold, method): + """Apply sharpening effect to images using optimized tensor operations""" + # Work directly with tensors for better performance + device = images.device + dtype = images.dtype + + # Process all images in batch + result = self._apply_sharpening_method_tensor( + images, strength, radius, threshold, method, device, dtype + ) + + return (result,) + + def _apply_sharpening_method_tensor(self, images, strength, radius, threshold, method, device, dtype): + """Apply the specified sharpening method using tensor operations""" + if method == "unsharp_mask": + return self._unsharp_mask_tensor(images, strength, radius, threshold) + elif method == "high_pass": + return self._high_pass_sharpen_tensor(images, strength, radius) + elif method == "edge_enhance": + return self._edge_enhance_tensor(images, strength) + + return images + + def _unsharp_mask_tensor(self, images, strength, radius, threshold): + """Apply unsharp mask sharpening using tensor operations""" + # Create blurred version using gaussian blur + blurred = self._gaussian_blur_tensor(images, radius) + + # Calculate the mask (difference between original and blurred) + mask = images - blurred + + # Apply threshold if specified + if threshold > 0: + threshold_norm = threshold / 255.0 + threshold_mask = torch.abs(mask) > threshold_norm + mask = torch.where(threshold_mask, mask, torch.zeros_like(mask)) + + # Apply the unsharp mask + sharpened = images + (mask * strength) + + # Clamp values to valid range + return torch.clamp(sharpened, 0, 1) + + def _high_pass_sharpen_tensor(self, images, strength, radius): + """Apply high-pass filter sharpening using tensor operations""" + # Create heavily blurred version for high-pass + blurred = self._gaussian_blur_tensor(images, radius * 2) + + # Create high-pass filter (original - blurred + 0.5) + high_pass = images - blurred + 0.5 + high_pass = torch.clamp(high_pass, 0, 1) + + # Blend with original using overlay mode + mask = images < 0.5 + result = torch.where( + mask, + 2 * images * high_pass, + 1 - 2 * (1 - images) * (1 - high_pass) + ) + + # Apply strength + result = images + (result - images) * strength + return torch.clamp(result, 0, 1) + + def _edge_enhance_tensor(self, images, strength): + """Apply edge enhancement sharpening using tensor operations""" + # Create a simple sharpening kernel + kernel = torch.tensor([ + [0, -1, 0], + [-1, 5, -1], + [0, -1, 0] + ], dtype=images.dtype, device=images.device).view(1, 1, 3, 3) + + # Expand kernel for all channels + kernel = kernel.expand(3, 1, 3, 3) + + # Apply convolution + images_conv = images.permute(0, 3, 1, 2) # [batch, channels, height, width] + sharpened_conv = F.conv2d(images_conv, kernel, padding=1, groups=3) + sharpened = sharpened_conv.permute(0, 2, 3, 1) # [batch, height, width, channels] + + # Blend with original + result = images + (sharpened - images) * strength + return torch.clamp(result, 0, 1) + + def _gaussian_blur_tensor(self, images, radius): + """Apply gaussian blur using torch operations (shared with bloom)""" + # Convert radius to sigma (approximation) + sigma = radius / 3.0 + + # Create gaussian kernel + kernel_size = int(2 * radius + 1) + if kernel_size % 2 == 0: + kernel_size += 1 + + # Create 1D gaussian kernel + x = torch.arange(kernel_size, dtype=images.dtype, device=images.device) + x = x - kernel_size // 2 + kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2) + kernel_1d = kernel_1d / kernel_1d.sum() + + # Apply separable gaussian blur + # Reshape for conv operations: [batch, channels, height, width] + images_conv = images.permute(0, 3, 1, 2) + + # Horizontal blur + kernel_h = kernel_1d.view(1, 1, 1, -1).expand(3, 1, 1, -1) + blurred_h = F.conv2d(images_conv, kernel_h, padding=(0, kernel_size//2), groups=3) + + # Vertical blur + kernel_v = kernel_1d.view(1, 1, -1, 1).expand(3, 1, -1, 1) + blurred = F.conv2d(blurred_h, kernel_v, padding=(kernel_size//2, 0), groups=3) + + # Reshape back to original format: [batch, height, width, channels] + return blurred.permute(0, 2, 3, 1) diff --git a/nodes/image_fx/split_toning.py b/nodes/image_fx/split_toning.py new file mode 100644 index 0000000..165e16a --- /dev/null +++ b/nodes/image_fx/split_toning.py @@ -0,0 +1,283 @@ +# APNext Split Toning Effect Node + +import torch +import numpy as np +from ...utils.constants import CUSTOM_CATEGORY + + +class APNextSplitToning: + """ + APNext Split Toning Effect Node + Applies different colors to highlights and shadows independently + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + # Highlight controls + "highlight_hue": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 1.0}), + "highlight_saturation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "highlight_luminance": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}), + # Shadow controls + "shadow_hue": ("FLOAT", {"default": 240.0, "min": 0.0, "max": 360.0, "step": 1.0}), + "shadow_saturation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "shadow_luminance": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + # Balance and blending + "balance": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), + "midtone_contrast": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), + "blend_mode": (["color", "soft_light", "overlay", "multiply"], {"default": "color"}), + "opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_split_toning" + CATEGORY = f"{CUSTOM_CATEGORY}/APNext FX Advanced" + + def apply_split_toning(self, images, highlight_hue, highlight_saturation, highlight_luminance, + shadow_hue, shadow_saturation, shadow_luminance, + balance, midtone_contrast, blend_mode, opacity): + """Apply split toning effect to images""" + device = images.device + dtype = images.dtype + + # Apply split toning + result = self._apply_split_toning_tensor( + images, highlight_hue, highlight_saturation, highlight_luminance, + shadow_hue, shadow_saturation, shadow_luminance, + balance, midtone_contrast, blend_mode, opacity + ) + + return (result,) + + def _apply_split_toning_tensor(self, images, h_hue, h_sat, h_lum, s_hue, s_sat, s_lum, + balance, midtone_contrast, blend_mode, opacity): + """Apply split toning using tensor operations""" + batch_size, height, width, channels = images.shape + device = images.device + dtype = images.dtype + + # Convert to HSV for easier manipulation + hsv_images = self._rgb_to_hsv_tensor(images) + + # Extract luminance for masking + luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2] + luminance = luminance.unsqueeze(-1) + + # Adjust luminance threshold based on balance + adjusted_h_lum = h_lum + balance * 0.2 + adjusted_s_lum = s_lum - balance * 0.2 + adjusted_h_lum = torch.clamp(torch.tensor(adjusted_h_lum), 0.0, 1.0) + adjusted_s_lum = torch.clamp(torch.tensor(adjusted_s_lum), 0.0, 1.0) + + # Create smooth masks for highlights and shadows + highlight_mask = self._create_luminance_mask(luminance, adjusted_h_lum, softness=0.2) + shadow_mask = self._create_luminance_mask(luminance, adjusted_s_lum, invert=True, softness=0.2) + + # Create toned colors + highlight_color = self._create_toned_color(h_hue, h_sat, device, dtype) + shadow_color = self._create_toned_color(s_hue, s_sat, device, dtype) + + # Apply toning based on blend mode + if blend_mode == "color": + result = self._apply_color_blend(images, highlight_color, shadow_color, + highlight_mask, shadow_mask) + elif blend_mode == "soft_light": + result = self._apply_soft_light_blend(images, highlight_color, shadow_color, + highlight_mask, shadow_mask) + elif blend_mode == "overlay": + result = self._apply_overlay_blend(images, highlight_color, shadow_color, + highlight_mask, shadow_mask) + elif blend_mode == "multiply": + result = self._apply_multiply_blend(images, highlight_color, shadow_color, + highlight_mask, shadow_mask) + + # Apply midtone contrast + if midtone_contrast != 0.0: + result = self._apply_midtone_contrast_tensor(result, midtone_contrast) + + # Blend with original based on opacity + result = images * (1 - opacity) + result * opacity + + # Clamp values + result = torch.clamp(result, 0, 1) + return result + + def _create_luminance_mask(self, luminance, threshold, invert=False, softness=0.1): + """Create smooth luminance-based mask""" + if invert: + mask = torch.sigmoid((threshold - luminance) / softness) + else: + mask = torch.sigmoid((luminance - threshold) / softness) + return mask + + def _create_toned_color(self, hue, saturation, device, dtype): + """Create a color from hue and saturation""" + # Convert hue from degrees to 0-1 range + h = (hue % 360.0) / 360.0 + s = saturation + v = 1.0 # Full brightness + + # Convert HSV to RGB + rgb = self._hsv_to_rgb_single(h, s, v, device, dtype) + return rgb + + def _hsv_to_rgb_single(self, h, s, v, device, dtype): + """Convert single HSV values to RGB""" + h = torch.tensor(h, device=device, dtype=dtype) + s = torch.tensor(s, device=device, dtype=dtype) + v = torch.tensor(v, device=device, dtype=dtype) + + c = v * s + x = c * (1 - torch.abs((h * 6) % 2 - 1)) + m = v - c + + h_i = torch.floor(h * 6).long() + + # Create RGB based on hue sector + if h_i == 0: + r, g, b = c, x, 0 + elif h_i == 1: + r, g, b = x, c, 0 + elif h_i == 2: + r, g, b = 0, c, x + elif h_i == 3: + r, g, b = 0, x, c + elif h_i == 4: + r, g, b = x, 0, c + else: + r, g, b = c, 0, x + + return torch.stack([r + m, g + m, b + m]) + + def _apply_color_blend(self, images, highlight_color, shadow_color, h_mask, s_mask): + """Apply color blend mode for split toning""" + # Convert to HSV + hsv = self._rgb_to_hsv_tensor(images) + + # Replace hue and saturation while preserving luminance + h_color_expanded = highlight_color.view(1, 1, 1, 3).expand_as(images) + s_color_expanded = shadow_color.view(1, 1, 1, 3).expand_as(images) + + # Blend colors + toned = images.clone() + toned = toned * (1 - h_mask) + h_color_expanded * h_mask + toned = toned * (1 - s_mask) + s_color_expanded * s_mask + + return toned + + def _apply_soft_light_blend(self, images, highlight_color, shadow_color, h_mask, s_mask): + """Apply soft light blend mode""" + h_color_expanded = highlight_color.view(1, 1, 1, 3).expand_as(images) + s_color_expanded = shadow_color.view(1, 1, 1, 3).expand_as(images) + + # Soft light formula + h_blend = torch.where( + h_color_expanded <= 0.5, + images - (1 - 2 * h_color_expanded) * images * (1 - images), + images + (2 * h_color_expanded - 1) * (torch.sqrt(images) - images) + ) + + s_blend = torch.where( + s_color_expanded <= 0.5, + images - (1 - 2 * s_color_expanded) * images * (1 - images), + images + (2 * s_color_expanded - 1) * (torch.sqrt(images) - images) + ) + + result = images * (1 - h_mask) + h_blend * h_mask + result = result * (1 - s_mask) + s_blend * s_mask + + return result + + def _apply_overlay_blend(self, images, highlight_color, shadow_color, h_mask, s_mask): + """Apply overlay blend mode""" + h_color_expanded = highlight_color.view(1, 1, 1, 3).expand_as(images) + s_color_expanded = shadow_color.view(1, 1, 1, 3).expand_as(images) + + # Overlay formula + h_blend = torch.where( + images <= 0.5, + 2 * images * h_color_expanded, + 1 - 2 * (1 - images) * (1 - h_color_expanded) + ) + + s_blend = torch.where( + images <= 0.5, + 2 * images * s_color_expanded, + 1 - 2 * (1 - images) * (1 - s_color_expanded) + ) + + result = images * (1 - h_mask) + h_blend * h_mask + result = result * (1 - s_mask) + s_blend * s_mask + + return result + + def _apply_multiply_blend(self, images, highlight_color, shadow_color, h_mask, s_mask): + """Apply multiply blend mode""" + h_color_expanded = highlight_color.view(1, 1, 1, 3).expand_as(images) + s_color_expanded = shadow_color.view(1, 1, 1, 3).expand_as(images) + + h_blend = images * h_color_expanded + s_blend = images * s_color_expanded + + result = images * (1 - h_mask) + h_blend * h_mask + result = result * (1 - s_mask) + s_blend * s_mask + + return result + + def _apply_midtone_contrast_tensor(self, images, contrast): + """Apply contrast specifically to midtones""" + # Create midtone mask + luminance = 0.299 * images[:, :, :, 0] + 0.587 * images[:, :, :, 1] + 0.114 * images[:, :, :, 2] + luminance = luminance.unsqueeze(-1) + + # Gaussian-like curve centered at 0.5 + midtone_mask = torch.exp(-((luminance - 0.5) ** 2) / (2 * 0.2 ** 2)) + + # Apply contrast + contrast_factor = 1.0 + contrast + mid_point = 0.5 + contrasted = (images - mid_point) * contrast_factor + mid_point + + # Blend based on midtone mask + result = images * (1 - midtone_mask) + contrasted * midtone_mask + return result + + def _rgb_to_hsv_tensor(self, rgb): + """Convert RGB to HSV color space""" + r, g, b = rgb[:, :, :, 0], rgb[:, :, :, 1], rgb[:, :, :, 2] + + max_val, max_idx = torch.max(rgb, dim=3) + min_val, _ = torch.min(rgb, dim=3) + + delta = max_val - min_val + + # Hue calculation + hue = torch.zeros_like(max_val) + mask = delta != 0 + + # Red is max + red_max = (max_idx == 0) & mask + hue[red_max] = ((g[red_max] - b[red_max]) / delta[red_max]) % 6 + + # Green is max + green_max = (max_idx == 1) & mask + hue[green_max] = (b[green_max] - r[green_max]) / delta[green_max] + 2 + + # Blue is max + blue_max = (max_idx == 2) & mask + hue[blue_max] = (r[blue_max] - g[blue_max]) / delta[blue_max] + 4 + + hue = hue / 6.0 + + # Saturation + saturation = torch.zeros_like(max_val) + saturation[max_val != 0] = delta[max_val != 0] / max_val[max_val != 0] + + # Value + value = max_val + + return torch.stack([hue, saturation, value], dim=3)