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>
261 lines
6.4 KiB
Python
261 lines
6.4 KiB
Python
"""Cel/toon shading effect."""
|
|
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from typing import Tuple, Optional
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class CelShadingConfig:
|
|
"""Configuration for cel shading effect."""
|
|
levels: int = 4 # Number of shading levels
|
|
outline_width: float = 1.0
|
|
outline_color: Tuple[float, float, float] = (0.0, 0.0, 0.0)
|
|
ambient: float = 0.3
|
|
highlight_threshold: float = 0.9
|
|
highlight_color: Optional[Tuple[float, float, float]] = None
|
|
|
|
|
|
def apply_cel_shading(
|
|
image: torch.Tensor,
|
|
config: CelShadingConfig
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply cel/toon shading effect to rendered image.
|
|
|
|
Args:
|
|
image: Input image [H, W, 3] or [N, H, W, 3]
|
|
config: Cel shading configuration
|
|
|
|
Returns:
|
|
Cel-shaded image
|
|
"""
|
|
was_batch = len(image.shape) == 4
|
|
if not was_batch:
|
|
image = image.unsqueeze(0)
|
|
|
|
output = []
|
|
for frame in image:
|
|
# Quantize colors to levels
|
|
quantized = quantize_colors(frame, config.levels)
|
|
|
|
# Add highlights if configured
|
|
if config.highlight_color is not None and config.highlight_threshold < 1.0:
|
|
quantized = add_highlights(
|
|
quantized, frame, config.highlight_threshold, config.highlight_color
|
|
)
|
|
|
|
# Add outline
|
|
if config.outline_width > 0:
|
|
quantized = add_outline(
|
|
quantized, config.outline_width, config.outline_color
|
|
)
|
|
|
|
output.append(quantized)
|
|
|
|
result = torch.stack(output)
|
|
|
|
if not was_batch:
|
|
result = result.squeeze(0)
|
|
|
|
return result
|
|
|
|
|
|
def quantize_colors(
|
|
image: torch.Tensor,
|
|
levels: int
|
|
) -> torch.Tensor:
|
|
"""
|
|
Quantize colors to discrete levels.
|
|
|
|
Args:
|
|
image: Input image [H, W, 3]
|
|
levels: Number of quantization levels
|
|
|
|
Returns:
|
|
Quantized image
|
|
"""
|
|
# Quantize each channel
|
|
quantized = torch.floor(image * levels) / levels
|
|
|
|
# Add a small amount back to avoid pure black
|
|
quantized = quantized * 0.9 + 0.1
|
|
|
|
return quantized
|
|
|
|
|
|
def add_highlights(
|
|
image: torch.Tensor,
|
|
original: torch.Tensor,
|
|
threshold: float,
|
|
highlight_color: Tuple[float, float, float]
|
|
) -> torch.Tensor:
|
|
"""
|
|
Add specular highlights to cel-shaded image.
|
|
|
|
Args:
|
|
image: Cel-shaded image [H, W, 3]
|
|
original: Original image for brightness reference [H, W, 3]
|
|
threshold: Brightness threshold for highlights
|
|
highlight_color: RGB color for highlights
|
|
|
|
Returns:
|
|
Image with highlights
|
|
"""
|
|
device = image.device
|
|
|
|
# Compute brightness
|
|
brightness = original.mean(dim=-1)
|
|
|
|
# Create highlight mask
|
|
highlight_mask = (brightness > threshold).unsqueeze(-1).float()
|
|
|
|
# Apply highlight color
|
|
highlight = torch.tensor(highlight_color, dtype=image.dtype, device=device)
|
|
highlight = highlight.view(1, 1, 3)
|
|
|
|
output = image * (1 - highlight_mask) + highlight * highlight_mask
|
|
|
|
return output
|
|
|
|
|
|
def add_outline(
|
|
image: torch.Tensor,
|
|
width: float,
|
|
color: Tuple[float, float, float]
|
|
) -> torch.Tensor:
|
|
"""
|
|
Add edge outline to image using Sobel edge detection.
|
|
|
|
Args:
|
|
image: Input image [H, W, 3]
|
|
width: Outline width
|
|
color: Outline color
|
|
|
|
Returns:
|
|
Image with outline
|
|
"""
|
|
device = image.device
|
|
|
|
# Convert to grayscale for edge detection
|
|
gray = image.mean(dim=-1, keepdim=True)
|
|
|
|
# Sobel kernels
|
|
sobel_x = torch.tensor([
|
|
[-1, 0, 1],
|
|
[-2, 0, 2],
|
|
[-1, 0, 1]
|
|
], dtype=image.dtype, device=device).view(1, 1, 3, 3)
|
|
|
|
sobel_y = torch.tensor([
|
|
[-1, -2, -1],
|
|
[0, 0, 0],
|
|
[1, 2, 1]
|
|
], dtype=image.dtype, device=device).view(1, 1, 3, 3)
|
|
|
|
# Prepare for convolution
|
|
gray_4d = gray.permute(2, 0, 1).unsqueeze(0) # [1, 1, H, W]
|
|
|
|
# Pad
|
|
padded = torch.nn.functional.pad(gray_4d, (1, 1, 1, 1), mode='replicate')
|
|
|
|
# Apply Sobel
|
|
gx = torch.nn.functional.conv2d(padded, sobel_x)
|
|
gy = torch.nn.functional.conv2d(padded, sobel_y)
|
|
|
|
# Edge magnitude
|
|
edges = torch.sqrt(gx ** 2 + gy ** 2)
|
|
edges = edges.squeeze(0).squeeze(0) # [H, W]
|
|
|
|
# Threshold for outline
|
|
threshold = 0.1 * width
|
|
outline_mask = (edges > threshold).unsqueeze(-1).float()
|
|
|
|
# Apply outline color
|
|
outline_color_t = torch.tensor(color, dtype=image.dtype, device=device)
|
|
outline_color_t = outline_color_t.view(1, 1, 3)
|
|
|
|
output = image * (1 - outline_mask) + outline_color_t * outline_mask
|
|
|
|
return output
|
|
|
|
|
|
def compute_toon_lighting(
|
|
normals: torch.Tensor,
|
|
light_dir: torch.Tensor,
|
|
levels: int = 4,
|
|
ambient: float = 0.3
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute quantized toon lighting.
|
|
|
|
Args:
|
|
normals: Surface normals [H, W, 3]
|
|
light_dir: Light direction [3]
|
|
levels: Number of shading levels
|
|
ambient: Ambient light contribution
|
|
|
|
Returns:
|
|
Lighting intensity [H, W, 1]
|
|
"""
|
|
# Normalize
|
|
n = normals / (torch.norm(normals, dim=-1, keepdim=True) + 1e-8)
|
|
l = light_dir / (torch.norm(light_dir) + 1e-8)
|
|
|
|
# Diffuse
|
|
ndotl = torch.sum(n * l.view(1, 1, 3), dim=-1, keepdim=True)
|
|
ndotl = ndotl.clamp(0, 1)
|
|
|
|
# Quantize
|
|
quantized = torch.floor(ndotl * levels) / levels
|
|
|
|
# Add ambient
|
|
lighting = ambient + (1 - ambient) * quantized
|
|
|
|
return lighting
|
|
|
|
|
|
def apply_rim_lighting(
|
|
image: torch.Tensor,
|
|
normals: torch.Tensor,
|
|
view_dir: torch.Tensor,
|
|
rim_color: Tuple[float, float, float] = (1.0, 1.0, 1.0),
|
|
rim_power: float = 2.0,
|
|
rim_intensity: float = 0.5
|
|
) -> torch.Tensor:
|
|
"""
|
|
Add rim lighting effect (common in cel shading).
|
|
|
|
Args:
|
|
image: Input image [H, W, 3]
|
|
normals: Surface normals [H, W, 3]
|
|
view_dir: View direction [3]
|
|
rim_color: Rim light color
|
|
rim_power: Rim falloff power
|
|
rim_intensity: Rim light intensity
|
|
|
|
Returns:
|
|
Image with rim lighting
|
|
"""
|
|
device = image.device
|
|
|
|
# Normalize
|
|
n = normals / (torch.norm(normals, dim=-1, keepdim=True) + 1e-8)
|
|
v = view_dir / (torch.norm(view_dir) + 1e-8)
|
|
|
|
# Rim factor (1 - dot(n, v))
|
|
ndotv = torch.sum(n * v.view(1, 1, 3), dim=-1, keepdim=True)
|
|
rim = torch.pow(1 - ndotv.clamp(0, 1), rim_power)
|
|
|
|
# Apply rim
|
|
rim_color_t = torch.tensor(rim_color, dtype=image.dtype, device=device)
|
|
rim_color_t = rim_color_t.view(1, 1, 3)
|
|
|
|
output = image + rim * rim_color_t * rim_intensity
|
|
|
|
return output.clamp(0, 1)
|