styletransferandClaude Opus 4.5
5d0cc8aa0a
Features: - 8 ComfyUI nodes for TRELLIS2 video-to-3D processing - Onion skinning with configurable frame interval - Multiple rendering modes (PBR, wireframe, cel, normal, depth) - 4 lighting presets (studio, outdoor, dramatic, flat) - Post-processing effects (bloom, DoF, motion blur, outlines) - Progress bars and auto model download - 5 example workflows 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
256 lines
6.6 KiB
Python
256 lines
6.6 KiB
Python
"""Lighting presets and utilities."""
|
|
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Tuple, Dict
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class Light:
|
|
"""Single light source."""
|
|
direction: Tuple[float, float, float]
|
|
color: Tuple[float, float, float] = (1.0, 1.0, 1.0)
|
|
intensity: float = 1.0
|
|
type: str = "directional" # "directional", "point", "spot"
|
|
|
|
|
|
@dataclass
|
|
class LightingConfig:
|
|
"""Complete lighting configuration."""
|
|
ambient: float = 0.3
|
|
ambient_color: Tuple[float, float, float] = (1.0, 1.0, 1.0)
|
|
lights: List[Light] = field(default_factory=list)
|
|
|
|
|
|
# Preset configurations
|
|
STUDIO_LIGHTING = LightingConfig(
|
|
ambient=0.3,
|
|
ambient_color=(1.0, 1.0, 1.0),
|
|
lights=[
|
|
Light(direction=(1.0, 1.0, 1.0), color=(1.0, 1.0, 1.0), intensity=0.6), # Key
|
|
Light(direction=(-1.0, 0.5, 0.5), color=(0.8, 0.9, 1.0), intensity=0.3), # Fill
|
|
Light(direction=(0.0, 1.0, -1.0), color=(1.0, 0.9, 0.8), intensity=0.2), # Rim
|
|
]
|
|
)
|
|
|
|
OUTDOOR_LIGHTING = LightingConfig(
|
|
ambient=0.4,
|
|
ambient_color=(0.7, 0.8, 1.0), # Sky blue ambient
|
|
lights=[
|
|
Light(direction=(0.5, 1.0, 0.3), color=(1.0, 0.95, 0.9), intensity=0.8), # Sun
|
|
Light(direction=(0.0, -1.0, 0.0), color=(0.6, 0.7, 1.0), intensity=0.2), # Sky fill
|
|
]
|
|
)
|
|
|
|
DRAMATIC_LIGHTING = LightingConfig(
|
|
ambient=0.1,
|
|
ambient_color=(0.5, 0.5, 0.6),
|
|
lights=[
|
|
Light(direction=(1.0, 0.5, 0.0), color=(1.0, 0.9, 0.7), intensity=1.0),
|
|
]
|
|
)
|
|
|
|
FLAT_LIGHTING = LightingConfig(
|
|
ambient=0.8,
|
|
ambient_color=(1.0, 1.0, 1.0),
|
|
lights=[
|
|
Light(direction=(0.0, 0.0, 1.0), color=(1.0, 1.0, 1.0), intensity=0.2),
|
|
]
|
|
)
|
|
|
|
PRESETS: Dict[str, LightingConfig] = {
|
|
"studio": STUDIO_LIGHTING,
|
|
"outdoor": OUTDOOR_LIGHTING,
|
|
"dramatic": DRAMATIC_LIGHTING,
|
|
"flat": FLAT_LIGHTING,
|
|
}
|
|
|
|
|
|
def create_studio_lighting() -> LightingConfig:
|
|
"""Create studio 3-point lighting setup."""
|
|
return STUDIO_LIGHTING
|
|
|
|
|
|
def create_outdoor_lighting() -> LightingConfig:
|
|
"""Create outdoor HDR-style lighting."""
|
|
return OUTDOOR_LIGHTING
|
|
|
|
|
|
def create_dramatic_lighting() -> LightingConfig:
|
|
"""Create high-contrast dramatic lighting."""
|
|
return DRAMATIC_LIGHTING
|
|
|
|
|
|
def create_flat_lighting() -> LightingConfig:
|
|
"""Create flat, even lighting for cel shading."""
|
|
return FLAT_LIGHTING
|
|
|
|
|
|
def apply_lighting_preset(
|
|
colors: torch.Tensor,
|
|
normals: torch.Tensor,
|
|
preset_name: str = "studio",
|
|
view_matrix: torch.Tensor = None
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply a lighting preset to colors based on normals.
|
|
|
|
Args:
|
|
colors: Base colors [H, W, 3]
|
|
normals: Surface normals [H, W, 3]
|
|
preset_name: Name of lighting preset
|
|
view_matrix: Optional view matrix for specular
|
|
|
|
Returns:
|
|
Lit colors [H, W, 3]
|
|
"""
|
|
config = PRESETS.get(preset_name, STUDIO_LIGHTING)
|
|
return apply_lighting(colors, normals, config, view_matrix)
|
|
|
|
|
|
def apply_lighting(
|
|
colors: torch.Tensor,
|
|
normals: torch.Tensor,
|
|
config: LightingConfig,
|
|
view_matrix: torch.Tensor = None
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply lighting configuration to colors.
|
|
|
|
Args:
|
|
colors: Base colors [H, W, 3]
|
|
normals: Surface normals [H, W, 3]
|
|
config: Lighting configuration
|
|
view_matrix: Optional view matrix for specular
|
|
|
|
Returns:
|
|
Lit colors [H, W, 3]
|
|
"""
|
|
device = colors.device
|
|
|
|
# Normalize normals
|
|
n = normals / (torch.norm(normals, dim=-1, keepdim=True) + 1e-8)
|
|
|
|
# Start with ambient
|
|
ambient_color = torch.tensor(config.ambient_color, dtype=colors.dtype, device=device)
|
|
lit = colors * config.ambient * ambient_color.view(1, 1, 3)
|
|
|
|
# Add each light
|
|
for light in config.lights:
|
|
# Light direction
|
|
light_dir = torch.tensor(light.direction, dtype=colors.dtype, device=device)
|
|
light_dir = light_dir / (torch.norm(light_dir) + 1e-8)
|
|
|
|
# Light color
|
|
light_color = torch.tensor(light.color, dtype=colors.dtype, device=device)
|
|
|
|
# Diffuse
|
|
ndotl = torch.sum(n * light_dir.view(1, 1, 3), dim=-1, keepdim=True)
|
|
ndotl = ndotl.clamp(0, 1)
|
|
|
|
diffuse = colors * light_color.view(1, 1, 3) * ndotl * light.intensity
|
|
lit = lit + diffuse
|
|
|
|
return lit.clamp(0, 1)
|
|
|
|
|
|
def compute_shadow_factor(
|
|
position: torch.Tensor,
|
|
normal: torch.Tensor,
|
|
light_dir: torch.Tensor,
|
|
shadow_softness: float = 0.0
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute a simple shadow factor based on normal orientation.
|
|
|
|
Args:
|
|
position: World position (unused for directional lights)
|
|
normal: Surface normal [H, W, 3]
|
|
light_dir: Light direction [3]
|
|
shadow_softness: Softness of shadow edge
|
|
|
|
Returns:
|
|
Shadow factor [H, W, 1] where 1 = fully lit, 0 = in shadow
|
|
"""
|
|
device = normal.device
|
|
|
|
n = normal / (torch.norm(normal, dim=-1, keepdim=True) + 1e-8)
|
|
l = light_dir / (torch.norm(light_dir) + 1e-8)
|
|
|
|
ndotl = torch.sum(n * l.view(1, 1, 3), dim=-1, keepdim=True)
|
|
|
|
if shadow_softness > 0:
|
|
# Soft shadow transition
|
|
shadow = torch.sigmoid((ndotl + shadow_softness) / shadow_softness)
|
|
else:
|
|
# Hard shadow
|
|
shadow = (ndotl > 0).float()
|
|
|
|
return shadow
|
|
|
|
|
|
def create_environment_light(
|
|
resolution: int = 64
|
|
) -> torch.Tensor:
|
|
"""
|
|
Create a simple environment light map.
|
|
|
|
Args:
|
|
resolution: Resolution of environment map
|
|
|
|
Returns:
|
|
Environment map [resolution, resolution*2, 3]
|
|
"""
|
|
# Simple gradient sky
|
|
env = torch.zeros(resolution, resolution * 2, 3)
|
|
|
|
for y in range(resolution):
|
|
# Vertical gradient (sky blue at top, white at horizon)
|
|
t = y / (resolution - 1)
|
|
sky_color = torch.tensor([0.5 + 0.5 * t, 0.7 + 0.3 * t, 1.0])
|
|
env[y, :] = sky_color
|
|
|
|
return env
|
|
|
|
|
|
def sample_environment(
|
|
env_map: torch.Tensor,
|
|
directions: torch.Tensor
|
|
) -> torch.Tensor:
|
|
"""
|
|
Sample environment map given directions.
|
|
|
|
Args:
|
|
env_map: Environment map [H, W, 3]
|
|
directions: Sample directions [N, 3]
|
|
|
|
Returns:
|
|
Sampled colors [N, 3]
|
|
"""
|
|
H, W, _ = env_map.shape
|
|
|
|
# Convert direction to spherical coordinates
|
|
x, y, z = directions[:, 0], directions[:, 1], directions[:, 2]
|
|
|
|
# Azimuth angle
|
|
phi = torch.atan2(x, z) # [-pi, pi]
|
|
# Elevation angle
|
|
theta = torch.asin(y.clamp(-1, 1)) # [-pi/2, pi/2]
|
|
|
|
# Convert to UV
|
|
u = (phi / (2 * 3.14159) + 0.5) * (W - 1)
|
|
v = (0.5 - theta / 3.14159) * (H - 1)
|
|
|
|
# Clamp to valid range
|
|
u = u.clamp(0, W - 1).long()
|
|
v = v.clamp(0, H - 1).long()
|
|
|
|
# Sample
|
|
colors = env_map[v, u]
|
|
|
|
return colors
|