1554 lines
60 KiB
Python
1554 lines
60 KiB
Python
"""
|
||
Custom ComfyUI Node: PBR Material Processor
|
||
Simplified version that takes Marigold/Lotus outputs and processes them
|
||
with unified brightness/gamma controls and seed management
|
||
"""
|
||
|
||
import torch
|
||
|
||
|
||
class PBRMaterialProcessor:
|
||
"""
|
||
Takes raw outputs from Marigold IID and processes them into PBR maps
|
||
With unified brightness controls for roughness and metallic
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
# Marigold outputs (these are 3-batch images typically)
|
||
"marigold_appearance": ("IMAGE",),
|
||
"marigold_lighting": ("IMAGE",),
|
||
|
||
# Brightness controls
|
||
"roughness_brightness": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Adjust roughness map brightness"
|
||
}),
|
||
"metallic_brightness": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Adjust metallic map brightness"
|
||
}),
|
||
"albedo_gamma": ("FLOAT", {
|
||
"default": 0.45,
|
||
"min": 0.1,
|
||
"max": 2.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Gamma correction for albedo"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE")
|
||
RETURN_NAMES = ("albedo", "roughness", "metallic", "ao")
|
||
FUNCTION = "process_materials"
|
||
CATEGORY = "Texture Alchemist/Materials"
|
||
|
||
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 adjust_brightness(self, image: torch.Tensor, brightness: float) -> torch.Tensor:
|
||
"""Adjust image brightness"""
|
||
return torch.clamp(image * brightness, 0.0, 1.0)
|
||
|
||
def grayscale(self, image: torch.Tensor) -> torch.Tensor:
|
||
"""Convert to grayscale"""
|
||
if image.shape[-1] == 1:
|
||
return image
|
||
# Luminance weights
|
||
weights = torch.tensor([0.2989, 0.5870, 0.1140], device=image.device, dtype=image.dtype)
|
||
return torch.sum(image * weights, dim=-1, keepdim=True)
|
||
|
||
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 process_materials(self, marigold_appearance, marigold_lighting,
|
||
roughness_brightness, metallic_brightness, albedo_gamma):
|
||
"""
|
||
Process Marigold outputs into clean PBR maps
|
||
|
||
Marigold IID Appearance typically outputs 3 images:
|
||
- [0]: Base appearance
|
||
- [1]: Material channels (R=roughness, G=metallic)
|
||
- [2]: Additional data
|
||
|
||
Marigold IID Lighting typically outputs 3 images:
|
||
- [0]: Albedo/diffuse color
|
||
- [1]: Ambient occlusion / lighting
|
||
- [2]: Additional lighting data
|
||
"""
|
||
|
||
# Process Appearance output
|
||
# Apply gamma first
|
||
appearance = self.apply_gamma(marigold_appearance, albedo_gamma)
|
||
|
||
# Extract material channels (typically batch index 1)
|
||
if appearance.shape[0] >= 2:
|
||
material_channels = appearance[1:2]
|
||
else:
|
||
material_channels = appearance[0:1]
|
||
|
||
# Extract roughness (red channel)
|
||
roughness = material_channels[:, :, :, 0:1]
|
||
roughness = self.channel_to_rgb(roughness)
|
||
roughness = self.adjust_brightness(roughness, roughness_brightness)
|
||
|
||
# Extract metallic (green channel)
|
||
metallic = material_channels[:, :, :, 1:2]
|
||
metallic = self.channel_to_rgb(metallic)
|
||
# Metallic uses inverse gamma effect
|
||
metallic_gamma = max(0.1, 2.0 - metallic_brightness)
|
||
metallic = self.apply_gamma(metallic, metallic_gamma)
|
||
|
||
# Process Lighting output
|
||
lighting = self.apply_gamma(marigold_lighting, 0.45)
|
||
|
||
# Extract albedo (batch index 0)
|
||
if lighting.shape[0] >= 1:
|
||
albedo = lighting[0:1]
|
||
else:
|
||
albedo = lighting
|
||
|
||
# Extract or generate AO (batch index 1 if available)
|
||
if lighting.shape[0] >= 2:
|
||
ao = lighting[1:2]
|
||
else:
|
||
ao = albedo
|
||
|
||
# Convert AO to grayscale
|
||
ao = self.grayscale(ao)
|
||
ao = self.channel_to_rgb(ao)
|
||
|
||
return (albedo, roughness, metallic, ao)
|
||
|
||
|
||
class PBRNormalProcessor:
|
||
"""
|
||
Process Lotus normal map outputs with flip control
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"normal_map": ("IMAGE",),
|
||
"flip_green": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Flip green channel (DirectX vs OpenGL)"
|
||
}),
|
||
"normalize": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Normalize to -1 to 1 range"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("normal",)
|
||
FUNCTION = "process_normal"
|
||
CATEGORY = "Texture Alchemist/Materials"
|
||
|
||
def process_normal(self, normal_map, flip_green, normalize):
|
||
"""Process normal map with optional green channel flip"""
|
||
|
||
normal = normal_map.clone()
|
||
|
||
# Flip green channel if requested
|
||
if flip_green and normal.shape[-1] >= 3:
|
||
normal[:, :, :, 1] = 1.0 - normal[:, :, :, 1]
|
||
|
||
# Normalize if requested
|
||
if normalize:
|
||
# Convert from 0-1 to -1 to 1
|
||
normal = (normal * 2.0) - 1.0
|
||
# Normalize vectors
|
||
norm = torch.sqrt(torch.sum(normal ** 2, dim=-1, keepdim=True))
|
||
norm = torch.clamp(norm, min=1e-8)
|
||
normal = normal / norm
|
||
# Convert back to 0-1
|
||
normal = (normal + 1.0) / 2.0
|
||
|
||
return (normal,)
|
||
|
||
|
||
class PBRHeightProcessor:
|
||
"""
|
||
Process Lotus depth/height map with remap control
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"depth_map": ("IMAGE",),
|
||
"invert": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": "Invert height values"
|
||
}),
|
||
"normalize": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Remap to 0-1 range"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("height",)
|
||
FUNCTION = "process_height"
|
||
CATEGORY = "Texture Alchemist/Materials"
|
||
|
||
def process_height(self, depth_map, invert, normalize):
|
||
"""Process height/depth map"""
|
||
|
||
height = depth_map.clone()
|
||
|
||
# Convert to grayscale if needed
|
||
if height.shape[-1] > 1:
|
||
weights = torch.tensor([0.2989, 0.5870, 0.1140],
|
||
device=height.device, dtype=height.dtype)
|
||
height = torch.sum(height * weights, dim=-1, keepdim=True)
|
||
height = height.repeat(1, 1, 1, 3)
|
||
|
||
# Normalize to 0-1
|
||
if normalize:
|
||
h_min = height.min()
|
||
h_max = height.max()
|
||
if h_max - h_min > 1e-6:
|
||
height = (height - h_min) / (h_max - h_min)
|
||
|
||
# Invert if requested
|
||
if invert:
|
||
height = 1.0 - height
|
||
|
||
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
|
||
- Green Channel: Height
|
||
- Blue Channel: Roughness
|
||
|
||
Outputs grayscale, height, roughness, original image, and PBR pipe
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"frankenmap": ("IMAGE",),
|
||
|
||
# Gamma controls for each channel
|
||
"albedo_gamma": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.1,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Gamma correction for grayscale (from red channel)"
|
||
}),
|
||
"height_gamma": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.1,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Gamma correction for height (from green channel)"
|
||
}),
|
||
"roughness_gamma": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.1,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Gamma correction for roughness (from blue channel)"
|
||
}),
|
||
|
||
# Brightness/contrast controls
|
||
"albedo_brightness": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Grayscale brightness multiplier"
|
||
}),
|
||
"height_contrast": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Height contrast multiplier"
|
||
}),
|
||
"roughness_brightness": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 3.0,
|
||
"step": 0.01,
|
||
"display": "slider",
|
||
"tooltip": "Roughness brightness multiplier"
|
||
}),
|
||
|
||
# Inversion controls
|
||
"invert_height": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": "Invert height values (bumps become indents)"
|
||
}),
|
||
"invert_roughness": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": "Invert roughness values (rough becomes smooth)"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "PBR_PIPE")
|
||
RETURN_NAMES = ("grayscale", "height", "roughness", "original_image", "pbr_pipe")
|
||
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 adjust_brightness(self, image: torch.Tensor, brightness: float) -> torch.Tensor:
|
||
"""Adjust image brightness"""
|
||
return torch.clamp(image * brightness, 0.0, 1.0)
|
||
|
||
def adjust_contrast(self, image: torch.Tensor, contrast: float) -> torch.Tensor:
|
||
"""Adjust image contrast around midpoint"""
|
||
midpoint = 0.5
|
||
return torch.clamp((image - midpoint) * contrast + midpoint, 0.0, 1.0)
|
||
|
||
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, frankenmap, albedo_gamma, height_gamma, roughness_gamma,
|
||
albedo_brightness, height_contrast, roughness_brightness,
|
||
invert_height, invert_roughness):
|
||
"""Extract PBR maps from FrankenMap"""
|
||
|
||
print("\n" + "="*60)
|
||
print("FrankenMap PBR Extractor")
|
||
print("="*60)
|
||
print(f"Input shape: {frankenmap.shape}")
|
||
|
||
# Validate input
|
||
if frankenmap.shape[-1] < 3:
|
||
print("⚠ Warning: FrankenMap should have 3 channels (RGB)")
|
||
print(" Padding with zeros for missing channels")
|
||
# Pad with zeros if needed
|
||
missing = 3 - frankenmap.shape[-1]
|
||
padding = torch.zeros(
|
||
(frankenmap.shape[0], frankenmap.shape[1], frankenmap.shape[2], missing),
|
||
device=frankenmap.device, dtype=frankenmap.dtype
|
||
)
|
||
frankenmap = torch.cat([frankenmap, padding], dim=-1)
|
||
|
||
# Extract channels
|
||
# Red = Albedo/Grayscale
|
||
red_channel = frankenmap[:, :, :, 0:1]
|
||
# Green = Height
|
||
green_channel = frankenmap[:, :, :, 1:2]
|
||
# Blue = Roughness
|
||
blue_channel = frankenmap[:, :, :, 2:3]
|
||
|
||
print(f"\n📊 CHANNEL EXTRACTION:")
|
||
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 Grayscale (from red channel)
|
||
print(f"\n🎨 PROCESSING GRAYSCALE:")
|
||
print(f" Gamma: {albedo_gamma}")
|
||
print(f" Brightness: {albedo_brightness}")
|
||
|
||
albedo = red_channel
|
||
albedo = self.apply_gamma(albedo, albedo_gamma)
|
||
albedo = self.adjust_brightness(albedo, albedo_brightness)
|
||
albedo = self.channel_to_rgb(albedo)
|
||
|
||
print(f" Output range: [{albedo.min():.3f}, {albedo.max():.3f}]")
|
||
|
||
# Process Height (from green channel)
|
||
print(f"\n⛰️ PROCESSING HEIGHT:")
|
||
print(f" Gamma: {height_gamma}")
|
||
print(f" Contrast: {height_contrast}")
|
||
print(f" Invert: {invert_height}")
|
||
|
||
height = green_channel
|
||
height = self.apply_gamma(height, height_gamma)
|
||
height = self.adjust_contrast(height, height_contrast)
|
||
|
||
if invert_height:
|
||
height = 1.0 - height
|
||
print(f" ✓ Height inverted")
|
||
|
||
height = self.channel_to_rgb(height)
|
||
|
||
print(f" Output range: [{height.min():.3f}, {height.max():.3f}]")
|
||
|
||
# Process Roughness (from blue channel)
|
||
print(f"\n🔨 PROCESSING ROUGHNESS:")
|
||
print(f" Gamma: {roughness_gamma}")
|
||
print(f" Brightness: {roughness_brightness}")
|
||
print(f" Invert: {invert_roughness}")
|
||
|
||
roughness = blue_channel
|
||
roughness = self.apply_gamma(roughness, roughness_gamma)
|
||
roughness = self.adjust_brightness(roughness, roughness_brightness)
|
||
|
||
if invert_roughness:
|
||
roughness = 1.0 - roughness
|
||
print(f" ✓ Roughness inverted")
|
||
|
||
roughness = self.channel_to_rgb(roughness)
|
||
|
||
print(f" Output range: [{roughness.min():.3f}, {roughness.max():.3f}]")
|
||
|
||
# Create PBR pipe
|
||
pbr_pipe = {
|
||
"albedo": albedo,
|
||
"normal": None,
|
||
"ao": None,
|
||
"height": height,
|
||
"roughness": roughness,
|
||
"metallic": None,
|
||
"transparency": None,
|
||
"emission": None,
|
||
"image": frankenmap, # Store original input
|
||
}
|
||
|
||
print(f"\n✓ FrankenMap extraction complete")
|
||
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, 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
|
||
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,
|
||
}
|
||
|
||
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)",
|
||
}
|