""" Relight Utilities Image relighting based on normal map dot product - inspired by Blender's Normal node """ import torch import numpy as np class NormalMapRelighter: """ Relight an image using a normal map and light direction Uses dot product between normal vectors and light direction to calculate lighting Similar to Blender's Normal node shader """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "normal_map": ("IMAGE",), # Light direction (normalized vector) "light_x": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, "tooltip": "Light direction X component" }), "light_y": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, "tooltip": "Light direction Y component" }), "light_z": ("FLOAT", { "default": 1.0, "min": -1.0, "max": 1.0, "step": 0.01, "tooltip": "Light direction Z component (pointing toward camera)" }), # Light properties "light_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Light intensity multiplier" }), "ambient_light": ("FLOAT", { "default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Ambient light level (minimum brightness)" }), # Normal map format "normal_format": (["OpenGL", "DirectX"], { "default": "OpenGL", "tooltip": "Normal map format (affects Y channel)" }), } } RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("relit_image", "lighting_mask") FUNCTION = "relight" CATEGORY = "Texture Alchemist/Lighting" def decode_normal_map(self, normal_map: torch.Tensor, format: str) -> torch.Tensor: """ Decode normal map from [0,1] color space to [-1,1] vector space """ # Convert from 0-1 to -1 to 1 normals = normal_map * 2.0 - 1.0 # Flip Y for DirectX format if format == "DirectX": normals[:, :, :, 1] = -normals[:, :, :, 1] # Normalize the vectors length = torch.sqrt(torch.sum(normals ** 2, dim=-1, keepdim=True)) length = torch.clamp(length, min=1e-8) normals = normals / length return normals def calculate_lighting(self, normals: torch.Tensor, light_dir: torch.Tensor, intensity: float, ambient: float) -> torch.Tensor: """ Calculate lighting using dot product between normals and light direction """ # Normalize light direction light_length = torch.sqrt(torch.sum(light_dir ** 2)) if light_length > 1e-8: light_dir = light_dir / light_length # Calculate dot product (N · L) # normals shape: [B, H, W, 3] # light_dir shape: [3] dot_product = torch.sum(normals * light_dir, dim=-1, keepdim=True) # Clamp dot product to [0, 1] (only front-facing surfaces are lit) dot_product = torch.clamp(dot_product, 0.0, 1.0) # Apply intensity and ambient light lighting = dot_product * intensity + ambient # Clamp final lighting lighting = torch.clamp(lighting, 0.0, 1.0) return lighting def relight(self, image, normal_map, light_x, light_y, light_z, light_intensity, ambient_light, normal_format): """ Relight image using normal map and light direction """ print("\n" + "="*60) print("Normal Map Relighter") print("="*60) print(f"Image shape: {image.shape}") print(f"Normal map shape: {normal_map.shape}") print(f"Light direction: ({light_x:.3f}, {light_y:.3f}, {light_z:.3f})") print(f"Light intensity: {light_intensity}") print(f"Ambient light: {ambient_light}") print(f"Normal format: {normal_format}") # Decode normal map print(f"\n🗺️ Decoding normal map...") normals = self.decode_normal_map(normal_map, normal_format) print(f" Normal range: [{normals.min():.3f}, {normals.max():.3f}]") # Create light direction vector light_dir = torch.tensor([light_x, light_y, light_z], device=image.device, dtype=image.dtype) # Calculate lighting print(f"\n💡 Calculating lighting...") lighting_mask = self.calculate_lighting(normals, light_dir, light_intensity, ambient_light) print(f" Lighting mask range: [{lighting_mask.min():.3f}, {lighting_mask.max():.3f}]") # Apply lighting to image # Expand lighting mask to match image channels if image.shape[-1] > 1: lighting_expanded = lighting_mask.repeat(1, 1, 1, image.shape[-1]) else: lighting_expanded = lighting_mask relit_image = image * lighting_expanded # Convert lighting mask to RGB for output lighting_mask_rgb = lighting_mask.repeat(1, 1, 1, 3) print(f"\n✓ Relighting complete") print(f" Relit image range: [{relit_image.min():.3f}, {relit_image.max():.3f}]") print("="*60 + "\n") return (relit_image, lighting_mask_rgb) class NormalMapRelighterAdvanced: """ Advanced relighting with multiple light types and properties Supports directional, point, and ambient lighting """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "normal_map": ("IMAGE",), # Light direction (normalized vector) "light_x": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, "tooltip": "Light direction X component" }), "light_y": ("FLOAT", { "default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, "tooltip": "Light direction Y component" }), "light_z": ("FLOAT", { "default": 1.0, "min": -1.0, "max": 1.0, "step": 0.01, "tooltip": "Light direction Z component" }), # Light properties "diffuse_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Diffuse light intensity" }), "specular_intensity": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Specular highlight intensity" }), "specular_power": ("FLOAT", { "default": 32.0, "min": 1.0, "max": 128.0, "step": 1.0, "tooltip": "Specular power (shininess)" }), "ambient_light": ("FLOAT", { "default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Ambient light level" }), # Light color "light_color_r": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Light color red channel" }), "light_color_g": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Light color green channel" }), "light_color_b": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Light color blue channel" }), # Normal map format "normal_format": (["OpenGL", "DirectX"], { "default": "OpenGL", "tooltip": "Normal map format" }), } } RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE") RETURN_NAMES = ("relit_image", "diffuse_lighting", "specular_lighting") FUNCTION = "relight" CATEGORY = "Texture Alchemist/Lighting" def decode_normal_map(self, normal_map: torch.Tensor, format: str) -> torch.Tensor: """Decode normal map from color space to vector space""" normals = normal_map * 2.0 - 1.0 if format == "DirectX": normals[:, :, :, 1] = -normals[:, :, :, 1] length = torch.sqrt(torch.sum(normals ** 2, dim=-1, keepdim=True)) length = torch.clamp(length, min=1e-8) normals = normals / length return normals def calculate_diffuse(self, normals: torch.Tensor, light_dir: torch.Tensor, intensity: float) -> torch.Tensor: """Calculate diffuse lighting (Lambertian)""" dot_product = torch.sum(normals * light_dir, dim=-1, keepdim=True) dot_product = torch.clamp(dot_product, 0.0, 1.0) return dot_product * intensity def calculate_specular(self, normals: torch.Tensor, light_dir: torch.Tensor, view_dir: torch.Tensor, intensity: float, power: float) -> torch.Tensor: """Calculate specular lighting (Blinn-Phong)""" # Calculate halfway vector half_vector = light_dir + view_dir half_length = torch.sqrt(torch.sum(half_vector ** 2)) if half_length > 1e-8: half_vector = half_vector / half_length # Calculate specular term spec_dot = torch.sum(normals * half_vector, dim=-1, keepdim=True) spec_dot = torch.clamp(spec_dot, 0.0, 1.0) specular = torch.pow(spec_dot, power) * intensity return specular def relight(self, image, normal_map, light_x, light_y, light_z, diffuse_intensity, specular_intensity, specular_power, ambient_light, light_color_r, light_color_g, light_color_b, normal_format): """Advanced relighting with diffuse and specular components""" print("\n" + "="*60) print("Advanced Normal Map Relighter") print("="*60) print(f"Light direction: ({light_x:.3f}, {light_y:.3f}, {light_z:.3f})") print(f"Light color: ({light_color_r:.3f}, {light_color_g:.3f}, {light_color_b:.3f})") # Decode normal map normals = self.decode_normal_map(normal_map, normal_format) # Normalize light and view directions light_dir = torch.tensor([light_x, light_y, light_z], device=image.device, dtype=image.dtype) light_length = torch.sqrt(torch.sum(light_dir ** 2)) if light_length > 1e-8: light_dir = light_dir / light_length # View direction (camera looking at surface) view_dir = torch.tensor([0.0, 0.0, 1.0], device=image.device, dtype=image.dtype) # Calculate lighting components diffuse = self.calculate_diffuse(normals, light_dir, diffuse_intensity) specular = self.calculate_specular(normals, light_dir, view_dir, specular_intensity, specular_power) # Create light color tensor light_color = torch.tensor([light_color_r, light_color_g, light_color_b], device=image.device, dtype=image.dtype) # Combine lighting components total_lighting = diffuse + ambient_light total_lighting = torch.clamp(total_lighting, 0.0, 1.0) # Expand to RGB total_lighting_rgb = total_lighting.repeat(1, 1, 1, 3) * light_color specular_rgb = specular.repeat(1, 1, 1, 3) * light_color # Apply to image relit_image = image * total_lighting_rgb + specular_rgb relit_image = torch.clamp(relit_image, 0.0, 1.0) # Output masks diffuse_rgb = total_lighting.repeat(1, 1, 1, 3) specular_output = specular.repeat(1, 1, 1, 3) print(f"\n✓ Advanced relighting complete") print(f" Diffuse range: [{diffuse.min():.3f}, {diffuse.max():.3f}]") print(f" Specular range: [{specular.min():.3f}, {specular.max():.3f}]") print("="*60 + "\n") return (relit_image, diffuse_rgb, specular_output) class NormalMapRelighterIBL: """ Image-Based Lighting (IBL) using equirectangular environment maps Samples environment map colors based on normal directions and reflection vectors Perfect for realistic HDRI lighting """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "normal_map": ("IMAGE",), "environment_map": ("IMAGE",), # Environment rotation "environment_rotation": ("FLOAT", { "default": 0.0, "min": -180.0, "max": 180.0, "step": 1.0, "tooltip": "Rotate environment map (degrees)" }), # Lighting properties "diffuse_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Diffuse lighting intensity" }), "specular_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Specular reflection intensity" }), "specular_power": ("FLOAT", { "default": 32.0, "min": 1.0, "max": 128.0, "step": 1.0, "tooltip": "Specular power (blur amount - lower = blurrier)" }), "ambient_light": ("FLOAT", { "default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Ambient light level" }), # Normal map format "normal_format": (["OpenGL", "DirectX"], { "default": "OpenGL", "tooltip": "Normal map format" }), } } RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE") RETURN_NAMES = ("relit_image", "diffuse_lighting", "specular_lighting") FUNCTION = "relight_ibl" CATEGORY = "Texture Alchemist/Lighting" def decode_normal_map(self, normal_map: torch.Tensor, format: str) -> torch.Tensor: """Decode normal map from color space to vector space""" normals = normal_map * 2.0 - 1.0 if format == "DirectX": normals[:, :, :, 1] = -normals[:, :, :, 1] length = torch.sqrt(torch.sum(normals ** 2, dim=-1, keepdim=True)) length = torch.clamp(length, min=1e-8) normals = normals / length return normals def direction_to_equirect_uv(self, direction: torch.Tensor, rotation_radians: float) -> torch.Tensor: """ Convert 3D direction vector to equirectangular UV coordinates direction: [B, H, W, 3] normalized vectors returns: [B, H, W, 2] UV coordinates in range [0, 1] """ x = direction[:, :, :, 0] y = direction[:, :, :, 1] z = direction[:, :, :, 2] # Calculate latitude (theta) and longitude (phi) # theta = arcsin(y) - vertical angle # phi = atan2(x, z) - horizontal angle theta = torch.asin(torch.clamp(y, -1.0, 1.0)) phi = torch.atan2(x, z) # Apply rotation to phi phi = phi + rotation_radians # Convert to UV coordinates [0, 1] u = (phi / (2 * np.pi)) + 0.5 v = (theta / np.pi) + 0.5 # Wrap U coordinate u = u % 1.0 # Clamp V coordinate v = torch.clamp(v, 0.0, 1.0) # Stack to create UV tensor uv = torch.stack([u, v], dim=-1) return uv def sample_equirect(self, env_map: torch.Tensor, uv: torch.Tensor) -> torch.Tensor: """ Sample equirectangular map using UV coordinates with bilinear interpolation env_map: [B, H, W, C] uv: [B, H, W, 2] returns: [B, H, W, C] """ batch_env, height_env, width_env, channels = env_map.shape batch_uv, height_uv, width_uv, _ = uv.shape # Convert UV to pixel coordinates # U maps to width, V maps to height px = uv[:, :, :, 0] * (width_env - 1) py = uv[:, :, :, 1] * (height_env - 1) # Bilinear interpolation px0 = torch.floor(px).long() py0 = torch.floor(py).long() px1 = px0 + 1 py1 = py0 + 1 # Clamp coordinates px0 = torch.clamp(px0, 0, width_env - 1) px1 = torch.clamp(px1, 0, width_env - 1) py0 = torch.clamp(py0, 0, height_env - 1) py1 = torch.clamp(py1, 0, height_env - 1) # Get fractional parts fx = px - px0.float() fy = py - py0.float() # Expand dimensions for broadcasting with channels fx = fx.unsqueeze(-1) # [B, H, W, 1] fy = fy.unsqueeze(-1) # [B, H, W, 1] # Sample four corners using advanced indexing # We need to create batch indices that match the spatial dimensions batch_idx = torch.arange(batch_env, device=env_map.device).view(batch_env, 1, 1).expand(batch_uv, height_uv, width_uv) # Sample all four corners c00 = env_map[batch_idx, py0, px0] # [B, H, W, C] c01 = env_map[batch_idx, py0, px1] # [B, H, W, C] c10 = env_map[batch_idx, py1, px0] # [B, H, W, C] c11 = env_map[batch_idx, py1, px1] # [B, H, W, C] # Bilinear interpolation c0 = c00 * (1 - fx) + c01 * fx c1 = c10 * (1 - fx) + c11 * fx result = c0 * (1 - fy) + c1 * fy return result def calculate_reflection(self, normals: torch.Tensor, view_dir: torch.Tensor) -> torch.Tensor: """ Calculate reflection vector: R = 2(N·V)N - V normals: [B, H, W, 3] view_dir: [3] viewing direction returns: [B, H, W, 3] reflection vectors """ # Calculate N·V dot_nv = torch.sum(normals * view_dir, dim=-1, keepdim=True) # R = 2(N·V)N - V reflection = 2.0 * dot_nv * normals - view_dir # Normalize length = torch.sqrt(torch.sum(reflection ** 2, dim=-1, keepdim=True)) length = torch.clamp(length, min=1e-8) reflection = reflection / length return reflection def blur_environment_sample(self, env_map: torch.Tensor, uv: torch.Tensor, blur_amount: float) -> torch.Tensor: """ Sample environment map with simple blur for roughness simulation Lower specular_power = more blur """ # For now, just sample directly # In a more advanced version, could implement mipmap-based sampling result = self.sample_equirect(env_map, uv) # Simple blur approximation by averaging nearby samples if blur_amount > 1.0: # This is a simplified version - proper IBL would use prefiltered mipmaps pass return result def relight_ibl(self, image, normal_map, environment_map, environment_rotation, diffuse_intensity, specular_intensity, specular_power, ambient_light, normal_format): """ Relight image using Image-Based Lighting with equirectangular environment map """ print("\n" + "="*60) print("Image-Based Lighting (IBL) Relighter") print("="*60) print(f"Image shape: {image.shape}") print(f"Normal map shape: {normal_map.shape}") print(f"Environment map shape: {environment_map.shape}") print(f"Environment rotation: {environment_rotation}°") # Convert rotation to radians rotation_rad = environment_rotation * (np.pi / 180.0) # Decode normal map print(f"\n🗺️ Decoding normal map...") normals = self.decode_normal_map(normal_map, normal_format) # View direction (camera looking at surface) view_dir = torch.tensor([0.0, 0.0, 1.0], device=image.device, dtype=image.dtype) # Calculate diffuse lighting from environment print(f"\n💡 Calculating diffuse lighting from environment...") diffuse_uv = self.direction_to_equirect_uv(normals, rotation_rad) diffuse_colors = self.sample_equirect(environment_map, diffuse_uv) # Apply diffuse intensity diffuse_lighting = diffuse_colors * diffuse_intensity + ambient_light diffuse_lighting = torch.clamp(diffuse_lighting, 0.0, 1.0) print(f" Diffuse range: [{diffuse_lighting.min():.3f}, {diffuse_lighting.max():.3f}]") # Calculate specular lighting from reflections print(f"\n✨ Calculating specular reflections...") reflection_vectors = self.calculate_reflection(normals, view_dir) specular_uv = self.direction_to_equirect_uv(reflection_vectors, rotation_rad) # Sample with blur based on specular power (inverse relationship) blur_factor = 128.0 / max(specular_power, 1.0) specular_colors = self.blur_environment_sample(environment_map, specular_uv, blur_factor) # Calculate Fresnel-like falloff for specular view_dot_normal = torch.sum(normals * view_dir, dim=-1, keepdim=True) view_dot_normal = torch.clamp(view_dot_normal, 0.0, 1.0) # Simple Fresnel approximation: more reflection at grazing angles fresnel = torch.pow(1.0 - view_dot_normal, 5.0) fresnel = fresnel * 0.5 + 0.5 # Soften the effect specular_lighting = specular_colors * specular_intensity * fresnel specular_lighting = torch.clamp(specular_lighting, 0.0, 1.0) print(f" Specular range: [{specular_lighting.min():.3f}, {specular_lighting.max():.3f}]") # Combine diffuse and specular relit_image = image * diffuse_lighting + specular_lighting relit_image = torch.clamp(relit_image, 0.0, 1.0) print(f"\n✓ IBL relighting complete") print(f" Final image range: [{relit_image.min():.3f}, {relit_image.max():.3f}]") print("="*60 + "\n") return (relit_image, diffuse_lighting, specular_lighting) class NormalMapRelighterIBL_PBR: """ PBR-aware Image-Based Lighting with roughness and metallic control Uses PBR material properties for physically accurate relighting Non-metallic surfaces won't look metallic! """ @classmethod def INPUT_TYPES(cls): return { "required": { "diffuse_image": ("IMAGE", {"tooltip": "Albedo/diffuse color map"}), "normal_map": ("IMAGE",), "environment_map": ("IMAGE",), # Environment rotation "environment_rotation": ("FLOAT", { "default": 0.0, "min": -180.0, "max": 180.0, "step": 1.0, "tooltip": "Rotate environment map (degrees)" }), # Lighting properties "diffuse_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Diffuse lighting intensity" }), "specular_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Specular reflection intensity" }), "ambient_light": ("FLOAT", { "default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Ambient light level" }), # Roughness control "roughness_multiplier": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Global roughness multiplier (affects both diffuse and specular blur)" }), "roughness_affects_diffuse": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "How much roughness affects diffuse detail (0=uniform, 1=fully controlled by roughness)" }), "diffuse_contrast": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Diffuse softness control (0=soft/compressed highlights, 1=normal, 2=enhanced contrast). Lower values reduce metallic look" }), # Diffuse color control "use_colored_diffuse": ("BOOLEAN", { "default": False, "tooltip": "Use environment colors for diffuse (False = grayscale shading only)" }), "diffuse_quality": (["Flat Ambient (No Direction)", "Pure Matte (No Environment)", "Soft (Heavy Blur)", "Detailed (Light Blur)", "Hybrid (Lambert + Env Colors)", "True Hemisphere Sampling"], { "default": "Flat Ambient (No Direction)", "tooltip": "Diffuse quality - Flat Ambient = pure even lighting (no reflection look)" }), # Normal map format "normal_format": (["OpenGL", "DirectX"], { "default": "OpenGL", "tooltip": "Normal map format" }), }, "optional": { "roughness_map": ("IMAGE",), "metallic_map": ("IMAGE",), "irradiance_map": ("IMAGE",), } } RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE") RETURN_NAMES = ("relit_image", "diffuse_lighting", "specular_lighting", "diffuse_color", "key_light_direction") FUNCTION = "relight_ibl_pbr" CATEGORY = "Texture Alchemist/Lighting" def decode_normal_map(self, normal_map: torch.Tensor, format: str) -> torch.Tensor: """Decode normal map from color space to vector space""" normals = normal_map * 2.0 - 1.0 if format == "DirectX": normals[:, :, :, 1] = -normals[:, :, :, 1] length = torch.sqrt(torch.sum(normals ** 2, dim=-1, keepdim=True)) length = torch.clamp(length, min=1e-8) normals = normals / length return normals def prepare_pbr_map(self, pbr_map: torch.Tensor, target_shape: tuple) -> torch.Tensor: """Convert PBR map to grayscale and match target shape""" if pbr_map.shape[-1] > 1: # Convert to grayscale weights = torch.tensor([0.299, 0.587, 0.114], device=pbr_map.device, dtype=pbr_map.dtype) pbr_map = torch.sum(pbr_map[..., :3] * weights, dim=-1, keepdim=True) # Expand to match target channels if needed if pbr_map.shape[-1] != target_shape[-1]: pbr_map = pbr_map.repeat(1, 1, 1, target_shape[-1]) return pbr_map def direction_to_equirect_uv(self, direction: torch.Tensor, rotation_radians: float) -> torch.Tensor: """Convert 3D direction vector to equirectangular UV coordinates""" x = direction[:, :, :, 0] y = direction[:, :, :, 1] z = direction[:, :, :, 2] theta = torch.asin(torch.clamp(y, -1.0, 1.0)) phi = torch.atan2(x, z) phi = phi + rotation_radians u = (phi / (2 * np.pi)) + 0.5 v = (theta / np.pi) + 0.5 u = u % 1.0 v = torch.clamp(v, 0.0, 1.0) uv = torch.stack([u, v], dim=-1) return uv def sample_equirect(self, env_map: torch.Tensor, uv: torch.Tensor) -> torch.Tensor: """Sample equirectangular map using UV coordinates with bilinear interpolation""" batch_env, height_env, width_env, channels = env_map.shape batch_uv, height_uv, width_uv, _ = uv.shape px = uv[:, :, :, 0] * (width_env - 1) py = uv[:, :, :, 1] * (height_env - 1) px0 = torch.floor(px).long() py0 = torch.floor(py).long() px1 = px0 + 1 py1 = py0 + 1 px0 = torch.clamp(px0, 0, width_env - 1) px1 = torch.clamp(px1, 0, width_env - 1) py0 = torch.clamp(py0, 0, height_env - 1) py1 = torch.clamp(py1, 0, height_env - 1) fx = px - px0.float() fy = py - py0.float() fx = fx.unsqueeze(-1) fy = fy.unsqueeze(-1) batch_idx = torch.arange(batch_env, device=env_map.device).view(batch_env, 1, 1).expand(batch_uv, height_uv, width_uv) c00 = env_map[batch_idx, py0, px0] c01 = env_map[batch_idx, py0, px1] c10 = env_map[batch_idx, py1, px0] c11 = env_map[batch_idx, py1, px1] c0 = c00 * (1 - fx) + c01 * fx c1 = c10 * (1 - fx) + c11 * fx result = c0 * (1 - fy) + c1 * fy return result def calculate_reflection(self, normals: torch.Tensor, view_dir: torch.Tensor) -> torch.Tensor: """Calculate reflection vector: R = 2(N·V)N - V""" dot_nv = torch.sum(normals * view_dir, dim=-1, keepdim=True) reflection = 2.0 * dot_nv * normals - view_dir length = torch.sqrt(torch.sum(reflection ** 2, dim=-1, keepdim=True)) length = torch.clamp(length, min=1e-8) reflection = reflection / length return reflection def calculate_fresnel_schlick(self, view_dot_normal: torch.Tensor, f0: torch.Tensor) -> torch.Tensor: """ Fresnel-Schlick approximation f0: base reflectivity (0.04 for dielectrics, albedo for metals) """ return f0 + (1.0 - f0) * torch.pow(1.0 - view_dot_normal, 5.0) def sample_equirect_blurred(self, env_map: torch.Tensor, uv: torch.Tensor, roughness: torch.Tensor) -> torch.Tensor: """ Sample environment map with blur based on roughness roughness: [B, H, W, C] where 0=smooth, 1=rough """ # Convert roughness to average grayscale roughness_gray = torch.mean(roughness, dim=-1, keepdim=True) # [B, H, W, 1] # For each pixel, sample multiple nearby UV coordinates based on roughness # This simulates blurred reflections for rough surfaces batch_uv, height_uv, width_uv, _ = uv.shape # Sample center point center_sample = self.sample_equirect(env_map, uv) # For rough surfaces, sample additional points in a pattern # Roughness determines blur radius blur_radius = roughness_gray * 0.05 # Max 5% of image size blur # Create offset pattern (8 samples around center) offsets = [ (1, 0), (-1, 0), (0, 1), (0, -1), (0.707, 0.707), (-0.707, 0.707), (0.707, -0.707), (-0.707, -0.707) ] # Accumulate samples result = center_sample weight = 1.0 for ox, oy in offsets: # Create offset UV offset_uv = uv.clone() offset_uv[:, :, :, 0] += ox * blur_radius.squeeze(-1) offset_uv[:, :, :, 1] += oy * blur_radius.squeeze(-1) # Wrap/clamp UV offset_uv[:, :, :, 0] = offset_uv[:, :, :, 0] % 1.0 offset_uv[:, :, :, 1] = torch.clamp(offset_uv[:, :, :, 1], 0.0, 1.0) # Sample and accumulate sample = self.sample_equirect(env_map, offset_uv) sample_weight = roughness_gray # Weight by roughness result += sample * sample_weight weight += sample_weight # Average all samples result = result / weight return result def relight_ibl_pbr(self, diffuse_image, normal_map, environment_map, environment_rotation, diffuse_intensity, specular_intensity, ambient_light, roughness_multiplier, roughness_affects_diffuse, diffuse_contrast, use_colored_diffuse, diffuse_quality, normal_format, roughness_map=None, metallic_map=None, irradiance_map=None): """ PBR-aware IBL relighting with roughness and metallic control """ print("\n" + "="*60) print("PBR Image-Based Lighting Relighter") print("="*60) print(f"Diffuse image shape: {diffuse_image.shape}") print(f"Normal map shape: {normal_map.shape}") print(f"Environment map shape: {environment_map.shape}") # Convert rotation to radians rotation_rad = environment_rotation * (np.pi / 180.0) # Decode normal map normals = self.decode_normal_map(normal_map, normal_format) # Process PBR maps if roughness_map is not None: roughness = self.prepare_pbr_map(roughness_map, diffuse_image.shape) print(f"Roughness map: [{roughness.min():.3f}, {roughness.max():.3f}]") else: # Default roughness (semi-glossy) roughness = torch.ones_like(diffuse_image) * 0.5 print(f"Roughness map: Using default (0.5)") # Apply roughness multiplier roughness = roughness * roughness_multiplier roughness = torch.clamp(roughness, 0.0, 1.0) print(f"Roughness (after multiplier {roughness_multiplier}): [{roughness.min():.3f}, {roughness.max():.3f}]") if metallic_map is not None: metallic = self.prepare_pbr_map(metallic_map, diffuse_image.shape) print(f"Metallic map: [{metallic.min():.3f}, {metallic.max():.3f}]") else: # Default non-metallic metallic = torch.zeros_like(diffuse_image) print(f"Metallic map: Using default (0.0 - non-metallic)") # View direction view_dir = torch.tensor([0.0, 0.0, 1.0], device=diffuse_image.device, dtype=diffuse_image.dtype) # Calculate diffuse lighting from environment print(f"\n💡 Calculating PBR diffuse lighting...") print(f" Diffuse quality mode: {diffuse_quality}") print(f" Roughness affects diffuse: {roughness_affects_diffuse:.2f}") # Check if we have a pre-computed irradiance map if irradiance_map is not None: print(f" Using provided irradiance map (pre-computed)") diffuse_uv = self.direction_to_equirect_uv(normals, rotation_rad) # Even with irradiance, we can apply additional roughness-based blur if roughness_affects_diffuse > 0.0: # Base irradiance + roughness-modulated extra blur base_blur = roughness * roughness_affects_diffuse * 0.3 # 0 to 0.3 extra blur diffuse_colors = self.sample_equirect_blurred(irradiance_map, diffuse_uv, base_blur) print(f" Applied roughness-based blur to irradiance (range: [{base_blur.min():.3f}, {base_blur.max():.3f}])") else: diffuse_colors = self.sample_equirect(irradiance_map, diffuse_uv) elif diffuse_quality == "Flat Ambient (No Direction)": # PURE AMBIENT: No directional lighting at all, just average environment color # This ensures ZERO reflection/refraction look - pure flat lighting print(f" 🎨 Using Flat Ambient mode (pure non-directional)") # Get average environment color avg_env_color = torch.mean(environment_map, dim=(1, 2), keepdim=True) # Broadcast to image dimensions (same color everywhere) batch, height, width, _ = normals.shape diffuse_colors = avg_env_color.expand(batch, height, width, 3) print(f" Average environment color: RGB({avg_env_color[0,0,0,0]:.3f}, {avg_env_color[0,0,0,1]:.3f}, {avg_env_color[0,0,0,2]:.3f})") print(f" Result: Completely flat/even lighting (no directionality)") elif diffuse_quality == "Pure Matte (No Environment)": # Pure matte diffuse - simple directional lighting with no environment sampling light_dir = torch.tensor([0.2, 0.3, 0.9], device=diffuse_image.device, dtype=diffuse_image.dtype) light_dir = light_dir / torch.sqrt(torch.sum(light_dir ** 2)) # Calculate N·L (Lambert) dot_nl = torch.sum(normals * light_dir, dim=-1, keepdim=True) dot_nl = torch.clamp(dot_nl, 0.0, 1.0) # Pure matte shading (grayscale) diffuse_colors = dot_nl.repeat(1, 1, 1, 3) print(f" Using pure matte shading (no environment, clay-like)") elif diffuse_quality == "Hybrid (Lambert + Env Colors)": # SWITCHLIGHT-STYLE: Lambert shading with environment colors # This creates pure diffuse shading without looking metallic print(f" 🎬 Using Hybrid mode (Switchlight-style)") # Step 1: Calculate pure Lambert shading (black & white) # Use top hemisphere average as light direction light_dir = torch.tensor([0.0, 0.0, 1.0], device=diffuse_image.device, dtype=diffuse_image.dtype) dot_nl = torch.sum(normals * light_dir, dim=-1, keepdim=True) dot_nl = torch.clamp(dot_nl, 0.0, 1.0) print(f" Lambert shading range: [{dot_nl.min():.3f}, {dot_nl.max():.3f}]") # Step 2: If colored diffuse is enabled, map to environment colors if use_colored_diffuse: # Get environment colors (using mean-based approach, not extreme quantiles) # Convert environment to luminance env_lum_weights = torch.tensor([0.299, 0.587, 0.114], device=environment_map.device, dtype=environment_map.dtype) env_luminance = torch.sum(environment_map * env_lum_weights, dim=-1, keepdim=True) # Find dark and light colors (less extreme than 10th/90th percentile) dark_color = torch.quantile(environment_map.flatten(0, 2), 0.2, dim=0, keepdim=True) # 20th percentile (less extreme) light_color = torch.quantile(environment_map.flatten(0, 2), 0.8, dim=0, keepdim=True) # 80th percentile (less extreme) # Reshape for broadcasting dark_color = dark_color.reshape(1, 1, 1, 3) light_color = light_color.reshape(1, 1, 1, 3) # Reduce color intensity to avoid looking specular # Mix with grayscale to reduce saturation avg_color = (dark_color + light_color) / 2.0 dark_color = dark_color * 0.5 + avg_color * 0.5 # 50% desaturate light_color = light_color * 0.5 + avg_color * 0.5 # Map Lambert shading to color range diffuse_colors = dark_color + (light_color - dark_color) * dot_nl print(f" Dark color: [{dark_color[0,0,0,0]:.3f}, {dark_color[0,0,0,1]:.3f}, {dark_color[0,0,0,2]:.3f}]") print(f" Light color: [{light_color[0,0,0,0]:.3f}, {light_color[0,0,0,1]:.3f}, {light_color[0,0,0,2]:.3f}]") else: # Pure grayscale Lambert (no color from environment) diffuse_colors = dot_nl.repeat(1, 1, 1, 3) print(f" Using pure grayscale Lambert (no environment color)") elif diffuse_quality == "True Hemisphere Sampling": # TRUE PHYSICALLY-BASED: Integrate over hemisphere # This is the most accurate but also slowest print(f" 🌐 Using True Hemisphere Sampling (slow but accurate)") batch, height, width, _ = normals.shape num_samples = 64 # Can be adjusted for quality vs speed tradeoff # Initialize accumulator diffuse_colors = torch.zeros_like(diffuse_image) total_weight = torch.zeros_like(diffuse_image[:, :, :, 0:1]) # Track total weights print(f" Sampling {num_samples} directions per pixel...") # Generate sample directions using Fibonacci sphere (hemisphere only) indices = torch.arange(0, num_samples, dtype=torch.float32, device=diffuse_image.device) phi = torch.acos(1 - (indices + 0.5) / num_samples) # Only upper hemisphere (0 to π/2) theta = np.pi * (1 + 5**0.5) * indices # Convert to Cartesian coordinates sample_dirs = torch.stack([ torch.sin(phi) * torch.cos(theta), torch.sin(phi) * torch.sin(theta), torch.cos(phi) # z always positive (hemisphere) ], dim=-1) # [num_samples, 3] # For each sample direction for i in range(num_samples): sample_dir = sample_dirs[i:i+1] # [1, 3] # Calculate N·L for this sample dot_nl = torch.sum(normals * sample_dir, dim=-1, keepdim=True) dot_nl = torch.clamp(dot_nl, 0.0, 1.0) # Convert direction to UV coordinates # Expand sample_dir to match batch dimensions sample_dir_expanded = sample_dir.expand(batch, height, width, 3) sample_uv = self.direction_to_equirect_uv(sample_dir_expanded, rotation_rad) # Sample environment env_sample = self.sample_equirect(environment_map, sample_uv) # Accumulate: weight by N·L (Lambert's cosine law) diffuse_colors += env_sample * dot_nl total_weight += dot_nl if (i + 1) % 16 == 0: progress = (i + 1) / num_samples * 100 print(f" Progress: {progress:.0f}%") # Average by total weight (proper Monte Carlo integration) # Avoid division by zero total_weight = torch.clamp(total_weight, min=1e-8) diffuse_colors = diffuse_colors / total_weight print(f" Hemisphere sampling complete!") print(f" Diffuse range: [{diffuse_colors.min():.3f}, {diffuse_colors.max():.3f}]") elif diffuse_quality == "Soft (Heavy Blur)": # Heavy environment blur for soft diffuse, modulated by roughness diffuse_uv = self.direction_to_equirect_uv(normals, rotation_rad) # Base: 0.95, add roughness influence base_blur = 0.95 roughness_influence = roughness * roughness_affects_diffuse * 0.05 diffuse_blur = base_blur + roughness_influence diffuse_colors = self.sample_equirect_blurred(environment_map, diffuse_uv, diffuse_blur) print(f" Using heavy blur with roughness modulation") print(f" Blur range: [{diffuse_blur.min():.3f}, {diffuse_blur.max():.3f}]") else: # "Detailed (Light Blur)" # Lighter blur - more environment detail, strongly affected by roughness diffuse_uv = self.direction_to_equirect_uv(normals, rotation_rad) # Base: 0.5-0.7, roughness adds 0 to 0.3 base_blur = 0.5 + (1.0 - roughness_affects_diffuse) * 0.2 roughness_influence = roughness * roughness_affects_diffuse * 0.3 diffuse_blur = base_blur + roughness_influence diffuse_colors = self.sample_equirect_blurred(environment_map, diffuse_uv, diffuse_blur) print(f" Using detailed blur with roughness control") print(f" Blur range: [{diffuse_blur.min():.3f}, {diffuse_blur.max():.3f}]") # Convert diffuse to grayscale (luminance) for non-colored shading # Colored lighting looks metallic, grayscale looks like proper shading # NOTE: Hybrid mode handles this internally, so skip for Hybrid if not use_colored_diffuse and diffuse_quality != "Hybrid (Lambert + Env Colors)": # Convert to luminance for natural light/shadow appearance weights = torch.tensor([0.299, 0.587, 0.114], device=diffuse_colors.device, dtype=diffuse_colors.dtype) diffuse_luminance = torch.sum(diffuse_colors * weights, dim=-1, keepdim=True) diffuse_colors = diffuse_luminance.repeat(1, 1, 1, 3) print(f" Converting to grayscale (natural light/shadow appearance)") elif diffuse_quality != "Hybrid (Lambert + Env Colors)": print(f" Using colored diffuse (environment colors)") # Apply diffuse contrast control to reduce metallic look # Use a soft compression curve instead of linear blend to grey if diffuse_contrast != 1.0: if diffuse_contrast < 1.0: # Compress highlights using a soft curve (reduces metallic look) # This applies a power curve that compresses bright values while preserving gradients # contrast = 0 → strong compression (very matte) # contrast = 1 → no change power = 1.0 + (1.0 - diffuse_contrast) * 2.0 # Range: 1.0 to 3.0 diffuse_colors = torch.pow(diffuse_colors, 1.0 / power) print(f" Applied highlight compression: {diffuse_contrast:.2f} (power: {power:.2f})") else: # Enhance contrast (values > 1.0) # Apply inverse power curve to increase contrast power = 1.0 / diffuse_contrast diffuse_colors = torch.pow(diffuse_colors, power) print(f" Enhanced contrast: {diffuse_contrast:.2f} (power: {power:.2f})") diffuse_colors = torch.clamp(diffuse_colors, 0.0, 1.0) print(f" Diffuse range after adjustment: [{diffuse_colors.min():.3f}, {diffuse_colors.max():.3f}]") # Calculate ambient light as average environment color (GI-style) # Use irradiance map if available, otherwise average the environment if irradiance_map is not None and ambient_light > 0.0: # Average the irradiance map for ambient ambient_color = torch.mean(irradiance_map, dim=(1, 2), keepdim=True) print(f" Ambient from irradiance: RGB({ambient_color[0,0,0,0]:.3f}, {ambient_color[0,0,0,1]:.3f}, {ambient_color[0,0,0,2]:.3f})") elif ambient_light > 0.0: # Average the environment map for ambient (GI-style) ambient_color = torch.mean(environment_map, dim=(1, 2), keepdim=True) print(f" Ambient from environment avg: RGB({ambient_color[0,0,0,0]:.3f}, {ambient_color[0,0,0,1]:.3f}, {ambient_color[0,0,0,2]:.3f})") else: # No ambient ambient_color = torch.zeros_like(diffuse_image[:, 0:1, 0:1, :]) print(f" No ambient light") # Apply ambient light intensity ambient_contribution = ambient_color * ambient_light # Create pure diffuse shading (without metallic masking) for output preview diffuse_shading_pure = diffuse_colors * diffuse_intensity + ambient_contribution diffuse_shading_pure = torch.clamp(diffuse_shading_pure, 0.0, 1.0) # Create diffuse for actual lighting calculation (with metallic masking) # In PBR, metals don't have diffuse reflection diffuse_for_lighting = diffuse_colors * diffuse_intensity * (1.0 - metallic) + ambient_contribution diffuse_for_lighting = torch.clamp(diffuse_for_lighting, 0.0, 1.0) print(f" Final diffuse shading range: [{diffuse_shading_pure.min():.3f}, {diffuse_shading_pure.max():.3f}]") # Calculate specular reflections print(f"\n✨ Calculating PBR specular reflections...") reflection_vectors = self.calculate_reflection(normals, view_dir) specular_uv = self.direction_to_equirect_uv(reflection_vectors, rotation_rad) # Sample environment with roughness-based blur # Black roughness (0) = sharp reflections, White roughness (1) = blurred reflections specular_colors = self.sample_equirect_blurred(environment_map, specular_uv, roughness) print(f" Applied roughness-based blur to specular reflections") # Calculate Fresnel view_dot_normal = torch.sum(normals * view_dir, dim=-1, keepdim=True) view_dot_normal = torch.clamp(view_dot_normal, 0.0, 1.0) # F0: Base reflectivity # Dielectrics: ~0.04 (4% reflection) # Metals: Use albedo color f0_dielectric = torch.ones_like(diffuse_image) * 0.04 f0_metal = diffuse_image # Metals reflect colored light f0 = f0_dielectric * (1.0 - metallic) + f0_metal * metallic # Apply Fresnel fresnel = self.calculate_fresnel_schlick(view_dot_normal, f0) # Specular intensity: Metals get more specular, non-metals less # White metallic (1) = full specular, Black metallic (0) = minimal specular metal_boost = 1.0 + metallic * 2.0 # Metals get 3x boost specular_lighting = specular_colors * specular_intensity * fresnel * metal_boost specular_lighting = torch.clamp(specular_lighting, 0.0, 1.0) print(f" Specular range: [{specular_lighting.min():.3f}, {specular_lighting.max():.3f}]") print(f" Fresnel applied: grazing angles get more reflection") print(f" Roughness applied: white (rough) = blurred, black (smooth) = sharp") print(f" Metallic applied: white (metal) = strong colored reflection, black (non-metal) = weak") # Combine lighting # For metals: specular is the main component (colored by albedo) # For dielectrics: diffuse is main, specular is additive relit_image = diffuse_image * diffuse_for_lighting + specular_lighting relit_image = torch.clamp(relit_image, 0.0, 1.0) # Create additional outputs print(f"\n📤 Creating additional outputs...") # Diffuse Color: Input albedo with ONLY ambient/GI applied (no directional diffuse, no specular) # This is: Albedo × Ambient GI only - for Switchlight-style compositing diffuse_color = diffuse_image * ambient_contribution diffuse_color = torch.clamp(diffuse_color, 0.0, 1.0) print(f" Diffuse Color: Albedo × Ambient GI only (no directional lighting)") print(f" Range: [{diffuse_color.min():.3f}, {diffuse_color.max():.3f}]") print(f" Usage: Base image with ambient lift, add diffuse_lighting + specular separately") # Key Light Direction: Pure Lambert shading that rotates with environment # The light direction rotates with the environment rotation # Start with light pointing from above/forward (typical key light position) base_light_dir = torch.tensor([0.3, 0.5, 0.8], device=diffuse_image.device, dtype=diffuse_image.dtype) base_light_dir = base_light_dir / torch.sqrt(torch.sum(base_light_dir ** 2)) # Normalize # Rotate light direction by environment rotation (around Y axis) cos_rot = torch.cos(torch.tensor(rotation_rad, device=diffuse_image.device, dtype=diffuse_image.dtype)) sin_rot = torch.sin(torch.tensor(rotation_rad, device=diffuse_image.device, dtype=diffuse_image.dtype)) # Rotation matrix around Y axis rotated_light_x = base_light_dir[0] * cos_rot - base_light_dir[2] * sin_rot rotated_light_y = base_light_dir[1] # Y doesn't change rotated_light_z = base_light_dir[0] * sin_rot + base_light_dir[2] * cos_rot key_light_dir = torch.stack([rotated_light_x, rotated_light_y, rotated_light_z]) # Calculate Lambert dot product with rotated light key_light_dot = torch.sum(normals * key_light_dir, dim=-1, keepdim=True) key_light_dot = torch.clamp(key_light_dot, 0.0, 1.0) key_light_direction = key_light_dot.repeat(1, 1, 1, 3) # Convert to RGB for output print(f" Key Light Direction: Pure Lambert mask (rotates with environment)") print(f" Light direction: [{key_light_dir[0]:.3f}, {key_light_dir[1]:.3f}, {key_light_dir[2]:.3f}]") print(f" Environment rotation: {environment_rotation}°") print(f"\n✓ PBR IBL relighting complete") print(f" Final image range: [{relit_image.min():.3f}, {relit_image.max():.3f}]") print(f" Material interpretation:") print(f" Roughness: Black=Shiny/Sharp, White=Rough/Blurred") print(f" Metallic: Black=Non-Metal, White=Metal") print("="*60 + "\n") return (relit_image, diffuse_shading_pure, specular_lighting, diffuse_color, key_light_direction) class PBR_AOV_Generator: """ PBR AOV (Arbitrary Output Variable) Generator Takes PBR material inputs and generates proper render passes for compositing: - Ambient Light - GI Light Pass - Diffuse (direct lighting) - Reflection (specular) """ @classmethod def INPUT_TYPES(cls): return { "required": { "albedo": ("IMAGE", {"tooltip": "Albedo/diffuse color map"}), "normal_map": ("IMAGE",), "environment_map": ("IMAGE", {"tooltip": "HDRI or irradiance map"}), # Environment rotation "environment_rotation": ("FLOAT", { "default": 0.0, "min": -180.0, "max": 180.0, "step": 1.0, "tooltip": "Rotate environment map (degrees)" }), # Light intensities "ambient_intensity": ("FLOAT", { "default": 0.1, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "Ambient fill light intensity" }), "gi_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Global illumination (indirect) intensity" }), "diffuse_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Direct diffuse lighting intensity" }), "reflection_intensity": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Reflection/specular intensity" }), # Normal map format "normal_format": (["OpenGL", "DirectX"], { "default": "OpenGL", "tooltip": "Normal map format" }), }, "optional": { "roughness_map": ("IMAGE",), "metallic_map": ("IMAGE",), } } RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE") RETURN_NAMES = ("ambient_light", "gi_light", "diffuse", "reflection") FUNCTION = "generate_aov_passes" CATEGORY = "Texture Alchemist/Lighting" def decode_normal_map(self, normal_map, format): """Decode normal map from color space to vector space""" normals = normal_map * 2.0 - 1.0 if format == "DirectX": normals[:, :, :, 1] = -normals[:, :, :, 1] length = torch.sqrt(torch.sum(normals ** 2, dim=-1, keepdim=True)) length = torch.clamp(length, min=1e-8) normals = normals / length return normals def prepare_pbr_map(self, pbr_map, target_shape): """Ensure PBR map matches target dimensions""" if pbr_map.shape[1:3] != target_shape[1:3]: import torch.nn.functional as F pbr_map = F.interpolate(pbr_map.permute(0, 3, 1, 2), size=(target_shape[1], target_shape[2]), mode='bilinear', align_corners=False).permute(0, 2, 3, 1) if pbr_map.shape[-1] == 1: pbr_map = pbr_map.repeat(1, 1, 1, 3) return pbr_map def direction_to_equirect_uv(self, directions, rotation_offset=0.0): """Convert 3D direction vectors to equirectangular UV coordinates""" x, y, z = directions[:, :, :, 0], directions[:, :, :, 1], directions[:, :, :, 2] theta = torch.atan2(x, z) + rotation_offset phi = torch.acos(torch.clamp(y, -1.0, 1.0)) u = (theta / (2 * np.pi)) % 1.0 v = phi / np.pi uv = torch.stack([u, v], dim=-1) return uv def sample_equirect(self, env_map, uv): """Sample equirectangular environment map using UV coordinates""" batch, height, width, _ = uv.shape env_height, env_width = env_map.shape[1], env_map.shape[2] u = uv[:, :, :, 0] * (env_width - 1) v = uv[:, :, :, 1] * (env_height - 1) u0 = torch.floor(u).long() u1 = torch.ceil(u).long() v0 = torch.floor(v).long() v1 = torch.ceil(v).long() u0 = torch.clamp(u0, 0, env_width - 1) u1 = torch.clamp(u1, 0, env_width - 1) v0 = torch.clamp(v0, 0, env_height - 1) v1 = torch.clamp(v1, 0, env_height - 1) wu = (u - u0.float()).unsqueeze(-1) wv = (v - v0.float()).unsqueeze(-1) batch_idx = torch.arange(batch, device=env_map.device).view(-1, 1, 1).expand(batch, height, width) c00 = env_map[batch_idx, v0, u0] c01 = env_map[batch_idx, v0, u1] c10 = env_map[batch_idx, v1, u0] c11 = env_map[batch_idx, v1, u1] c0 = c00 * (1 - wu) + c01 * wu c1 = c10 * (1 - wu) + c11 * wu result = c0 * (1 - wv) + c1 * wv return result def calculate_reflection(self, normals, view_dir): """Calculate reflection vector for environment sampling""" dot_nv = torch.sum(normals * view_dir, dim=-1, keepdim=True) reflection = 2.0 * dot_nv * normals - view_dir length = torch.sqrt(torch.sum(reflection ** 2, dim=-1, keepdim=True)) length = torch.clamp(length, min=1e-8) reflection = reflection / length return reflection def calculate_fresnel_schlick(self, cos_theta, f0): """Fresnel-Schlick approximation""" return f0 + (1.0 - f0) * torch.pow(1.0 - cos_theta, 5.0) def generate_aov_passes(self, albedo, normal_map, environment_map, environment_rotation, ambient_intensity, gi_intensity, diffuse_intensity, reflection_intensity, normal_format, roughness_map=None, metallic_map=None): """ Generate proper AOV render passes for compositing """ print("\n" + "="*60) print("PBR AOV Generator - Render Passes") print("="*60) print(f"Albedo shape: {albedo.shape}") print(f"Normal map shape: {normal_map.shape}") print(f"Environment shape: {environment_map.shape}") # Convert rotation to radians rotation_rad = environment_rotation * (np.pi / 180.0) # Decode normal map normals = self.decode_normal_map(normal_map, normal_format) # Process PBR maps if roughness_map is not None: roughness = self.prepare_pbr_map(roughness_map, albedo.shape) print(f"Roughness map: [{roughness.min():.3f}, {roughness.max():.3f}]") else: roughness = torch.ones_like(albedo) * 0.5 print(f"Roughness: Using default (0.5)") if metallic_map is not None: metallic = self.prepare_pbr_map(metallic_map, albedo.shape) print(f"Metallic map: [{metallic.min():.3f}, {metallic.max():.3f}]") else: metallic = torch.zeros_like(albedo) print(f"Metallic: Using default (0.0)") # View direction view_dir = torch.tensor([0.0, 0.0, 1.0], device=albedo.device, dtype=albedo.dtype) # ===================================================================== # AOV 1: AMBIENT LIGHT # Pure flat ambient fill (no directionality) # ===================================================================== print(f"\n💡 Generating AOV 1: Ambient Light...") # Average environment color ambient_color = torch.mean(environment_map, dim=(1, 2), keepdim=True) batch, height, width, _ = albedo.shape ambient_fill = ambient_color.expand(batch, height, width, 3) ambient_light = albedo * ambient_fill * ambient_intensity ambient_light = torch.clamp(ambient_light, 0.0, 1.0) print(f" Ambient color: RGB({ambient_color[0,0,0,0]:.3f}, {ambient_color[0,0,0,1]:.3f}, {ambient_color[0,0,0,2]:.3f})") print(f" Result: Flat ambient fill × {ambient_intensity}") # ===================================================================== # AOV 2: GI LIGHT PASS (Global Illumination - Indirect) # Soft hemisphere-averaged environment lighting # ===================================================================== print(f"\n🌐 Generating AOV 2: GI Light Pass...") # Sample environment based on normals but heavily blurred (indirect) gi_uv = self.direction_to_equirect_uv(normals, rotation_rad) # Heavy blur for GI (indirect lighting is soft) # Simple multi-sample blur gi_colors = self.sample_equirect(environment_map, gi_uv) # Apply to albedo (non-metals only for GI) gi_light = albedo * gi_colors * gi_intensity * (1.0 - metallic) gi_light = torch.clamp(gi_light, 0.0, 1.0) print(f" Indirect lighting from environment") print(f" Masked by metallic (metals don't have diffuse GI)") # ===================================================================== # AOV 3: DIFFUSE (Direct Lighting) # Lambert-based direct light from environment # ===================================================================== print(f"\n☀️ Generating AOV 3: Diffuse (Direct)...") # Key light direction (from environment dominant direction) # For now, use upward direction as key light key_light_dir = torch.tensor([0.3, 0.5, 0.8], device=albedo.device, dtype=albedo.dtype) key_light_dir = key_light_dir / torch.sqrt(torch.sum(key_light_dir ** 2)) # Rotate with environment cos_rot = torch.cos(torch.tensor(rotation_rad, device=albedo.device, dtype=albedo.dtype)) sin_rot = torch.sin(torch.tensor(rotation_rad, device=albedo.device, dtype=albedo.dtype)) rotated_x = key_light_dir[0] * cos_rot - key_light_dir[2] * sin_rot rotated_y = key_light_dir[1] rotated_z = key_light_dir[0] * sin_rot + key_light_dir[2] * cos_rot key_light = torch.stack([rotated_x, rotated_y, rotated_z]) # Lambert shading dot_nl = torch.sum(normals * key_light, dim=-1, keepdim=True) dot_nl = torch.clamp(dot_nl, 0.0, 1.0) # Sample environment at key light direction for color key_light_uv = self.direction_to_equirect_uv( key_light.unsqueeze(0).unsqueeze(0).expand(batch, height, width, 3), rotation_rad ) key_light_color = self.sample_equirect(environment_map, key_light_uv) # Direct diffuse (non-metals only) diffuse = albedo * key_light_color * dot_nl * diffuse_intensity * (1.0 - metallic) diffuse = torch.clamp(diffuse, 0.0, 1.0) print(f" Direct Lambert lighting from key light") print(f" Key light direction: [{key_light[0]:.3f}, {key_light[1]:.3f}, {key_light[2]:.3f}]") # ===================================================================== # AOV 4: REFLECTION (Specular) # Environment reflections with PBR # ===================================================================== print(f"\n✨ Generating AOV 4: Reflection...") # Calculate reflection vectors reflection_vectors = self.calculate_reflection(normals, view_dir) reflection_uv = self.direction_to_equirect_uv(reflection_vectors, rotation_rad) # Sample environment (would need blur based on roughness for proper PBR) reflection_colors = self.sample_equirect(environment_map, reflection_uv) # Calculate Fresnel view_dot_normal = torch.sum(normals * view_dir, dim=-1, keepdim=True) view_dot_normal = torch.clamp(view_dot_normal, 0.0, 1.0) f0_dielectric = torch.ones_like(albedo) * 0.04 f0_metal = albedo f0 = f0_dielectric * (1.0 - metallic) + f0_metal * metallic fresnel = self.calculate_fresnel_schlick(view_dot_normal, f0) # Metal boost for reflections metal_boost = 1.0 + metallic * 2.0 reflection = reflection_colors * fresnel * reflection_intensity * metal_boost reflection = torch.clamp(reflection, 0.0, 1.0) print(f" Environment reflections with Fresnel") print(f" Boosted for metallic surfaces") print(f"\n✓ AOV Generation Complete") print(f" 4 render passes ready for compositing") print("="*60 + "\n") return (ambient_light, gi_light, diffuse, reflection) class Marigold_AOV_Extractor: """ Extract render passes (AOVs) from Marigold IId outputs Based on Octane render passes: https://help.otoy.com/hc/en-us/articles/30487201925019 Extracts: - Diffuse (Albedo from Lighting) - Lighting/Shading (Direct + Indirect illumination) - Reflection/Specular (Residual highlights) - Ambient Occlusion (from shading analysis) - Emission (bright residual areas) """ @classmethod def INPUT_TYPES(cls): return { "required": { "marigold_lighting": ("IMAGE", { "tooltip": "Marigold IId Lighting output (contains albedo + shading)" }), "marigold_appearance": ("IMAGE", { "tooltip": "Marigold IId Appearance output (contains residual/specular)" }), # Extraction controls "diffuse_gamma": ("FLOAT", { "default": 1.0, "min": 0.1, "max": 3.0, "step": 0.1, "tooltip": "Gamma correction for diffuse/albedo" }), "shading_strength": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 3.0, "step": 0.1, "tooltip": "Lighting/shading intensity" }), "specular_threshold": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Threshold for extracting specular from residual" }), "specular_strength": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 3.0, "step": 0.1, "tooltip": "Specular/reflection intensity" }), "ao_strength": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Ambient occlusion extraction strength" }), "emission_threshold": ("FLOAT", { "default": 0.8, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Threshold for extracting emission from bright areas" }), }, "optional": { "normal_map": ("IMAGE", { "tooltip": "Optional normal map for enhanced AO extraction" }), } } RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE", "IMAGE") RETURN_NAMES = ("diffuse_albedo", "lighting_direct", "lighting_indirect", "reflection_specular", "ambient_occlusion", "emission", "beauty_composite") FUNCTION = "extract_render_passes" CATEGORY = "Texture Alchemist/Lighting" def extract_render_passes(self, marigold_lighting, marigold_appearance, diffuse_gamma, shading_strength, specular_threshold, specular_strength, ao_strength, emission_threshold, normal_map=None): """ Extract render passes from Marigold IId outputs Marigold IId Lighting decomposition: - Albedo: Diffuse surface color (no lighting) - Shading: Direct + Indirect illumination Marigold IId Appearance: - Residual: Information not captured by albedo/shading (specular, reflections) """ print("\n" + "="*60) print("Marigold AOV Extractor - Render Pass Decomposition") print("="*60) print(f"Lighting input shape: {marigold_lighting.shape}") print(f"Appearance input shape: {marigold_appearance.shape}") # Ensure both inputs are RGB (3 channels) # Marigold sometimes outputs grayscale or 2-channel images def ensure_rgb(tensor, name): """Convert grayscale/2-channel to RGB""" if tensor.shape[-1] == 1: # Grayscale to RGB print(f" Converting {name} from grayscale to RGB") return tensor.repeat(1, 1, 1, 3) elif tensor.shape[-1] == 2: # 2-channel to RGB (repeat first channel for blue) print(f" Converting {name} from 2-channel to RGB") return torch.cat([tensor, tensor[:, :, :, 0:1]], dim=-1) elif tensor.shape[-1] == 3: # Already RGB return tensor else: # More than 3 channels, take first 3 print(f" Warning: {name} has {tensor.shape[-1]} channels, using first 3") return tensor[:, :, :, :3] marigold_lighting = ensure_rgb(marigold_lighting, "Lighting") marigold_appearance = ensure_rgb(marigold_appearance, "Appearance") # Ensure batch sizes match (take minimum) if marigold_lighting.shape[0] != marigold_appearance.shape[0]: min_batch = min(marigold_lighting.shape[0], marigold_appearance.shape[0]) print(f" Warning: Batch size mismatch!") print(f" Lighting batch: {marigold_lighting.shape[0]}, Appearance batch: {marigold_appearance.shape[0]}") print(f" Using first {min_batch} images from each") marigold_lighting = marigold_lighting[:min_batch] marigold_appearance = marigold_appearance[:min_batch] # Ensure inputs match spatial dimensions if marigold_lighting.shape[1:3] != marigold_appearance.shape[1:3]: import torch.nn.functional as F target_shape = marigold_lighting.shape[1:3] marigold_appearance = F.interpolate( marigold_appearance.permute(0, 3, 1, 2), size=target_shape, mode='bilinear', align_corners=False ).permute(0, 2, 3, 1) print(f" Resized appearance to match lighting: {marigold_appearance.shape}") print(f"Final shapes - Lighting: {marigold_lighting.shape}, Appearance: {marigold_appearance.shape}") # ===================================================================== # PASS 1: DIFFUSE ALBEDO # Pure surface color without lighting # ===================================================================== print(f"\n📦 Pass 1: Diffuse Albedo...") # Marigold Lighting contains the albedo component # We need to extract it by removing the shading # Simple approach: The lighting image is albedo × shading # We can approximate albedo by analyzing the image # For now, use the lighting output as diffuse (it's already close) diffuse_albedo = marigold_lighting.clone() # Apply gamma correction if diffuse_gamma != 1.0: diffuse_albedo = torch.pow(diffuse_albedo, diffuse_gamma) diffuse_albedo = torch.clamp(diffuse_albedo, 0.0, 1.0) print(f" Extracted diffuse albedo with gamma={diffuse_gamma}") # ===================================================================== # PASS 2 & 3: LIGHTING (Direct + Indirect) # ===================================================================== print(f"\n💡 Pass 2 & 3: Lighting (Direct + Indirect)...") # The "shading" in Marigold represents the lighting # We approximate this by dividing the lit image by albedo # shading = lit_image / albedo # To extract shading, we need the ratio # Avoid division by zero albedo_safe = torch.clamp(marigold_lighting, min=0.01) # The appearance contains the residual, so we can use lighting as base # Convert to grayscale for shading intensity lighting_intensity = torch.mean(marigold_lighting, dim=-1, keepdim=True) lighting_intensity = lighting_intensity.repeat(1, 1, 1, 3) # Direct lighting: Brighter areas (peaks) lighting_direct = torch.clamp(lighting_intensity * shading_strength, 0.0, 1.0) # Indirect lighting: Softer, ambient component # GI should be VERY soft and diffuse (no specular) # Apply heavy blur to remove specular highlights import torch.nn.functional as F # Convert to BCHW format for conv2d lighting_for_blur = lighting_intensity.permute(0, 3, 1, 2) # Create large gaussian-like blur kernel for soft GI kernel_size = 21 # Larger kernel = softer, more diffuse sigma = kernel_size / 3.0 # Simple box blur (faster than gaussian) # Multiple passes of small blur = larger blur for _ in range(3): # 3 passes for soft result lighting_for_blur = F.avg_pool2d( lighting_for_blur, kernel_size=5, stride=1, padding=2, count_include_pad=False ) # Convert back to BHWC lighting_indirect = lighting_for_blur.permute(0, 2, 3, 1) # Further reduce contrast for GI (very soft) lighting_indirect = (lighting_indirect - 0.5) * 0.3 + 0.5 # Much softer lighting_indirect = torch.clamp(lighting_indirect * shading_strength * 0.4, 0.0, 1.0) print(f" Extracted direct lighting (strength={shading_strength})") print(f" Extracted indirect lighting (heavily blurred, matte)") # ===================================================================== # PASS 4: REFLECTION/SPECULAR # Extract from residual (appearance - lighting) # ===================================================================== print(f"\n✨ Pass 4: Reflection/Specular...") # The residual is what appearance captures that lighting doesn't # This includes specular highlights and reflections residual = marigold_appearance - marigold_lighting residual = torch.clamp(residual, 0.0, 1.0) # Extract bright spots as specular residual_intensity = torch.mean(residual, dim=-1, keepdim=True) # Threshold to get specular highlights specular_mask = (residual_intensity > specular_threshold).float() # Apply mask and boost reflection_specular = residual * specular_mask * specular_strength reflection_specular = torch.clamp(reflection_specular, 0.0, 1.0) print(f" Extracted specular with threshold={specular_threshold}") print(f" Specular pixels: {(specular_mask.sum() / specular_mask.numel() * 100):.1f}%") # ===================================================================== # PASS 5: AMBIENT OCCLUSION # Extract from shadow analysis # ===================================================================== print(f"\n🌑 Pass 5: Ambient Occlusion...") # AO is typically the darkest areas in shading # Convert lighting to grayscale and invert dark areas lighting_gray = torch.mean(marigold_lighting, dim=-1, keepdim=True) # Invert and boost shadows ao_raw = 1.0 - lighting_gray # Apply contrast to enhance ao_raw = torch.pow(ao_raw, 1.0 / ao_strength) if ao_strength > 0 else ao_raw # Normalize ao_min = ao_raw.min() ao_max = ao_raw.max() if ao_max > ao_min: ambient_occlusion = (ao_raw - ao_min) / (ao_max - ao_min) else: ambient_occlusion = ao_raw # Convert to RGB for display ambient_occlusion = ambient_occlusion.repeat(1, 1, 1, 3) ambient_occlusion = torch.clamp(ambient_occlusion, 0.0, 1.0) print(f" Extracted AO from shadow analysis (strength={ao_strength})") # Enhance with normal map if provided if normal_map is not None: print(f" Enhancing AO with normal map...") # Use normal map to refine AO (darker in crevices) normal_gray = torch.mean(normal_map, dim=-1, keepdim=True) normal_factor = 1.0 - torch.abs(normal_gray - 0.5) * 2.0 ambient_occlusion = ambient_occlusion * (0.7 + normal_factor * 0.3) ambient_occlusion = torch.clamp(ambient_occlusion, 0.0, 1.0) # ===================================================================== # PASS 6: EMISSION # Extract very bright areas as self-illumination # ===================================================================== print(f"\n💥 Pass 6: Emission...") # Emission is extracted from very bright areas in appearance appearance_intensity = torch.mean(marigold_appearance, dim=-1, keepdim=True) # Threshold for emission emission_mask = (appearance_intensity > emission_threshold).float() # Extract emission emission = marigold_appearance * emission_mask # Boost emission emission = torch.clamp(emission * 1.5, 0.0, 1.0) emission_percent = (emission_mask.sum() / emission_mask.numel() * 100) print(f" Extracted emission with threshold={emission_threshold}") print(f" Emissive pixels: {emission_percent:.1f}%") # ===================================================================== # PASS 7: BEAUTY COMPOSITE # Combine all passes for reference # ===================================================================== print(f"\n🎨 Pass 7: Beauty Composite...") # Beauty = (Diffuse × Lighting_Direct) + Lighting_Indirect + Reflection + Emission beauty_composite = ( diffuse_albedo * lighting_direct + lighting_indirect + reflection_specular + emission ) beauty_composite = torch.clamp(beauty_composite, 0.0, 1.0) print(f" Composed beauty pass from all AOVs") # ===================================================================== # Summary # ===================================================================== print(f"\n✓ Render Pass Extraction Complete") print(f" 7 passes ready for compositing:") print(f" 1. Diffuse Albedo - Base surface color") print(f" 2. Lighting Direct - Direct illumination") print(f" 3. Lighting Indirect - Ambient/bounce light") print(f" 4. Reflection - Specular highlights") print(f" 5. Ambient Occlusion - Shadow/cavity AO") print(f" 6. Emission - Self-illumination") print(f" 7. Beauty Composite - Final combined result") print("="*60 + "\n") return (diffuse_albedo, lighting_direct, lighting_indirect, reflection_specular, ambient_occlusion, emission, beauty_composite) # Node registration NODE_CLASS_MAPPINGS = { "NormalMapRelighter": NormalMapRelighter, "NormalMapRelighterAdvanced": NormalMapRelighterAdvanced, "NormalMapRelighterIBL": NormalMapRelighterIBL, "NormalMapRelighterIBL_PBR": NormalMapRelighterIBL_PBR, "PBR_AOV_Generator": PBR_AOV_Generator, "Marigold_AOV_Extractor": Marigold_AOV_Extractor, } NODE_DISPLAY_NAME_MAPPINGS = { "NormalMapRelighter": "Normal Map Relighter", "NormalMapRelighterAdvanced": "Normal Map Relighter (Advanced)", "NormalMapRelighterIBL": "Normal Map Relighter (IBL)", "NormalMapRelighterIBL_PBR": "Normal Map Relighter (IBL + PBR)", "PBR_AOV_Generator": "🎬 PBR AOV Generator (Render Passes)", "Marigold_AOV_Extractor": "📦 Marigold AOV Extractor (Render Passes)", }