252 lines
8.4 KiB
Python
252 lines
8.4 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,)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PBRMaterialProcessor": PBRMaterialProcessor,
|
|
"PBRNormalProcessor": PBRNormalProcessor,
|
|
"PBRHeightProcessor": PBRHeightProcessor,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PBRMaterialProcessor": "PBR Material Processor (Marigold)",
|
|
"PBRNormalProcessor": "PBR Normal Processor (Lotus)",
|
|
"PBRHeightProcessor": "PBR Height Processor (Lotus)",
|
|
}
|