460 lines
16 KiB
Python
460 lines
16 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 FrankenMapExtractor:
|
|
"""
|
|
Extract PBR maps from FrankenMap format
|
|
FrankenMap channel layout:
|
|
- Red Channel: Grayscale/Albedo
|
|
- Green Channel: Height
|
|
- Blue Channel: Roughness
|
|
"""
|
|
|
|
@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 albedo (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": "Albedo 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", "PBR_PIPE")
|
|
RETURN_NAMES = ("albedo", "height", "roughness", "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 (Albedo): range [{red_channel.min():.3f}, {red_channel.max():.3f}]")
|
|
print(f" Green (Height): range [{green_channel.min():.3f}, {green_channel.max():.3f}]")
|
|
print(f" Blue (Roughness): range [{blue_channel.min():.3f}, {blue_channel.max():.3f}]")
|
|
|
|
# Process Albedo (from red channel)
|
|
print(f"\n🎨 PROCESSING ALBEDO:")
|
|
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,
|
|
}
|
|
|
|
print(f"\n✓ FrankenMap extraction complete")
|
|
print(f" Extracted: Albedo, Height, Roughness")
|
|
print(f" PBR Pipe created")
|
|
print("="*60 + "\n")
|
|
|
|
return (albedo, height, roughness, pbr_pipe)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PBRMaterialProcessor": PBRMaterialProcessor,
|
|
"PBRNormalProcessor": PBRNormalProcessor,
|
|
"PBRHeightProcessor": PBRHeightProcessor,
|
|
"FrankenMapExtractor": FrankenMapExtractor,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PBRMaterialProcessor": "PBR Material Processor (Marigold)",
|
|
"PBRNormalProcessor": "PBR Normal Processor (Lotus)",
|
|
"PBRHeightProcessor": "PBR Height Processor (Lotus)",
|
|
"FrankenMapExtractor": "PBR Extractor (FrankenMap)",
|
|
}
|