453 lines
16 KiB
Python
453 lines
16 KiB
Python
"""
|
|
Effect Utilities
|
|
Curvature, wear, detail blending, and masking tools
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class CurvatureGenerator:
|
|
"""
|
|
Generate curvature map from normal map or height map
|
|
Detects edges and crevices for wear/dirt masks
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"input_map": ("IMAGE",),
|
|
"input_type": (["normal", "height"], {
|
|
"default": "normal",
|
|
"tooltip": "Type of input map"
|
|
}),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"display": "number",
|
|
"tooltip": "Curvature detection strength"
|
|
}),
|
|
"blur_radius": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"display": "number",
|
|
"tooltip": "Smoothing radius"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("curvature",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Texture Alchemist/Effects"
|
|
|
|
def generate(self, input_map, input_type, strength, blur_radius):
|
|
"""Generate curvature map"""
|
|
|
|
print("\n" + "="*60)
|
|
print("Curvature Generator")
|
|
print("="*60)
|
|
print(f"Input shape: {input_map.shape}")
|
|
print(f"Input type: {input_type}")
|
|
print(f"Strength: {strength}")
|
|
|
|
# Convert to grayscale
|
|
if input_map.shape[-1] > 1:
|
|
if input_type == "normal":
|
|
# For normals, use Z channel (blue)
|
|
gray = input_map[:, :, :, 2:3]
|
|
else:
|
|
# For height, use luminance
|
|
weights = torch.tensor([0.2989, 0.5870, 0.1140],
|
|
device=input_map.device, dtype=input_map.dtype)
|
|
gray = torch.sum(input_map * weights, dim=-1, keepdim=True)
|
|
else:
|
|
gray = input_map
|
|
|
|
# Compute second derivative (curvature)
|
|
gray_bchw = gray.permute(0, 3, 1, 2)
|
|
|
|
# Laplacian kernel (detects edges and curvature)
|
|
laplacian = torch.tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]],
|
|
device=input_map.device, dtype=input_map.dtype).view(1, 1, 3, 3) / 4.0
|
|
|
|
curvature = F.conv2d(gray_bchw, laplacian, padding=1)
|
|
|
|
# Apply strength
|
|
curvature = curvature * strength
|
|
|
|
# Normalize and enhance
|
|
curvature = torch.abs(curvature)
|
|
|
|
# Apply blur if requested
|
|
if blur_radius > 0:
|
|
kernel_size = int(blur_radius * 4) + 1
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
sigma = blur_radius
|
|
|
|
# Gaussian blur
|
|
x = torch.arange(kernel_size, device=input_map.device, dtype=input_map.dtype)
|
|
kernel = torch.exp(-0.5 * ((x - kernel_size // 2) / sigma) ** 2)
|
|
kernel = kernel / kernel.sum()
|
|
|
|
kernel_h = kernel.view(1, 1, 1, -1)
|
|
kernel_v = kernel.view(1, 1, -1, 1)
|
|
|
|
padding = kernel_size // 2
|
|
curvature = F.pad(curvature, (padding, padding, 0, 0), mode='replicate')
|
|
curvature = F.conv2d(curvature, kernel_h)
|
|
curvature = F.pad(curvature, (0, 0, padding, padding), mode='replicate')
|
|
curvature = F.conv2d(curvature, kernel_v)
|
|
|
|
# Normalize to 0-1
|
|
curvature = (curvature - curvature.min()) / (curvature.max() - curvature.min() + 1e-8)
|
|
|
|
# Convert to RGB
|
|
curvature = curvature.permute(0, 2, 3, 1).repeat(1, 1, 1, 3)
|
|
|
|
print(f"✓ Curvature map generated")
|
|
print(f" Output shape: {curvature.shape}")
|
|
print("="*60 + "\n")
|
|
|
|
return (curvature,)
|
|
|
|
|
|
class DetailMapBlender:
|
|
"""
|
|
Blend detail maps onto base maps without washing out
|
|
Uses overlay/detail blending for normals and roughness
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base": ("IMAGE",),
|
|
"detail": ("IMAGE",),
|
|
"map_type": (["normal", "roughness", "generic"], {
|
|
"default": "generic",
|
|
"tooltip": "Map type determines blend method"
|
|
}),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Detail strength (0=no detail, 1=full)"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"mask": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("blended",)
|
|
FUNCTION = "blend"
|
|
CATEGORY = "Texture Alchemist/Effects"
|
|
|
|
def blend(self, base, detail, map_type, strength, mask=None):
|
|
"""Blend detail map onto base"""
|
|
|
|
print("\n" + "="*60)
|
|
print("Detail Map Blender")
|
|
print("="*60)
|
|
print(f"Base shape: {base.shape}")
|
|
print(f"Detail shape: {detail.shape}")
|
|
print(f"Map type: {map_type}")
|
|
print(f"Strength: {strength}")
|
|
|
|
# Resize detail to match base if needed
|
|
if detail.shape[1:3] != base.shape[1:3]:
|
|
detail = self._resize_to_match(detail, base)
|
|
|
|
# Apply blending based on map type
|
|
if map_type == "normal":
|
|
# RNM (Reoriented Normal Mapping) for normals
|
|
result = self._blend_normals_rnm(base, detail, strength)
|
|
elif map_type == "roughness":
|
|
# Multiply/overlay for roughness
|
|
result = base * (1.0 + (detail - 0.5) * strength)
|
|
else:
|
|
# Generic overlay blend
|
|
result = base + (detail - 0.5) * strength
|
|
|
|
# Apply mask if provided
|
|
if mask is not None:
|
|
if mask.shape[1:3] != base.shape[1:3]:
|
|
mask = self._resize_to_match(mask, base)
|
|
if mask.shape[-1] != 1:
|
|
mask_weights = torch.tensor([0.2989, 0.5870, 0.1140],
|
|
device=mask.device, dtype=mask.dtype)
|
|
mask = torch.sum(mask * mask_weights, dim=-1, keepdim=True)
|
|
|
|
result = base * (1.0 - mask) + result * mask
|
|
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
print(f"✓ Detail blended")
|
|
print("="*60 + "\n")
|
|
|
|
return (result,)
|
|
|
|
def _blend_normals_rnm(self, base, detail, strength):
|
|
"""Reoriented Normal Mapping blend"""
|
|
# Convert from [0,1] to [-1,1]
|
|
base_n = base * 2.0 - 1.0
|
|
detail_n = detail * 2.0 - 1.0
|
|
|
|
# Apply strength to detail
|
|
detail_n = detail_n * strength
|
|
|
|
# RNM blend
|
|
result_n = torch.zeros_like(base_n)
|
|
result_n[:, :, :, 0] = base_n[:, :, :, 0] + detail_n[:, :, :, 0]
|
|
result_n[:, :, :, 1] = base_n[:, :, :, 1] + detail_n[:, :, :, 1]
|
|
if base_n.shape[-1] >= 3:
|
|
result_n[:, :, :, 2] = base_n[:, :, :, 2] * detail_n[:, :, :, 2]
|
|
|
|
# Normalize
|
|
length = torch.sqrt(torch.sum(result_n ** 2, dim=-1, keepdim=True)) + 1e-8
|
|
result_n = result_n / length
|
|
|
|
# Convert back to [0,1]
|
|
result = (result_n + 1.0) / 2.0
|
|
|
|
return result
|
|
|
|
def _resize_to_match(self, source, target):
|
|
"""Resize source to match target"""
|
|
target_h, target_w = target.shape[1:3]
|
|
source_bchw = source.permute(0, 3, 1, 2)
|
|
resized = F.interpolate(source_bchw, size=(target_h, target_w),
|
|
mode='bilinear', align_corners=False)
|
|
return resized.permute(0, 2, 3, 1)
|
|
|
|
|
|
class WearGenerator:
|
|
"""
|
|
Generate wear and edge damage for materials
|
|
Uses curvature and AO for realistic wear patterns
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"albedo": ("IMAGE",),
|
|
"wear_strength": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Overall wear intensity"
|
|
}),
|
|
"edge_wear": ("FLOAT", {
|
|
"default": 0.7,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Wear on edges (needs curvature or height)"
|
|
}),
|
|
"dirt_strength": ("FLOAT", {
|
|
"default": 0.3,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Dirt accumulation in crevices"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"curvature": ("IMAGE",),
|
|
"ao": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("worn_albedo", "wear_mask")
|
|
FUNCTION = "generate_wear"
|
|
CATEGORY = "Texture Alchemist/Effects"
|
|
|
|
def generate_wear(self, albedo, wear_strength, edge_wear, dirt_strength,
|
|
curvature=None, ao=None):
|
|
"""Generate wear effects"""
|
|
|
|
print("\n" + "="*60)
|
|
print("Wear Generator")
|
|
print("="*60)
|
|
print(f"Albedo shape: {albedo.shape}")
|
|
print(f"Wear strength: {wear_strength}")
|
|
|
|
batch, height, width, channels = albedo.shape
|
|
device = albedo.device
|
|
dtype = albedo.dtype
|
|
|
|
# Create wear mask
|
|
wear_mask = torch.zeros((batch, height, width, 1), device=device, dtype=dtype)
|
|
|
|
# Add edge wear from curvature
|
|
if curvature is not None and edge_wear > 0:
|
|
curv = self._to_grayscale(curvature)
|
|
if curv.shape[1:3] != albedo.shape[1:3]:
|
|
curv = self._resize_to_match(curv, albedo)
|
|
wear_mask = wear_mask + curv * edge_wear
|
|
|
|
# Add dirt from AO (inverted)
|
|
if ao is not None and dirt_strength > 0:
|
|
ao_gray = self._to_grayscale(ao)
|
|
if ao_gray.shape[1:3] != albedo.shape[1:3]:
|
|
ao_gray = self._resize_to_match(ao_gray, albedo)
|
|
dirt = (1.0 - ao_gray) * dirt_strength
|
|
wear_mask = wear_mask + dirt
|
|
|
|
# Normalize and apply strength
|
|
wear_mask = torch.clamp(wear_mask * wear_strength, 0.0, 1.0)
|
|
|
|
# Apply wear to albedo (darken)
|
|
worn_albedo = albedo * (1.0 - wear_mask * 0.5)
|
|
|
|
# Convert mask to RGB for preview
|
|
wear_mask_rgb = wear_mask.repeat(1, 1, 1, 3)
|
|
|
|
print(f"✓ Wear generated")
|
|
print(f" Worn albedo: {worn_albedo.shape}")
|
|
print(f" Wear mask: {wear_mask_rgb.shape}")
|
|
print("="*60 + "\n")
|
|
|
|
return (worn_albedo, wear_mask_rgb)
|
|
|
|
def _to_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)
|
|
return torch.sum(image * weights, dim=-1, keepdim=True)
|
|
|
|
def _resize_to_match(self, source, target):
|
|
"""Resize source to match target"""
|
|
target_h, target_w = target.shape[1:3]
|
|
source_bchw = source.permute(0, 3, 1, 2)
|
|
resized = F.interpolate(source_bchw, size=(target_h, target_w),
|
|
mode='bilinear', align_corners=False)
|
|
return resized.permute(0, 2, 3, 1)
|
|
|
|
|
|
class GradientMap:
|
|
"""
|
|
Create masks from gradients (like Color Ramp but outputs grayscale masks)
|
|
Useful for creating selection masks from height/AO/curvature
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"input_range_min": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Values below this become black"
|
|
}),
|
|
"input_range_max": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Values above this become white"
|
|
}),
|
|
"invert": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Invert the mask"
|
|
}),
|
|
"smoothness": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 0.5,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Edge smoothness (0=hard, >0=soft)"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("mask",)
|
|
FUNCTION = "create_mask"
|
|
CATEGORY = "Texture Alchemist/Effects"
|
|
|
|
def create_mask(self, image, input_range_min, input_range_max, invert, smoothness):
|
|
"""Create gradient mask"""
|
|
|
|
print("\n" + "="*60)
|
|
print("Gradient Map")
|
|
print("="*60)
|
|
print(f"Input shape: {image.shape}")
|
|
print(f"Range: [{input_range_min:.2f}, {input_range_max:.2f}]")
|
|
|
|
# Convert to grayscale
|
|
if image.shape[-1] > 1:
|
|
weights = torch.tensor([0.2989, 0.5870, 0.1140],
|
|
device=image.device, dtype=image.dtype)
|
|
gray = torch.sum(image * weights, dim=-1, keepdim=True)
|
|
else:
|
|
gray = image
|
|
|
|
# Remap range
|
|
mask = (gray - input_range_min) / (input_range_max - input_range_min + 1e-8)
|
|
mask = torch.clamp(mask, 0.0, 1.0)
|
|
|
|
# Apply smoothness (ease in/out)
|
|
if smoothness > 0:
|
|
# Smoothstep function
|
|
t = mask
|
|
mask = t * t * (3.0 - 2.0 * t)
|
|
# Additional smoothing passes
|
|
smooth_passes = int(smoothness * 10)
|
|
for _ in range(smooth_passes):
|
|
mask = mask * mask * (3.0 - 2.0 * mask)
|
|
|
|
# Invert if requested
|
|
if invert:
|
|
mask = 1.0 - mask
|
|
|
|
# Convert to RGB
|
|
mask_rgb = mask.repeat(1, 1, 1, 3)
|
|
|
|
print(f"✓ Gradient mask created")
|
|
print(f" Output range: [{mask_rgb.min():.3f}, {mask_rgb.max():.3f}]")
|
|
print("="*60 + "\n")
|
|
|
|
return (mask_rgb,)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CurvatureGenerator": CurvatureGenerator,
|
|
"DetailMapBlender": DetailMapBlender,
|
|
"WearGenerator": WearGenerator,
|
|
"GradientMap": GradientMap,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CurvatureGenerator": "Curvature Map Generator",
|
|
"DetailMapBlender": "Detail Map Blender",
|
|
"WearGenerator": "Wear & Edge Damage Generator",
|
|
"GradientMap": "Gradient Map (Mask)",
|
|
}
|
|
|