366 lines
14 KiB
Python
366 lines
14 KiB
Python
"""
|
|
PBR Core Nodes
|
|
Extraction and basic adjustment of PBR maps from Marigold outputs
|
|
"""
|
|
|
|
import torch
|
|
|
|
|
|
class PBRExtractor:
|
|
"""
|
|
Extracts PBR maps from Marigold IID Appearance and Lighting outputs
|
|
|
|
Inputs:
|
|
- Marigold Appearance output (3 batches: 0=albedo, 1=rough/metal, 2=unused)
|
|
- Marigold Lighting output (3 batches: 0=albedo, 1=AO/Lighting, 2=unused)
|
|
|
|
Outputs: PBR_PIPE, Albedo, Roughness, Metallic, Lighting
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
# Marigold outputs
|
|
"marigold_appearance": ("IMAGE",),
|
|
"marigold_lighting": ("IMAGE",),
|
|
|
|
# Albedo choice
|
|
"albedo_source": (["appearance", "lighting"], {
|
|
"default": "lighting",
|
|
"tooltip": "Which Marigold output to use for albedo"
|
|
}),
|
|
|
|
# Gamma correction - separate controls
|
|
"gamma_albedo": ("FLOAT", {
|
|
"default": 0.45,
|
|
"min": 0.1,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Gamma correction for albedo (applied to both appearance and lighting sources)"
|
|
}),
|
|
"gamma_metal_rough": ("FLOAT", {
|
|
"default": 2.2,
|
|
"min": 0.1,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Gamma correction for metallic and roughness maps"
|
|
}),
|
|
"gamma_lighting_ao": ("FLOAT", {
|
|
"default": 0.45,
|
|
"min": 0.1,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Gamma correction for lighting and AO maps"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("PBR_PIPE",)
|
|
RETURN_NAMES = ("pbr_pipe",)
|
|
FUNCTION = "extract"
|
|
CATEGORY = "Texture Alchemist"
|
|
|
|
def apply_gamma(self, image, gamma):
|
|
"""Apply gamma correction"""
|
|
return torch.pow(torch.clamp(image, 0.0, 1.0), gamma)
|
|
|
|
def channel_to_rgb(self, channel):
|
|
"""Convert single channel to RGB"""
|
|
if channel.shape[-1] == 1:
|
|
return channel.repeat(1, 1, 1, 3)
|
|
return channel
|
|
|
|
def grayscale(self, image):
|
|
"""Convert to grayscale"""
|
|
if image.shape[-1] == 1:
|
|
return image
|
|
weights = torch.tensor([0.2989, 0.5870, 0.1140],
|
|
device=image.device, dtype=image.dtype)
|
|
gray = torch.sum(image * weights, dim=-1, keepdim=True)
|
|
return self.channel_to_rgb(gray)
|
|
|
|
def extract(self, marigold_appearance, marigold_lighting, albedo_source, gamma_albedo, gamma_metal_rough, gamma_lighting_ao):
|
|
"""Extract PBR maps"""
|
|
|
|
print("\n" + "="*60)
|
|
print("PBR Extractor")
|
|
print("="*60)
|
|
print(f"Appearance shape: {marigold_appearance.shape}")
|
|
print(f"Lighting shape: {marigold_lighting.shape}")
|
|
print(f"Albedo source: {albedo_source}")
|
|
print(f"Gamma - Albedo: {gamma_albedo}, Metal/Rough: {gamma_metal_rough}, Lighting/AO: {gamma_lighting_ao}")
|
|
|
|
# ===== ALBEDO =====
|
|
if albedo_source == "appearance":
|
|
# Get index 0 from appearance
|
|
albedo = marigold_appearance[0:1] if marigold_appearance.shape[0] > 0 else marigold_appearance
|
|
print("\n✓ Albedo from Appearance (index 0)")
|
|
else:
|
|
# Get index 0 from lighting
|
|
albedo = marigold_lighting[0:1] if marigold_lighting.shape[0] > 0 else marigold_lighting
|
|
print("\n✓ Albedo from Lighting (index 0)")
|
|
|
|
albedo = self.apply_gamma(albedo, gamma_albedo)
|
|
print(f" Gamma: {gamma_albedo}")
|
|
print(f" Range: [{albedo.min():.3f}, {albedo.max():.3f}]")
|
|
|
|
# ===== ROUGHNESS (Red channel from Appearance index 1) =====
|
|
if marigold_appearance.shape[0] >= 2:
|
|
material_batch = marigold_appearance[1:2]
|
|
roughness = material_batch[:, :, :, 0:1] # Red channel
|
|
roughness = self.channel_to_rgb(roughness)
|
|
roughness = self.apply_gamma(roughness, gamma_metal_rough)
|
|
print("\n✓ Roughness from Appearance index 1, RED channel")
|
|
print(f" Gamma: {gamma_metal_rough}")
|
|
print(f" Range: [{roughness.min():.3f}, {roughness.max():.3f}]")
|
|
else:
|
|
print("\n✗ WARNING: Appearance doesn't have index 1, using zeros")
|
|
roughness = torch.zeros_like(albedo)
|
|
|
|
# ===== METALLIC (Green channel from Appearance index 1) =====
|
|
if marigold_appearance.shape[0] >= 2:
|
|
material_batch = marigold_appearance[1:2]
|
|
metallic = material_batch[:, :, :, 1:2] # Green channel
|
|
metallic = self.channel_to_rgb(metallic)
|
|
metallic = self.apply_gamma(metallic, gamma_metal_rough)
|
|
print("\n✓ Metallic from Appearance index 1, GREEN channel")
|
|
print(f" Gamma: {gamma_metal_rough}")
|
|
print(f" Range: [{metallic.min():.3f}, {metallic.max():.3f}]")
|
|
else:
|
|
print("\n✗ WARNING: Appearance doesn't have index 1, using zeros")
|
|
metallic = torch.zeros_like(albedo)
|
|
|
|
# ===== LIGHTING / AO (Lighting index 1) =====
|
|
if marigold_lighting.shape[0] >= 2:
|
|
lighting = marigold_lighting[1:2]
|
|
lighting = self.grayscale(lighting)
|
|
lighting = self.apply_gamma(lighting, gamma_lighting_ao)
|
|
print("\n✓ Lighting/AO from Lighting index 1")
|
|
print(f" Gamma: {gamma_lighting_ao}")
|
|
print(f" Range: [{lighting.min():.3f}, {lighting.max():.3f}]")
|
|
else:
|
|
print("\n✗ WARNING: Lighting doesn't have index 1, using albedo luminance")
|
|
lighting = self.grayscale(albedo)
|
|
|
|
# ===== CREATE PBR PIPE =====
|
|
pbr_pipe = {
|
|
"albedo": albedo,
|
|
"normal": None,
|
|
"ao": lighting, # Store as 'ao' in pipe for compatibility
|
|
"height": None,
|
|
"roughness": roughness,
|
|
"metallic": metallic,
|
|
"transparency": None,
|
|
}
|
|
|
|
print("\n" + "="*60)
|
|
print("✓ Extraction Complete")
|
|
print("✓ PBR Pipeline created")
|
|
print("="*60 + "\n")
|
|
|
|
return (pbr_pipe,)
|
|
|
|
|
|
class PBRAdjuster:
|
|
"""
|
|
Adjust brightness, contrast, and invert for each PBR map
|
|
|
|
Takes: Albedo, AO, Roughness, Metallic
|
|
Outputs: Adjusted versions
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
# Inputs
|
|
"albedo": ("IMAGE",),
|
|
"ao": ("IMAGE",),
|
|
"roughness": ("IMAGE",),
|
|
"metallic": ("IMAGE",),
|
|
|
|
# Albedo adjustments
|
|
"albedo_brightness": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"albedo_contrast": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"albedo_invert": ("BOOLEAN", {"default": False}),
|
|
|
|
# AO adjustments
|
|
"ao_brightness": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"ao_contrast": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"ao_invert": ("BOOLEAN", {"default": False}),
|
|
|
|
# Roughness adjustments
|
|
"roughness_brightness": ("FLOAT", {
|
|
"default": 1.5,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"roughness_contrast": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"roughness_invert": ("BOOLEAN", {"default": False}),
|
|
|
|
# Metallic adjustments
|
|
"metallic_brightness": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"metallic_contrast": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"metallic_invert": ("BOOLEAN", {"default": False}),
|
|
|
|
# Embed transparency option
|
|
"embed_transparency": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Embed transparency map into albedo's alpha channel"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"transparency": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("albedo", "ao", "roughness", "metallic")
|
|
FUNCTION = "adjust"
|
|
CATEGORY = "Texture Alchemist"
|
|
|
|
def adjust_map(self, image, brightness, contrast, invert):
|
|
"""Apply brightness, contrast, and invert to a map"""
|
|
# Apply brightness
|
|
adjusted = image * brightness
|
|
|
|
# Apply contrast (around midpoint 0.5)
|
|
adjusted = (adjusted - 0.5) * contrast + 0.5
|
|
|
|
# Clamp
|
|
adjusted = torch.clamp(adjusted, 0.0, 1.0)
|
|
|
|
# Invert if requested
|
|
if invert:
|
|
adjusted = 1.0 - adjusted
|
|
|
|
return adjusted
|
|
|
|
def adjust(self, albedo, ao, roughness, metallic,
|
|
albedo_brightness, albedo_contrast, albedo_invert,
|
|
ao_brightness, ao_contrast, ao_invert,
|
|
roughness_brightness, roughness_contrast, roughness_invert,
|
|
metallic_brightness, metallic_contrast, metallic_invert,
|
|
embed_transparency, transparency=None):
|
|
"""Adjust all PBR maps"""
|
|
|
|
print("\n" + "="*60)
|
|
print("PBR Adjuster")
|
|
print("="*60)
|
|
|
|
# Adjust Albedo
|
|
albedo_adj = self.adjust_map(albedo, albedo_brightness, albedo_contrast, albedo_invert)
|
|
print(f"Albedo: brightness={albedo_brightness}, contrast={albedo_contrast}, invert={albedo_invert}")
|
|
print(f" Range: [{albedo_adj.min():.3f}, {albedo_adj.max():.3f}]")
|
|
|
|
# Adjust AO
|
|
ao_adj = self.adjust_map(ao, ao_brightness, ao_contrast, ao_invert)
|
|
print(f"\nAO: brightness={ao_brightness}, contrast={ao_contrast}, invert={ao_invert}")
|
|
print(f" Range: [{ao_adj.min():.3f}, {ao_adj.max():.3f}]")
|
|
|
|
# Adjust Roughness
|
|
roughness_adj = self.adjust_map(roughness, roughness_brightness, roughness_contrast, roughness_invert)
|
|
print(f"\nRoughness: brightness={roughness_brightness}, contrast={roughness_contrast}, invert={roughness_invert}")
|
|
print(f" Range: [{roughness_adj.min():.3f}, {roughness_adj.max():.3f}]")
|
|
|
|
# Adjust Metallic
|
|
metallic_adj = self.adjust_map(metallic, metallic_brightness, metallic_contrast, metallic_invert)
|
|
print(f"\nMetallic: brightness={metallic_brightness}, contrast={metallic_contrast}, invert={metallic_invert}")
|
|
print(f" Range: [{metallic_adj.min():.3f}, {metallic_adj.max():.3f}]")
|
|
|
|
# ===== EMBED TRANSPARENCY INTO ALBEDO ALPHA =====
|
|
if transparency is not None and embed_transparency:
|
|
# Convert transparency to grayscale if needed
|
|
trans_gray = transparency
|
|
if trans_gray.shape[-1] == 3:
|
|
weights = torch.tensor([0.299, 0.587, 0.114],
|
|
device=trans_gray.device, dtype=trans_gray.dtype)
|
|
trans_gray = torch.sum(trans_gray * weights, dim=-1, keepdim=True)
|
|
|
|
# Resize transparency to match albedo if needed
|
|
if trans_gray.shape[1:3] != albedo_adj.shape[1:3]:
|
|
trans_gray = torch.nn.functional.interpolate(
|
|
trans_gray.permute(0, 3, 1, 2),
|
|
size=albedo_adj.shape[1:3],
|
|
mode='bilinear',
|
|
align_corners=False
|
|
).permute(0, 2, 3, 1)
|
|
|
|
# Ensure albedo is RGB (3 channels)
|
|
if albedo_adj.shape[-1] == 4:
|
|
albedo_adj = albedo_adj[:, :, :, :3]
|
|
elif albedo_adj.shape[-1] == 1:
|
|
albedo_adj = albedo_adj.repeat(1, 1, 1, 3)
|
|
|
|
# Concatenate transparency as alpha channel
|
|
albedo_adj = torch.cat([albedo_adj, trans_gray], dim=-1)
|
|
print("\n✓ Transparency embedded into albedo alpha channel")
|
|
print(f" Albedo shape: {albedo_adj.shape}")
|
|
|
|
print("\n" + "="*60)
|
|
print("✓ Adjustments Complete")
|
|
print("="*60 + "\n")
|
|
|
|
return (albedo_adj, ao_adj, roughness_adj, metallic_adj)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PBRExtractor": PBRExtractor,
|
|
"PBRAdjuster": PBRAdjuster,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PBRExtractor": "PBR Extractor (Marigold)",
|
|
"PBRAdjuster": "PBR Adjuster",
|
|
}
|
|
|