diff --git a/pbr_extractor_node.py b/pbr_extractor_node.py index 0c70ee9..eab6694 100644 --- a/pbr_extractor_node.py +++ b/pbr_extractor_node.py @@ -237,13 +237,246 @@ class PBRHeightProcessor: return (height,) +class LightingExtractor: + """ + Extract lighting components from Marigold Lighting output + Marigold Lighting outputs 3 images in a batch: + - Image 0: Albedo (0-1, linear space) + - Image 1: Diffuse shading (0-1, linear space) + - Image 2: Non-diffuse residual (0-1, linear space) + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "marigold_lighting": ("IMAGE",), + + # Gamma controls for each output + "albedo_gamma": ("FLOAT", { + "default": 0.4545, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma correction for albedo (linear space)" + }), + "diffuse_shading_gamma": ("FLOAT", { + "default": 0.4545, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma correction for diffuse shading (linear space)" + }), + "non_diffuse_gamma": ("FLOAT", { + "default": 0.4545, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma correction for non-diffuse residual (linear space)" + }), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("albedo", "diffuse_shading", "non_diffuse_residual") + FUNCTION = "extract" + CATEGORY = "Texture Alchemist/Extractors" + + def apply_gamma(self, image: torch.Tensor, gamma: float) -> torch.Tensor: + """Apply gamma correction""" + return torch.pow(torch.clamp(image, 0.0, 1.0), gamma) + + def extract(self, marigold_lighting, albedo_gamma, diffuse_shading_gamma, non_diffuse_gamma): + """Extract lighting components from Marigold batch output""" + + print("\n" + "="*60) + print("Lighting Extractor (Marigold)") + print("="*60) + print(f"Input shape: {marigold_lighting.shape}") + print(f"Batch size: {marigold_lighting.shape[0]}") + + # Extract each image from the batch + if marigold_lighting.shape[0] >= 1: + albedo = marigold_lighting[0:1] + print(f"\n🎨 Albedo (Image 0):") + print(f" Range: [{albedo.min():.3f}, {albedo.max():.3f}]") + print(f" Gamma: {albedo_gamma}") + albedo = self.apply_gamma(albedo, albedo_gamma) + print(f" Output range: [{albedo.min():.3f}, {albedo.max():.3f}]") + else: + print("⚠ Warning: No albedo image found in batch") + albedo = marigold_lighting[0:1] + + if marigold_lighting.shape[0] >= 2: + diffuse_shading = marigold_lighting[1:2] + print(f"\n💡 Diffuse Shading (Image 1):") + print(f" Range: [{diffuse_shading.min():.3f}, {diffuse_shading.max():.3f}]") + print(f" Gamma: {diffuse_shading_gamma}") + diffuse_shading = self.apply_gamma(diffuse_shading, diffuse_shading_gamma) + print(f" Output range: [{diffuse_shading.min():.3f}, {diffuse_shading.max():.3f}]") + else: + print("⚠ Warning: No diffuse shading image found in batch, using albedo") + diffuse_shading = albedo + + if marigold_lighting.shape[0] >= 3: + non_diffuse = marigold_lighting[2:3] + print(f"\n✨ Non-Diffuse Residual (Image 2):") + print(f" Range: [{non_diffuse.min():.3f}, {non_diffuse.max():.3f}]") + print(f" Gamma: {non_diffuse_gamma}") + non_diffuse = self.apply_gamma(non_diffuse, non_diffuse_gamma) + print(f" Output range: [{non_diffuse.min():.3f}, {non_diffuse.max():.3f}]") + else: + print("⚠ Warning: No non-diffuse residual found in batch, using albedo") + non_diffuse = albedo + + print(f"\n✓ Lighting extraction complete") + print("="*60 + "\n") + + return (albedo, diffuse_shading, non_diffuse) + + +class AppearanceExtractor: + """ + Extract appearance components from Marigold Appearance output + Marigold Appearance outputs 3 images in a batch: + - Image 0: Albedo (0-1, sRGB space) + - Image 1: Roughness and Metallicity (0-1, linear space) + - Red channel: Roughness + - Green channel: Metallicity + - Image 2: Uncertainty maps (optional, for ensemble size > 2) + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "marigold_appearance": ("IMAGE",), + + # Gamma controls for each output + "albedo_gamma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma correction for albedo (sRGB space)" + }), + "roughness_gamma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma correction for roughness (linear space)" + }), + "metallic_gamma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma correction for metallic (linear space)" + }), + "uncertainty_gamma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma correction for uncertainty maps" + }), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("albedo", "roughness", "metallic", "uncertainty") + FUNCTION = "extract" + CATEGORY = "Texture Alchemist/Extractors" + + def apply_gamma(self, image: torch.Tensor, gamma: float) -> torch.Tensor: + """Apply gamma correction""" + return torch.pow(torch.clamp(image, 0.0, 1.0), gamma) + + def channel_to_rgb(self, channel: torch.Tensor) -> torch.Tensor: + """Convert single channel to RGB""" + if channel.shape[-1] == 1: + return channel.repeat(1, 1, 1, 3) + return channel + + def extract(self, marigold_appearance, albedo_gamma, roughness_gamma, metallic_gamma, uncertainty_gamma): + """Extract appearance components from Marigold batch output""" + + print("\n" + "="*60) + print("Appearance Extractor (Marigold)") + print("="*60) + print(f"Input shape: {marigold_appearance.shape}") + print(f"Batch size: {marigold_appearance.shape[0]}") + + # Extract albedo (Image 0 - sRGB space) + if marigold_appearance.shape[0] >= 1: + albedo = marigold_appearance[0:1] + print(f"\n🎨 Albedo (Image 0 - sRGB):") + print(f" Range: [{albedo.min():.3f}, {albedo.max():.3f}]") + print(f" Gamma: {albedo_gamma}") + albedo = self.apply_gamma(albedo, albedo_gamma) + print(f" Output range: [{albedo.min():.3f}, {albedo.max():.3f}]") + else: + print("⚠ Warning: No albedo image found in batch") + albedo = marigold_appearance[0:1] + + # Extract material channels (Image 1 - linear space) + if marigold_appearance.shape[0] >= 2: + material_channels = marigold_appearance[1:2] + print(f"\n🔨 Material Channels (Image 1 - linear):") + print(f" Shape: {material_channels.shape}") + print(f" Range: [{material_channels.min():.3f}, {material_channels.max():.3f}]") + + # Extract roughness (red channel) + roughness = material_channels[:, :, :, 0:1] + print(f"\n 🟥 Roughness (Red Channel):") + print(f" Range: [{roughness.min():.3f}, {roughness.max():.3f}]") + print(f" Gamma: {roughness_gamma}") + roughness = self.apply_gamma(roughness, roughness_gamma) + roughness = self.channel_to_rgb(roughness) + print(f" Output range: [{roughness.min():.3f}, {roughness.max():.3f}]") + + # Extract metallic (green channel) + metallic = material_channels[:, :, :, 1:2] + print(f"\n 🟩 Metallicity (Green Channel):") + print(f" Range: [{metallic.min():.3f}, {metallic.max():.3f}]") + print(f" Gamma: {metallic_gamma}") + metallic = self.apply_gamma(metallic, metallic_gamma) + metallic = self.channel_to_rgb(metallic) + print(f" Output range: [{metallic.min():.3f}, {metallic.max():.3f}]") + else: + print("⚠ Warning: No material channels found in batch, using albedo") + roughness = albedo + metallic = albedo + + # Extract uncertainty maps (Image 2 - optional) + if marigold_appearance.shape[0] >= 3: + uncertainty = marigold_appearance[2:3] + print(f"\n📊 Uncertainty Maps (Image 2):") + print(f" Range: [{uncertainty.min():.3f}, {uncertainty.max():.3f}]") + print(f" Gamma: {uncertainty_gamma}") + uncertainty = self.apply_gamma(uncertainty, uncertainty_gamma) + print(f" Output range: [{uncertainty.min():.3f}, {uncertainty.max():.3f}]") + else: + print("⚠ Note: No uncertainty maps in batch (requires ensemble size > 2)") + # Create a zero/neutral uncertainty map + uncertainty = torch.zeros_like(albedo) + + print(f"\n✓ Appearance extraction complete") + print("="*60 + "\n") + + return (albedo, roughness, metallic, uncertainty) + + class FrankenMapExtractor: """ Extract PBR maps from FrankenMap format FrankenMap channel layout: - - Red Channel: Grayscale/Albedo + - Red Channel: Grayscale - Green Channel: Height - Blue Channel: Roughness + + Outputs grayscale, height, roughness, original image, and PBR pipe """ @classmethod @@ -259,7 +492,7 @@ class FrankenMapExtractor: "max": 3.0, "step": 0.01, "display": "slider", - "tooltip": "Gamma correction for albedo (from red channel)" + "tooltip": "Gamma correction for grayscale (from red channel)" }), "height_gamma": ("FLOAT", { "default": 1.0, @@ -285,7 +518,7 @@ class FrankenMapExtractor: "max": 3.0, "step": 0.01, "display": "slider", - "tooltip": "Albedo brightness multiplier" + "tooltip": "Grayscale brightness multiplier" }), "height_contrast": ("FLOAT", { "default": 1.0, @@ -316,8 +549,8 @@ class FrankenMapExtractor: } } - RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "PBR_PIPE") - RETURN_NAMES = ("albedo", "height", "roughness", "pbr_pipe") + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "PBR_PIPE") + RETURN_NAMES = ("grayscale", "height", "roughness", "original_image", "pbr_pipe") FUNCTION = "extract" CATEGORY = "Texture Alchemist/Extractors" @@ -371,12 +604,12 @@ class FrankenMapExtractor: blue_channel = frankenmap[:, :, :, 2:3] print(f"\n📊 CHANNEL EXTRACTION:") - print(f" Red (Albedo): range [{red_channel.min():.3f}, {red_channel.max():.3f}]") + print(f" Red (Grayscale): range [{red_channel.min():.3f}, {red_channel.max():.3f}]") print(f" Green (Height): range [{green_channel.min():.3f}, {green_channel.max():.3f}]") print(f" Blue (Roughness): range [{blue_channel.min():.3f}, {blue_channel.max():.3f}]") - # Process Albedo (from red channel) - print(f"\n🎨 PROCESSING ALBEDO:") + # Process Grayscale (from red channel) + print(f"\n🎨 PROCESSING GRAYSCALE:") print(f" Gamma: {albedo_gamma}") print(f" Brightness: {albedo_brightness}") @@ -433,14 +666,855 @@ class FrankenMapExtractor: "metallic": None, "transparency": None, "emission": None, + "image": frankenmap, # Store original input } print(f"\n✓ FrankenMap extraction complete") - print(f" Extracted: Albedo, Height, Roughness") + print(f" Extracted: Grayscale, Height, Roughness") + print(f" Original image stored in pipe") print(f" PBR Pipe created") print("="*60 + "\n") - return (albedo, height, roughness, pbr_pipe) + return (albedo, height, roughness, frankenmap, pbr_pipe) + + +class BasicLightingBuilder: + """ + Combine lighting components into final composite image + Formula: Composite = Albedo * Shading + Residual + + Perfect for recombining outputs from Lighting Extractor (Marigold) + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "albedo": ("IMAGE",), + "shading": ("IMAGE",), + "residual": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("composite",) + FUNCTION = "build" + CATEGORY = "Texture Alchemist/Materials" + + def build(self, albedo, shading, residual): + """Combine lighting components using multiplicative and additive blending""" + + print("\n" + "="*60) + print("Basic Lighting Builder") + print("="*60) + print(f"Albedo shape: {albedo.shape}") + print(f"Shading shape: {shading.shape}") + print(f"Residual shape: {residual.shape}") + + # Step 1: Multiply shading over albedo + print(f"\n📐 Step 1: Multiply Shading over Albedo") + print(f" Albedo range: [{albedo.min():.3f}, {albedo.max():.3f}]") + print(f" Shading range: [{shading.min():.3f}, {shading.max():.3f}]") + + base = albedo * shading + print(f" Result range: [{base.min():.3f}, {base.max():.3f}]") + + # Step 2: Add residual + print(f"\n➕ Step 2: Add Residual") + print(f" Base range: [{base.min():.3f}, {base.max():.3f}]") + print(f" Residual range: [{residual.min():.3f}, {residual.max():.3f}]") + + composite = base + residual + print(f" Pre-clamp range: [{composite.min():.3f}, {composite.max():.3f}]") + + # Clamp to valid range + composite = torch.clamp(composite, 0.0, 1.0) + + print(f"\n✓ Lighting composite complete") + print(f" Final range: [{composite.min():.3f}, {composite.max():.3f}]") + print(f" Formula: Composite = Albedo * Shading + Residual") + print("="*60 + "\n") + + return (composite,) + + +class CompositeMaskAdjuster: + """ + Apply mask-based adjustments to all three lighting passes simultaneously + Allows targeted artistic changes that affect Albedo, Shading, and Residual consistently + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("IMAGE",), + + # Adjustment controls + "brightness": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 3.0, + "step": 0.01, + "tooltip": "Brightness multiplier in masked areas" + }), + "contrast": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 3.0, + "step": 0.01, + "tooltip": "Contrast adjustment in masked areas" + }), + "gamma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma correction in masked areas" + }), + "saturation": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 3.0, + "step": 0.01, + "tooltip": "Saturation adjustment in masked areas (affects albedo only)" + }), + }, + "optional": { + "albedo": ("IMAGE",), + "shading": ("IMAGE",), + "residual": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("albedo", "shading", "residual") + FUNCTION = "adjust" + CATEGORY = "Texture Alchemist/Materials" + + def apply_brightness(self, image: torch.Tensor, brightness: float) -> torch.Tensor: + """Apply brightness adjustment""" + return image * brightness + + def apply_contrast(self, image: torch.Tensor, contrast: float) -> torch.Tensor: + """Apply contrast adjustment around midpoint""" + midpoint = 0.5 + return (image - midpoint) * contrast + midpoint + + def apply_gamma(self, image: torch.Tensor, gamma: float) -> torch.Tensor: + """Apply gamma correction""" + return torch.pow(torch.clamp(image, 0.0, 1.0), gamma) + + def apply_saturation(self, image: torch.Tensor, saturation: float) -> torch.Tensor: + """Apply saturation adjustment (luminance-based)""" + # Calculate luminance + if image.shape[-1] >= 3: + weights = torch.tensor([0.2989, 0.5870, 0.1140], device=image.device, dtype=image.dtype) + luminance = torch.sum(image[..., :3] * weights, dim=-1, keepdim=True) + + # Blend between luminance (desaturated) and original (saturated) + result = luminance + (image[..., :3] - luminance) * saturation + + # If there are extra channels (alpha, etc.), preserve them + if image.shape[-1] > 3: + result = torch.cat([result, image[..., 3:]], dim=-1) + + return result + else: + # Grayscale image, saturation has no effect + return image + + def prepare_mask(self, mask: torch.Tensor, target_shape: tuple) -> torch.Tensor: + """Prepare mask to match target image shape""" + # Convert mask to grayscale if needed + if mask.shape[-1] > 1: + weights = torch.tensor([0.299, 0.587, 0.114], device=mask.device, dtype=mask.dtype) + mask = torch.sum(mask[..., :3] * weights, dim=-1, keepdim=True) + + # Expand mask to match target channels + if mask.shape[-1] != target_shape[-1]: + mask = mask.repeat(1, 1, 1, target_shape[-1]) + + # Clamp mask to [0, 1] + mask = torch.clamp(mask, 0.0, 1.0) + + return mask + + def adjust(self, mask, brightness, contrast, gamma, saturation, albedo=None, shading=None, residual=None): + """Apply masked adjustments to all three passes""" + + print("\n" + "="*60) + print("Composite Mask Adjuster") + print("="*60) + print(f"Mask shape: {mask.shape}") + if albedo is not None: + print(f"Albedo shape: {albedo.shape}") + if shading is not None: + print(f"Shading shape: {shading.shape}") + if residual is not None: + print(f"Residual shape: {residual.shape}") + print(f"\nAdjustments:") + print(f" Brightness: {brightness}") + print(f" Contrast: {contrast}") + print(f" Gamma: {gamma}") + print(f" Saturation: {saturation}") + + print(f"\nMask range: [{mask.min():.3f}, {mask.max():.3f}]") + + # Process Albedo (if provided) + if albedo is not None: + mask_albedo = self.prepare_mask(mask, albedo.shape) + + print(f"\n🎨 Processing Albedo:") + print(f" Input range: [{albedo.min():.3f}, {albedo.max():.3f}]") + + albedo_adjusted = albedo.clone() + albedo_adjusted = self.apply_brightness(albedo_adjusted, brightness) + albedo_adjusted = self.apply_contrast(albedo_adjusted, contrast) + albedo_adjusted = self.apply_gamma(albedo_adjusted, gamma) + albedo_adjusted = self.apply_saturation(albedo_adjusted, saturation) + albedo_adjusted = torch.clamp(albedo_adjusted, 0.0, 1.0) + + # Blend based on mask + albedo_final = albedo * (1.0 - mask_albedo) + albedo_adjusted * mask_albedo + print(f" Output range: [{albedo_final.min():.3f}, {albedo_final.max():.3f}]") + else: + albedo_final = None + print(f"\n🎨 Albedo: Not provided (skipped)") + + # Process Shading (if provided) + if shading is not None: + mask_shading = self.prepare_mask(mask, shading.shape) + + print(f"\n💡 Processing Shading:") + print(f" Input range: [{shading.min():.3f}, {shading.max():.3f}]") + + shading_adjusted = shading.clone() + shading_adjusted = self.apply_brightness(shading_adjusted, brightness) + shading_adjusted = self.apply_contrast(shading_adjusted, contrast) + shading_adjusted = self.apply_gamma(shading_adjusted, gamma) + shading_adjusted = torch.clamp(shading_adjusted, 0.0, 1.0) + + # Blend based on mask + shading_final = shading * (1.0 - mask_shading) + shading_adjusted * mask_shading + print(f" Output range: [{shading_final.min():.3f}, {shading_final.max():.3f}]") + else: + shading_final = None + print(f"\n💡 Shading: Not provided (skipped)") + + # Process Residual (if provided) + if residual is not None: + mask_residual = self.prepare_mask(mask, residual.shape) + + print(f"\n✨ Processing Residual:") + print(f" Input range: [{residual.min():.3f}, {residual.max():.3f}]") + + residual_adjusted = residual.clone() + residual_adjusted = self.apply_brightness(residual_adjusted, brightness) + residual_adjusted = self.apply_contrast(residual_adjusted, contrast) + residual_adjusted = self.apply_gamma(residual_adjusted, gamma) + residual_adjusted = torch.clamp(residual_adjusted, 0.0, 1.0) + + # Blend based on mask + residual_final = residual * (1.0 - mask_residual) + residual_adjusted * mask_residual + print(f" Output range: [{residual_final.min():.3f}, {residual_final.max():.3f}]") + else: + residual_final = None + print(f"\n✨ Residual: Not provided (skipped)") + + print(f"\n✓ Composite mask adjustment complete") + print("="*60 + "\n") + + return (albedo_final, shading_final, residual_final) + + +class CompositeMaskExtractor: + """ + Extract masks from lighting passes based on various criteria + Useful for creating targeted adjustments based on lighting properties + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + # Extraction method + "extract_from": (["Albedo", "Shading", "Residual", "Combined"], { + "default": "Shading", + "tooltip": "Which pass to extract mask from" + }), + + "method": (["Luminance", "Threshold", "Range", "Inverse Luminance"], { + "default": "Luminance", + "tooltip": "Method for mask extraction" + }), + + # Threshold controls + "threshold": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Threshold value for Threshold method" + }), + + "range_min": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Minimum value for Range method" + }), + + "range_max": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Maximum value for Range method" + }), + + # Post-processing + "feather": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 0.5, + "step": 0.01, + "tooltip": "Soften mask edges (0 = hard edges)" + }), + + "gamma": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 3.0, + "step": 0.01, + "tooltip": "Gamma adjustment for mask" + }), + }, + "optional": { + "albedo": ("IMAGE",), + "shading": ("IMAGE",), + "residual": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("mask",) + FUNCTION = "extract" + CATEGORY = "Texture Alchemist/Materials" + + def get_luminance(self, image: torch.Tensor) -> torch.Tensor: + """Calculate luminance from RGB image""" + if image.shape[-1] >= 3: + weights = torch.tensor([0.2989, 0.5870, 0.1140], device=image.device, dtype=image.dtype) + return torch.sum(image[..., :3] * weights, dim=-1, keepdim=True) + else: + return image[..., :1] + + def apply_feather(self, mask: torch.Tensor, feather: float) -> torch.Tensor: + """Soften mask edges""" + if feather <= 0.0: + return mask + + # Simple feathering using contrast reduction around threshold + feather_strength = 1.0 - feather * 2.0 # 0.0-0.5 maps to 1.0-0.0 + midpoint = 0.5 + return torch.clamp((mask - midpoint) * feather_strength + midpoint, 0.0, 1.0) + + def extract(self, extract_from, method, threshold, range_min, range_max, feather, gamma, + albedo=None, shading=None, residual=None): + """Extract mask from lighting passes""" + + print("\n" + "="*60) + print("Composite Mask Extractor") + print("="*60) + print(f"Extract from: {extract_from}") + print(f"Method: {method}") + + # Select source image + if extract_from == "Albedo": + if albedo is None: + raise ValueError("Albedo input required when extract_from is 'Albedo'") + source = albedo + print(f"Source: Albedo, range [{source.min():.3f}, {source.max():.3f}]") + elif extract_from == "Shading": + if shading is None: + raise ValueError("Shading input required when extract_from is 'Shading'") + source = shading + print(f"Source: Shading, range [{source.min():.3f}, {source.max():.3f}]") + elif extract_from == "Residual": + if residual is None: + raise ValueError("Residual input required when extract_from is 'Residual'") + source = residual + print(f"Source: Residual, range [{source.min():.3f}, {source.max():.3f}]") + else: # Combined + # Average all available passes + available_passes = [] + if albedo is not None: + available_passes.append(albedo) + if shading is not None: + available_passes.append(shading) + if residual is not None: + available_passes.append(residual) + + if not available_passes: + raise ValueError("At least one pass required when extract_from is 'Combined'") + + source = sum(available_passes) / len(available_passes) + print(f"Source: Combined ({len(available_passes)} passes), range [{source.min():.3f}, {source.max():.3f}]") + + # Extract mask based on method + if method == "Luminance": + print(f"\n🔍 Extracting by Luminance") + mask = self.get_luminance(source) + + elif method == "Inverse Luminance": + print(f"\n🔍 Extracting by Inverse Luminance") + mask = 1.0 - self.get_luminance(source) + + elif method == "Threshold": + print(f"\n🔍 Extracting by Threshold") + print(f" Threshold: {threshold}") + luminance = self.get_luminance(source) + mask = (luminance >= threshold).float() + + elif method == "Range": + print(f"\n🔍 Extracting by Range") + print(f" Range: [{range_min}, {range_max}]") + luminance = self.get_luminance(source) + mask = ((luminance >= range_min) & (luminance <= range_max)).float() + + print(f" Mask range (pre-processing): [{mask.min():.3f}, {mask.max():.3f}]") + + # Apply feathering + if feather > 0.0: + print(f"\n🌫️ Applying feather: {feather}") + mask = self.apply_feather(mask, feather) + print(f" Mask range (post-feather): [{mask.min():.3f}, {mask.max():.3f}]") + + # Apply gamma + if gamma != 1.0: + print(f"\n🎚️ Applying gamma: {gamma}") + mask = torch.pow(torch.clamp(mask, 0.0, 1.0), gamma) + print(f" Mask range (post-gamma): [{mask.min():.3f}, {mask.max():.3f}]") + + # Convert to RGB for compatibility + if mask.shape[-1] == 1: + mask = mask.repeat(1, 1, 1, 3) + + print(f"\n✓ Mask extraction complete") + print(f" Final mask range: [{mask.min():.3f}, {mask.max():.3f}]") + print("="*60 + "\n") + + return (mask,) + + +class LightingPassMaskApplicator: + """ + Apply a mask to lighting passes (albedo, shading, residual) + Outputs the masked-out version of each input (mask defines what to keep) + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("IMAGE",), + }, + "optional": { + "albedo": ("IMAGE",), + "shading": ("IMAGE",), + "residual": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("albedo", "shading", "residual") + FUNCTION = "apply_mask" + CATEGORY = "Texture Alchemist/Materials" + + def prepare_mask(self, mask: torch.Tensor, target_shape: tuple) -> torch.Tensor: + """Prepare mask to match target image shape""" + # Convert mask to grayscale if needed + if mask.shape[-1] > 1: + weights = torch.tensor([0.299, 0.587, 0.114], device=mask.device, dtype=mask.dtype) + mask = torch.sum(mask[..., :3] * weights, dim=-1, keepdim=True) + + # Expand mask to match target channels + if mask.shape[-1] != target_shape[-1]: + mask = mask.repeat(1, 1, 1, target_shape[-1]) + + # Clamp mask to [0, 1] + mask = torch.clamp(mask, 0.0, 1.0) + + return mask + + def apply_mask(self, mask, albedo=None, shading=None, residual=None): + """Apply mask to all provided passes""" + + print("\n" + "="*60) + print("Lighting Pass Mask Applicator") + print("="*60) + print(f"Mask shape: {mask.shape}") + print(f"Mask range: [{mask.min():.3f}, {mask.max():.3f}]") + + # Apply mask to albedo if provided + if albedo is not None: + mask_albedo = self.prepare_mask(mask, albedo.shape) + albedo_masked = albedo * mask_albedo + print(f"\n🎨 Albedo:") + print(f" Input range: [{albedo.min():.3f}, {albedo.max():.3f}]") + print(f" Output range: [{albedo_masked.min():.3f}, {albedo_masked.max():.3f}]") + else: + albedo_masked = None + print(f"\n🎨 Albedo: Not provided") + + # Apply mask to shading if provided + if shading is not None: + mask_shading = self.prepare_mask(mask, shading.shape) + shading_masked = shading * mask_shading + print(f"\n💡 Shading:") + print(f" Input range: [{shading.min():.3f}, {shading.max():.3f}]") + print(f" Output range: [{shading_masked.min():.3f}, {shading_masked.max():.3f}]") + else: + shading_masked = None + print(f"\n💡 Shading: Not provided") + + # Apply mask to residual if provided + if residual is not None: + mask_residual = self.prepare_mask(mask, residual.shape) + residual_masked = residual * mask_residual + print(f"\n✨ Residual:") + print(f" Input range: [{residual.min():.3f}, {residual.max():.3f}]") + print(f" Output range: [{residual_masked.min():.3f}, {residual_masked.max():.3f}]") + else: + residual_masked = None + print(f"\n✨ Residual: Not provided") + + print(f"\n✓ Mask applied to all passes") + print("="*60 + "\n") + + return (albedo_masked, shading_masked, residual_masked) + + +class PassMaskApplicator: + """ + Apply a mask to up to 6 images + Outputs the masked-out version of each input (mask defines what to keep) + General purpose mask applicator for any images + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("IMAGE",), + }, + "optional": { + "image_1": ("IMAGE",), + "image_2": ("IMAGE",), + "image_3": ("IMAGE",), + "image_4": ("IMAGE",), + "image_5": ("IMAGE",), + "image_6": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("image_1", "image_2", "image_3", "image_4", "image_5", "image_6") + FUNCTION = "apply_mask" + CATEGORY = "Texture Alchemist/Adjustment" + + def prepare_mask(self, mask: torch.Tensor, target_shape: tuple) -> torch.Tensor: + """Prepare mask to match target image shape""" + # Convert mask to grayscale if needed + if mask.shape[-1] > 1: + weights = torch.tensor([0.299, 0.587, 0.114], device=mask.device, dtype=mask.dtype) + mask = torch.sum(mask[..., :3] * weights, dim=-1, keepdim=True) + + # Expand mask to match target channels + if mask.shape[-1] != target_shape[-1]: + mask = mask.repeat(1, 1, 1, target_shape[-1]) + + # Clamp mask to [0, 1] + mask = torch.clamp(mask, 0.0, 1.0) + + return mask + + def apply_mask(self, mask, image_1=None, image_2=None, image_3=None, + image_4=None, image_5=None, image_6=None): + """Apply mask to all provided images""" + + print("\n" + "="*60) + print("Pass Mask Applicator") + print("="*60) + print(f"Mask shape: {mask.shape}") + print(f"Mask range: [{mask.min():.3f}, {mask.max():.3f}]") + + results = [] + images = [image_1, image_2, image_3, image_4, image_5, image_6] + + for i, image in enumerate(images, 1): + if image is not None: + mask_prepared = self.prepare_mask(mask, image.shape) + masked_image = image * mask_prepared + results.append(masked_image) + print(f"\n📷 Image {i}:") + print(f" Input range: [{image.min():.3f}, {image.max():.3f}]") + print(f" Output range: [{masked_image.min():.3f}, {masked_image.max():.3f}]") + else: + results.append(None) + print(f"\n📷 Image {i}: Not provided") + + print(f"\n✓ Mask applied to all images") + print("="*60 + "\n") + + return tuple(results) + + +class MultimattePassApplicator: + """ + Apply a specific channel from a multimatte mask to up to 6 images + Extracts Red, Green, or Blue channel and uses it as a mask + Perfect for VFX workflows with channel-packed masks + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "multimatte_mask": ("IMAGE",), + "channel": (["Red", "Green", "Blue", "Luminance"], { + "default": "Red", + "tooltip": "Which channel to use as the mask" + }), + }, + "optional": { + "image_1": ("IMAGE",), + "image_2": ("IMAGE",), + "image_3": ("IMAGE",), + "image_4": ("IMAGE",), + "image_5": ("IMAGE",), + "image_6": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE") + RETURN_NAMES = ("image_1", "image_2", "image_3", "image_4", "image_5", "image_6") + FUNCTION = "apply_mask" + CATEGORY = "Texture Alchemist/Adjustment" + + def extract_channel(self, multimatte: torch.Tensor, channel: str) -> torch.Tensor: + """Extract specific channel from multimatte""" + if channel == "Luminance": + # Calculate luminance from RGB + if multimatte.shape[-1] >= 3: + weights = torch.tensor([0.2989, 0.5870, 0.1140], device=multimatte.device, dtype=multimatte.dtype) + return torch.sum(multimatte[..., :3] * weights, dim=-1, keepdim=True) + else: + return multimatte[..., 0:1] + + if multimatte.shape[-1] < 3 and channel != "Luminance": + print(f"⚠ Warning: Multimatte has fewer than 3 channels, using first channel") + return multimatte[..., 0:1] + + if channel == "Red": + return multimatte[..., 0:1] + elif channel == "Green": + return multimatte[..., 1:2] + elif channel == "Blue": + return multimatte[..., 2:3] + else: + return multimatte[..., 0:1] + + def prepare_mask(self, mask: torch.Tensor, target_shape: tuple) -> torch.Tensor: + """Prepare mask to match target image shape""" + # Mask is already single channel from extract_channel + + # Expand mask to match target channels + if mask.shape[-1] != target_shape[-1]: + mask = mask.repeat(1, 1, 1, target_shape[-1]) + + # Clamp mask to [0, 1] + mask = torch.clamp(mask, 0.0, 1.0) + + return mask + + def apply_mask(self, multimatte_mask, channel, image_1=None, image_2=None, image_3=None, + image_4=None, image_5=None, image_6=None): + """Apply selected channel mask to all provided images""" + + print("\n" + "="*60) + print("Multimatte Pass Applicator") + print("="*60) + print(f"Multimatte shape: {multimatte_mask.shape}") + print(f"Selected channel: {channel}") + + # Extract the selected channel + mask = self.extract_channel(multimatte_mask, channel) + print(f"Extracted mask range: [{mask.min():.3f}, {mask.max():.3f}]") + + results = [] + images = [image_1, image_2, image_3, image_4, image_5, image_6] + + for i, image in enumerate(images, 1): + if image is not None: + mask_prepared = self.prepare_mask(mask, image.shape) + masked_image = image * mask_prepared + results.append(masked_image) + print(f"\n📷 Image {i}:") + print(f" Input range: [{image.min():.3f}, {image.max():.3f}]") + print(f" Output range: [{masked_image.min():.3f}, {masked_image.max():.3f}]") + else: + results.append(None) + print(f"\n📷 Image {i}: Not provided") + + print(f"\n✓ {channel} channel mask applied to all images") + print("="*60 + "\n") + + return tuple(results) + + +class SinglePassMaskApplicator: + """ + Apply a mask to a single image + Simple mask applicator for one image at a time + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mask": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_mask" + CATEGORY = "Texture Alchemist/Adjustment" + + def prepare_mask(self, mask: torch.Tensor, target_shape: tuple) -> torch.Tensor: + """Prepare mask to match target image shape""" + # Convert mask to grayscale if needed + if mask.shape[-1] > 1: + weights = torch.tensor([0.299, 0.587, 0.114], device=mask.device, dtype=mask.dtype) + mask = torch.sum(mask[..., :3] * weights, dim=-1, keepdim=True) + + # Expand mask to match target channels + if mask.shape[-1] != target_shape[-1]: + mask = mask.repeat(1, 1, 1, target_shape[-1]) + + # Clamp mask to [0, 1] + mask = torch.clamp(mask, 0.0, 1.0) + + return mask + + def apply_mask(self, image, mask): + """Apply mask to image""" + + print("\n" + "="*60) + print("Single Pass Mask Applicator") + print("="*60) + print(f"Image shape: {image.shape}") + print(f"Mask shape: {mask.shape}") + print(f"Mask range: [{mask.min():.3f}, {mask.max():.3f}]") + + mask_prepared = self.prepare_mask(mask, image.shape) + masked_image = image * mask_prepared + + print(f"\nInput range: [{image.min():.3f}, {image.max():.3f}]") + print(f"Output range: [{masked_image.min():.3f}, {masked_image.max():.3f}]") + print(f"\n✓ Mask applied") + print("="*60 + "\n") + + return (masked_image,) + + +class SingleMultimattePassApplicator: + """ + Apply a specific channel from a multimatte mask to a single image + Extracts Red, Green, Blue, or Luminance channel and uses it as a mask + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "multimatte_mask": ("IMAGE",), + "channel": (["Red", "Green", "Blue", "Luminance"], { + "default": "Red", + "tooltip": "Which channel to use as the mask" + }), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "apply_mask" + CATEGORY = "Texture Alchemist/Adjustment" + + def extract_channel(self, multimatte: torch.Tensor, channel: str) -> torch.Tensor: + """Extract specific channel from multimatte""" + if channel == "Luminance": + # Calculate luminance from RGB + if multimatte.shape[-1] >= 3: + weights = torch.tensor([0.2989, 0.5870, 0.1140], device=multimatte.device, dtype=multimatte.dtype) + return torch.sum(multimatte[..., :3] * weights, dim=-1, keepdim=True) + else: + return multimatte[..., 0:1] + + if multimatte.shape[-1] < 3 and channel != "Luminance": + print(f"⚠ Warning: Multimatte has fewer than 3 channels, using first channel") + return multimatte[..., 0:1] + + if channel == "Red": + return multimatte[..., 0:1] + elif channel == "Green": + return multimatte[..., 1:2] + elif channel == "Blue": + return multimatte[..., 2:3] + else: + return multimatte[..., 0:1] + + def prepare_mask(self, mask: torch.Tensor, target_shape: tuple) -> torch.Tensor: + """Prepare mask to match target image shape""" + # Mask is already single channel from extract_channel + + # Expand mask to match target channels + if mask.shape[-1] != target_shape[-1]: + mask = mask.repeat(1, 1, 1, target_shape[-1]) + + # Clamp mask to [0, 1] + mask = torch.clamp(mask, 0.0, 1.0) + + return mask + + def apply_mask(self, image, multimatte_mask, channel): + """Apply selected channel mask to image""" + + print("\n" + "="*60) + print("Single Multimatte Pass Applicator") + print("="*60) + print(f"Image shape: {image.shape}") + print(f"Multimatte shape: {multimatte_mask.shape}") + print(f"Selected channel: {channel}") + + # Extract the selected channel + mask = self.extract_channel(multimatte_mask, channel) + print(f"Extracted mask range: [{mask.min():.3f}, {mask.max():.3f}]") + + mask_prepared = self.prepare_mask(mask, image.shape) + masked_image = image * mask_prepared + + print(f"\nInput range: [{image.min():.3f}, {image.max():.3f}]") + print(f"Output range: [{masked_image.min():.3f}, {masked_image.max():.3f}]") + print(f"\n✓ {channel} channel mask applied") + print("="*60 + "\n") + + return (masked_image,) # Node registration @@ -448,6 +1522,16 @@ NODE_CLASS_MAPPINGS = { "PBRMaterialProcessor": PBRMaterialProcessor, "PBRNormalProcessor": PBRNormalProcessor, "PBRHeightProcessor": PBRHeightProcessor, + "LightingExtractor": LightingExtractor, + "AppearanceExtractor": AppearanceExtractor, + "BasicLightingBuilder": BasicLightingBuilder, + "CompositeMaskAdjuster": CompositeMaskAdjuster, + "CompositeMaskExtractor": CompositeMaskExtractor, + "LightingPassMaskApplicator": LightingPassMaskApplicator, + "PassMaskApplicator": PassMaskApplicator, + "MultimattePassApplicator": MultimattePassApplicator, + "SinglePassMaskApplicator": SinglePassMaskApplicator, + "SingleMultimattePassApplicator": SingleMultimattePassApplicator, "FrankenMapExtractor": FrankenMapExtractor, } @@ -455,5 +1539,15 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PBRMaterialProcessor": "PBR Material Processor (Marigold)", "PBRNormalProcessor": "PBR Normal Processor (Lotus)", "PBRHeightProcessor": "PBR Height Processor (Lotus)", + "LightingExtractor": "Lighting Extractor (Marigold)", + "AppearanceExtractor": "Appearance Extractor (Marigold)", + "BasicLightingBuilder": "Basic Lighting Builder", + "CompositeMaskAdjuster": "Composite Mask Adjuster", + "CompositeMaskExtractor": "Composite Mask Extractor", + "LightingPassMaskApplicator": "Lighting Pass Mask Applicator", + "PassMaskApplicator": "Pass Mask Applicator", + "MultimattePassApplicator": "Multimatte Pass Applicator", + "SinglePassMaskApplicator": "Single Pass Mask Applicator", + "SingleMultimattePassApplicator": "Single Multimatte Pass Applicator", "FrankenMapExtractor": "PBR Extractor (FrankenMap)", }