266 lines
9.3 KiB
Python
266 lines
9.3 KiB
Python
"""
|
|
Map Utilities
|
|
Height and AO map processing
|
|
"""
|
|
|
|
import torch
|
|
|
|
|
|
class LotusHeightProcessor:
|
|
"""
|
|
Process Lotus depth/height map output
|
|
Connect: LotusSampler → VAEDecode → This node
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"lotus_depth": ("IMAGE",),
|
|
"invert": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Invert height values"
|
|
}),
|
|
"bit_depth": (["8-bit", "16-bit", "32-bit"], {
|
|
"default": "16-bit",
|
|
"tooltip": "Output bit depth (save as EXR for 16/32-bit)"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("height",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "Texture Alchemist/Maps"
|
|
|
|
def process(self, lotus_depth, invert, bit_depth):
|
|
"""Process height map"""
|
|
|
|
height = lotus_depth.clone()
|
|
|
|
# Convert to grayscale
|
|
if height.shape[-1] == 3:
|
|
weights = torch.tensor([0.299, 0.587, 0.114],
|
|
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
|
|
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
|
|
|
|
print(f"✓ Height processed ({bit_depth})")
|
|
if bit_depth in ["16-bit", "32-bit"]:
|
|
print(f" Note: Save as OpenEXR to preserve {bit_depth} precision")
|
|
|
|
return (height,)
|
|
|
|
|
|
class AOApproximator:
|
|
"""
|
|
Generate Ambient Occlusion map from height and normal maps
|
|
Uses height sampling to detect concave areas that would be occluded
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"height": ("IMAGE",),
|
|
"radius": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 64,
|
|
"step": 1,
|
|
"display": "number",
|
|
"tooltip": "Sampling radius in pixels (larger = broader occlusion)"
|
|
}),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "AO intensity multiplier"
|
|
}),
|
|
"samples": ("INT", {
|
|
"default": 16,
|
|
"min": 4,
|
|
"max": 64,
|
|
"step": 1,
|
|
"display": "number",
|
|
"tooltip": "Number of sampling directions (more = better quality, slower)"
|
|
}),
|
|
"contrast": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Contrast adjustment for AO"
|
|
}),
|
|
"use_normal": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Use normal map to improve AO quality (if provided)"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"normal": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("ao",)
|
|
FUNCTION = "generate_ao"
|
|
CATEGORY = "Texture Alchemist/Maps"
|
|
|
|
def generate_ao(self, height, radius, strength, samples, contrast, use_normal, normal=None):
|
|
"""Generate AO from height map"""
|
|
|
|
print("\n" + "="*60)
|
|
print("AO Approximator")
|
|
print("="*60)
|
|
print(f"Height shape: {height.shape}")
|
|
if normal is not None:
|
|
print(f"Normal shape: {normal.shape}")
|
|
print(f"Radius: {radius}, Samples: {samples}, Strength: {strength}, Contrast: {contrast}")
|
|
print(f"Use normal: {use_normal}")
|
|
|
|
# Convert height to grayscale if needed
|
|
height_map = height.clone()
|
|
if height_map.shape[-1] == 3:
|
|
weights = torch.tensor([0.299, 0.587, 0.114],
|
|
device=height_map.device, dtype=height_map.dtype)
|
|
height_map = torch.sum(height_map * weights, dim=-1, keepdim=True)
|
|
elif height_map.shape[-1] != 1:
|
|
height_map = height_map[..., 0:1]
|
|
|
|
batch, h, w, channels = height.shape
|
|
|
|
# Initialize AO map (start fully lit)
|
|
ao = torch.ones((batch, h, w, 1), device=height.device, dtype=height.dtype)
|
|
|
|
# Generate sampling angles
|
|
import math
|
|
angles = [2.0 * math.pi * i / samples for i in range(samples)]
|
|
|
|
print(f"\n⚙ Computing AO with {samples} samples at radius {radius}...")
|
|
|
|
# For each sampling direction
|
|
occlusion_sum = torch.zeros_like(ao)
|
|
|
|
for angle in angles:
|
|
# Calculate offset
|
|
dx = int(round(math.cos(angle) * radius))
|
|
dy = int(round(math.sin(angle) * radius))
|
|
|
|
# Skip if no offset
|
|
if dx == 0 and dy == 0:
|
|
continue
|
|
|
|
# Sample height at offset position (with boundary handling)
|
|
# Pad the height map to handle boundaries
|
|
pad_h = abs(dy) if dy != 0 else 0
|
|
pad_w = abs(dx) if dx != 0 else 0
|
|
|
|
# Use reflection padding
|
|
height_padded = torch.nn.functional.pad(
|
|
height_map.permute(0, 3, 1, 2), # BHWC -> BCHW
|
|
(pad_w, pad_w, pad_h, pad_h),
|
|
mode='replicate'
|
|
).permute(0, 2, 3, 1) # BCHW -> BHWC
|
|
|
|
# Calculate the sampling positions
|
|
y_start = max(0, dy) + pad_h
|
|
y_end = y_start + h
|
|
x_start = max(0, dx) + pad_w
|
|
x_end = x_start + w
|
|
|
|
# Center position
|
|
center_y = pad_h
|
|
center_x = pad_w
|
|
|
|
# Sample neighbor and center
|
|
neighbor_height = height_padded[:, y_start:y_end, x_start:x_end, :]
|
|
center_height = height_padded[:, center_y:center_y+h, center_x:center_x+w, :]
|
|
|
|
# Calculate height difference (positive if neighbor is higher = occlusion)
|
|
height_diff = neighbor_height - center_height
|
|
|
|
# Convert to occlusion (only positive differences contribute)
|
|
occlusion = torch.clamp(height_diff, 0.0, 1.0)
|
|
|
|
occlusion_sum += occlusion
|
|
|
|
# Average the occlusion
|
|
if samples > 0:
|
|
occlusion_avg = occlusion_sum / samples
|
|
else:
|
|
occlusion_avg = occlusion_sum
|
|
|
|
# Apply strength
|
|
occlusion_avg = occlusion_avg * strength
|
|
|
|
# Convert to AO (1.0 = no occlusion, 0.0 = full occlusion)
|
|
ao = 1.0 - torch.clamp(occlusion_avg, 0.0, 1.0)
|
|
|
|
# Apply contrast
|
|
if contrast != 1.0:
|
|
ao = (ao - 0.5) * contrast + 0.5
|
|
ao = torch.clamp(ao, 0.0, 1.0)
|
|
|
|
# Optional: Use normal map to bias AO
|
|
if use_normal and normal is not None:
|
|
print("✓ Applying normal-based bias")
|
|
|
|
# Convert normal to grayscale or use blue channel (up direction)
|
|
if normal.shape[-1] >= 3:
|
|
# Use blue channel (Z/up direction)
|
|
# Areas facing up get less AO, areas facing down/sideways get more
|
|
normal_z = normal[..., 2:3]
|
|
|
|
# Resize to match AO if needed
|
|
if normal_z.shape[1:3] != ao.shape[1:3]:
|
|
normal_z = torch.nn.functional.interpolate(
|
|
normal_z.permute(0, 3, 1, 2),
|
|
size=ao.shape[1:3],
|
|
mode='bilinear',
|
|
align_corners=False
|
|
).permute(0, 2, 3, 1)
|
|
|
|
# Blend AO with normal bias
|
|
# Faces pointing up (normal_z close to 1.0) = less AO
|
|
# Faces pointing sideways/down (normal_z close to 0.5 or less) = more AO
|
|
normal_bias = torch.clamp((1.0 - normal_z) * 0.5, 0.0, 0.5)
|
|
ao = ao * (1.0 - normal_bias)
|
|
|
|
# Convert to RGB
|
|
ao_rgb = ao.repeat(1, 1, 1, 3)
|
|
|
|
print(f"\n✓ AO generated")
|
|
print(f" Range: [{ao_rgb.min():.3f}, {ao_rgb.max():.3f}]")
|
|
print("\n" + "="*60)
|
|
print("✓ AO Approximation Complete")
|
|
print("="*60 + "\n")
|
|
|
|
return (ao_rgb,)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LotusHeightProcessor": LotusHeightProcessor,
|
|
"AOApproximator": AOApproximator,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LotusHeightProcessor": "Height Processor (Lotus)",
|
|
"AOApproximator": "AO Approximator",
|
|
}
|
|
|