296 lines
9.9 KiB
Python
296 lines
9.9 KiB
Python
"""
|
|
Advanced Height/Displacement Tools for TextureAlchemy
|
|
Height amplifier, combiner, displacement to vector
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class HeightAmplifier:
|
|
"""
|
|
Amplify or compress height range
|
|
Preserves mid-tones while increasing/decreasing contrast
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"height_map": ("IMAGE",),
|
|
"amount": ("FLOAT", {
|
|
"default": 2.0,
|
|
"min": 0.1,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"display": "number",
|
|
"tooltip": "Amplification amount (1.0 = no change, >1 = amplify, <1 = compress)"
|
|
}),
|
|
"center_point": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Midpoint to preserve"
|
|
}),
|
|
"method": (["power", "linear", "smooth"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("height",)
|
|
FUNCTION = "amplify"
|
|
CATEGORY = "Texture Alchemist/Height"
|
|
|
|
def amplify(self, height_map, amount, center_point, method):
|
|
"""Amplify height map"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Height Amplifier - {method.upper()}")
|
|
print(f"{'='*60}")
|
|
|
|
# Convert to grayscale if needed
|
|
if height_map.shape[3] > 1:
|
|
height = height_map.mean(dim=3, keepdim=True)
|
|
else:
|
|
height = height_map
|
|
|
|
if method == "power":
|
|
# Power curve
|
|
centered = height - center_point
|
|
amplified = torch.sign(centered) * torch.pow(torch.abs(centered), 1.0 / amount)
|
|
result = amplified + center_point
|
|
|
|
elif method == "linear":
|
|
# Linear amplification around center
|
|
result = (height - center_point) * amount + center_point
|
|
|
|
else: # smooth
|
|
# Smooth S-curve
|
|
centered = height - center_point
|
|
amplified = centered * amount
|
|
# Apply smooth step
|
|
result = amplified / torch.sqrt(1.0 + amplified * amplified) + center_point
|
|
|
|
# Convert back to 3-channel
|
|
result = result.repeat(1, 1, 1, 3)
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
print(f"✓ Height amplified")
|
|
print(f" Amount: {amount:.2f}x")
|
|
print(f" Method: {method}")
|
|
print(f" Center: {center_point:.2f}")
|
|
print(f" Range: [{result.min():.3f}, {result.max():.3f}]")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (result,)
|
|
|
|
|
|
class HeightCombiner:
|
|
"""
|
|
Combine multiple height maps intelligently
|
|
Layer macro and micro detail
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"layer1": ("IMAGE",),
|
|
"layer1_mode": (["add", "max", "min", "multiply", "screen", "overlay"], {"default": "add"}),
|
|
"layer1_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01, "display": "number"}),
|
|
|
|
"layer2": ("IMAGE",),
|
|
"layer2_mode": (["add", "max", "min", "multiply", "screen", "overlay"], {"default": "overlay"}),
|
|
"layer2_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01, "display": "number"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("height",)
|
|
FUNCTION = "combine"
|
|
CATEGORY = "Texture Alchemist/Height"
|
|
|
|
def combine(self, base, **kwargs):
|
|
"""Combine height maps"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Height Combiner")
|
|
print(f"{'='*60}")
|
|
|
|
result = base.clone()
|
|
if result.shape[3] > 1:
|
|
result = result.mean(dim=3, keepdim=True)
|
|
|
|
layers_applied = 0
|
|
|
|
for i in range(1, 3):
|
|
layer_key = f"layer{i}"
|
|
mode_key = f"layer{i}_mode"
|
|
strength_key = f"layer{i}_strength"
|
|
|
|
if layer_key in kwargs and kwargs[layer_key] is not None:
|
|
layer = kwargs[layer_key]
|
|
mode = kwargs.get(mode_key, "add")
|
|
strength = kwargs.get(strength_key, 1.0)
|
|
|
|
# Resize if needed
|
|
if layer.shape[1:3] != result.shape[1:3]:
|
|
layer = F.interpolate(
|
|
layer.permute(0, 3, 1, 2),
|
|
size=(result.shape[1], result.shape[2]),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
).permute(0, 2, 3, 1)
|
|
|
|
# Convert to grayscale
|
|
if layer.shape[3] > 1:
|
|
layer = layer.mean(dim=3, keepdim=True)
|
|
|
|
# Apply blend mode
|
|
if mode == "add":
|
|
blended = result + layer - 0.5
|
|
elif mode == "max":
|
|
blended = torch.max(result, layer)
|
|
elif mode == "min":
|
|
blended = torch.min(result, layer)
|
|
elif mode == "multiply":
|
|
blended = result * layer
|
|
elif mode == "screen":
|
|
blended = 1.0 - (1.0 - result) * (1.0 - layer)
|
|
else: # overlay
|
|
blended = torch.where(
|
|
result < 0.5,
|
|
2.0 * result * layer,
|
|
1.0 - 2.0 * (1.0 - result) * (1.0 - layer)
|
|
)
|
|
|
|
# Apply strength
|
|
result = result * (1.0 - strength) + blended * strength
|
|
layers_applied += 1
|
|
|
|
print(f" Layer {i}: {mode} (strength {strength:.2f})")
|
|
|
|
# Convert to 3-channel
|
|
result = result.repeat(1, 1, 1, 3)
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
print(f"\n✓ Combined {layers_applied} height layers")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (result,)
|
|
|
|
|
|
class DisplacementToVector:
|
|
"""
|
|
Convert height map to XYZ displacement vector map
|
|
For advanced displacement mapping
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"height_map": ("IMAGE",),
|
|
"magnitude": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"display": "number",
|
|
"tooltip": "Displacement magnitude"
|
|
}),
|
|
"direction_x": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -1.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"direction_y": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": -1.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"direction_z": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": -1.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Default is pure vertical displacement"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("vector_displacement",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "Texture Alchemist/Height"
|
|
|
|
def convert(self, height_map, magnitude, direction_x, direction_y, direction_z):
|
|
"""Convert to vector displacement"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"Displacement to Vector")
|
|
print(f"{'='*60}")
|
|
|
|
# Convert to grayscale
|
|
if height_map.shape[3] > 1:
|
|
height = height_map.mean(dim=3, keepdim=True)
|
|
else:
|
|
height = height_map
|
|
|
|
# Normalize direction vector
|
|
length = torch.sqrt(direction_x**2 + direction_y**2 + direction_z**2)
|
|
if length < 1e-7:
|
|
direction_x, direction_y, direction_z = 0.0, 0.0, 1.0
|
|
length = 1.0
|
|
|
|
dir_x = direction_x / length
|
|
dir_y = direction_y / length
|
|
dir_z = direction_z / length
|
|
|
|
# Create vector displacement
|
|
displacement = (height - 0.5) * magnitude # Center around 0
|
|
|
|
vec_x = displacement * dir_x
|
|
vec_y = displacement * dir_y
|
|
vec_z = displacement * dir_z
|
|
|
|
# Normalize to 0-1 range for output
|
|
vec_x = (vec_x + magnitude/2) / magnitude
|
|
vec_y = (vec_y + magnitude/2) / magnitude
|
|
vec_z = (vec_z + magnitude/2) / magnitude
|
|
|
|
result = torch.cat([vec_x, vec_y, vec_z], dim=3)
|
|
result = torch.clamp(result, 0.0, 1.0)
|
|
|
|
print(f"✓ Vector displacement created")
|
|
print(f" Magnitude: {magnitude:.2f}")
|
|
print(f" Direction: ({dir_x:.2f}, {dir_y:.2f}, {dir_z:.2f})")
|
|
print(f"{'='*60}\n")
|
|
|
|
return (result,)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"HeightAmplifier": HeightAmplifier,
|
|
"HeightCombiner": HeightCombiner,
|
|
"DisplacementToVector": DisplacementToVector,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"HeightAmplifier": "Height Amplifier",
|
|
"HeightCombiner": "Height Combiner",
|
|
"DisplacementToVector": "Displacement to Vector",
|
|
}
|
|
|