411 lines
15 KiB
Python
411 lines
15 KiB
Python
"""
|
|
Advanced Material Generation for TextureAlchemy
|
|
SSS, Anisotropy, Translucency, Emission maps
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class SSSMapGenerator:
|
|
"""
|
|
Generate Subsurface Scattering maps from albedo and thickness
|
|
For skin, wax, jade, marble materials
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"albedo": ("IMAGE",),
|
|
"scatter_color": ("STRING", {
|
|
"default": "#FFCCAA",
|
|
"tooltip": "SSS tint color (e.g., #FFCCAA for skin)"
|
|
}),
|
|
"scatter_strength": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"depth_influence": ("FLOAT", {
|
|
"default": 0.3,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "How much darker areas scatter more"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"thickness": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("sss_map",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Texture Alchemist/Materials"
|
|
|
|
def generate(self, albedo, scatter_color, scatter_strength, depth_influence, thickness=None):
|
|
"""Generate SSS map"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"SSS Map Generator")
|
|
print(f"{'='*60}")
|
|
|
|
# Parse scatter color
|
|
scatter_color = scatter_color.lstrip('#')
|
|
r = int(scatter_color[0:2], 16) / 255.0
|
|
g = int(scatter_color[2:4], 16) / 255.0
|
|
b = int(scatter_color[4:6], 16) / 255.0
|
|
|
|
color_tint = torch.tensor([r, g, b], device=albedo.device, dtype=albedo.dtype)
|
|
|
|
# Base SSS from albedo luminosity
|
|
lum = 0.299 * albedo[:, :, :, 0] + 0.587 * albedo[:, :, :, 1] + 0.114 * albedo[:, :, :, 2]
|
|
|
|
# Darker areas scatter more
|
|
sss_strength = (1.0 - lum) * depth_influence + (1.0 - depth_influence)
|
|
sss_strength = sss_strength.unsqueeze(3)
|
|
|
|
# Apply thickness if provided
|
|
if thickness is not None:
|
|
if thickness.shape[1:3] != albedo.shape[1:3]:
|
|
thickness = F.interpolate(
|
|
thickness.permute(0, 3, 1, 2),
|
|
size=(albedo.shape[1], albedo.shape[2]),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
).permute(0, 2, 3, 1)
|
|
|
|
thick_val = thickness.mean(dim=3, keepdim=True)
|
|
sss_strength = sss_strength * thick_val
|
|
|
|
# Apply scatter color
|
|
sss_map = color_tint.view(1, 1, 1, 3) * sss_strength * scatter_strength
|
|
|
|
# Blend with albedo
|
|
result = albedo * (1.0 - scatter_strength) + sss_map
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
print(f"✓ SSS map generated")
|
|
print(f" Scatter color: #{scatter_color}")
|
|
print(f" Strength: {scatter_strength:.2f}")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (result,)
|
|
|
|
|
|
class AnisotropyMapGenerator:
|
|
"""
|
|
Generate anisotropy maps for brushed metal, hair, fabric
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}),
|
|
"pattern_type": (["brushed_horizontal", "brushed_vertical", "brushed_circular", "hair_flow", "custom_angle"],),
|
|
"angle": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 360.0,
|
|
"step": 1.0,
|
|
"display": "number",
|
|
"tooltip": "Brush/flow direction in degrees"
|
|
}),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"variation": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Random variation in direction"
|
|
}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("anisotropy",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Texture Alchemist/Materials"
|
|
|
|
def generate(self, width, height, pattern_type, angle, strength, variation, seed):
|
|
"""Generate anisotropy map"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Anisotropy Map Generator - {pattern_type.upper()}")
|
|
print(f"{'='*60}")
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
dtype = torch.float32
|
|
|
|
torch.manual_seed(seed)
|
|
|
|
# Create coordinate grid
|
|
y_coords = torch.linspace(-1, 1, height, device=device, dtype=dtype).view(-1, 1).repeat(1, width)
|
|
x_coords = torch.linspace(-1, 1, width, device=device, dtype=dtype).view(1, -1).repeat(height, 1)
|
|
|
|
# Generate direction map based on pattern
|
|
if pattern_type == "brushed_horizontal":
|
|
direction = torch.zeros_like(x_coords)
|
|
elif pattern_type == "brushed_vertical":
|
|
direction = torch.ones_like(x_coords) * 90.0
|
|
elif pattern_type == "brushed_circular":
|
|
direction = torch.atan2(y_coords, x_coords) * 180.0 / 3.14159
|
|
elif pattern_type == "hair_flow":
|
|
# Wavy flow pattern
|
|
direction = x_coords * 45.0 + torch.sin(y_coords * 5.0) * 15.0
|
|
else: # custom_angle
|
|
direction = torch.ones_like(x_coords) * angle
|
|
|
|
# Add variation
|
|
if variation > 0.0:
|
|
noise = torch.rand((height, width), device=device, dtype=dtype) * 2.0 - 1.0
|
|
direction = direction + noise * variation * 45.0
|
|
|
|
# Convert direction to tangent vectors (normalized)
|
|
# Store as RGB where RG = tangent XY, B = strength
|
|
import math
|
|
dir_rad = direction * math.pi / 180.0
|
|
|
|
tangent_x = torch.cos(dir_rad) * strength
|
|
tangent_y = torch.sin(dir_rad) * strength
|
|
|
|
# Normalize to 0-1 range
|
|
tangent_x = (tangent_x + 1.0) * 0.5
|
|
tangent_y = (tangent_y + 1.0) * 0.5
|
|
|
|
aniso_map = torch.stack([tangent_x, tangent_y, torch.ones_like(tangent_x) * strength], dim=2)
|
|
aniso_map = aniso_map.unsqueeze(0)
|
|
aniso_map = torch.clamp(aniso_map, 0.0, 1.0)
|
|
|
|
print(f"✓ Anisotropy map generated")
|
|
print(f" Pattern: {pattern_type}")
|
|
print(f" Angle: {angle:.1f}°")
|
|
print(f" Strength: {strength:.2f}")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (aniso_map,)
|
|
|
|
|
|
class TranslucencyMapGenerator:
|
|
"""
|
|
Generate translucency maps for leaves, paper, thin fabrics
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"albedo": ("IMAGE",),
|
|
"translucency_color": ("STRING", {
|
|
"default": "#88FF88",
|
|
"tooltip": "Color when backlit (e.g., green for leaves)"
|
|
}),
|
|
"strength": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"thickness": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("translucency",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Texture Alchemist/Materials"
|
|
|
|
def generate(self, albedo, translucency_color, strength, thickness=None):
|
|
"""Generate translucency map"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Translucency Map Generator")
|
|
print(f"{'='*60}")
|
|
|
|
# Parse color
|
|
translucency_color = translucency_color.lstrip('#')
|
|
r = int(translucency_color[0:2], 16) / 255.0
|
|
g = int(translucency_color[2:4], 16) / 255.0
|
|
b = int(translucency_color[4:6], 16) / 255.0
|
|
|
|
trans_color = torch.tensor([r, g, b], device=albedo.device, dtype=albedo.dtype)
|
|
|
|
# Base translucency from albedo brightness
|
|
lum = 0.299 * albedo[:, :, :, 0] + 0.587 * albedo[:, :, :, 1] + 0.114 * albedo[:, :, :, 2]
|
|
trans_strength = lum.unsqueeze(3) * strength
|
|
|
|
# Apply thickness if provided
|
|
if thickness is not None:
|
|
if thickness.shape[1:3] != albedo.shape[1:3]:
|
|
thickness = F.interpolate(
|
|
thickness.permute(0, 3, 1, 2),
|
|
size=(albedo.shape[1], albedo.shape[2]),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
).permute(0, 2, 3, 1)
|
|
|
|
thick_val = 1.0 - thickness.mean(dim=3, keepdim=True) # Thinner = more translucent
|
|
trans_strength = trans_strength * thick_val
|
|
|
|
# Apply translucency color
|
|
result = trans_color.view(1, 1, 1, 3) * trans_strength
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
print(f"✓ Translucency map generated")
|
|
print(f" Color: #{translucency_color}")
|
|
print(f" Strength: {strength:.2f}")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (result,)
|
|
|
|
|
|
class EmissionMaskGenerator:
|
|
"""
|
|
Generate emission masks from brightness threshold or patterns
|
|
For glowing elements, lights, screens, magic effects
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mode": (["brightness_threshold", "color_select", "procedural"],),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"display": "number",
|
|
"tooltip": "Emission intensity"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"source_image": ("IMAGE",),
|
|
"threshold": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number"}),
|
|
"emission_color": ("STRING", {"default": "#FFFFFF"}),
|
|
"width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}),
|
|
"pattern_scale": ("FLOAT", {"default": 20.0, "min": 1.0, "max": 100.0, "step": 1.0, "display": "number"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("emission",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Texture Alchemist/Materials"
|
|
|
|
def generate(self, mode, strength, source_image=None, threshold=0.7, emission_color="#FFFFFF",
|
|
width=512, height=512, pattern_scale=20.0, seed=0):
|
|
"""Generate emission mask"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Emission Mask Generator - {mode.upper()}")
|
|
print(f"{'='*60}")
|
|
|
|
if mode == "brightness_threshold" and source_image is not None:
|
|
# Create emission from bright areas
|
|
lum = 0.299 * source_image[:, :, :, 0] + 0.587 * source_image[:, :, :, 1] + 0.114 * source_image[:, :, :, 2]
|
|
mask = torch.clamp((lum - threshold) / (1.0 - threshold), 0.0, 1.0)
|
|
emission = source_image * mask.unsqueeze(3) * strength
|
|
|
|
elif mode == "color_select" and source_image is not None:
|
|
# Use specific color as emission
|
|
color = emission_color.lstrip('#')
|
|
r = int(color[0:2], 16) / 255.0
|
|
g = int(color[2:4], 16) / 255.0
|
|
b = int(color[4:6], 16) / 255.0
|
|
|
|
target = torch.tensor([r, g, b], device=source_image.device, dtype=source_image.dtype)
|
|
|
|
# Color distance
|
|
diff = source_image - target.view(1, 1, 1, 3)
|
|
dist = torch.sqrt((diff * diff).sum(dim=3, keepdim=True))
|
|
|
|
# Create mask from similar colors
|
|
mask = torch.clamp(1.0 - dist / threshold, 0.0, 1.0)
|
|
emission = source_image * mask * strength
|
|
|
|
else: # procedural
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
dtype = torch.float32
|
|
|
|
torch.manual_seed(seed)
|
|
|
|
# Generate procedural emission pattern
|
|
noise = torch.rand((1, height, width, 1), device=device, dtype=dtype)
|
|
|
|
# Multi-scale
|
|
for octave in range(3):
|
|
freq = pattern_scale * (2 ** octave)
|
|
amp = 1.0 / (2 ** octave)
|
|
|
|
grid_h = max(int(height / freq), 2) + 2
|
|
grid_w = max(int(width / freq), 2) + 2
|
|
|
|
grid = torch.rand((1, grid_h, grid_w, 1), device=device, dtype=dtype)
|
|
upsampled = F.interpolate(
|
|
grid.permute(0, 3, 1, 2),
|
|
size=(height, width),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
).permute(0, 2, 3, 1)
|
|
|
|
noise = noise + upsampled * amp
|
|
|
|
noise = (noise - noise.min()) / (noise.max() - noise.min() + 1e-7)
|
|
|
|
# Apply color
|
|
color = emission_color.lstrip('#')
|
|
r = int(color[0:2], 16) / 255.0
|
|
g = int(color[2:4], 16) / 255.0
|
|
b = int(color[4:6], 16) / 255.0
|
|
|
|
color_vec = torch.tensor([r, g, b], device=device, dtype=dtype)
|
|
|
|
emission = noise * color_vec.view(1, 1, 1, 3) * strength
|
|
|
|
emission = torch.clamp(emission, 0.0, 1.0)
|
|
|
|
print(f"✓ Emission mask generated")
|
|
print(f" Mode: {mode}")
|
|
print(f" Strength: {strength:.2f}")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (emission,)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SSSMapGenerator": SSSMapGenerator,
|
|
"AnisotropyMapGenerator": AnisotropyMapGenerator,
|
|
"TranslucencyMapGenerator": TranslucencyMapGenerator,
|
|
"EmissionMaskGenerator": EmissionMaskGenerator,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SSSMapGenerator": "SSS Map Generator",
|
|
"AnisotropyMapGenerator": "Anisotropy Map Generator",
|
|
"TranslucencyMapGenerator": "Translucency Map Generator",
|
|
"EmissionMaskGenerator": "Emission Mask Generator",
|
|
}
|
|
|