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>
322 lines
8.4 KiB
Python
322 lines
8.4 KiB
Python
"""Post-processing effects for rendered images."""
|
|
|
|
import logging
|
|
from typing import Tuple, Optional
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def apply_bloom(
|
|
images: torch.Tensor,
|
|
intensity: float = 0.5,
|
|
threshold: float = 0.8,
|
|
radius: int = 20
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply bloom/glow effect to bright areas.
|
|
|
|
Args:
|
|
images: Input images [N, H, W, C] or [H, W, C]
|
|
intensity: Bloom intensity
|
|
threshold: Brightness threshold
|
|
radius: Blur radius for bloom
|
|
|
|
Returns:
|
|
Images with bloom applied
|
|
"""
|
|
was_3d = len(images.shape) == 3
|
|
if was_3d:
|
|
images = images.unsqueeze(0)
|
|
|
|
device = images.device
|
|
N, H, W, C = images.shape
|
|
|
|
output = []
|
|
for frame in images:
|
|
# Extract bright areas
|
|
brightness = frame.mean(dim=-1, keepdim=True)
|
|
bright_mask = (brightness > threshold).float()
|
|
bright_areas = frame * bright_mask
|
|
|
|
# Blur the bright areas
|
|
# Convert to NCHW for convolution
|
|
bright_nchw = bright_areas.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W]
|
|
|
|
# Create Gaussian kernel
|
|
kernel_size = radius * 2 + 1
|
|
sigma = radius / 3
|
|
kernel_1d = _gaussian_kernel_1d(kernel_size, sigma, device)
|
|
|
|
# Separable blur
|
|
padded = F.pad(bright_nchw, (radius, radius, radius, radius), mode='reflect')
|
|
|
|
# Horizontal blur
|
|
blurred_h = F.conv2d(
|
|
padded,
|
|
kernel_1d.view(1, 1, 1, -1).expand(C, 1, 1, -1),
|
|
groups=C
|
|
)
|
|
|
|
# Vertical blur
|
|
blurred = F.conv2d(
|
|
blurred_h,
|
|
kernel_1d.view(1, 1, -1, 1).expand(C, 1, -1, 1),
|
|
groups=C
|
|
)
|
|
|
|
# Convert back to HWC
|
|
bloom = blurred.squeeze(0).permute(1, 2, 0)
|
|
|
|
# Composite
|
|
result = frame + bloom * intensity
|
|
|
|
output.append(result.clamp(0, 1))
|
|
|
|
result = torch.stack(output)
|
|
|
|
if was_3d:
|
|
result = result.squeeze(0)
|
|
|
|
return result
|
|
|
|
|
|
def apply_depth_of_field(
|
|
images: torch.Tensor,
|
|
depth_map: torch.Tensor,
|
|
focal_distance: float = 0.5,
|
|
aperture: float = 0.1,
|
|
max_blur: int = 20
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply depth of field blur effect.
|
|
|
|
Args:
|
|
images: Input images [N, H, W, C]
|
|
depth_map: Depth values [N, H, W, 1] (0 = close, 1 = far)
|
|
focal_distance: Distance in focus (0-1)
|
|
aperture: Aperture size (larger = more blur)
|
|
max_blur: Maximum blur radius
|
|
|
|
Returns:
|
|
Images with DoF applied
|
|
"""
|
|
was_3d = len(images.shape) == 3
|
|
if was_3d:
|
|
images = images.unsqueeze(0)
|
|
depth_map = depth_map.unsqueeze(0)
|
|
|
|
device = images.device
|
|
N, H, W, C = images.shape
|
|
|
|
output = []
|
|
for i in range(N):
|
|
frame = images[i]
|
|
depth = depth_map[i]
|
|
|
|
# Compute blur amount based on distance from focal plane
|
|
blur_amount = torch.abs(depth - focal_distance) * aperture * max_blur
|
|
blur_amount = blur_amount.squeeze(-1) # [H, W]
|
|
|
|
# For simplicity, use uniform blur based on average
|
|
# (true DoF would need per-pixel blur which is expensive)
|
|
avg_blur = int(blur_amount.mean().item())
|
|
|
|
if avg_blur > 0:
|
|
# Convert for convolution
|
|
frame_nchw = frame.permute(2, 0, 1).unsqueeze(0)
|
|
|
|
# Gaussian blur
|
|
kernel_size = avg_blur * 2 + 1
|
|
sigma = avg_blur / 3
|
|
kernel = _gaussian_kernel_2d(kernel_size, sigma, device)
|
|
|
|
padded = F.pad(frame_nchw, (avg_blur, avg_blur, avg_blur, avg_blur), mode='reflect')
|
|
blurred = F.conv2d(padded, kernel.unsqueeze(0).unsqueeze(0).expand(C, 1, -1, -1), groups=C)
|
|
|
|
blurred_hwc = blurred.squeeze(0).permute(1, 2, 0)
|
|
|
|
# Blend based on blur amount
|
|
blend = (blur_amount / max_blur).unsqueeze(-1).clamp(0, 1)
|
|
result = frame * (1 - blend) + blurred_hwc * blend
|
|
else:
|
|
result = frame
|
|
|
|
output.append(result)
|
|
|
|
result = torch.stack(output)
|
|
|
|
if was_3d:
|
|
result = result.squeeze(0)
|
|
|
|
return result
|
|
|
|
|
|
def apply_motion_blur(
|
|
images: torch.Tensor,
|
|
intensity: float = 0.5
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply simple motion blur between frames.
|
|
|
|
Args:
|
|
images: Video frames [N, H, W, C]
|
|
intensity: Blur intensity (0-1)
|
|
|
|
Returns:
|
|
Motion blurred frames
|
|
"""
|
|
if len(images.shape) != 4 or images.shape[0] < 2:
|
|
return images
|
|
|
|
N, H, W, C = images.shape
|
|
output = images.clone()
|
|
|
|
for i in range(1, N):
|
|
# Blend with previous frame
|
|
prev = images[i - 1]
|
|
curr = images[i]
|
|
output[i] = curr * (1 - intensity) + prev * intensity
|
|
|
|
return output
|
|
|
|
|
|
def apply_outline(
|
|
images: torch.Tensor,
|
|
width: float = 1.0,
|
|
color: Tuple[float, float, float] = (0.0, 0.0, 0.0),
|
|
threshold: float = 0.1
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply edge outline effect.
|
|
|
|
Args:
|
|
images: Input images [N, H, W, C] or [H, W, C]
|
|
width: Outline width
|
|
color: Outline color
|
|
threshold: Edge detection threshold
|
|
|
|
Returns:
|
|
Images with outlines
|
|
"""
|
|
was_3d = len(images.shape) == 3
|
|
if was_3d:
|
|
images = images.unsqueeze(0)
|
|
|
|
device = images.device
|
|
N, H, W, C = images.shape
|
|
|
|
# Sobel kernels
|
|
sobel_x = torch.tensor([
|
|
[-1, 0, 1],
|
|
[-2, 0, 2],
|
|
[-1, 0, 1]
|
|
], dtype=images.dtype, device=device)
|
|
|
|
sobel_y = torch.tensor([
|
|
[-1, -2, -1],
|
|
[0, 0, 0],
|
|
[1, 2, 1]
|
|
], dtype=images.dtype, device=device)
|
|
|
|
output = []
|
|
for frame in images:
|
|
# Convert to grayscale
|
|
gray = frame.mean(dim=-1)
|
|
|
|
# Pad for convolution
|
|
gray_4d = gray.unsqueeze(0).unsqueeze(0) # [1, 1, H, W]
|
|
padded = F.pad(gray_4d, (1, 1, 1, 1), mode='replicate')
|
|
|
|
# Apply Sobel
|
|
gx = F.conv2d(padded, sobel_x.view(1, 1, 3, 3))
|
|
gy = F.conv2d(padded, sobel_y.view(1, 1, 3, 3))
|
|
|
|
# Edge magnitude
|
|
edges = torch.sqrt(gx ** 2 + gy ** 2).squeeze()
|
|
|
|
# Dilate edges for width
|
|
if width > 1:
|
|
dilation = int(width)
|
|
edges_4d = edges.unsqueeze(0).unsqueeze(0)
|
|
kernel = torch.ones(1, 1, dilation * 2 + 1, dilation * 2 + 1, device=device)
|
|
edges_dilated = F.conv2d(
|
|
F.pad(edges_4d, (dilation, dilation, dilation, dilation), mode='replicate'),
|
|
kernel
|
|
).squeeze()
|
|
edges = (edges_dilated > 0).float()
|
|
else:
|
|
edges = (edges > threshold * width).float()
|
|
|
|
# Apply outline color
|
|
outline_mask = edges.unsqueeze(-1)
|
|
outline_color = torch.tensor(color, dtype=frame.dtype, device=device)
|
|
result = frame * (1 - outline_mask) + outline_color.view(1, 1, 3) * outline_mask
|
|
|
|
output.append(result)
|
|
|
|
result = torch.stack(output)
|
|
|
|
if was_3d:
|
|
result = result.squeeze(0)
|
|
|
|
return result
|
|
|
|
|
|
def apply_vignette(
|
|
images: torch.Tensor,
|
|
intensity: float = 0.5,
|
|
radius: float = 0.8
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply vignette effect (darkening at edges).
|
|
|
|
Args:
|
|
images: Input images [N, H, W, C] or [H, W, C]
|
|
intensity: Vignette intensity
|
|
radius: Vignette radius (1 = no effect at center)
|
|
|
|
Returns:
|
|
Images with vignette
|
|
"""
|
|
was_3d = len(images.shape) == 3
|
|
if was_3d:
|
|
images = images.unsqueeze(0)
|
|
|
|
device = images.device
|
|
N, H, W, C = images.shape
|
|
|
|
# Create vignette mask
|
|
y, x = torch.meshgrid(
|
|
torch.linspace(-1, 1, H, device=device),
|
|
torch.linspace(-1, 1, W, device=device),
|
|
indexing='ij'
|
|
)
|
|
|
|
dist = torch.sqrt(x ** 2 + y ** 2)
|
|
vignette = 1 - torch.clamp((dist - radius) / (1 - radius), 0, 1) * intensity
|
|
vignette = vignette.unsqueeze(-1) # [H, W, 1]
|
|
|
|
output = images * vignette.unsqueeze(0)
|
|
|
|
if was_3d:
|
|
output = output.squeeze(0)
|
|
|
|
return output
|
|
|
|
|
|
def _gaussian_kernel_1d(size: int, sigma: float, device: torch.device) -> torch.Tensor:
|
|
"""Create 1D Gaussian kernel."""
|
|
x = torch.arange(size, dtype=torch.float32, device=device) - size // 2
|
|
kernel = torch.exp(-x ** 2 / (2 * sigma ** 2))
|
|
return kernel / kernel.sum()
|
|
|
|
|
|
def _gaussian_kernel_2d(size: int, sigma: float, device: torch.device) -> torch.Tensor:
|
|
"""Create 2D Gaussian kernel."""
|
|
kernel_1d = _gaussian_kernel_1d(size, sigma, device)
|
|
kernel_2d = kernel_1d.unsqueeze(1) @ kernel_1d.unsqueeze(0)
|
|
return kernel_2d
|