nodes/styletransfer-ComfyUI-TRELLIS2_Motion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit: ComfyUI-TRELLIS2_Motion node package

Browse Source 
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>
This commit is contained in:
styletransfer
2025-12-20 13:30:04 -05:00
co-authored by Claude Opus 4.5
commit 5d0cc8aa0a
31 changed files with 6414 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+71
View File
@@ -0,0 +1,71 @@
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
# Virtual environments
venv/
ENV/
env/
.venv/
# IDE
.idea/
.vscode/
*.swp
*.swo
*~
# OS
.DS_Store
Thumbs.db
# Logs
*.log
logs/
# Temporary files
tmp/
temp/
*.tmp
# Model files (large)
*.pt
*.pth
*.ckpt
*.safetensors
*.bin
# Cache
.cache/
*.cache
# Jupyter
.ipynb_checkpoints/
# Testing
.pytest_cache/
.coverage
htmlcov/
# Output files
output/
outputs/
results/
LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2024
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
README.md
+181
View File
@@ -0,0 +1,181 @@
# ComfyUI-TRELLIS2_Motion
A comprehensive ComfyUI node package that uses Microsoft's TRELLIS2 to transform video footage into 3D mesh reconstructions with advanced rendering effects.
![TRELLIS2 Motion](https://img.shields.io/badge/ComfyUI-Custom_Nodes-blue)
![Python](https://img.shields.io/badge/Python-3.10+-green)
![CUDA](https://img.shields.io/badge/CUDA-Required-red)
## Features
- **Image to 3D Mesh**: Convert single images to high-quality 3D meshes with PBR materials
- **Video Processing**: Batch process video frames through TRELLIS2
- **Multiple Rendering Modes**: PBR, wireframe, cel/toon shading, normal map, depth map
- **Onion Skinning**: Motion ghosting effect with configurable interval (every 5, 10, 30 frames)
- **Camera Control**: Orbit, dolly, pan trajectories with animation support
- **Lighting Presets**: Studio, outdoor, dramatic, flat lighting configurations
- **Post-Processing**: Bloom, depth of field, motion blur, outlines
- **Progress Bars**: Visual feedback on all long-running operations
- **Auto Model Download**: Automatically downloads TRELLIS2 model (~8GB) on first use
## Nodes
| Node | Description |
|------|-------------|
| **Trellis2Predict** | Convert single image to 3D mesh with PBR materials |
| **VideoToMeshes** | Batch process video frames through TRELLIS2 |
| **CameraTrajectory** | Define camera movement (orbit, dolly, pan, static) |
| **MultiAngleCamera** | Generate multiple viewing angles for grid preview |
| **MeshRenderer** | Render mesh from camera viewpoint with effects |
| **MeshSequenceRenderer** | Render video sequence with consistent camera |
| **OnionSkinning** | Add motion ghosting effect to rendered frames |
| **RenderingEffects** | Apply bloom, DoF, cel shading, outlines |
| **Trellis2VideoAngleShift** | All-in-one convenience node |
## Installation
### 1. Clone the Repository
```bash
cd ComfyUI/custom_nodes
git clone https://github.com/yourusername/ComfyUI-TRELLIS2_Motion.git
cd ComfyUI-TRELLIS2_Motion
```
### 2. Install Dependencies
```bash
python install.py
```
### 3. Install TRELLIS2 (Required)
```bash
pip install git+https://github.com/microsoft/TRELLIS.git
```
### 4. Install nvdiffrast (Required for Rendering)
```bash
pip install git+https://github.com/NVlabs/nvdiffrast.git
```
### 5. Restart ComfyUI
The nodes will appear under the `TRELLIS2` category.
## Requirements
- **Python**: 3.10+
- **CUDA**: Required (tested on CUDA 11.8+)
- **GPU**: NVIDIA GPU with 24GB+ VRAM recommended (A100, RTX 4090)
- **ComfyUI**: Latest version
## Workflows
Pre-built workflows are included in the `workflows/` folder:
| Workflow | Description |
|----------|-------------|
| `single_image_novel_view.json` | Basic image to 3D with camera control |
| `video_angle_shift.json` | Full video processing pipeline |
| `onion_skinning_demo.json` | Motion ghosting showcase |
| `effects_showcase.json` | Cel shading and post-processing demo |
| `multiview_preview.json` | 8-angle grid preview |
## Usage Examples
### Single Image Novel View
```
LoadImage → Trellis2Predict → MeshRenderer → SaveImage
↑
CameraTrajectory
```
### Video Angle Shift
```
VHS_LoadVideo → VideoToMeshes → MeshSequenceRenderer → VHS_VideoCombine
↑
CameraTrajectory
```
### With Onion Skinning
```
VHS_LoadVideo → VideoToMeshes → MeshSequenceRenderer → OnionSkinning → VHS_VideoCombine
```
## Onion Skinning Parameters
| Parameter | Description | Default |
|-----------|-------------|---------|
| `ghost_count` | Number of ghost frames | 3 |
| `frame_interval` | Frames between ghosts | 5 |
| `ghost_opacity` | Base opacity of ghosts | 0.3 |
| `opacity_falloff` | linear, exponential, constant | linear |
| `ghost_color_mode` | original, tinted, silhouette | tinted |
| `past_tint` | Color for past frames | Blue (#0066FF) |
| `future_tint` | Color for future frames | Orange (#FF6600) |
## Rendering Modes
- **PBR**: Full physically-based rendering with materials
- **Wireframe**: Edge visualization
- **Cel/Toon**: Quantized shading with outlines
- **Normal**: RGB-encoded surface normals
- **Depth**: Grayscale depth visualization
## Lighting Presets
- **Studio**: Professional 3-point lighting
- **Outdoor**: HDR sky with sun
- **Dramatic**: High contrast single light
- **Flat**: Even lighting for cel shading
## Performance
| Resolution | TRELLIS2 Time | Render Time | Total/Frame |
|------------|---------------|-------------|-------------|
| 512x512 | ~3s | ~0.1s | ~3.1s |
| 1024x1024 | ~10s | ~0.2s | ~10.2s |
*Tested on NVIDIA A100 80GB*
## Comparison with SHARP
| Aspect | TRELLIS2 | SHARP |
|--------|----------|-------|
| Output | PBR Mesh + Materials | Gaussian Splats |
| Processing Time | 3-60s per frame | <1s per frame |
| Materials | Full PBR | SH coefficients |
| Rendering | nvdiffrast | gsplat |
| Quality | Higher fidelity meshes | Real-time splatting |
## Troubleshooting
### "CUDA out of memory"
- Reduce resolution from 1024 to 512
- Increase frame_skip to process fewer frames
- Use a GPU with more VRAM
### "nvdiffrast not found"
```bash
pip install git+https://github.com/NVlabs/nvdiffrast.git
```
### "TRELLIS not found"
```bash
pip install git+https://github.com/microsoft/TRELLIS.git
```
## Credits
- [Microsoft TRELLIS](https://github.com/microsoft/TRELLIS) - Core 3D reconstruction model
- [nvdiffrast](https://github.com/NVlabs/nvdiffrast) - CUDA mesh rendering
- [ComfyUI](https://github.com/comfyanonymous/ComfyUI) - Node-based UI framework
## License
MIT License - See LICENSE file for details.
__init__.py
+68
View File
@@ -0,0 +1,68 @@
"""
ComfyUI-TRELLIS2_Motion - Video to 3D Mesh Reconstruction with Effects
A comprehensive ComfyUI node package that uses Microsoft's TRELLIS2 to transform
video footage into 3D mesh reconstructions with advanced rendering effects.
Features:
- Single image to 3D mesh conversion
- Video batch processing
- Camera trajectory control (orbit, dolly, pan)
- Multiple rendering modes (PBR, wireframe, cel shading)
- Onion skinning for motion visualization
- Post-processing effects (bloom, DoF, motion blur)
"""
import logging
logger = logging.getLogger(__name__)
# Node imports
from .nodes.trellis_predict import Trellis2Predict
from .nodes.video_to_meshes import VideoToMeshes, MergeMeshSequences
from .nodes.camera_trajectory import CameraTrajectory
from .nodes.mesh_renderer import MeshRenderer
from .nodes.mesh_sequence_renderer import MeshSequenceRenderer
from .nodes.onion_skinning import OnionSkinning
from .nodes.rendering_effects import RenderingEffects
from .nodes.combined_node import Trellis2VideoAngleShift
# Collect all node mappings
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
# Core nodes
NODE_CLASS_MAPPINGS["Trellis2Predict"] = Trellis2Predict
NODE_DISPLAY_NAME_MAPPINGS["Trellis2Predict"] = "TRELLIS2 Predict (Image to 3D Mesh)"
NODE_CLASS_MAPPINGS["VideoToMeshes"] = VideoToMeshes
NODE_DISPLAY_NAME_MAPPINGS["VideoToMeshes"] = "Video to Meshes (Batch TRELLIS2)"
NODE_CLASS_MAPPINGS["MergeMeshSequences"] = MergeMeshSequences
NODE_DISPLAY_NAME_MAPPINGS["MergeMeshSequences"] = "Merge Mesh Sequences"
# Camera
NODE_CLASS_MAPPINGS["CameraTrajectory_T2"] = CameraTrajectory
NODE_DISPLAY_NAME_MAPPINGS["CameraTrajectory_T2"] = "Camera Trajectory (TRELLIS2)"
# Rendering
NODE_CLASS_MAPPINGS["MeshRenderer"] = MeshRenderer
NODE_DISPLAY_NAME_MAPPINGS["MeshRenderer"] = "Mesh Renderer (3D to Image)"
NODE_CLASS_MAPPINGS["MeshSequenceRenderer"] = MeshSequenceRenderer
NODE_DISPLAY_NAME_MAPPINGS["MeshSequenceRenderer"] = "Mesh Sequence Renderer"
# Effects
NODE_CLASS_MAPPINGS["OnionSkinning"] = OnionSkinning
NODE_DISPLAY_NAME_MAPPINGS["OnionSkinning"] = "Onion Skinning (Motion Ghosts)"
NODE_CLASS_MAPPINGS["RenderingEffects"] = RenderingEffects
NODE_DISPLAY_NAME_MAPPINGS["RenderingEffects"] = "Rendering Effects"
# Combined convenience node
NODE_CLASS_MAPPINGS["Trellis2VideoAngleShift"] = Trellis2VideoAngleShift
NODE_DISPLAY_NAME_MAPPINGS["Trellis2VideoAngleShift"] = "TRELLIS2 Video Angle Shift"
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
logger.info(f"ComfyUI-TRELLIS2_Motion loaded {len(NODE_CLASS_MAPPINGS)} nodes")
effects/__init__.py
+41
View File
@@ -0,0 +1,41 @@
"""Visual effects modules for ComfyUI-TRELLIS2_Motion."""
from .onion_skin import apply_onion_skinning, OnionSkinConfig
from .wireframe import render_wireframe, wireframe_overlay
from .cel_shading import apply_cel_shading, CelShadingConfig
from .lighting import (
apply_lighting_preset,
create_studio_lighting,
create_outdoor_lighting,
create_dramatic_lighting,
create_flat_lighting,
)
from .post_processing import (
apply_bloom,
apply_depth_of_field,
apply_motion_blur,
apply_outline,
)
__all__ = [
# Onion skinning
"apply_onion_skinning",
"OnionSkinConfig",
# Wireframe
"render_wireframe",
"wireframe_overlay",
# Cel shading
"apply_cel_shading",
"CelShadingConfig",
# Lighting
"apply_lighting_preset",
"create_studio_lighting",
"create_outdoor_lighting",
"create_dramatic_lighting",
"create_flat_lighting",
# Post-processing
"apply_bloom",
"apply_depth_of_field",
"apply_motion_blur",
"apply_outline",
]
effects/cel_shading.py
+260
View File
@@ -0,0 +1,260 @@
"""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)
effects/lighting.py
+255
View File
@@ -0,0 +1,255 @@
"""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
effects/onion_skin.py
+176
View File
@@ -0,0 +1,176 @@
"""Onion skinning effect - motion ghosting for animation visualization."""
import logging
from dataclasses import dataclass
from typing import List, Tuple
import torch
logger = logging.getLogger(__name__)
@dataclass
class OnionSkinConfig:
"""Configuration for onion skinning effect."""
ghost_count: int = 3
frame_interval: int = 5
ghost_opacity: float = 0.3
opacity_falloff: str = "linear" # "linear", "exponential", "constant"
color_mode: str = "tinted" # "original", "tinted", "silhouette"
past_tint: Tuple[float, float, float] = (0.0, 0.4, 1.0) # Blue
future_tint: Tuple[float, float, float] = (1.0, 0.4, 0.0) # Orange
include_future: bool = False
blend_mode: str = "over" # "over", "add", "screen"
def apply_onion_skinning(
frames: torch.Tensor,
config: OnionSkinConfig
) -> torch.Tensor:
"""
Apply onion skinning effect to a sequence of frames.
Args:
frames: Video frames [N, H, W, C]
config: Onion skinning configuration
Returns:
Frames with onion skinning applied [N, H, W, C]
"""
num_frames = frames.shape[0]
output = frames.clone()
device = frames.device
for i in range(num_frames):
current_frame = frames[i].clone()
composite = current_frame.clone()
# Add past ghosts
for g in range(1, config.ghost_count + 1):
ghost_idx = i - (g * config.frame_interval)
if ghost_idx < 0:
continue
ghost = frames[ghost_idx].clone()
opacity = _compute_opacity(g, config.ghost_count, config.ghost_opacity, config.opacity_falloff)
# Apply color treatment
if config.color_mode == "tinted":
tint = torch.tensor(config.past_tint, dtype=ghost.dtype, device=device)
ghost = _apply_tint(ghost, tint, 0.5)
elif config.color_mode == "silhouette":
tint = torch.tensor(config.past_tint, dtype=ghost.dtype, device=device)
ghost = tint.view(1, 1, 3).expand_as(ghost)
# Blend
composite = _blend_frames(composite, ghost, opacity, config.blend_mode)
# Add future ghosts (if enabled)
if config.include_future:
for g in range(1, config.ghost_count + 1):
ghost_idx = i + (g * config.frame_interval)
if ghost_idx >= num_frames:
continue
ghost = frames[ghost_idx].clone()
opacity = _compute_opacity(g, config.ghost_count, config.ghost_opacity, config.opacity_falloff)
if config.color_mode == "tinted":
tint = torch.tensor(config.future_tint, dtype=ghost.dtype, device=device)
ghost = _apply_tint(ghost, tint, 0.5)
elif config.color_mode == "silhouette":
tint = torch.tensor(config.future_tint, dtype=ghost.dtype, device=device)
ghost = tint.view(1, 1, 3).expand_as(ghost)
composite = _blend_frames(composite, ghost, opacity, config.blend_mode)
# Blend composite with current frame (current on top)
# Current frame is always fully opaque on top
output[i] = composite
logger.info(
f"OnionSkinning: Applied {config.ghost_count} ghosts at interval {config.frame_interval}"
)
return output
def _compute_opacity(
ghost_index: int,
total_ghosts: int,
base_opacity: float,
falloff: str
) -> float:
"""Compute opacity for a ghost based on its distance."""
if falloff == "constant":
return base_opacity
elif falloff == "exponential":
decay = 0.5 ** ghost_index
return base_opacity * decay
else: # linear
linear_factor = 1.0 - (ghost_index - 1) / total_ghosts
return base_opacity * max(0.1, linear_factor)
def _apply_tint(
image: torch.Tensor,
tint: torch.Tensor,
strength: float
) -> torch.Tensor:
"""Apply color tint to image."""
# Lerp between original and tint
tinted = image * (1 - strength) + tint.view(1, 1, 3) * strength
return tinted
def _blend_frames(
base: torch.Tensor,
overlay: torch.Tensor,
opacity: float,
mode: str
) -> torch.Tensor:
"""Blend overlay onto base with given opacity and mode."""
if mode == "add":
return base + overlay * opacity
elif mode == "screen":
# Screen blend: 1 - (1 - base) * (1 - overlay * opacity)
return 1 - (1 - base) * (1 - overlay * opacity)
else: # "over" (standard alpha blend)
return base * (1 - opacity) + overlay * opacity
def create_onion_skin_preview(
frames: torch.Tensor,
ghost_count: int = 3,
frame_interval: int = 5
) -> torch.Tensor:
"""
Create a quick preview of onion skinning with default settings.
Args:
frames: Video frames [N, H, W, C]
ghost_count: Number of ghosts
frame_interval: Frames between ghosts
Returns:
Single frame showing onion skin effect at middle of sequence
"""
config = OnionSkinConfig(
ghost_count=ghost_count,
frame_interval=frame_interval,
ghost_opacity=0.3,
opacity_falloff="linear",
color_mode="tinted",
)
# Get middle frame
mid_idx = len(frames) // 2
start_idx = max(0, mid_idx - ghost_count * frame_interval)
end_idx = min(len(frames), mid_idx + ghost_count * frame_interval + 1)
subset = frames[start_idx:end_idx]
result = apply_onion_skinning(subset, config)
# Return the middle frame of the result
result_mid = (mid_idx - start_idx)
return result[result_mid].unsqueeze(0)
effects/post_processing.py
+321
View File
@@ -0,0 +1,321 @@
"""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
effects/wireframe.py
+240
View File
@@ -0,0 +1,240 @@
"""Wireframe rendering effect."""
import logging
from typing import Tuple, Optional
import torch
logger = logging.getLogger(__name__)
def render_wireframe(
vertices: torch.Tensor,
faces: torch.Tensor,
mvp: torch.Tensor,
width: int,
height: int,
line_width: float = 1.0,
line_color: Tuple[float, float, float] = (1.0, 1.0, 1.0),
background_color: Tuple[float, float, float] = (0.0, 0.0, 0.0),
fill_color: Optional[Tuple[float, float, float]] = None,
) -> torch.Tensor:
"""
Render a mesh as wireframe.
Args:
vertices: Mesh vertices [N, 3]
faces: Face indices [F, 3]
mvp: Model-view-projection matrix [4, 4]
width: Output width
height: Output height
line_width: Width of wireframe lines
line_color: RGB color for lines
background_color: RGB background color
fill_color: Optional fill color for faces
Returns:
Rendered wireframe image [H, W, 3]
"""
device = vertices.device
# Transform vertices to clip space
verts_homo = torch.cat([
vertices,
torch.ones(vertices.shape[0], 1, device=device)
], dim=1)
verts_clip = (mvp @ verts_homo.T).T
# Perspective divide to NDC
verts_ndc = verts_clip[:, :3] / (verts_clip[:, 3:4] + 1e-8)
# Convert to screen space
verts_screen = torch.zeros_like(verts_ndc[:, :2])
verts_screen[:, 0] = (verts_ndc[:, 0] + 1) * 0.5 * width
verts_screen[:, 1] = (1 - verts_ndc[:, 1]) * 0.5 * height # Flip Y
# Initialize output
output = torch.zeros(height, width, 3, dtype=torch.float32, device=device)
output[:, :, 0] = background_color[0]
output[:, :, 1] = background_color[1]
output[:, :, 2] = background_color[2]
# Draw edges
line_color_t = torch.tensor(line_color, dtype=torch.float32, device=device)
for face in faces:
for j in range(3):
v0 = verts_screen[face[j]]
v1 = verts_screen[face[(j + 1) % 3]]
# Draw line between v0 and v1
output = _draw_line(output, v0, v1, line_color_t, line_width)
return output
def wireframe_overlay(
base_image: torch.Tensor,
wireframe: torch.Tensor,
opacity: float = 0.5,
blend_mode: str = "over"
) -> torch.Tensor:
"""
Overlay wireframe on top of base image.
Args:
base_image: Base rendered image [H, W, 3]
wireframe: Wireframe image [H, W, 3]
opacity: Wireframe opacity
blend_mode: "over", "add", or "screen"
Returns:
Composited image [H, W, 3]
"""
if blend_mode == "add":
return (base_image + wireframe * opacity).clamp(0, 1)
elif blend_mode == "screen":
return (1 - (1 - base_image) * (1 - wireframe * opacity)).clamp(0, 1)
else: # over
# Only overlay where wireframe has color
mask = (wireframe.sum(dim=-1, keepdim=True) > 0).float()
return base_image * (1 - mask * opacity) + wireframe * mask * opacity
def _draw_line(
image: torch.Tensor,
p0: torch.Tensor,
p1: torch.Tensor,
color: torch.Tensor,
width: float = 1.0
) -> torch.Tensor:
"""
Draw a line on an image using Bresenham's algorithm.
Args:
image: Image tensor [H, W, 3]
p0: Start point [2]
p1: End point [2]
color: Line color [3]
width: Line width
Returns:
Image with line drawn
"""
H, W, _ = image.shape
x0, y0 = int(p0[0].item()), int(p0[1].item())
x1, y1 = int(p1[0].item()), int(p1[1].item())
# Bresenham's line algorithm
dx = abs(x1 - x0)
dy = abs(y1 - y0)
sx = 1 if x0 < x1 else -1
sy = 1 if y0 < y1 else -1
err = dx - dy
half_w = int(width / 2)
while True:
# Draw pixel with width
for wx in range(-half_w, half_w + 1):
for wy in range(-half_w, half_w + 1):
px, py = x0 + wx, y0 + wy
if 0 <= px < W and 0 <= py < H:
image[py, px] = color
if x0 == x1 and y0 == y1:
break
e2 = 2 * err
if e2 > -dy:
err -= dy
x0 += sx
if e2 < dx:
err += dx
y0 += sy
return image
def extract_edges(
faces: torch.Tensor
) -> torch.Tensor:
"""
Extract unique edges from face indices.
Args:
faces: Face indices [F, 3]
Returns:
Edge indices [E, 2]
"""
# Get all edges (including duplicates)
edges = []
for i in range(3):
j = (i + 1) % 3
e = torch.stack([faces[:, i], faces[:, j]], dim=1)
edges.append(e)
all_edges = torch.cat(edges, dim=0)
# Sort each edge so smaller index comes first
sorted_edges = torch.sort(all_edges, dim=1).values
# Remove duplicates
unique_edges = torch.unique(sorted_edges, dim=0)
return unique_edges
def compute_edge_visibility(
vertices: torch.Tensor,
faces: torch.Tensor,
viewmat: torch.Tensor,
edges: Optional[torch.Tensor] = None
) -> torch.Tensor:
"""
Compute which edges are visible (silhouette edges).
An edge is a silhouette edge if its two adjacent faces have
different front/back facing orientations.
Args:
vertices: Vertex positions [N, 3]
faces: Face indices [F, 3]
viewmat: View matrix [4, 4]
edges: Optional pre-computed edges [E, 2]
Returns:
Boolean mask of visible edges [E]
"""
device = vertices.device
if edges is None:
edges = extract_edges(faces)
# Compute face normals
v0 = vertices[faces[:, 0]]
v1 = vertices[faces[:, 1]]
v2 = vertices[faces[:, 2]]
face_normals = torch.cross(v1 - v0, v2 - v0)
face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8)
# Face centers
face_centers = (v0 + v1 + v2) / 3
# View direction (from face center to camera)
cam_pos = -viewmat[:3, :3].T @ viewmat[:3, 3]
view_dirs = cam_pos.unsqueeze(0) - face_centers
view_dirs = view_dirs / (torch.norm(view_dirs, dim=1, keepdim=True) + 1e-8)
# Front-facing test
dots = (face_normals * view_dirs).sum(dim=1)
front_facing = dots > 0
# Build edge-to-face mapping (simplified: mark edge visible if any adjacent face is front-facing)
# For proper silhouette, need to check if edge is between front and back facing faces
visible = torch.ones(len(edges), dtype=torch.bool, device=device)
return visible
install.py
+105
View File
@@ -0,0 +1,105 @@
"""
Installation script for ComfyUI-TRELLIS2_Motion dependencies.
Run with: python install.py
"""
import subprocess
import sys
import os
def run_pip(*args):
"""Run pip with the given arguments."""
subprocess.check_call([sys.executable, "-m", "pip", *args])
def install_dependencies():
"""Install all required dependencies."""
print("Installing ComfyUI-TRELLIS2_Motion dependencies...")
# Core dependencies
print("\n[1/5] Installing core dependencies...")
run_pip("install", "torch>=2.0.0", "numpy>=1.21.0", "pillow>=9.0.0")
# HuggingFace and transformers
print("\n[2/5] Installing HuggingFace dependencies...")
run_pip("install", "huggingface_hub>=0.20.0", "transformers>=4.30.0")
# Mesh processing
print("\n[3/5] Installing mesh processing libraries...")
run_pip("install", "trimesh>=3.0.0", "xatlas")
# Background removal
print("\n[4/5] Installing background removal...")
run_pip("install", "rembg")
# Image processing
print("\n[5/5] Installing image processing libraries...")
run_pip("install", "kornia>=0.7.0", "opencv-python-headless>=4.8.0")
run_pip("install", "scipy>=1.10.0", "einops>=0.6.0")
# CUDA-specific packages (may need manual installation)
print("\n[Note] The following packages may need manual installation:")
print(" - nvdiffrast: pip install git+https://github.com/NVlabs/nvdiffrast.git")
print(" - spconv: pip install spconv-cu118 (or cu121 for CUDA 12.1)")
print(" - flash-attn: pip install flash-attn --no-build-isolation")
try:
run_pip("install", "nvdiffrast")
except Exception as e:
print(f" Warning: nvdiffrast installation failed: {e}")
print(" Try: pip install git+https://github.com/NVlabs/nvdiffrast.git")
print("\n" + "="*50)
print("Installation complete!")
print("="*50)
print("\nTo use TRELLIS2, you also need to install the TRELLIS package:")
print(" pip install git+https://github.com/microsoft/TRELLIS.git")
print("\nOr follow the instructions at:")
print(" https://github.com/microsoft/TRELLIS")
def check_dependencies():
"""Check which dependencies are installed."""
deps = {
"torch": "torch",
"numpy": "numpy",
"PIL": "pillow",
"huggingface_hub": "huggingface_hub",
"transformers": "transformers",
"trimesh": "trimesh",
"kornia": "kornia",
"cv2": "opencv-python",
"nvdiffrast": "nvdiffrast",
"rembg": "rembg",
}
print("Checking dependencies...")
print("-" * 40)
for import_name, package_name in deps.items():
try:
__import__(import_name)
print(f" [OK] {package_name}")
except ImportError:
print(f" [MISSING] {package_name}")
# Check CUDA
try:
import torch
if torch.cuda.is_available():
print(f"\n CUDA: Available ({torch.cuda.get_device_name(0)})")
print(f" CUDA Version: {torch.version.cuda}")
else:
print("\n CUDA: Not available (CPU only)")
except Exception as e:
print(f"\n CUDA check failed: {e}")
if __name__ == "__main__":
if len(sys.argv) > 1 and sys.argv[1] == "check":
check_dependencies()
else:
install_dependencies()
print("\nRun 'python install.py check' to verify installation.")
nodes/__init__.py
+1
View File
@@ -0,0 +1 @@
"""TRELLIS2 Motion nodes for ComfyUI."""
nodes/camera_trajectory.py
+251
View File
@@ -0,0 +1,251 @@
"""Camera Trajectory node - define camera movement parameters."""
import logging
from typing import Any, Dict, Tuple
logger = logging.getLogger(__name__)
class CameraTrajectory:
"""
Define camera trajectory parameters for 3D mesh rendering.
Creates camera path definitions that can be used with MeshRenderer
to render from specific viewpoints or animated paths.
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"trajectory_type": (["orbit", "dolly", "pan", "static"], {
"default": "static",
"tooltip": "Type of camera movement"
}),
"angle_x": ("FLOAT", {
"default": 0.0,
"min": -89.0,
"max": 89.0,
"step": 0.5,
"display": "slider",
"tooltip": "Vertical angle in degrees (pitch) - positive looks from above"
}),
"angle_y": ("FLOAT", {
"default": 0.0,
"min": -180.0,
"max": 180.0,
"step": 0.5,
"display": "slider",
"tooltip": "Horizontal angle in degrees (yaw) - positive rotates right"
}),
},
"optional": {
"angle_z": ("FLOAT", {
"default": 0.0,
"min": -180.0,
"max": 180.0,
"step": 0.5,
"display": "slider",
"tooltip": "Roll angle in degrees"
}),
"distance": ("FLOAT", {
"default": 2.5,
"min": 0.5,
"max": 20.0,
"step": 0.1,
"tooltip": "Distance from scene center"
}),
"fov": ("FLOAT", {
"default": 60.0,
"min": 10.0,
"max": 120.0,
"step": 1.0,
"tooltip": "Field of view in degrees"
}),
"center_x": ("FLOAT", {
"default": 0.0,
"min": -10.0,
"max": 10.0,
"step": 0.1,
"tooltip": "Scene center X coordinate"
}),
"center_y": ("FLOAT", {
"default": 0.0,
"min": -10.0,
"max": 10.0,
"step": 0.1,
"tooltip": "Scene center Y coordinate"
}),
"center_z": ("FLOAT", {
"default": 0.0,
"min": -10.0,
"max": 10.0,
"step": 0.1,
"tooltip": "Scene center Z coordinate"
}),
"animate": ("BOOLEAN", {
"default": False,
"tooltip": "Enable camera animation over sequence"
}),
"animation_range": ("FLOAT", {
"default": 30.0,
"min": 5.0,
"max": 360.0,
"step": 5.0,
"tooltip": "Range of animation in degrees (orbit/pan) or distance (dolly)"
}),
"animation_axis": (["y", "x"], {
"default": "y",
"tooltip": "Axis to animate around (y=horizontal, x=vertical)"
}),
}
}
RETURN_TYPES = ("CAMERA_TRAJECTORY",)
RETURN_NAMES = ("camera_trajectory",)
FUNCTION = "create_trajectory"
CATEGORY = "TRELLIS2/Camera"
def create_trajectory(
self,
trajectory_type: str,
angle_x: float,
angle_y: float,
angle_z: float = 0.0,
distance: float = 2.5,
fov: float = 60.0,
center_x: float = 0.0,
center_y: float = 0.0,
center_z: float = 0.0,
animate: bool = False,
animation_range: float = 30.0,
animation_axis: str = "y",
) -> Tuple[Dict[str, Any]]:
"""
Create camera trajectory parameters.
Returns:
Dictionary containing camera trajectory configuration
"""
trajectory = {
"trajectory_type": trajectory_type,
"angle_x": angle_x,
"angle_y": angle_y,
"angle_z": angle_z,
"distance": distance,
"fov": fov,
"center": (center_x, center_y, center_z),
"animate": animate,
"animation_range": animation_range,
"animation_axis": animation_axis,
}
logger.debug(
f"CameraTrajectory: type={trajectory_type}, "
f"angles=({angle_x}, {angle_y}, {angle_z}), "
f"distance={distance}, fov={fov}, "
f"animate={animate}"
)
return (trajectory,)
class MultiAngleCamera:
"""
Generate multiple camera angles for batch rendering.
Creates a grid of camera positions for multi-view rendering.
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"num_angles": ("INT", {
"default": 8,
"min": 2,
"max": 36,
"step": 1,
"tooltip": "Number of viewing angles to generate"
}),
},
"optional": {
"angle_x": ("FLOAT", {
"default": 15.0,
"min": -60.0,
"max": 60.0,
"step": 1.0,
"tooltip": "Fixed vertical angle for all views"
}),
"distance": ("FLOAT", {
"default": 2.5,
"min": 0.5,
"max": 20.0,
"step": 0.1,
"tooltip": "Distance from scene center"
}),
"fov": ("FLOAT", {
"default": 60.0,
"min": 10.0,
"max": 120.0,
"step": 1.0,
"tooltip": "Field of view in degrees"
}),
"start_angle": ("FLOAT", {
"default": 0.0,
"min": 0.0,
"max": 360.0,
"step": 1.0,
"tooltip": "Starting horizontal angle"
}),
}
}
RETURN_TYPES = ("CAMERA_TRAJECTORY",)
RETURN_NAMES = ("camera_trajectory",)
FUNCTION = "create_multi_angle"
CATEGORY = "TRELLIS2/Camera"
def create_multi_angle(
self,
num_angles: int,
angle_x: float = 15.0,
distance: float = 2.5,
fov: float = 60.0,
start_angle: float = 0.0,
) -> Tuple[Dict[str, Any]]:
"""
Create multi-angle camera trajectory.
Returns:
Dictionary with list of camera positions
"""
# Generate evenly spaced angles
angle_step = 360.0 / num_angles
angles = [(start_angle + i * angle_step) % 360.0 - 180.0 for i in range(num_angles)]
trajectory = {
"trajectory_type": "multi_angle",
"angles_y": angles,
"angle_x": angle_x,
"distance": distance,
"fov": fov,
"center": (0.0, 0.0, 0.0),
"num_angles": num_angles,
}
logger.info(f"MultiAngleCamera: Generated {num_angles} viewing angles")
return (trajectory,)
# Node registration
NODE_CLASS_MAPPINGS = {
"CameraTrajectory_T2": CameraTrajectory,
"MultiAngleCamera": MultiAngleCamera,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"CameraTrajectory_T2": "Camera Trajectory (TRELLIS2)",
"MultiAngleCamera": "Multi-Angle Camera",
}
nodes/combined_node.py
+501
View File
@@ -0,0 +1,501 @@
"""Combined all-in-one TRELLIS2 Video Angle Shift node."""
import logging
from typing import Any, Dict, Tuple
import numpy as np
import torch
from PIL import Image
logger = logging.getLogger(__name__)
from ..utils.camera_utils import (
create_orbit_camera,
create_intrinsic_matrix,
estimate_mesh_center,
estimate_mesh_distance,
)
from ..utils.video_utils import extract_frames, interpolate_missing_frames
from ..utils.download_utils import ensure_trellis_model, get_progress_bar
# Check for TRELLIS
TRELLIS_AVAILABLE = False
try:
from trellis.pipelines import TrslIs2Img
TRELLIS_AVAILABLE = True
except ImportError:
pass
# Check for rembg
REMBG_AVAILABLE = False
try:
from rembg import remove as remove_background
REMBG_AVAILABLE = True
except ImportError:
pass
# Check for nvdiffrast
NVDIFFRAST_AVAILABLE = False
try:
import nvdiffrast.torch as dr
NVDIFFRAST_AVAILABLE = True
except ImportError:
pass
class Trellis2VideoAngleShift:
"""
All-in-one TRELLIS2 video angle shift node.
This convenience node combines the entire pipeline:
1. Extract video frames
2. Process each frame through TRELLIS2 to create 3D mesh
3. Render all frames from a shifted camera angle
4. Optionally apply onion skinning effect
5. Optionally interpolate skipped frames
6. Return the shifted video
For more control, use the individual nodes (VideoToMeshes, CameraTrajectory,
MeshSequenceRenderer) instead.
Requirements:
- TRELLIS2 (pip install git+https://github.com/microsoft/TRELLIS.git)
- nvdiffrast (pip install git+https://github.com/NVlabs/nvdiffrast.git)
- CUDA GPU
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"images": ("IMAGE", {
"tooltip": "Video frames as image batch [N, H, W, C]"
}),
"angle_x": ("FLOAT", {
"default": 15.0,
"min": -60.0,
"max": 60.0,
"step": 0.5,
"display": "slider",
"tooltip": "Vertical camera angle shift in degrees (pitch)"
}),
"angle_y": ("FLOAT", {
"default": 0.0,
"min": -180.0,
"max": 180.0,
"step": 0.5,
"display": "slider",
"tooltip": "Horizontal camera angle shift in degrees (yaw)"
}),
},
"optional": {
"frame_skip": ("INT", {
"default": 5,
"min": 1,
"max": 30,
"step": 1,
"tooltip": "Process every Nth frame"
}),
"resolution": (["512", "1024"], {
"default": "1024",
"tooltip": "TRELLIS2 processing resolution"
}),
"output_width": ("INT", {
"default": 0,
"min": 0,
"max": 4096,
"step": 64,
"tooltip": "Output width (0 = match input)"
}),
"output_height": ("INT", {
"default": 0,
"min": 0,
"max": 4096,
"step": 64,
"tooltip": "Output height (0 = match input)"
}),
"fov": ("FLOAT", {
"default": 60.0,
"min": 30.0,
"max": 120.0,
"step": 1.0,
"tooltip": "Field of view in degrees"
}),
"render_mode": (["pbr", "wireframe", "cel"], {
"default": "pbr",
"tooltip": "Rendering style"
}),
"lighting_preset": (["studio", "outdoor", "dramatic", "flat"], {
"default": "studio",
"tooltip": "Lighting configuration"
}),
"animate_camera": ("BOOLEAN", {
"default": False,
"tooltip": "Animate camera movement over video duration"
}),
"animation_range": ("FLOAT", {
"default": 30.0,
"min": 5.0,
"max": 90.0,
"step": 1.0,
"tooltip": "Range of camera animation in degrees"
}),
"enable_onion_skin": ("BOOLEAN", {
"default": False,
"tooltip": "Enable onion skinning effect"
}),
"onion_skin_interval": ("INT", {
"default": 5,
"min": 1,
"max": 30,
"step": 1,
"tooltip": "Frames between onion skin ghosts"
}),
"interpolate_frames": ("BOOLEAN", {
"default": True,
"tooltip": "Interpolate skipped frames"
}),
"background_color": ("COLOR", {
"default": "#000000",
"tooltip": "Background color"
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("shifted_video",)
FUNCTION = "process"
CATEGORY = "TRELLIS2"
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self._pipeline = None
self._glctx = None
def process(
self,
images: torch.Tensor,
angle_x: float,
angle_y: float,
frame_skip: int = 5,
resolution: str = "1024",
output_width: int = 0,
output_height: int = 0,
fov: float = 60.0,
render_mode: str = "pbr",
lighting_preset: str = "studio",
animate_camera: bool = False,
animation_range: float = 30.0,
enable_onion_skin: bool = False,
onion_skin_interval: int = 5,
interpolate_frames: bool = True,
background_color: str = "#000000",
) -> Tuple[torch.Tensor]:
"""
Process video through TRELLIS2 and render from shifted angle.
Args:
images: Video frames [N, H, W, C]
angle_x: Vertical camera angle
angle_y: Horizontal camera angle
frame_skip: Process every Nth frame
resolution: TRELLIS2 resolution
output_width: Output width (0 = match input)
output_height: Output height (0 = match input)
fov: Field of view
render_mode: Rendering style
lighting_preset: Lighting preset
animate_camera: Animate camera over video
animation_range: Animation range in degrees
enable_onion_skin: Apply onion skinning
onion_skin_interval: Frames between ghosts
interpolate_frames: Interpolate skipped frames
background_color: Background color
Returns:
Shifted video tensor
"""
# Validate dependencies
if not TRELLIS_AVAILABLE:
raise RuntimeError(
"Trellis2VideoAngleShift: TRELLIS2 not available. "
"Install with: pip install git+https://github.com/microsoft/TRELLIS.git"
)
if not NVDIFFRAST_AVAILABLE:
raise RuntimeError(
"Trellis2VideoAngleShift: nvdiffrast not available. "
"Install with: pip install git+https://github.com/NVlabs/nvdiffrast.git"
)
if not torch.cuda.is_available():
raise RuntimeError("Trellis2VideoAngleShift: CUDA required")
# Get dimensions
total_frames, H, W, C = images.shape
out_w = output_width if output_width > 0 else W
out_h = output_height if output_height > 0 else H
# Extract frames
frames, indices = extract_frames(images, frame_skip)
num_to_process = len(frames)
# Progress: model + predict per frame + render per frame + finalize
total_steps = num_to_process * 2 + 3
pbar = get_progress_bar(total_steps)
pbar.update_absolute(0, total_steps, "Loading TRELLIS2")
# Load model
model_path = ensure_trellis_model()
if self._pipeline is None:
self._pipeline = TrslIs2Img.from_pretrained(model_path)
self._pipeline.to("cuda")
if self._glctx is None:
self._glctx = dr.RasterizeCudaContext()
logger.info(
f"Trellis2VideoAngleShift: Processing {num_to_process}/{total_frames} frames "
f"at {out_w}x{out_h}, angle=({angle_x}, {angle_y})"
)
pbar.update_absolute(1, total_steps, "Starting processing")
rendered_frames = []
bg_color = self._parse_color(background_color)
background = torch.tensor(bg_color, dtype=torch.float32, device=self.device)
for i, (frame, orig_idx) in enumerate(zip(frames, indices)):
# Step 1: TRELLIS2 prediction
pbar.update_absolute(
2 + i * 2, total_steps,
f"TRELLIS2 predict frame {i+1}/{num_to_process}"
)
try:
# Convert to PIL
img_np = frame.cpu().numpy()
img_np = np.clip(img_np, 0, 1)
img_np = (img_np * 255).astype(np.uint8)
pil_image = Image.fromarray(img_np, mode="RGB")
# Remove background
if REMBG_AVAILABLE:
pil_image = remove_background(pil_image)
if pil_image.mode == "RGBA":
bg = Image.new("RGB", pil_image.size, (255, 255, 255))
bg.paste(pil_image, mask=pil_image.split()[3])
pil_image = bg
# Run TRELLIS2
outputs = self._pipeline.run(pil_image, seed=i)
mesh_output = outputs['mesh'][0] if isinstance(outputs['mesh'], list) else outputs['mesh']
# Convert mesh
mesh = self._convert_mesh(mesh_output)
except Exception as e:
logger.error(f"TRELLIS2 failed for frame {i}: {e}")
blank = torch.zeros(out_h, out_w, 3, dtype=torch.float32)
rendered_frames.append(blank)
continue
# Step 2: Render
pbar.update_absolute(
3 + i * 2, total_steps,
f"Rendering frame {i+1}/{num_to_process}"
)
# Compute camera angle
if animate_camera and total_frames > 1:
t = orig_idx / max(total_frames - 1, 1)
current_angle_y = angle_y + (t - 0.5) * animation_range
else:
current_angle_y = angle_y
# Estimate center/distance
center = estimate_mesh_center(mesh["vertices"])
distance = estimate_mesh_distance(mesh["vertices"], center)
# Create camera
viewmat = create_orbit_camera(angle_x, current_angle_y, distance, center)
K = create_intrinsic_matrix(fov, out_w, out_h)
# Render
try:
rendered = self._render_mesh(
mesh, viewmat, out_w, out_h,
background, render_mode, lighting_preset
)
rendered_frames.append(rendered.cpu())
except Exception as e:
logger.error(f"Render failed for frame {i}: {e}")
blank = torch.zeros(out_h, out_w, 3, dtype=torch.float32)
rendered_frames.append(blank)
# Clean up
del mesh
torch.cuda.empty_cache()
# Interpolate if needed
pbar.update_absolute(total_steps - 1, total_steps, "Finalizing")
if interpolate_frames and frame_skip > 1:
output = interpolate_missing_frames(rendered_frames, indices, total_frames)
else:
output = torch.stack(rendered_frames)
# Apply onion skinning
if enable_onion_skin:
from ..effects.onion_skin import apply_onion_skinning, OnionSkinConfig
config = OnionSkinConfig(
ghost_count=3,
frame_interval=onion_skin_interval,
ghost_opacity=0.3,
opacity_falloff="linear",
color_mode="tinted",
past_tint=(0.0, 0.4, 1.0),
future_tint=(1.0, 0.4, 0.0),
)
output = apply_onion_skinning(output, config)
pbar.update_absolute(total_steps, total_steps, "Complete")
logger.info(
f"Trellis2VideoAngleShift: Completed {output.shape[0]} frames "
f"at {out_w}x{out_h}"
)
return (output,)
def _convert_mesh(self, mesh_output) -> Dict[str, torch.Tensor]:
"""Convert TRELLIS mesh to dict format."""
mesh = {}
if hasattr(mesh_output, 'v_pos'):
mesh["vertices"] = mesh_output.v_pos.to(self.device)
elif hasattr(mesh_output, 'vertices'):
mesh["vertices"] = torch.tensor(
mesh_output.vertices, dtype=torch.float32, device=self.device
)
if hasattr(mesh_output, 't_pos_idx'):
mesh["faces"] = mesh_output.t_pos_idx.to(self.device)
elif hasattr(mesh_output, 'faces'):
mesh["faces"] = torch.tensor(
mesh_output.faces, dtype=torch.int64, device=self.device
)
if hasattr(mesh_output, 'v_nrm') and mesh_output.v_nrm is not None:
mesh["normals"] = mesh_output.v_nrm.to(self.device)
else:
mesh["normals"] = self._compute_normals(mesh["vertices"], mesh["faces"])
mesh["colors"] = torch.ones(
mesh["vertices"].shape[0], 3,
dtype=torch.float32, device=self.device
) * 0.8
return mesh
def _compute_normals(self, vertices, faces):
"""Compute vertex normals."""
v0 = vertices[faces[:, 0]]
v1 = vertices[faces[:, 1]]
v2 = vertices[faces[:, 2]]
face_normals = torch.cross(v1 - v0, v2 - v0)
face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8)
vertex_normals = torch.zeros_like(vertices)
for j in range(3):
vertex_normals.index_add_(0, faces[:, j], face_normals)
return vertex_normals / (torch.norm(vertex_normals, dim=1, keepdim=True) + 1e-8)
def _render_mesh(
self,
mesh: Dict[str, torch.Tensor],
viewmat: np.ndarray,
width: int,
height: int,
background: torch.Tensor,
render_mode: str,
lighting_preset: str,
) -> torch.Tensor:
"""Render mesh using nvdiffrast."""
from ..utils.camera_utils import create_projection_matrix
vertices = mesh["vertices"]
faces = mesh["faces"]
normals = mesh["normals"]
colors = mesh["colors"]
# Create MVP matrix
proj = create_projection_matrix(60.0, width / height)
mvp = torch.tensor(proj @ viewmat, dtype=torch.float32, device=self.device)
viewmat_t = torch.tensor(viewmat, dtype=torch.float32, device=self.device)
# Transform vertices
verts_homo = torch.cat([
vertices,
torch.ones(vertices.shape[0], 1, device=self.device)
], dim=1)
verts_clip = (mvp @ verts_homo.T).T
# Rasterize
rast_out, _ = dr.rasterize(
self._glctx, verts_clip.unsqueeze(0), faces.int(),
resolution=[height, width]
)
# Interpolate
normals_interp, _ = dr.interpolate(
normals.unsqueeze(0), rast_out, faces.int()
)
colors_interp, _ = dr.interpolate(
colors.unsqueeze(0), rast_out, faces.int()
)
# Simple lighting
light_dir = torch.tensor([1, 1, 1], dtype=torch.float32, device=self.device)
light_dir = light_dir / torch.norm(light_dir)
n = normals_interp[0]
n = n / (torch.norm(n, dim=2, keepdim=True) + 1e-8)
ndotl = torch.sum(n * light_dir.view(1, 1, 3), dim=2, keepdim=True)
ndotl = ndotl.clamp(0, 1)
lit_color = colors_interp[0] * (0.3 + 0.7 * ndotl)
# Cel shading
if render_mode == "cel":
levels = 4
lit_color = torch.floor(lit_color * levels) / levels
# Apply background
mask = (rast_out[0, :, :, 3:4] > 0).float()
output = lit_color * mask + background.view(1, 1, 3) * (1 - mask)
return output.clamp(0, 1)
def _parse_color(self, color_str: str) -> list:
"""Parse hex color."""
color_str = color_str.lstrip("#")
if len(color_str) == 6:
r = int(color_str[0:2], 16) / 255.0
g = int(color_str[2:4], 16) / 255.0
b = int(color_str[4:6], 16) / 255.0
return [r, g, b]
return [0.0, 0.0, 0.0]
# Node registration
NODE_CLASS_MAPPINGS = {
"Trellis2VideoAngleShift": Trellis2VideoAngleShift,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Trellis2VideoAngleShift": "TRELLIS2 Video Angle Shift",
}
nodes/mesh_renderer.py
+638
View File
@@ -0,0 +1,638 @@
"""Mesh Renderer node - render 3D meshes from camera viewpoints."""
import logging
from typing import Any, Dict, Tuple, Optional
import numpy as np
import torch
logger = logging.getLogger(__name__)
from ..utils.camera_utils import (
create_orbit_camera,
create_intrinsic_matrix,
create_projection_matrix,
estimate_mesh_center,
estimate_mesh_distance,
)
# Check for nvdiffrast
NVDIFFRAST_AVAILABLE = False
try:
import nvdiffrast.torch as dr
NVDIFFRAST_AVAILABLE = True
except ImportError:
logger.warning(
"nvdiffrast not available. Install with: "
"pip install git+https://github.com/NVlabs/nvdiffrast.git"
)
# Progress bar with fallback
try:
import comfy.utils
COMFY_AVAILABLE = True
except ImportError:
COMFY_AVAILABLE = False
# Lighting presets
LIGHTING_PRESETS = {
"studio": {
"ambient": 0.3,
"lights": [
{"direction": [1, 1, 1], "color": [1.0, 1.0, 1.0], "intensity": 0.6}, # Key
{"direction": [-1, 0.5, 0.5], "color": [0.8, 0.9, 1.0], "intensity": 0.3}, # Fill
{"direction": [0, 1, -1], "color": [1.0, 0.9, 0.8], "intensity": 0.2}, # Rim
]
},
"outdoor": {
"ambient": 0.4,
"lights": [
{"direction": [0.5, 1, 0.3], "color": [1.0, 0.95, 0.9], "intensity": 0.8}, # Sun
{"direction": [0, -1, 0], "color": [0.6, 0.7, 1.0], "intensity": 0.2}, # Sky
]
},
"dramatic": {
"ambient": 0.1,
"lights": [
{"direction": [1, 0.5, 0], "color": [1.0, 0.9, 0.7], "intensity": 1.0},
]
},
"flat": {
"ambient": 0.8,
"lights": [
{"direction": [0, 0, 1], "color": [1.0, 1.0, 1.0], "intensity": 0.2},
]
},
}
class MeshRenderer:
"""
Render a 3D mesh from a specified camera viewpoint.
Supports multiple rendering modes (PBR, wireframe, cel shading, etc.)
and lighting presets.
"""
_glctx = None # Class-level OpenGL context
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"mesh": ("TRELLIS_MESH", {
"tooltip": "3D mesh from TRELLIS2"
}),
"camera_trajectory": ("CAMERA_TRAJECTORY", {
"tooltip": "Camera parameters from CameraTrajectory node"
}),
},
"optional": {
"output_width": ("INT", {
"default": 1920,
"min": 256,
"max": 4096,
"step": 64,
"tooltip": "Output image width"
}),
"output_height": ("INT", {
"default": 1080,
"min": 256,
"max": 4096,
"step": 64,
"tooltip": "Output image height"
}),
"render_mode": (["pbr", "wireframe", "cel", "normal", "depth"], {
"default": "pbr",
"tooltip": "Rendering style"
}),
"background_color": ("COLOR", {
"default": "#000000",
"tooltip": "Background color for rendered image"
}),
"lighting_preset": (["studio", "outdoor", "dramatic", "flat"], {
"default": "studio",
"tooltip": "Lighting configuration"
}),
"num_frames": ("INT", {
"default": 1,
"min": 1,
"max": 300,
"step": 1,
"tooltip": "Number of frames to render (for animated trajectories)"
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("rendered_image",)
FUNCTION = "render"
CATEGORY = "TRELLIS2/Render"
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
@classmethod
def get_glctx(cls):
"""Get or create OpenGL context (singleton)."""
if cls._glctx is None and NVDIFFRAST_AVAILABLE:
cls._glctx = dr.RasterizeCudaContext()
return cls._glctx
def render(
self,
mesh: Dict[str, Any],
camera_trajectory: Dict[str, Any],
output_width: int = 1920,
output_height: int = 1080,
render_mode: str = "pbr",
background_color: str = "#000000",
lighting_preset: str = "studio",
num_frames: int = 1,
) -> Tuple[torch.Tensor]:
"""
Render mesh from camera viewpoint(s).
Args:
mesh: Dictionary of mesh data from TRELLIS2
camera_trajectory: Camera parameters
output_width: Output image width
output_height: Output image height
render_mode: Rendering style
background_color: Background hex color
lighting_preset: Lighting configuration name
num_frames: Number of frames to render
Returns:
Rendered image tensor [N, H, W, 3]
"""
if not NVDIFFRAST_AVAILABLE:
raise RuntimeError(
"MeshRenderer: nvdiffrast is not available. "
"Install with: pip install git+https://github.com/NVlabs/nvdiffrast.git"
)
if not torch.cuda.is_available():
raise RuntimeError(
"MeshRenderer: CUDA is required for nvdiffrast rendering."
)
# Parse background color
bg_color = self._parse_color(background_color)
background = torch.tensor(bg_color, dtype=torch.float32, device=self.device)
# Get mesh data
vertices = mesh["vertices"].to(self.device)
faces = mesh["faces"].to(self.device)
# Get or compute normals
if "normals" in mesh:
normals = mesh["normals"].to(self.device)
else:
normals = self._compute_normals(vertices, faces)
# Get colors/textures
if "colors" in mesh:
colors = mesh["colors"].to(self.device)
else:
colors = torch.ones(vertices.shape[0], 3, dtype=torch.float32, device=self.device)
# Estimate scene center and distance if not in mesh data
center = camera_trajectory.get("center", (0.0, 0.0, 0.0))
if center == (0.0, 0.0, 0.0) and "auto_center" not in mesh:
center = estimate_mesh_center(vertices)
mesh["auto_center"] = center
base_distance = camera_trajectory.get("distance", 2.5)
if base_distance == 2.5:
auto_dist = estimate_mesh_distance(vertices, center)
if auto_dist > 0:
base_distance = auto_dist
# Create camera matrices
cameras = self._create_cameras(
camera_trajectory, num_frames, output_width, output_height,
center, base_distance
)
# Get lighting
lighting = LIGHTING_PRESETS.get(lighting_preset, LIGHTING_PRESETS["studio"])
# Render each frame
rendered_frames = []
glctx = self.get_glctx()
if COMFY_AVAILABLE:
pbar = comfy.utils.ProgressBar(num_frames)
for i, (mvp, viewmat, K) in enumerate(cameras):
if render_mode == "pbr":
rendered = self._render_pbr(
glctx, vertices, faces, normals, colors,
mvp, viewmat, output_width, output_height,
background, lighting, mesh.get("textures")
)
elif render_mode == "wireframe":
rendered = self._render_wireframe(
glctx, vertices, faces, mvp,
output_width, output_height, background
)
elif render_mode == "cel":
rendered = self._render_cel(
glctx, vertices, faces, normals, colors,
mvp, viewmat, output_width, output_height,
background, lighting
)
elif render_mode == "normal":
rendered = self._render_normal(
glctx, vertices, faces, normals,
mvp, output_width, output_height, background
)
elif render_mode == "depth":
rendered = self._render_depth(
glctx, vertices, faces, mvp,
output_width, output_height, background
)
else:
rendered = self._render_pbr(
glctx, vertices, faces, normals, colors,
mvp, viewmat, output_width, output_height,
background, lighting, mesh.get("textures")
)
rendered_frames.append(rendered.cpu())
if COMFY_AVAILABLE:
pbar.update(1)
# Stack frames
output = torch.stack(rendered_frames)
logger.info(
f"MeshRenderer: Rendered {num_frames} frame(s) at {output_width}x{output_height} "
f"mode={render_mode}"
)
return (output,)
def _create_cameras(
self,
camera_trajectory: Dict[str, Any],
num_frames: int,
width: int,
height: int,
center: tuple,
distance: float,
) -> list:
"""Create camera matrices for each frame."""
from ..utils.camera_utils import create_camera_trajectory
trajectory_type = camera_trajectory.get("trajectory_type", "static")
angle_x = camera_trajectory.get("angle_x", 0.0)
angle_y = camera_trajectory.get("angle_y", 0.0)
fov = camera_trajectory.get("fov", 60.0)
animate = camera_trajectory.get("animate", False)
animation_range = camera_trajectory.get("animation_range", 30.0)
animation_axis = camera_trajectory.get("animation_axis", "y")
# Handle multi-angle trajectory
if trajectory_type == "multi_angle":
angles_y = camera_trajectory.get("angles_y", [0.0])
cameras = []
for ay in angles_y:
viewmat = create_orbit_camera(angle_x, ay, distance, center)
K = create_intrinsic_matrix(fov, width, height)
proj = create_projection_matrix(fov, width / height)
mvp = proj @ viewmat
cameras.append((
torch.tensor(mvp, dtype=torch.float32, device=self.device),
torch.tensor(viewmat, dtype=torch.float32, device=self.device),
torch.tensor(K, dtype=torch.float32, device=self.device)
))
return cameras
# Create trajectory
viewmats = create_camera_trajectory(
trajectory_type, angle_x, angle_y, distance, center,
num_frames, animate, animation_range, animation_axis
)
cameras = []
for viewmat in viewmats:
K = create_intrinsic_matrix(fov, width, height)
proj = create_projection_matrix(fov, width / height)
mvp = proj @ viewmat
cameras.append((
torch.tensor(mvp, dtype=torch.float32, device=self.device),
torch.tensor(viewmat, dtype=torch.float32, device=self.device),
torch.tensor(K, dtype=torch.float32, device=self.device)
))
return cameras
def _render_pbr(
self,
glctx,
vertices: torch.Tensor,
faces: torch.Tensor,
normals: torch.Tensor,
colors: torch.Tensor,
mvp: torch.Tensor,
viewmat: torch.Tensor,
width: int,
height: int,
background: torch.Tensor,
lighting: dict,
textures: Optional[dict] = None,
) -> torch.Tensor:
"""Render with PBR lighting."""
# Transform vertices to clip space
verts_homo = torch.cat([
vertices,
torch.ones(vertices.shape[0], 1, device=vertices.device)
], dim=1)
verts_clip = (mvp @ verts_homo.T).T
# Rasterize
rast_out, _ = dr.rasterize(
glctx, verts_clip.unsqueeze(0), faces.int(),
resolution=[height, width]
)
# Interpolate attributes
normals_interp, _ = dr.interpolate(
normals.unsqueeze(0), rast_out, faces.int()
)
colors_interp, _ = dr.interpolate(
colors.unsqueeze(0), rast_out, faces.int()
)
# Apply lighting
lit_color = self._apply_lighting(
colors_interp[0], normals_interp[0], viewmat, lighting
)
# Apply background
mask = (rast_out[0, :, :, 3:4] > 0).float()
output = lit_color * mask + background.view(1, 1, 3) * (1 - mask)
return output.clamp(0, 1)
def _render_wireframe(
self,
glctx,
vertices: torch.Tensor,
faces: torch.Tensor,
mvp: torch.Tensor,
width: int,
height: int,
background: torch.Tensor,
) -> torch.Tensor:
"""Render wireframe view."""
# Transform vertices
verts_homo = torch.cat([
vertices,
torch.ones(vertices.shape[0], 1, device=vertices.device)
], dim=1)
verts_clip = (mvp @ verts_homo.T).T
# Rasterize with edge detection
rast_out, _ = dr.rasterize(
glctx, verts_clip.unsqueeze(0), faces.int(),
resolution=[height, width]
)
# Get barycentric coordinates
mask = (rast_out[0, :, :, 3:4] > 0).float()
# Simple wireframe: detect triangle edges
bary = rast_out[0, :, :, :2]
edge_w = 0.02
edges = torch.min(bary, 1.0 - bary[:, :, 0:1] - bary[:, :, 1:2])
edges = torch.min(edges, dim=2, keepdim=True).values
wire = (edges < edge_w).float()
# White wireframe on background
wire_color = torch.tensor([1.0, 1.0, 1.0], device=vertices.device).view(1, 1, 3)
output = wire * wire_color * mask + background.view(1, 1, 3) * (1 - mask)
return output.clamp(0, 1)
def _render_cel(
self,
glctx,
vertices: torch.Tensor,
faces: torch.Tensor,
normals: torch.Tensor,
colors: torch.Tensor,
mvp: torch.Tensor,
viewmat: torch.Tensor,
width: int,
height: int,
background: torch.Tensor,
lighting: dict,
levels: int = 4,
) -> torch.Tensor:
"""Render with cel/toon shading."""
# Transform vertices
verts_homo = torch.cat([
vertices,
torch.ones(vertices.shape[0], 1, device=vertices.device)
], dim=1)
verts_clip = (mvp @ verts_homo.T).T
# Rasterize
rast_out, _ = dr.rasterize(
glctx, verts_clip.unsqueeze(0), faces.int(),
resolution=[height, width]
)
# Interpolate
normals_interp, _ = dr.interpolate(
normals.unsqueeze(0), rast_out, faces.int()
)
colors_interp, _ = dr.interpolate(
colors.unsqueeze(0), rast_out, faces.int()
)
# Compute lighting intensity
light_dir = torch.tensor(
lighting["lights"][0]["direction"],
dtype=torch.float32,
device=vertices.device
)
light_dir = light_dir / (torch.norm(light_dir) + 1e-8)
n = normals_interp[0]
n = n / (torch.norm(n, dim=2, keepdim=True) + 1e-8)
intensity = torch.sum(n * light_dir.view(1, 1, 3), dim=2, keepdim=True)
intensity = intensity.clamp(0, 1)
# Quantize to levels (cel shading)
intensity = torch.floor(intensity * levels) / levels
intensity = intensity * 0.7 + 0.3 # Add ambient
# Apply to colors
cel_color = colors_interp[0] * intensity
# Apply background
mask = (rast_out[0, :, :, 3:4] > 0).float()
output = cel_color * mask + background.view(1, 1, 3) * (1 - mask)
return output.clamp(0, 1)
def _render_normal(
self,
glctx,
vertices: torch.Tensor,
faces: torch.Tensor,
normals: torch.Tensor,
mvp: torch.Tensor,
width: int,
height: int,
background: torch.Tensor,
) -> torch.Tensor:
"""Render normal map view."""
verts_homo = torch.cat([
vertices,
torch.ones(vertices.shape[0], 1, device=vertices.device)
], dim=1)
verts_clip = (mvp @ verts_homo.T).T
rast_out, _ = dr.rasterize(
glctx, verts_clip.unsqueeze(0), faces.int(),
resolution=[height, width]
)
normals_interp, _ = dr.interpolate(
normals.unsqueeze(0), rast_out, faces.int()
)
# Convert normals to color ([-1,1] -> [0,1])
normal_color = (normals_interp[0] + 1.0) * 0.5
mask = (rast_out[0, :, :, 3:4] > 0).float()
output = normal_color * mask + background.view(1, 1, 3) * (1 - mask)
return output.clamp(0, 1)
def _render_depth(
self,
glctx,
vertices: torch.Tensor,
faces: torch.Tensor,
mvp: torch.Tensor,
width: int,
height: int,
background: torch.Tensor,
) -> torch.Tensor:
"""Render depth map."""
verts_homo = torch.cat([
vertices,
torch.ones(vertices.shape[0], 1, device=vertices.device)
], dim=1)
verts_clip = (mvp @ verts_homo.T).T
rast_out, _ = dr.rasterize(
glctx, verts_clip.unsqueeze(0), faces.int(),
resolution=[height, width]
)
# Get depth from clip-space z
depth = rast_out[0, :, :, 2:3]
# Normalize depth to [0, 1]
mask = (rast_out[0, :, :, 3:4] > 0).float()
depth_valid = depth * mask
if depth_valid.max() > depth_valid.min():
depth_norm = (depth - depth_valid.min()) / (depth_valid.max() - depth_valid.min() + 1e-8)
else:
depth_norm = torch.zeros_like(depth)
# Convert to grayscale RGB
depth_color = depth_norm.repeat(1, 1, 3)
output = depth_color * mask + background.view(1, 1, 3) * (1 - mask)
return output.clamp(0, 1)
def _apply_lighting(
self,
colors: torch.Tensor,
normals: torch.Tensor,
viewmat: torch.Tensor,
lighting: dict,
) -> torch.Tensor:
"""Apply lighting to colors."""
# Normalize normals
n = normals / (torch.norm(normals, dim=2, keepdim=True) + 1e-8)
# Start with ambient
lit = colors * lighting["ambient"]
# Add directional lights
for light in lighting["lights"]:
light_dir = torch.tensor(
light["direction"],
dtype=torch.float32,
device=colors.device
)
light_dir = light_dir / (torch.norm(light_dir) + 1e-8)
light_color = torch.tensor(
light["color"],
dtype=torch.float32,
device=colors.device
)
# Diffuse lighting
ndotl = torch.sum(n * light_dir.view(1, 1, 3), dim=2, keepdim=True)
ndotl = ndotl.clamp(0, 1)
lit = lit + colors * light_color.view(1, 1, 3) * ndotl * light["intensity"]
return lit
def _compute_normals(
self,
vertices: torch.Tensor,
faces: torch.Tensor
) -> torch.Tensor:
"""Compute vertex normals from faces."""
v0 = vertices[faces[:, 0]]
v1 = vertices[faces[:, 1]]
v2 = vertices[faces[:, 2]]
# Face normals
face_normals = torch.cross(v1 - v0, v2 - v0)
face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8)
# Accumulate to vertex normals
vertex_normals = torch.zeros_like(vertices)
for i in range(3):
vertex_normals.index_add_(0, faces[:, i], face_normals)
# Normalize
vertex_normals = vertex_normals / (torch.norm(vertex_normals, dim=1, keepdim=True) + 1e-8)
return vertex_normals
def _parse_color(self, color_str: str) -> list:
"""Parse hex color to RGB floats."""
color_str = color_str.lstrip("#")
if len(color_str) == 6:
r = int(color_str[0:2], 16) / 255.0
g = int(color_str[2:4], 16) / 255.0
b = int(color_str[4:6], 16) / 255.0
return [r, g, b]
return [0.0, 0.0, 0.0]
# Node registration
NODE_CLASS_MAPPINGS = {
"MeshRenderer": MeshRenderer,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MeshRenderer": "Mesh Renderer (3D to Image)",
}
nodes/mesh_sequence_renderer.py
+271
View File
@@ -0,0 +1,271 @@
"""Mesh Sequence Renderer node - render video sequence with consistent camera."""
import logging
from typing import Any, Dict, Tuple
import torch
logger = logging.getLogger(__name__)
from ..utils.camera_utils import (
create_orbit_camera,
create_intrinsic_matrix,
create_projection_matrix,
estimate_mesh_center,
estimate_mesh_distance,
)
# Progress bar
try:
import comfy.utils
COMFY_AVAILABLE = True
except ImportError:
COMFY_AVAILABLE = False
class MeshSequenceRenderer:
"""
Render a sequence of meshes from a consistent camera angle.
Designed to work with VideoToMeshes output, rendering each frame's
mesh from the same camera angle for video output.
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"mesh_sequence": ("TRELLIS_MESH_SEQUENCE", {
"tooltip": "Sequence of meshes from VideoToMeshes"
}),
"camera_trajectory": ("CAMERA_TRAJECTORY", {
"tooltip": "Camera parameters from CameraTrajectory node"
}),
},
"optional": {
"output_width": ("INT", {
"default": 1920,
"min": 256,
"max": 4096,
"step": 64,
"tooltip": "Output image width"
}),
"output_height": ("INT", {
"default": 1080,
"min": 256,
"max": 4096,
"step": 64,
"tooltip": "Output image height"
}),
"render_mode": (["pbr", "wireframe", "cel", "normal", "depth"], {
"default": "pbr",
"tooltip": "Rendering style"
}),
"background_color": ("COLOR", {
"default": "#000000",
"tooltip": "Background color"
}),
"lighting_preset": (["studio", "outdoor", "dramatic", "flat"], {
"default": "studio",
"tooltip": "Lighting configuration"
}),
"enable_onion_skin": ("BOOLEAN", {
"default": False,
"tooltip": "Enable onion skinning effect during rendering"
}),
"onion_skin_frames": ("INT", {
"default": 3,
"min": 1,
"max": 10,
"step": 1,
"tooltip": "Number of ghost frames"
}),
"onion_skin_interval": ("INT", {
"default": 5,
"min": 1,
"max": 30,
"step": 1,
"tooltip": "Frames between ghosts"
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("rendered_video",)
FUNCTION = "render_sequence"
CATEGORY = "TRELLIS2/Render"
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
def render_sequence(
self,
mesh_sequence: Dict[str, Any],
camera_trajectory: Dict[str, Any],
output_width: int = 1920,
output_height: int = 1080,
render_mode: str = "pbr",
background_color: str = "#000000",
lighting_preset: str = "studio",
enable_onion_skin: bool = False,
onion_skin_frames: int = 3,
onion_skin_interval: int = 5,
) -> Tuple[torch.Tensor]:
"""
Render a sequence of meshes.
Args:
mesh_sequence: Dictionary with meshes list and frame indices
camera_trajectory: Camera parameters
output_width: Output width
output_height: Output height
render_mode: Rendering style
background_color: Background hex color
lighting_preset: Lighting preset name
enable_onion_skin: Whether to apply onion skinning
onion_skin_frames: Number of ghost frames
onion_skin_interval: Frames between ghosts
Returns:
Rendered video tensor [N, H, W, 3]
"""
from .mesh_renderer import MeshRenderer
meshes_list = mesh_sequence.get("meshes", [])
frame_indices = mesh_sequence.get("frame_indices", [])
total_frames = mesh_sequence.get("total_frames", len(meshes_list))
if not meshes_list:
raise ValueError("MeshSequenceRenderer: Empty mesh sequence")
# Create renderer instance
renderer = MeshRenderer()
# Extract camera parameters
trajectory_type = camera_trajectory.get("trajectory_type", "static")
angle_x = camera_trajectory.get("angle_x", 0.0)
angle_y = camera_trajectory.get("angle_y", 0.0)
distance = camera_trajectory.get("distance", 2.5)
center = camera_trajectory.get("center", (0.0, 0.0, 0.0))
fov = camera_trajectory.get("fov", 60.0)
animate = camera_trajectory.get("animate", False)
animation_range = camera_trajectory.get("animation_range", 30.0)
animation_axis = camera_trajectory.get("animation_axis", "y")
logger.info(
f"MeshSequenceRenderer: Rendering {len(meshes_list)} meshes, "
f"camera: type={trajectory_type}, angles=({angle_x}, {angle_y})"
)
rendered_frames = []
if COMFY_AVAILABLE:
pbar = comfy.utils.ProgressBar(len(meshes_list))
for i, (idx, mesh) in enumerate(zip(frame_indices, meshes_list)):
# Compute per-frame camera angles if animating
if animate and total_frames > 1:
t = idx / max(total_frames - 1, 1)
if trajectory_type == "orbit":
if animation_axis == "y":
current_angle_y = angle_y + (t - 0.5) * animation_range
current_angle_x = angle_x
else:
current_angle_x = angle_x + (t - 0.5) * animation_range
current_angle_y = angle_y
current_distance = distance
current_center = center
elif trajectory_type == "dolly":
current_angle_x = angle_x
current_angle_y = angle_y
current_distance = max(0.5, distance + (t - 0.5) * animation_range)
current_center = center
elif trajectory_type == "pan":
pan_offset = (t - 0.5) * animation_range * 0.1
current_angle_x = angle_x
current_angle_y = angle_y
current_distance = distance
current_center = (center[0] + pan_offset, center[1], center[2])
else:
current_angle_x = angle_x
current_angle_y = angle_y
current_distance = distance
current_center = center
else:
current_angle_x = angle_x
current_angle_y = angle_y
current_distance = distance
current_center = center
# Auto-compute center/distance from mesh if needed
if current_center == (0.0, 0.0, 0.0):
current_center = estimate_mesh_center(mesh["vertices"])
if current_distance == 2.5:
auto_dist = estimate_mesh_distance(mesh["vertices"], current_center)
if auto_dist > 0:
current_distance = auto_dist
# Create frame-specific camera trajectory
frame_trajectory = {
"trajectory_type": "static",
"angle_x": current_angle_x,
"angle_y": current_angle_y,
"distance": current_distance,
"center": current_center,
"fov": fov,
"animate": False,
}
# Render single frame
(rendered,) = renderer.render(
mesh=mesh,
camera_trajectory=frame_trajectory,
output_width=output_width,
output_height=output_height,
render_mode=render_mode,
background_color=background_color,
lighting_preset=lighting_preset,
num_frames=1,
)
rendered_frames.append(rendered[0])
if COMFY_AVAILABLE:
pbar.update(1)
# Stack into batch
output = torch.stack(rendered_frames)
# Apply onion skinning if enabled
if enable_onion_skin:
from ..effects.onion_skin import apply_onion_skinning, OnionSkinConfig
config = OnionSkinConfig(
ghost_count=onion_skin_frames,
frame_interval=onion_skin_interval,
ghost_opacity=0.3,
opacity_falloff="linear",
color_mode="tinted",
past_tint=(0.0, 0.4, 1.0), # Blue
future_tint=(1.0, 0.4, 0.0), # Orange
)
output = apply_onion_skinning(output, config)
logger.info(
f"MeshSequenceRenderer: Rendered {len(meshes_list)} frames "
f"at {output_width}x{output_height}"
)
return (output,)
# Node registration
NODE_CLASS_MAPPINGS = {
"MeshSequenceRenderer": MeshSequenceRenderer,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MeshSequenceRenderer": "Mesh Sequence Renderer",
}
nodes/onion_skinning.py
+165
View File
@@ -0,0 +1,165 @@
"""Onion Skinning node - add motion ghosting effect to rendered frames."""
import logging
from typing import Any, Dict, Tuple
import torch
logger = logging.getLogger(__name__)
from ..effects.onion_skin import apply_onion_skinning, OnionSkinConfig
class OnionSkinning:
"""
Add motion ghosting effect to rendered video frames.
Creates a visual trail showing previous (and optionally future) frames
as transparent overlays, commonly used in animation to visualize motion.
Parameters allow control over:
- Number of ghost frames
- Frame interval between ghosts (every 5 frames, 10 frames, etc.)
- Ghost opacity and falloff
- Color tinting for past/future ghosts
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"images": ("IMAGE", {
"tooltip": "Rendered video frames [N, H, W, C]"
}),
},
"optional": {
"ghost_count": ("INT", {
"default": 3,
"min": 1,
"max": 10,
"step": 1,
"tooltip": "Number of ghost frames to show"
}),
"frame_interval": ("INT", {
"default": 5,
"min": 1,
"max": 30,
"step": 1,
"tooltip": "Frames between each ghost (5 = every 5th frame)"
}),
"ghost_opacity": ("FLOAT", {
"default": 0.3,
"min": 0.05,
"max": 0.8,
"step": 0.05,
"tooltip": "Base opacity of ghost frames"
}),
"opacity_falloff": (["linear", "exponential", "constant"], {
"default": "linear",
"tooltip": "How opacity decreases with distance"
}),
"ghost_color_mode": (["original", "tinted", "silhouette"], {
"default": "tinted",
"tooltip": "Color treatment for ghosts"
}),
"past_tint": ("COLOR", {
"default": "#0066FF",
"tooltip": "Tint color for past frames (blue)"
}),
"future_tint": ("COLOR", {
"default": "#FF6600",
"tooltip": "Tint color for future frames (orange)"
}),
"include_future": ("BOOLEAN", {
"default": False,
"tooltip": "Include future frames as well as past"
}),
"blend_mode": (["over", "add", "screen"], {
"default": "over",
"tooltip": "How ghosts blend with current frame"
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("onion_skinned",)
FUNCTION = "apply"
CATEGORY = "TRELLIS2/Effects"
def apply(
self,
images: torch.Tensor,
ghost_count: int = 3,
frame_interval: int = 5,
ghost_opacity: float = 0.3,
opacity_falloff: str = "linear",
ghost_color_mode: str = "tinted",
past_tint: str = "#0066FF",
future_tint: str = "#FF6600",
include_future: bool = False,
blend_mode: str = "over",
) -> Tuple[torch.Tensor]:
"""
Apply onion skinning effect to video frames.
Args:
images: Video frames [N, H, W, C]
ghost_count: Number of ghost frames
frame_interval: Frames between ghosts
ghost_opacity: Base ghost opacity
opacity_falloff: Opacity falloff type
ghost_color_mode: Color treatment
past_tint: Hex color for past frames
future_tint: Hex color for future frames
include_future: Include future frames
blend_mode: Blend mode for compositing
Returns:
Video tensor with onion skinning applied
"""
# Parse colors
past_rgb = self._parse_color(past_tint)
future_rgb = self._parse_color(future_tint)
# Create config
config = OnionSkinConfig(
ghost_count=ghost_count,
frame_interval=frame_interval,
ghost_opacity=ghost_opacity,
opacity_falloff=opacity_falloff,
color_mode=ghost_color_mode,
past_tint=past_rgb,
future_tint=future_rgb,
include_future=include_future,
blend_mode=blend_mode,
)
# Apply effect
output = apply_onion_skinning(images, config)
logger.info(
f"OnionSkinning: Applied {ghost_count} ghosts at interval {frame_interval} "
f"to {images.shape[0]} frames"
)
return (output,)
def _parse_color(self, color_str: str) -> Tuple[float, float, float]:
"""Parse hex color to RGB floats."""
color_str = color_str.lstrip("#")
if len(color_str) == 6:
r = int(color_str[0:2], 16) / 255.0
g = int(color_str[2:4], 16) / 255.0
b = int(color_str[4:6], 16) / 255.0
return (r, g, b)
return (0.0, 0.0, 0.0)
# Node registration
NODE_CLASS_MAPPINGS = {
"OnionSkinning": OnionSkinning,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"OnionSkinning": "Onion Skinning (Motion Ghosts)",
}
nodes/rendering_effects.py
+255
View File
@@ -0,0 +1,255 @@
"""Rendering Effects node - apply visual effects to rendered output."""
import logging
from typing import Any, Dict, Tuple
import torch
logger = logging.getLogger(__name__)
class RenderingEffects:
"""
Apply post-processing visual effects to rendered images.
Includes effects like:
- Bloom (glow on bright areas)
- Depth of field (blur based on depth)
- Motion blur
- Cel shading quantization
- Outline/edge detection
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"images": ("IMAGE", {
"tooltip": "Rendered images [N, H, W, C]"
}),
},
"optional": {
# Bloom
"bloom_intensity": ("FLOAT", {
"default": 0.0,
"min": 0.0,
"max": 2.0,
"step": 0.1,
"tooltip": "Bloom/glow intensity (0 = disabled)"
}),
"bloom_threshold": ("FLOAT", {
"default": 0.8,
"min": 0.0,
"max": 1.0,
"step": 0.05,
"tooltip": "Brightness threshold for bloom"
}),
"bloom_radius": ("INT", {
"default": 20,
"min": 5,
"max": 100,
"step": 5,
"tooltip": "Bloom blur radius"
}),
# Motion blur (for video)
"motion_blur": ("FLOAT", {
"default": 0.0,
"min": 0.0,
"max": 1.0,
"step": 0.1,
"tooltip": "Motion blur intensity (0 = disabled)"
}),
# Cel shading post-process
"cel_shading_levels": ("INT", {
"default": 0,
"min": 0,
"max": 8,
"step": 1,
"tooltip": "Color quantization levels (0 = disabled)"
}),
# Outline
"outline_width": ("FLOAT", {
"default": 0.0,
"min": 0.0,
"max": 5.0,
"step": 0.5,
"tooltip": "Edge outline width (0 = disabled)"
}),
"outline_color": ("COLOR", {
"default": "#000000",
"tooltip": "Outline color"
}),
"outline_threshold": ("FLOAT", {
"default": 0.1,
"min": 0.01,
"max": 0.5,
"step": 0.01,
"tooltip": "Edge detection sensitivity"
}),
# Color grading
"contrast": ("FLOAT", {
"default": 1.0,
"min": 0.5,
"max": 2.0,
"step": 0.05,
"tooltip": "Contrast adjustment"
}),
"saturation": ("FLOAT", {
"default": 1.0,
"min": 0.0,
"max": 2.0,
"step": 0.05,
"tooltip": "Color saturation"
}),
"brightness": ("FLOAT", {
"default": 1.0,
"min": 0.5,
"max": 2.0,
"step": 0.05,
"tooltip": "Brightness adjustment"
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("processed",)
FUNCTION = "apply_effects"
CATEGORY = "TRELLIS2/Effects"
def apply_effects(
self,
images: torch.Tensor,
bloom_intensity: float = 0.0,
bloom_threshold: float = 0.8,
bloom_radius: int = 20,
motion_blur: float = 0.0,
cel_shading_levels: int = 0,
outline_width: float = 0.0,
outline_color: str = "#000000",
outline_threshold: float = 0.1,
contrast: float = 1.0,
saturation: float = 1.0,
brightness: float = 1.0,
) -> Tuple[torch.Tensor]:
"""
Apply visual effects to images.
Args:
images: Input images [N, H, W, C]
Various effect parameters
Returns:
Processed images
"""
output = images.clone()
device = output.device
# Apply bloom
if bloom_intensity > 0:
output = self._apply_bloom(output, bloom_intensity, bloom_threshold, bloom_radius)
# Apply motion blur (for sequences)
if motion_blur > 0 and output.shape[0] > 1:
output = self._apply_motion_blur(output, motion_blur)
# Apply cel shading quantization
if cel_shading_levels > 0:
output = self._apply_cel_quantize(output, cel_shading_levels)
# Apply outline
if outline_width > 0:
outline_rgb = self._parse_color(outline_color)
output = self._apply_outline(output, outline_width, outline_rgb, outline_threshold)
# Apply color grading
if contrast != 1.0 or saturation != 1.0 or brightness != 1.0:
output = self._apply_color_grading(output, contrast, saturation, brightness)
# Clamp to valid range
output = output.clamp(0, 1)
logger.info(f"RenderingEffects: Applied effects to {output.shape[0]} frames")
return (output,)
def _apply_bloom(
self,
images: torch.Tensor,
intensity: float,
threshold: float,
radius: int,
) -> torch.Tensor:
"""Apply bloom effect."""
from ..effects.post_processing import apply_bloom
return apply_bloom(images, intensity, threshold, radius)
def _apply_motion_blur(
self,
images: torch.Tensor,
intensity: float,
) -> torch.Tensor:
"""Apply motion blur between frames."""
from ..effects.post_processing import apply_motion_blur
return apply_motion_blur(images, intensity)
def _apply_cel_quantize(
self,
images: torch.Tensor,
levels: int,
) -> torch.Tensor:
"""Quantize colors for cel shading look."""
# Quantize each channel
quantized = torch.floor(images * levels) / levels
return quantized
def _apply_outline(
self,
images: torch.Tensor,
width: float,
color: Tuple[float, float, float],
threshold: float,
) -> torch.Tensor:
"""Apply edge outline."""
from ..effects.post_processing import apply_outline
return apply_outline(images, width, color, threshold)
def _apply_color_grading(
self,
images: torch.Tensor,
contrast: float,
saturation: float,
brightness: float,
) -> torch.Tensor:
"""Apply color grading adjustments."""
# Brightness
output = images * brightness
# Contrast (around 0.5 midpoint)
output = (output - 0.5) * contrast + 0.5
# Saturation
if saturation != 1.0:
gray = output.mean(dim=-1, keepdim=True)
output = gray + saturation * (output - gray)
return output
def _parse_color(self, color_str: str) -> Tuple[float, float, float]:
"""Parse hex color to RGB tuple."""
color_str = color_str.lstrip("#")
if len(color_str) == 6:
r = int(color_str[0:2], 16) / 255.0
g = int(color_str[2:4], 16) / 255.0
b = int(color_str[4:6], 16) / 255.0
return (r, g, b)
return (0.0, 0.0, 0.0)
# Node registration
NODE_CLASS_MAPPINGS = {
"RenderingEffects": RenderingEffects,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"RenderingEffects": "Rendering Effects",
}
nodes/trellis_predict.py
+330
View File
@@ -0,0 +1,330 @@
"""TRELLIS2 Predict node - convert single image to 3D mesh."""
import logging
import os
import tempfile
from pathlib import Path
from typing import Any, Dict, Tuple, Optional
import numpy as np
import torch
from PIL import Image
logger = logging.getLogger(__name__)
# Progress bar with fallback
try:
import comfy.utils
COMFY_AVAILABLE = True
except ImportError:
COMFY_AVAILABLE = False
from ..utils.download_utils import ensure_trellis_model, get_progress_bar
# Check for TRELLIS availability
TRELLIS_AVAILABLE = False
try:
from trellis.pipelines import TrsEllis2Img
from trellis.utils import postprocessing_utils
TRELLIS_AVAILABLE = True
except ImportError:
logger.warning(
"TRELLIS2 not available. Install with: "
"pip install git+https://github.com/microsoft/TRELLIS.git"
)
# Check for rembg (background removal)
REMBG_AVAILABLE = False
try:
from rembg import remove as remove_background
REMBG_AVAILABLE = True
except ImportError:
logger.warning("rembg not available. Install with: pip install rembg")
class Trellis2Predict:
"""
Convert a single image to a 3D mesh with PBR materials using TRELLIS2.
TRELLIS2 (4B parameter model) processes a single image and produces a
high-quality 3D mesh with PBR materials (albedo, roughness, metallic).
Note: This node requires the TRELLIS2 model and CUDA GPU.
First run will download the model (~8GB).
"""
_pipeline = None # Class-level singleton
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"image": ("IMAGE", {
"tooltip": "Input image to convert to 3D mesh"
}),
},
"optional": {
"resolution": (["512", "1024", "1536"], {
"default": "1024",
"tooltip": "Processing resolution (higher = better quality, slower)"
}),
"seed": ("INT", {
"default": 0,
"min": 0,
"max": 2147483647,
"tooltip": "Random seed for reproducible results (0 = random)"
}),
"remove_background": ("BOOLEAN", {
"default": True,
"tooltip": "Automatically remove background before processing"
}),
"simplify_mesh": ("BOOLEAN", {
"default": False,
"tooltip": "Simplify mesh to reduce polygon count"
}),
"target_faces": ("INT", {
"default": 50000,
"min": 1000,
"max": 500000,
"step": 1000,
"tooltip": "Target face count when simplifying (only if simplify_mesh is True)"
}),
}
}
RETURN_TYPES = ("TRELLIS_MESH",)
RETURN_NAMES = ("mesh",)
FUNCTION = "predict"
CATEGORY = "TRELLIS2/Core"
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
@classmethod
def get_pipeline(cls):
"""Get or create the TRELLIS2 pipeline (singleton)."""
if cls._pipeline is None:
if not TRELLIS_AVAILABLE:
raise RuntimeError(
"TRELLIS2 not available. Install with: "
"pip install git+https://github.com/microsoft/TRELLIS.git"
)
logger.info("Loading TRELLIS2 pipeline...")
# Ensure model is downloaded
model_path = ensure_trellis_model()
# Initialize pipeline
cls._pipeline = TrslIs2Img.from_pretrained(model_path)
cls._pipeline.to("cuda")
logger.info("TRELLIS2 pipeline loaded")
return cls._pipeline
def predict(
self,
image: torch.Tensor,
resolution: str = "1024",
seed: int = 0,
remove_background: bool = True,
simplify_mesh: bool = False,
target_faces: int = 50000,
) -> Tuple[Dict[str, Any]]:
"""
Run TRELLIS2 prediction on a single image.
Args:
image: Input image tensor [1, H, W, C] or [H, W, C]
resolution: Processing resolution
seed: Random seed
remove_background: Whether to remove background
simplify_mesh: Whether to simplify the output mesh
target_faces: Target face count for simplification
Returns:
Tuple containing dictionary of mesh data
"""
if not torch.cuda.is_available():
raise RuntimeError(
"Trellis2Predict: CUDA is required for TRELLIS2. "
"This node cannot run on CPU or MPS."
)
# Initialize progress bar (5 steps)
pbar = get_progress_bar(5)
pbar.update_absolute(0, 5, "Initializing TRELLIS2")
# Handle batch dimension
if len(image.shape) == 4:
if image.shape[0] != 1:
logger.warning(
f"Trellis2Predict received batch of {image.shape[0]} images, "
"only processing first image"
)
image = image[0]
# Convert to PIL Image
pbar.update_absolute(1, 5, "Preprocessing image")
img_np = image.cpu().numpy()
img_np = np.clip(img_np, 0, 1)
img_np = (img_np * 255).astype(np.uint8)
pil_image = Image.fromarray(img_np, mode="RGB")
# Remove background if requested
if remove_background and REMBG_AVAILABLE:
logger.info("Removing background...")
pil_image = remove_background(pil_image)
# Convert RGBA to RGB with white background
if pil_image.mode == "RGBA":
background = Image.new("RGB", pil_image.size, (255, 255, 255))
background.paste(pil_image, mask=pil_image.split()[3])
pil_image = background
# Get resolution
res = int(resolution)
# Set seed
if seed == 0:
seed = torch.randint(0, 2147483647, (1,)).item()
torch.manual_seed(seed)
np.random.seed(seed)
# Run TRELLIS2
pbar.update_absolute(2, 5, "Running TRELLIS2 inference")
logger.info(f"Running TRELLIS2 at {res}x{res} resolution...")
try:
pipeline = self.get_pipeline()
# Run inference
outputs = pipeline.run(
pil_image,
seed=seed,
)
# Extract mesh
pbar.update_absolute(3, 5, "Extracting mesh")
mesh_output = outputs['mesh'][0] if isinstance(outputs['mesh'], list) else outputs['mesh']
except Exception as e:
logger.error(f"TRELLIS2 inference failed: {e}")
raise RuntimeError(f"TRELLIS2 inference failed: {e}")
# Convert to our mesh format
pbar.update_absolute(4, 5, "Converting mesh format")
mesh_data = self._convert_mesh(mesh_output)
# Optionally simplify
if simplify_mesh and target_faces < mesh_data["faces"].shape[0]:
logger.info(f"Simplifying mesh to {target_faces} faces...")
mesh_data = self._simplify_mesh(mesh_data, target_faces)
# Add metadata
mesh_data["source_resolution"] = res
mesh_data["seed"] = seed
mesh_data["source_image_shape"] = list(image.shape)
pbar.update_absolute(5, 5, "Complete")
logger.info(
f"Trellis2Predict: Generated mesh with "
f"{mesh_data['vertices'].shape[0]} vertices, "
f"{mesh_data['faces'].shape[0]} faces"
)
return (mesh_data,)
def _convert_mesh(self, mesh_output) -> Dict[str, Any]:
"""Convert TRELLIS mesh output to our format."""
mesh_data = {}
# Extract vertices and faces
if hasattr(mesh_output, 'v_pos'):
mesh_data["vertices"] = mesh_output.v_pos.to(self.device)
elif hasattr(mesh_output, 'vertices'):
mesh_data["vertices"] = torch.tensor(
mesh_output.vertices, dtype=torch.float32, device=self.device
)
if hasattr(mesh_output, 't_pos_idx'):
mesh_data["faces"] = mesh_output.t_pos_idx.to(self.device)
elif hasattr(mesh_output, 'faces'):
mesh_data["faces"] = torch.tensor(
mesh_output.faces, dtype=torch.int64, device=self.device
)
# Extract normals
if hasattr(mesh_output, 'v_nrm') and mesh_output.v_nrm is not None:
mesh_data["normals"] = mesh_output.v_nrm.to(self.device)
# Extract UVs
if hasattr(mesh_output, 'v_tex') and mesh_output.v_tex is not None:
mesh_data["uvs"] = mesh_output.v_tex.to(self.device)
# Extract textures/materials
textures = {}
if hasattr(mesh_output, 'albedo') and mesh_output.albedo is not None:
textures["albedo"] = mesh_output.albedo
if hasattr(mesh_output, 'roughness') and mesh_output.roughness is not None:
textures["roughness"] = mesh_output.roughness
if hasattr(mesh_output, 'metallic') and mesh_output.metallic is not None:
textures["metallic"] = mesh_output.metallic
if hasattr(mesh_output, 'normal') and mesh_output.normal is not None:
textures["normal_map"] = mesh_output.normal
if textures:
mesh_data["textures"] = textures
# Store original for advanced operations
mesh_data["_original"] = mesh_output
return mesh_data
def _simplify_mesh(
self,
mesh_data: Dict[str, Any],
target_faces: int
) -> Dict[str, Any]:
"""Simplify mesh to target face count."""
try:
import trimesh
vertices = mesh_data["vertices"].cpu().numpy()
faces = mesh_data["faces"].cpu().numpy()
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
# Simplify using quadric decimation
simplified = mesh.simplify_quadric_decimation(target_faces)
# Update mesh data
mesh_data["vertices"] = torch.tensor(
simplified.vertices, dtype=torch.float32, device=self.device
)
mesh_data["faces"] = torch.tensor(
simplified.faces, dtype=torch.int64, device=self.device
)
# Recompute normals
mesh_data["normals"] = torch.tensor(
simplified.vertex_normals, dtype=torch.float32, device=self.device
)
logger.info(f"Simplified mesh to {len(simplified.faces)} faces")
except ImportError:
logger.warning("trimesh not available for mesh simplification")
return mesh_data
# Node registration
NODE_CLASS_MAPPINGS = {
"Trellis2Predict": Trellis2Predict,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Trellis2Predict": "TRELLIS2 Predict (Image to 3D Mesh)",
}
nodes/video_to_meshes.py
+327
View File
@@ -0,0 +1,327 @@
"""Video to Meshes node - batch process video frames through TRELLIS2."""
import logging
from typing import Any, Dict, List, Tuple
import numpy as np
import torch
from PIL import Image
logger = logging.getLogger(__name__)
from ..utils.video_utils import extract_frames
from ..utils.download_utils import ensure_trellis_model, get_progress_bar
# Check for TRELLIS availability
TRELLIS_AVAILABLE = False
try:
from trellis.pipelines import TrslIs2Img
TRELLIS_AVAILABLE = True
except ImportError:
pass
# Check for rembg
REMBG_AVAILABLE = False
try:
from rembg import remove as remove_background
REMBG_AVAILABLE = True
except ImportError:
pass
# Progress bar
try:
import comfy.utils
COMFY_AVAILABLE = True
except ImportError:
COMFY_AVAILABLE = False
class VideoToMeshes:
"""
Batch process video frames through TRELLIS2 to create mesh sequence.
Takes a batch of video frames and processes each one (with optional frame skipping)
through TRELLIS2 to produce a sequence of 3D meshes.
Note: This is compute-intensive. Processing time scales with frame count.
Use frame_skip to reduce processing time for long videos.
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"images": ("IMAGE", {
"tooltip": "Video frames as image batch [N, H, W, C]"
}),
},
"optional": {
"frame_skip": ("INT", {
"default": 5,
"min": 1,
"max": 30,
"step": 1,
"tooltip": "Process every Nth frame (1 = all frames)"
}),
"max_frames": ("INT", {
"default": 0,
"min": 0,
"max": 1000,
"step": 1,
"tooltip": "Maximum frames to process (0 = no limit)"
}),
"resolution": (["512", "1024"], {
"default": "1024",
"tooltip": "TRELLIS2 processing resolution"
}),
"remove_background": ("BOOLEAN", {
"default": True,
"tooltip": "Remove background from each frame"
}),
"seed": ("INT", {
"default": 0,
"min": 0,
"max": 2147483647,
"tooltip": "Random seed (0 = different seed per frame)"
}),
}
}
RETURN_TYPES = ("TRELLIS_MESH_SEQUENCE",)
RETURN_NAMES = ("mesh_sequence",)
FUNCTION = "process_batch"
CATEGORY = "TRELLIS2/Video"
def __init__(self):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self._pipeline = None
def get_pipeline(self):
"""Get or create the TRELLIS2 pipeline."""
if self._pipeline is None:
if not TRELLIS_AVAILABLE:
raise RuntimeError(
"TRELLIS2 not available. Install with: "
"pip install git+https://github.com/microsoft/TRELLIS.git"
)
model_path = ensure_trellis_model()
self._pipeline = TrslIs2Img.from_pretrained(model_path)
self._pipeline.to("cuda")
return self._pipeline
def process_batch(
self,
images: torch.Tensor,
frame_skip: int = 5,
max_frames: int = 0,
resolution: str = "1024",
remove_background: bool = True,
seed: int = 0,
) -> Tuple[Dict[str, Any]]:
"""
Process video frames through TRELLIS2.
Args:
images: Video frames [N, H, W, C]
frame_skip: Process every Nth frame
max_frames: Maximum frames to process (0 = all)
resolution: TRELLIS2 processing resolution
remove_background: Whether to remove backgrounds
seed: Random seed
Returns:
Dictionary containing:
- meshes: List of mesh dictionaries
- frame_indices: List of original frame indices
- total_frames: Total number of frames in input
"""
if not torch.cuda.is_available():
raise RuntimeError("VideoToMeshes: CUDA is required for TRELLIS2")
total_frames = images.shape[0]
# Extract frames with skipping
frames, indices = extract_frames(images, frame_skip, max_frames)
num_to_process = len(frames)
logger.info(
f"VideoToMeshes: Processing {num_to_process} frames "
f"from {total_frames} total (skip={frame_skip})"
)
# Initialize progress bar
total_steps = num_to_process + 2 # model load + frames + complete
pbar = get_progress_bar(total_steps)
pbar.update_absolute(0, total_steps, "Loading TRELLIS2 model")
# Get pipeline
pipeline = self.get_pipeline()
res = int(resolution)
pbar.update_absolute(1, total_steps, f"Processing {num_to_process} frames")
meshes_list = []
for i, (frame, orig_idx) in enumerate(zip(frames, indices)):
pbar.update_absolute(
i + 2, total_steps,
f"Processing frame {i+1}/{num_to_process}"
)
try:
# Convert to PIL
img_np = frame.cpu().numpy()
img_np = np.clip(img_np, 0, 1)
img_np = (img_np * 255).astype(np.uint8)
pil_image = Image.fromarray(img_np, mode="RGB")
# Remove background
if remove_background and REMBG_AVAILABLE:
pil_image = remove_background(pil_image)
if pil_image.mode == "RGBA":
background = Image.new("RGB", pil_image.size, (255, 255, 255))
background.paste(pil_image, mask=pil_image.split()[3])
pil_image = background
# Set seed
frame_seed = seed if seed > 0 else torch.randint(0, 2147483647, (1,)).item()
torch.manual_seed(frame_seed)
# Run inference
outputs = pipeline.run(pil_image, seed=frame_seed)
# Extract mesh
mesh_output = outputs['mesh'][0] if isinstance(outputs['mesh'], list) else outputs['mesh']
# Convert to our format
mesh_data = self._convert_mesh(mesh_output)
mesh_data["frame_index"] = orig_idx
mesh_data["seed"] = frame_seed
meshes_list.append(mesh_data)
except Exception as e:
logger.error(f"VideoToMeshes: Failed to process frame {i}: {e}")
continue
# Clear CUDA cache periodically
if i % 10 == 0:
torch.cuda.empty_cache()
result = {
"meshes": meshes_list,
"frame_indices": [m["frame_index"] for m in meshes_list],
"total_frames": total_frames,
"frame_skip": frame_skip,
}
pbar.update_absolute(total_steps, total_steps, "Complete")
logger.info(f"VideoToMeshes: Successfully processed {len(meshes_list)} frames")
return (result,)
def _convert_mesh(self, mesh_output) -> Dict[str, Any]:
"""Convert TRELLIS mesh output to our format."""
mesh_data = {}
if hasattr(mesh_output, 'v_pos'):
mesh_data["vertices"] = mesh_output.v_pos.to(self.device)
elif hasattr(mesh_output, 'vertices'):
mesh_data["vertices"] = torch.tensor(
mesh_output.vertices, dtype=torch.float32, device=self.device
)
if hasattr(mesh_output, 't_pos_idx'):
mesh_data["faces"] = mesh_output.t_pos_idx.to(self.device)
elif hasattr(mesh_output, 'faces'):
mesh_data["faces"] = torch.tensor(
mesh_output.faces, dtype=torch.int64, device=self.device
)
if hasattr(mesh_output, 'v_nrm') and mesh_output.v_nrm is not None:
mesh_data["normals"] = mesh_output.v_nrm.to(self.device)
if hasattr(mesh_output, 'v_tex') and mesh_output.v_tex is not None:
mesh_data["uvs"] = mesh_output.v_tex.to(self.device)
textures = {}
if hasattr(mesh_output, 'albedo') and mesh_output.albedo is not None:
textures["albedo"] = mesh_output.albedo
if hasattr(mesh_output, 'roughness') and mesh_output.roughness is not None:
textures["roughness"] = mesh_output.roughness
if hasattr(mesh_output, 'metallic') and mesh_output.metallic is not None:
textures["metallic"] = mesh_output.metallic
if textures:
mesh_data["textures"] = textures
mesh_data["_original"] = mesh_output
return mesh_data
class MergeMeshSequences:
"""
Merge multiple mesh sequences into one.
Useful for combining sequences processed in batches.
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"sequence_1": ("TRELLIS_MESH_SEQUENCE", {
"tooltip": "First mesh sequence"
}),
"sequence_2": ("TRELLIS_MESH_SEQUENCE", {
"tooltip": "Second mesh sequence"
}),
},
}
RETURN_TYPES = ("TRELLIS_MESH_SEQUENCE",)
RETURN_NAMES = ("merged_sequence",)
FUNCTION = "merge"
CATEGORY = "TRELLIS2/Video"
def merge(
self,
sequence_1: Dict[str, Any],
sequence_2: Dict[str, Any]
) -> Tuple[Dict[str, Any]]:
"""Merge two mesh sequences."""
meshes = sequence_1["meshes"] + sequence_2["meshes"]
indices = sequence_1["frame_indices"] + sequence_2["frame_indices"]
# Sort by frame index
combined = sorted(zip(indices, meshes), key=lambda x: x[0])
sorted_indices = [x[0] for x in combined]
sorted_meshes = [x[1] for x in combined]
total_frames = max(
sequence_1.get("total_frames", max(sequence_1["frame_indices"]) + 1),
sequence_2.get("total_frames", max(sequence_2["frame_indices"]) + 1)
)
result = {
"meshes": sorted_meshes,
"frame_indices": sorted_indices,
"total_frames": total_frames,
}
logger.info(f"MergeMeshSequences: Merged {len(sorted_meshes)} total frames")
return (result,)
# Node registration
NODE_CLASS_MAPPINGS = {
"VideoToMeshes": VideoToMeshes,
"MergeMeshSequences": MergeMeshSequences,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VideoToMeshes": "Video to Meshes (Batch TRELLIS2)",
"MergeMeshSequences": "Merge Mesh Sequences",
}
requirements.txt
+29
View File
@@ -0,0 +1,29 @@
# ComfyUI-TRELLIS2_Motion dependencies
# Core
torch>=2.0.0
numpy>=1.21.0
pillow>=9.0.0
# TRELLIS2 dependencies
huggingface_hub>=0.20.0
transformers>=4.30.0
trimesh>=3.0.0
xatlas # UV unwrapping
rembg # Background removal
spconv-cu118 # Sparse convolutions (CUDA 11.8)
flash-attn # Efficient attention
# Mesh rendering
nvdiffrast # CUDA mesh rendering
kaolin # NVIDIA 3D deep learning library
# Image processing effects
kornia>=0.7.0 # Differentiable image processing
opencv-python-headless>=4.8.0
# Utilities
scipy>=1.10.0
tqdm>=4.65.0
requests>=2.28.0
einops>=0.6.0
utils/__init__.py
+60
View File
@@ -0,0 +1,60 @@
"""Utility modules for ComfyUI-TRELLIS2_Motion."""
from .camera_utils import (
create_rotation_matrix,
create_orbit_camera,
create_intrinsic_matrix,
interpolate_cameras,
create_camera_trajectory,
estimate_mesh_center,
estimate_mesh_distance,
)
from .video_utils import (
extract_frames,
interpolate_missing_frames,
batch_to_list,
list_to_batch,
resize_frames,
get_video_info,
)
from .download_utils import (
ensure_trellis_model,
is_trellis_model_downloaded,
get_progress_bar,
)
from .mesh_utils import (
load_mesh,
save_mesh,
mesh_to_tensors,
get_mesh_info,
)
__all__ = [
# Camera
"create_rotation_matrix",
"create_orbit_camera",
"create_intrinsic_matrix",
"interpolate_cameras",
"create_camera_trajectory",
"estimate_mesh_center",
"estimate_mesh_distance",
# Video
"extract_frames",
"interpolate_missing_frames",
"batch_to_list",
"list_to_batch",
"resize_frames",
"get_video_info",
# Download
"ensure_trellis_model",
"is_trellis_model_downloaded",
"get_progress_bar",
# Mesh
"load_mesh",
"save_mesh",
"mesh_to_tensors",
"get_mesh_info",
]
utils/camera_utils.py
+367
View File
@@ -0,0 +1,367 @@
"""Camera utilities for creating view matrices and trajectories."""
import logging
from typing import List, Optional, Tuple, Union
import numpy as np
import torch
logger = logging.getLogger(__name__)
def create_rotation_matrix(
angle_x: float,
angle_y: float,
angle_z: float = 0.0
) -> np.ndarray:
"""
Create a 3x3 rotation matrix from Euler angles.
Uses XYZ order (rotate around X, then Y, then Z).
Args:
angle_x: Rotation around X axis in degrees (pitch/tilt)
angle_y: Rotation around Y axis in degrees (yaw/pan)
angle_z: Rotation around Z axis in degrees (roll)
Returns:
3x3 rotation matrix as numpy array
"""
ax, ay, az = np.radians([angle_x, angle_y, angle_z])
# Rotation around X axis
Rx = np.array([
[1, 0, 0],
[0, np.cos(ax), -np.sin(ax)],
[0, np.sin(ax), np.cos(ax)]
], dtype=np.float32)
# Rotation around Y axis
Ry = np.array([
[np.cos(ay), 0, np.sin(ay)],
[0, 1, 0],
[-np.sin(ay), 0, np.cos(ay)]
], dtype=np.float32)
# Rotation around Z axis
Rz = np.array([
[np.cos(az), -np.sin(az), 0],
[np.sin(az), np.cos(az), 0],
[0, 0, 1]
], dtype=np.float32)
# Combined rotation: Rz @ Ry @ Rx
return Rz @ Ry @ Rx
def create_orbit_camera(
angle_x: float,
angle_y: float,
distance: float,
center: Tuple[float, float, float] = (0.0, 0.0, 0.0)
) -> np.ndarray:
"""
Create a 4x4 camera extrinsic matrix for an orbital view around a scene center.
Uses OpenGL camera convention (x-right, y-up, z-backward) which is standard
for mesh rendering with nvdiffrast.
Args:
angle_x: Vertical angle in degrees (positive = look from above)
angle_y: Horizontal angle in degrees (positive = rotate right)
distance: Distance from camera to scene center
center: Scene center point (x, y, z). TRELLIS2 meshes are typically
centered at origin (0, 0, 0).
Returns:
4x4 world-to-camera transformation matrix (extrinsic)
"""
# Convert angles to radians
theta = np.radians(angle_y) # Horizontal orbit (azimuth)
phi = np.radians(angle_x) # Vertical tilt (elevation)
# Calculate camera position on a sphere around the center
# Spherical to Cartesian conversion (OpenGL: y-up)
x = distance * np.sin(theta) * np.cos(phi)
y = distance * np.sin(phi)
z = distance * np.cos(theta) * np.cos(phi)
# Camera position in world coordinates
camera_pos = np.array([x, y, z], dtype=np.float32) + np.array(center, dtype=np.float32)
# Look-at direction (from camera to center)
forward = np.array(center, dtype=np.float32) - camera_pos
forward = forward / (np.linalg.norm(forward) + 1e-8)
# Up vector in OpenGL convention (y-up)
world_up = np.array([0.0, 1.0, 0.0], dtype=np.float32)
# Right vector (cross product of forward and up)
right = np.cross(forward, world_up)
right_norm = np.linalg.norm(right)
# Handle gimbal lock case (looking straight up or down)
if right_norm < 1e-6:
# Use a different up vector
world_up = np.array([0.0, 0.0, 1.0], dtype=np.float32)
right = np.cross(forward, world_up)
right_norm = np.linalg.norm(right)
right = right / (right_norm + 1e-8)
# Recalculate up to be orthogonal
up = np.cross(right, forward)
up = up / (np.linalg.norm(up) + 1e-8)
# Build rotation matrix (camera axes in world coordinates)
# Camera coordinate system: x=right, y=up, z=-forward (OpenGL)
R_cam_to_world = np.stack([right, up, -forward], axis=1)
# World-to-camera rotation
R = R_cam_to_world.T
# Translation: camera position in camera coordinates
t = -R @ camera_pos
# Build 4x4 extrinsic matrix
extrinsic = np.eye(4, dtype=np.float32)
extrinsic[:3, :3] = R
extrinsic[:3, 3] = t
return extrinsic
def create_intrinsic_matrix(
fov: float,
width: int,
height: int,
fov_type: str = "horizontal"
) -> np.ndarray:
"""
Create a 3x3 camera intrinsic matrix.
Args:
fov: Field of view in degrees
width: Image width in pixels
height: Image height in pixels
fov_type: "horizontal" or "vertical" FOV specification
Returns:
3x3 intrinsic matrix K
"""
fov_rad = np.radians(fov)
if fov_type == "horizontal":
focal = width / (2 * np.tan(fov_rad / 2))
else: # vertical
focal = height / (2 * np.tan(fov_rad / 2))
# Principal point at image center
cx = width / 2.0
cy = height / 2.0
K = np.array([
[focal, 0, cx],
[0, focal, cy],
[0, 0, 1]
], dtype=np.float32)
return K
def create_projection_matrix(
fov: float,
aspect: float,
near: float = 0.01,
far: float = 100.0
) -> np.ndarray:
"""
Create a 4x4 perspective projection matrix.
Args:
fov: Vertical field of view in degrees
aspect: Aspect ratio (width / height)
near: Near clipping plane
far: Far clipping plane
Returns:
4x4 projection matrix
"""
fov_rad = np.radians(fov)
f = 1.0 / np.tan(fov_rad / 2)
proj = np.zeros((4, 4), dtype=np.float32)
proj[0, 0] = f / aspect
proj[1, 1] = f
proj[2, 2] = (far + near) / (near - far)
proj[2, 3] = (2 * far * near) / (near - far)
proj[3, 2] = -1
return proj
def interpolate_cameras(
cam1: np.ndarray,
cam2: np.ndarray,
t: float
) -> np.ndarray:
"""
Interpolate between two camera matrices.
Uses linear interpolation for translation and SLERP approximation for rotation.
Args:
cam1: First 4x4 camera extrinsic matrix
cam2: Second 4x4 camera extrinsic matrix
t: Interpolation factor (0 = cam1, 1 = cam2)
Returns:
Interpolated 4x4 camera extrinsic matrix
"""
t = np.clip(t, 0.0, 1.0)
# Extract rotation and translation
R1, t1 = cam1[:3, :3], cam1[:3, 3]
R2, t2 = cam2[:3, :3], cam2[:3, 3]
# Linear interpolation for translation
t_interp = t1 * (1 - t) + t2 * t
# For rotation, we use a simple linear interpolation and re-orthogonalize
# (proper SLERP would require quaternion conversion)
R_interp = R1 * (1 - t) + R2 * t
# Re-orthogonalize using SVD
U, _, Vt = np.linalg.svd(R_interp)
R_interp = U @ Vt
# Ensure proper rotation (det = 1, not -1)
if np.linalg.det(R_interp) < 0:
R_interp = -R_interp
# Build interpolated matrix
result = np.eye(4, dtype=np.float32)
result[:3, :3] = R_interp
result[:3, 3] = t_interp
return result
def create_camera_trajectory(
trajectory_type: str,
angle_x: float,
angle_y: float,
distance: float,
center: Tuple[float, float, float],
num_frames: int,
animate: bool = False,
animation_range: float = 10.0,
animation_axis: str = "y"
) -> List[np.ndarray]:
"""
Create a sequence of camera matrices for a trajectory.
Args:
trajectory_type: One of "orbit", "dolly", "pan", "static"
angle_x: Base vertical angle in degrees
angle_y: Base horizontal angle in degrees
distance: Base distance from scene center
center: Scene center point
num_frames: Number of frames in the trajectory
animate: Whether to animate the camera over the sequence
animation_range: Range of animation in degrees (for orbit/pan) or distance (for dolly)
animation_axis: Axis to animate around ("x", "y", or "z")
Returns:
List of 4x4 camera extrinsic matrices, one per frame
"""
cameras = []
for i in range(num_frames):
t = i / max(num_frames - 1, 1) # Normalized time [0, 1]
if not animate:
# Static camera - same for all frames
cam = create_orbit_camera(angle_x, angle_y, distance, center)
else:
# Animated camera
if trajectory_type == "orbit":
# Orbit around the scene
if animation_axis == "y":
# Horizontal orbit
current_angle_y = angle_y + (t - 0.5) * animation_range
cam = create_orbit_camera(angle_x, current_angle_y, distance, center)
else: # x axis
# Vertical orbit
current_angle_x = angle_x + (t - 0.5) * animation_range
cam = create_orbit_camera(current_angle_x, angle_y, distance, center)
elif trajectory_type == "dolly":
# Move camera closer/farther
current_distance = distance + (t - 0.5) * animation_range
current_distance = max(0.5, current_distance) # Don't go too close
cam = create_orbit_camera(angle_x, angle_y, current_distance, center)
elif trajectory_type == "pan":
# Pan the camera (translate horizontally)
pan_offset = (t - 0.5) * animation_range * 0.1 # Scale down for translation
current_center = (
center[0] + pan_offset,
center[1],
center[2]
)
cam = create_orbit_camera(angle_x, angle_y, distance, current_center)
else: # static or unknown
cam = create_orbit_camera(angle_x, angle_y, distance, center)
cameras.append(cam)
return cameras
def estimate_mesh_center(vertices: Union[torch.Tensor, np.ndarray]) -> Tuple[float, float, float]:
"""
Estimate the scene center from mesh vertices.
Args:
vertices: [N, 3] tensor/array of vertex positions
Returns:
Tuple (x, y, z) of estimated scene center
"""
if isinstance(vertices, torch.Tensor):
center = vertices.mean(dim=0).cpu().numpy()
else:
center = vertices.mean(axis=0)
return tuple(center.tolist())
def estimate_mesh_distance(
vertices: Union[torch.Tensor, np.ndarray],
center: Optional[Tuple[float, float, float]] = None
) -> float:
"""
Estimate an appropriate camera distance for viewing the mesh.
Args:
vertices: [N, 3] tensor/array of vertex positions
center: Optional pre-computed scene center
Returns:
Recommended camera distance
"""
if center is None:
center = estimate_mesh_center(vertices)
if isinstance(vertices, torch.Tensor):
center_tensor = torch.tensor(center, device=vertices.device, dtype=vertices.dtype)
distances = torch.norm(vertices - center_tensor, dim=1)
scene_radius = torch.quantile(distances, 0.95).item()
else:
center_arr = np.array(center, dtype=vertices.dtype)
distances = np.linalg.norm(vertices - center_arr, axis=1)
scene_radius = np.percentile(distances, 95)
# Camera should be at roughly 2-3x the scene radius
return scene_radius * 2.5
utils/download_utils.py
+205
View File
@@ -0,0 +1,205 @@
"""Download utilities for TRELLIS2 model with progress tracking."""
import logging
import os
from pathlib import Path
from typing import Optional
logger = logging.getLogger(__name__)
# TRELLIS2 model configuration
TRELLIS_HF_REPO = "microsoft/TRELLIS-2-4B"
TRELLIS_CACHE_DIR = Path.home() / ".cache" / "huggingface" / "hub"
# Progress bar fallback for when ComfyUI is not available
try:
import comfy.utils
COMFY_AVAILABLE = True
except ImportError:
COMFY_AVAILABLE = False
class ProgressBar:
"""Fallback progress bar when ComfyUI is not available."""
def __init__(self, total: int):
self.total = total
self.current = 0
def update_absolute(self, current: int, total: Optional[int] = None, desc: str = ""):
self.current = current
total = total or self.total
pct = (current / total * 100) if total > 0 else 0
logger.info(f"{desc}: {pct:.1f}% ({current}/{total})")
def update(self, value: int):
self.current += value
pct = (self.current / self.total * 100) if self.total > 0 else 0
logger.info(f"Progress: {pct:.1f}%")
def get_progress_bar(total: int):
"""Get a progress bar instance (ComfyUI or fallback)."""
if COMFY_AVAILABLE:
return comfy.utils.ProgressBar(total)
return ProgressBar(total)
def get_trellis_cache_dir() -> Path:
"""Get the TRELLIS model cache directory."""
cache_dir = TRELLIS_CACHE_DIR
cache_dir.mkdir(parents=True, exist_ok=True)
return cache_dir
def is_trellis_model_downloaded() -> bool:
"""Check if the TRELLIS model is already downloaded."""
try:
from huggingface_hub import snapshot_download, HfFileSystem
fs = HfFileSystem()
# Check if model files exist in cache
cache_dir = get_trellis_cache_dir()
# Look for model files
model_dir = cache_dir / f"models--{TRELLIS_HF_REPO.replace('/', '--')}"
if model_dir.exists():
# Check for key files
snapshots = model_dir / "snapshots"
if snapshots.exists():
# Check any snapshot directory
for snapshot in snapshots.iterdir():
if (snapshot / "config.json").exists():
return True
return False
except Exception as e:
logger.warning(f"Could not check TRELLIS model status: {e}")
return False
def ensure_trellis_model(pbar: Optional[ProgressBar] = None) -> str:
"""
Ensure the TRELLIS model is available, downloading if necessary.
Args:
pbar: Optional progress bar for download progress
Returns:
Path to the TRELLIS model directory
"""
try:
from huggingface_hub import snapshot_download
except ImportError:
raise ImportError(
"huggingface_hub is required for downloading TRELLIS2 model. "
"Install with: pip install huggingface_hub"
)
if is_trellis_model_downloaded():
logger.info("TRELLIS2 model found in cache")
cache_dir = get_trellis_cache_dir()
model_dir = cache_dir / f"models--{TRELLIS_HF_REPO.replace('/', '--')}"
# Return the actual model path
snapshots = model_dir / "snapshots"
for snapshot in snapshots.iterdir():
if (snapshot / "config.json").exists():
return str(snapshot)
return str(model_dir)
logger.info(f"Downloading TRELLIS2 model from {TRELLIS_HF_REPO}...")
logger.info("This may take a while (~8GB download)")
if pbar:
pbar.update_absolute(0, 100, "Downloading TRELLIS2 model")
try:
model_path = snapshot_download(
TRELLIS_HF_REPO,
local_dir_use_symlinks=False,
)
if pbar:
pbar.update_absolute(100, 100, "TRELLIS2 download complete")
logger.info(f"TRELLIS2 model downloaded to {model_path}")
return model_path
except Exception as e:
logger.error(f"Failed to download TRELLIS2 model: {e}")
raise RuntimeError(
f"Failed to download TRELLIS2 model from {TRELLIS_HF_REPO}. "
f"Error: {e}\n"
"Try manually downloading with:\n"
f" huggingface-cli download {TRELLIS_HF_REPO}"
)
def check_trellis_dependencies() -> dict:
"""
Check if all TRELLIS2 dependencies are available.
Returns:
Dictionary with status of each dependency
"""
status = {
"trellis": False,
"model_downloaded": False,
"nvdiffrast": False,
"cuda": False,
"rembg": False,
}
# Check TRELLIS package
try:
import trellis
status["trellis"] = True
except ImportError:
pass
# Check model
status["model_downloaded"] = is_trellis_model_downloaded()
# Check nvdiffrast
try:
import nvdiffrast
status["nvdiffrast"] = True
except ImportError:
pass
# Check CUDA
try:
import torch
status["cuda"] = torch.cuda.is_available()
except ImportError:
pass
# Check rembg
try:
import rembg
status["rembg"] = True
except ImportError:
pass
return status
def print_dependency_status():
"""Print the status of all dependencies."""
status = check_trellis_dependencies()
print("\nTRELLIS2 Dependency Status:")
print("-" * 40)
for name, available in status.items():
icon = "OK" if available else "MISSING"
print(f" [{icon}] {name}")
if not all(status.values()):
print("\nTo install missing dependencies:")
if not status["trellis"]:
print(" pip install git+https://github.com/microsoft/TRELLIS.git")
if not status["nvdiffrast"]:
print(" pip install git+https://github.com/NVlabs/nvdiffrast.git")
if not status["rembg"]:
print(" pip install rembg")
if not status["model_downloaded"]:
print(f" huggingface-cli download {TRELLIS_HF_REPO}")
utils/mesh_utils.py
+314
View File
@@ -0,0 +1,314 @@
"""Mesh I/O utilities for TRELLIS2 meshes."""
import logging
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union
import numpy as np
import torch
logger = logging.getLogger(__name__)
# Try to import trimesh for mesh I/O
try:
import trimesh
TRIMESH_AVAILABLE = True
except ImportError:
TRIMESH_AVAILABLE = False
logger.warning("trimesh not available. Install with: pip install trimesh")
def load_mesh(
path: Union[str, Path],
device: str = "cuda"
) -> Dict[str, Any]:
"""
Load a mesh from file and convert to tensor format.
Args:
path: Path to mesh file (supports .obj, .ply, .glb, .gltf)
device: Device for output tensors
Returns:
Dictionary containing mesh data:
- vertices: [N, 3] tensor of vertex positions
- faces: [F, 3] tensor of face indices
- normals: [N, 3] tensor of vertex normals (if available)
- uvs: [N, 2] tensor of texture coordinates (if available)
- colors: [N, 3/4] tensor of vertex colors (if available)
- textures: dict of texture images (if available)
"""
if not TRIMESH_AVAILABLE:
raise ImportError("trimesh is required for mesh loading. Install with: pip install trimesh")
path = Path(path)
if not path.exists():
raise FileNotFoundError(f"Mesh file not found: {path}")
logger.info(f"Loading mesh from {path}")
# Load mesh
mesh = trimesh.load(str(path), force='mesh')
result = {}
# Vertices
vertices = torch.tensor(mesh.vertices, dtype=torch.float32, device=device)
result["vertices"] = vertices
# Faces
faces = torch.tensor(mesh.faces, dtype=torch.int64, device=device)
result["faces"] = faces
# Normals
if hasattr(mesh, 'vertex_normals') and mesh.vertex_normals is not None:
normals = torch.tensor(mesh.vertex_normals, dtype=torch.float32, device=device)
result["normals"] = normals
# UVs
if hasattr(mesh.visual, 'uv') and mesh.visual.uv is not None:
uvs = torch.tensor(mesh.visual.uv, dtype=torch.float32, device=device)
result["uvs"] = uvs
# Vertex colors
if hasattr(mesh.visual, 'vertex_colors') and mesh.visual.vertex_colors is not None:
colors = torch.tensor(
mesh.visual.vertex_colors[:, :3] / 255.0,
dtype=torch.float32,
device=device
)
result["colors"] = colors
# Textures (if PBR material)
if hasattr(mesh.visual, 'material'):
material = mesh.visual.material
textures = {}
if hasattr(material, 'baseColorTexture') and material.baseColorTexture is not None:
textures["albedo"] = np.array(material.baseColorTexture)
if hasattr(material, 'metallicRoughnessTexture') and material.metallicRoughnessTexture is not None:
textures["metallic_roughness"] = np.array(material.metallicRoughnessTexture)
if hasattr(material, 'normalTexture') and material.normalTexture is not None:
textures["normal"] = np.array(material.normalTexture)
if textures:
result["textures"] = textures
# Store original mesh reference for advanced operations
result["_trimesh"] = mesh
logger.info(
f"Loaded mesh: {len(vertices)} vertices, {len(faces)} faces"
)
return result
def save_mesh(
mesh_data: Dict[str, Any],
path: Union[str, Path],
file_format: Optional[str] = None
) -> None:
"""
Save mesh data to file.
Args:
mesh_data: Dictionary containing mesh tensors
path: Output file path
file_format: Optional format override (obj, ply, glb, gltf)
"""
if not TRIMESH_AVAILABLE:
raise ImportError("trimesh is required for mesh saving. Install with: pip install trimesh")
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
# If we have the original trimesh, use it
if "_trimesh" in mesh_data:
mesh = mesh_data["_trimesh"]
else:
# Reconstruct from tensors
vertices = mesh_data["vertices"].cpu().numpy()
faces = mesh_data["faces"].cpu().numpy()
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
if "normals" in mesh_data:
mesh.vertex_normals = mesh_data["normals"].cpu().numpy()
if "colors" in mesh_data:
colors = mesh_data["colors"].cpu().numpy()
# Convert to 0-255 range and add alpha
colors_255 = (colors * 255).astype(np.uint8)
if colors_255.shape[1] == 3:
alpha = np.full((len(colors_255), 1), 255, dtype=np.uint8)
colors_255 = np.concatenate([colors_255, alpha], axis=1)
mesh.visual.vertex_colors = colors_255
# Save
mesh.export(str(path), file_type=file_format)
logger.info(f"Saved mesh to {path}")
def mesh_to_tensors(
mesh: Any,
device: str = "cuda"
) -> Dict[str, torch.Tensor]:
"""
Convert a trimesh/TRELLIS mesh object to tensor dictionary.
Args:
mesh: Mesh object (trimesh.Trimesh or TRELLIS mesh)
device: Target device for tensors
Returns:
Dictionary of mesh tensors
"""
result = {}
# Handle different mesh types
if hasattr(mesh, 'vertices'):
# trimesh or similar
if isinstance(mesh.vertices, torch.Tensor):
result["vertices"] = mesh.vertices.to(device)
else:
result["vertices"] = torch.tensor(mesh.vertices, dtype=torch.float32, device=device)
if hasattr(mesh, 'faces'):
if isinstance(mesh.faces, torch.Tensor):
result["faces"] = mesh.faces.to(device)
else:
result["faces"] = torch.tensor(mesh.faces, dtype=torch.int64, device=device)
# Handle TRELLIS-specific mesh format
elif hasattr(mesh, 'v_pos'):
result["vertices"] = mesh.v_pos.to(device)
if hasattr(mesh, 't_pos_idx'):
result["faces"] = mesh.t_pos_idx.to(device)
# Handle dictionary input
elif isinstance(mesh, dict):
for key, value in mesh.items():
if isinstance(value, torch.Tensor):
result[key] = value.to(device)
elif isinstance(value, np.ndarray):
result[key] = torch.tensor(value, device=device)
return result
def get_mesh_info(mesh_data: Dict[str, Any]) -> Dict[str, Any]:
"""
Get information about a mesh.
Args:
mesh_data: Dictionary containing mesh tensors
Returns:
Dictionary with mesh properties
"""
info = {}
if "vertices" in mesh_data:
vertices = mesh_data["vertices"]
info["num_vertices"] = len(vertices)
info["vertex_min"] = vertices.min(dim=0).values.cpu().tolist()
info["vertex_max"] = vertices.max(dim=0).values.cpu().tolist()
info["vertex_center"] = vertices.mean(dim=0).cpu().tolist()
# Compute bounding box size
bbox_min = vertices.min(dim=0).values
bbox_max = vertices.max(dim=0).values
bbox_size = bbox_max - bbox_min
info["bbox_size"] = bbox_size.cpu().tolist()
if "faces" in mesh_data:
info["num_faces"] = len(mesh_data["faces"])
if "normals" in mesh_data:
info["has_normals"] = True
if "uvs" in mesh_data:
info["has_uvs"] = True
if "colors" in mesh_data:
info["has_colors"] = True
if "textures" in mesh_data:
info["texture_types"] = list(mesh_data["textures"].keys())
return info
def normalize_mesh(
mesh_data: Dict[str, torch.Tensor],
target_scale: float = 1.0
) -> Dict[str, torch.Tensor]:
"""
Normalize mesh to fit within a unit cube centered at origin.
Args:
mesh_data: Dictionary containing mesh tensors
target_scale: Scale factor for the normalized mesh
Returns:
Normalized mesh data
"""
vertices = mesh_data["vertices"]
# Compute center and scale
center = vertices.mean(dim=0, keepdim=True)
vertices_centered = vertices - center
max_extent = vertices_centered.abs().max()
scale = target_scale / max_extent
# Apply normalization
result = mesh_data.copy()
result["vertices"] = vertices_centered * scale
return result
def combine_meshes(
meshes: list,
device: str = "cuda"
) -> Dict[str, torch.Tensor]:
"""
Combine multiple meshes into a single mesh.
Args:
meshes: List of mesh dictionaries
device: Target device
Returns:
Combined mesh dictionary
"""
all_vertices = []
all_faces = []
vertex_offset = 0
for mesh in meshes:
vertices = mesh["vertices"]
faces = mesh["faces"]
all_vertices.append(vertices)
all_faces.append(faces + vertex_offset)
vertex_offset += len(vertices)
result = {
"vertices": torch.cat(all_vertices, dim=0).to(device),
"faces": torch.cat(all_faces, dim=0).to(device),
}
# Combine other attributes if present
if all(["normals" in m for m in meshes]):
result["normals"] = torch.cat([m["normals"] for m in meshes], dim=0).to(device)
if all(["colors" in m for m in meshes]):
result["colors"] = torch.cat([m["colors"] for m in meshes], dim=0).to(device)
return result
utils/video_utils.py
+219
View File
@@ -0,0 +1,219 @@
"""Video frame utilities for extraction and interpolation."""
import logging
from typing import List, Tuple
import torch
logger = logging.getLogger(__name__)
def extract_frames(
video_tensor: torch.Tensor,
skip: int = 1,
max_frames: int = 0
) -> Tuple[List[torch.Tensor], List[int]]:
"""
Extract frames from a video tensor with optional frame skipping.
Args:
video_tensor: Video frames as [N, H, W, C] tensor (ComfyUI format)
skip: Process every Nth frame (1 = all frames, 2 = every other, etc.)
max_frames: Maximum number of frames to extract (0 = no limit)
Returns:
Tuple of:
- frames: List of [H, W, C] tensors for each extracted frame
- indices: List of original frame indices that were extracted
"""
total_frames = video_tensor.shape[0]
# Generate indices with skipping
indices = list(range(0, total_frames, skip))
# Apply max_frames limit
if max_frames > 0:
indices = indices[:max_frames]
# Extract frames
frames = [video_tensor[i] for i in indices]
logger.info(
f"Extracted {len(frames)} frames from {total_frames} total "
f"(skip={skip}, max={max_frames})"
)
return frames, indices
def interpolate_missing_frames(
rendered_frames: List[torch.Tensor],
original_indices: List[int],
total_frames: int
) -> torch.Tensor:
"""
Interpolate frames that were skipped during processing.
Uses linear interpolation between the nearest processed frames to fill in gaps.
Args:
rendered_frames: List of rendered [H, W, C] tensors for processed frames
original_indices: The original frame indices that were processed
total_frames: Total number of frames in the original video
Returns:
Full video tensor [N, H, W, C] with all frames filled in
"""
if len(original_indices) == total_frames:
# No interpolation needed
return torch.stack(rendered_frames)
if len(rendered_frames) == 0:
raise ValueError("No rendered frames provided")
# Get shape from first frame
H, W, C = rendered_frames[0].shape
device = rendered_frames[0].device
dtype = rendered_frames[0].dtype
# Create output tensor
output = torch.zeros(total_frames, H, W, C, dtype=dtype, device=device)
# Create a set for faster lookup
processed_set = set(original_indices)
# Fill in processed frames
for i, idx in enumerate(original_indices):
output[idx] = rendered_frames[i]
# Interpolate gaps
for frame_idx in range(total_frames):
if frame_idx in processed_set:
continue
# Find nearest processed frames before and after
prev_idx = None
next_idx = None
for idx in original_indices:
if idx < frame_idx:
prev_idx = idx
elif idx > frame_idx:
next_idx = idx
break
# Handle edge cases
if prev_idx is None:
# Before first processed frame - use first frame
output[frame_idx] = output[original_indices[0]]
elif next_idx is None:
# After last processed frame - use last frame
output[frame_idx] = output[original_indices[-1]]
else:
# Interpolate between prev and next
# Find which indices in original_indices correspond to prev/next
prev_frame = output[prev_idx]
next_frame = output[next_idx]
# Linear interpolation factor
t = (frame_idx - prev_idx) / (next_idx - prev_idx)
# Interpolate pixel values
output[frame_idx] = prev_frame * (1 - t) + next_frame * t
logger.info(
f"Interpolated {total_frames - len(original_indices)} missing frames "
f"(total: {total_frames})"
)
return output
def batch_to_list(batch_tensor: torch.Tensor) -> List[torch.Tensor]:
"""
Convert a batch tensor to a list of individual tensors.
Args:
batch_tensor: [N, H, W, C] tensor
Returns:
List of N tensors, each [H, W, C]
"""
return [batch_tensor[i] for i in range(batch_tensor.shape[0])]
def list_to_batch(frame_list: List[torch.Tensor]) -> torch.Tensor:
"""
Convert a list of frame tensors to a batch tensor.
Args:
frame_list: List of [H, W, C] tensors
Returns:
Stacked tensor [N, H, W, C]
"""
return torch.stack(frame_list)
def resize_frames(
frames: torch.Tensor,
width: int,
height: int,
mode: str = "bilinear"
) -> torch.Tensor:
"""
Resize a batch of frames to a target resolution.
Args:
frames: [N, H, W, C] tensor in ComfyUI format
width: Target width
height: Target height
mode: Interpolation mode ("nearest", "bilinear", "bicubic")
Returns:
Resized [N, height, width, C] tensor
"""
# ComfyUI format: [N, H, W, C]
# PyTorch interpolate expects: [N, C, H, W]
n_frames = frames.shape[0]
# Convert to [N, C, H, W]
frames_nchw = frames.permute(0, 3, 1, 2)
# Resize
resized = torch.nn.functional.interpolate(
frames_nchw,
size=(height, width),
mode=mode,
align_corners=False if mode != "nearest" else None
)
# Convert back to [N, H, W, C]
resized_nhwc = resized.permute(0, 2, 3, 1)
return resized_nhwc
def get_video_info(video_tensor: torch.Tensor) -> dict:
"""
Get information about a video tensor.
Args:
video_tensor: [N, H, W, C] tensor
Returns:
Dictionary with video properties
"""
n, h, w, c = video_tensor.shape
return {
"num_frames": n,
"height": h,
"width": w,
"channels": c,
"dtype": str(video_tensor.dtype),
"device": str(video_tensor.device),
"min_value": video_tensor.min().item(),
"max_value": video_tensor.max().item(),
}
workflows/effects_showcase.json
+170
View File
@@ -0,0 +1,170 @@
{
"last_node_id": 7,
"last_link_id": 6,
"nodes": [
{
"id": 1,
"type": "LoadImage",
"pos": [100, 200],
"size": [300, 300],
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{"name": "IMAGE", "type": "IMAGE", "links": [1], "slot_index": 0},
{"name": "MASK", "type": "MASK", "links": null}
],
"properties": {"Node name for S&R": "LoadImage"},
"widgets_values": ["input_image.png", "image"]
},
{
"id": 2,
"type": "Trellis2Predict",
"pos": [500, 200],
"size": [280, 180],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{"name": "image", "type": "IMAGE", "link": 1}
],
"outputs": [
{"name": "mesh", "type": "TRELLIS_MESH", "links": [2], "slot_index": 0}
],
"properties": {"Node name for S&R": "Trellis2Predict"},
"widgets_values": ["1024", 0, true, false, 50000]
},
{
"id": 3,
"type": "CameraTrajectory_T2",
"pos": [500, 450],
"size": [280, 350],
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "links": [3], "slot_index": 0}
],
"properties": {"Node name for S&R": "CameraTrajectory_T2"},
"widgets_values": [
"static",
20.0,
-15.0,
0.0,
2.5,
60.0,
0.0,
0.0,
0.0,
false,
30.0,
"y"
]
},
{
"id": 4,
"type": "MeshRenderer",
"pos": [900, 280],
"size": [280, 220],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{"name": "mesh", "type": "TRELLIS_MESH", "link": 2},
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "link": 3}
],
"outputs": [
{"name": "rendered_image", "type": "IMAGE", "links": [4], "slot_index": 0}
],
"properties": {"Node name for S&R": "MeshRenderer"},
"widgets_values": [1920, 1080, "cel", "#000000", "dramatic", 1]
},
{
"id": 5,
"type": "RenderingEffects",
"pos": [1300, 200],
"size": [280, 400],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 4}
],
"outputs": [
{"name": "processed", "type": "IMAGE", "links": [5], "slot_index": 0}
],
"properties": {"Node name for S&R": "RenderingEffects"},
"widgets_values": [
0.3,
0.8,
20,
0.0,
0,
1.5,
"#000000",
0.1,
1.1,
1.1,
1.0
]
},
{
"id": 6,
"type": "SaveImage",
"pos": [1700, 280],
"size": [300, 300],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 5}
],
"properties": {"Node name for S&R": "SaveImage"},
"widgets_values": ["trellis2_effects"]
}
],
"links": [
[1, 1, 0, 2, 0, "IMAGE"],
[2, 2, 0, 4, 0, "TRELLIS_MESH"],
[3, 3, 0, 4, 1, "CAMERA_TRAJECTORY"],
[4, 4, 0, 5, 0, "IMAGE"],
[5, 5, 0, 6, 0, "IMAGE"]
],
"groups": [
{
"title": "1. Load Image",
"bounding": [50, 100, 400, 450],
"color": "#3f789e"
},
{
"title": "2. Convert to 3D",
"bounding": [460, 100, 380, 750],
"color": "#8f5e8e"
},
{
"title": "3. Render (Cel Shading)",
"bounding": [860, 180, 380, 400],
"color": "#5e8f6e"
},
{
"title": "4. Post-Processing Effects",
"bounding": [1260, 100, 380, 550],
"color": "#e67e22"
},
{
"title": "5. Save",
"bounding": [1660, 180, 380, 450],
"color": "#9b59b6"
}
],
"config": {},
"extra": {
"info": {
"name": "Effects Showcase (TRELLIS2)",
"description": "Demonstrate cel shading, bloom, and other rendering effects.",
"author": "ComfyUI-TRELLIS2_Motion",
"version": "1.0"
}
},
"version": 0.4
}
workflows/multiview_preview.json
+146
View File
@@ -0,0 +1,146 @@
{
"last_node_id": 6,
"last_link_id": 5,
"nodes": [
{
"id": 1,
"type": "LoadImage",
"pos": [100, 200],
"size": [300, 300],
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{"name": "IMAGE", "type": "IMAGE", "links": [1], "slot_index": 0},
{"name": "MASK", "type": "MASK", "links": null}
],
"properties": {"Node name for S&R": "LoadImage"},
"widgets_values": ["input_image.png", "image"]
},
{
"id": 2,
"type": "Trellis2Predict",
"pos": [500, 200],
"size": [280, 180],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{"name": "image", "type": "IMAGE", "link": 1}
],
"outputs": [
{"name": "mesh", "type": "TRELLIS_MESH", "links": [2], "slot_index": 0}
],
"properties": {"Node name for S&R": "Trellis2Predict"},
"widgets_values": ["1024", 0, true, false, 50000]
},
{
"id": 3,
"type": "MultiAngleCamera",
"pos": [500, 450],
"size": [280, 180],
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "links": [3], "slot_index": 0}
],
"properties": {"Node name for S&R": "MultiAngleCamera"},
"widgets_values": [8, 15.0, 2.5, 60.0, 0.0]
},
{
"id": 4,
"type": "MeshRenderer",
"pos": [900, 280],
"size": [280, 220],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{"name": "mesh", "type": "TRELLIS_MESH", "link": 2},
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "link": 3}
],
"outputs": [
{"name": "rendered_image", "type": "IMAGE", "links": [4], "slot_index": 0}
],
"properties": {"Node name for S&R": "MeshRenderer"},
"widgets_values": [512, 512, "pbr", "#000000", "studio", 8]
},
{
"id": 5,
"type": "ImageGridComposite",
"pos": [1300, 280],
"size": [280, 150],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 4}
],
"outputs": [
{"name": "IMAGE", "type": "IMAGE", "links": [5], "slot_index": 0}
],
"properties": {"Node name for S&R": "ImageGridComposite"},
"widgets_values": [4, 2, 5]
},
{
"id": 6,
"type": "SaveImage",
"pos": [1700, 280],
"size": [300, 300],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 5}
],
"properties": {"Node name for S&R": "SaveImage"},
"widgets_values": ["trellis2_multiview"]
}
],
"links": [
[1, 1, 0, 2, 0, "IMAGE"],
[2, 2, 0, 4, 0, "TRELLIS_MESH"],
[3, 3, 0, 4, 1, "CAMERA_TRAJECTORY"],
[4, 4, 0, 5, 0, "IMAGE"],
[5, 5, 0, 6, 0, "IMAGE"]
],
"groups": [
{
"title": "1. Load Image",
"bounding": [50, 100, 400, 450],
"color": "#3f789e"
},
{
"title": "2. Convert to 3D",
"bounding": [460, 100, 380, 580],
"color": "#8f5e8e"
},
{
"title": "3. Render Multi-View",
"bounding": [860, 180, 380, 400],
"color": "#5e8f6e"
},
{
"title": "4. Composite Grid",
"bounding": [1260, 180, 380, 300],
"color": "#e67e22"
},
{
"title": "5. Save",
"bounding": [1660, 180, 380, 450],
"color": "#9b59b6"
}
],
"config": {},
"extra": {
"info": {
"name": "Multi-View Preview (TRELLIS2)",
"description": "Generate a grid of 8 viewing angles from a single image.",
"author": "ComfyUI-TRELLIS2_Motion",
"version": "1.0",
"notes": "Requires an ImageGridComposite node or similar for the grid layout."
}
},
"version": 0.4
}
workflows/onion_skinning_demo.json
+162
View File
@@ -0,0 +1,162 @@
{
"last_node_id": 7,
"last_link_id": 6,
"nodes": [
{
"id": 1,
"type": "VHS_LoadVideo",
"pos": [100, 200],
"size": [300, 250],
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{"name": "IMAGE", "type": "IMAGE", "links": [1], "slot_index": 0},
{"name": "frame_count", "type": "INT", "links": null},
{"name": "audio", "type": "VHS_AUDIO", "links": null}
],
"properties": {"Node name for S&R": "VHS_LoadVideo"},
"widgets_values": ["input_video.mp4", "video", 0, false, true, 30]
},
{
"id": 2,
"type": "VideoToMeshes",
"pos": [500, 200],
"size": [280, 180],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 1}
],
"outputs": [
{"name": "mesh_sequence", "type": "TRELLIS_MESH_SEQUENCE", "links": [2], "slot_index": 0}
],
"properties": {"Node name for S&R": "VideoToMeshes"},
"widgets_values": [3, 0, "1024", true, 0]
},
{
"id": 3,
"type": "CameraTrajectory_T2",
"pos": [500, 450],
"size": [280, 350],
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "links": [3], "slot_index": 0}
],
"properties": {"Node name for S&R": "CameraTrajectory_T2"},
"widgets_values": [
"static",
20.0,
-15.0,
0.0,
2.5,
60.0,
0.0,
0.0,
0.0,
false,
30.0,
"y"
]
},
{
"id": 4,
"type": "MeshSequenceRenderer",
"pos": [900, 200],
"size": [280, 260],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{"name": "mesh_sequence", "type": "TRELLIS_MESH_SEQUENCE", "link": 2},
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "link": 3}
],
"outputs": [
{"name": "rendered_video", "type": "IMAGE", "links": [4], "slot_index": 0}
],
"properties": {"Node name for S&R": "MeshSequenceRenderer"},
"widgets_values": [1920, 1080, "pbr", "#000000", "studio", false, 3, 5]
},
{
"id": 5,
"type": "OnionSkinning",
"pos": [1300, 200],
"size": [280, 320],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 4}
],
"outputs": [
{"name": "onion_skinned", "type": "IMAGE", "links": [5], "slot_index": 0}
],
"properties": {"Node name for S&R": "OnionSkinning"},
"widgets_values": [3, 5, 0.3, "linear", "tinted", "#0066FF", "#FF6600", false, "over"]
},
{
"id": 6,
"type": "VHS_VideoCombine",
"pos": [1700, 200],
"size": [300, 300],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 5}
],
"outputs": [
{"name": "Filenames", "type": "VHS_FILENAMES", "links": null}
],
"properties": {"Node name for S&R": "VHS_VideoCombine"},
"widgets_values": ["trellis2_onion_skin", 30, "video/h264-mp4", true, true, ""]
}
],
"links": [
[1, 1, 0, 2, 0, "IMAGE"],
[2, 2, 0, 4, 0, "TRELLIS_MESH_SEQUENCE"],
[3, 3, 0, 4, 1, "CAMERA_TRAJECTORY"],
[4, 4, 0, 5, 0, "IMAGE"],
[5, 5, 0, 6, 0, "IMAGE"]
],
"groups": [
{
"title": "1. Load Video",
"bounding": [50, 100, 400, 400],
"color": "#3f789e"
},
{
"title": "2. Process Frames",
"bounding": [460, 100, 380, 750],
"color": "#8f5e8e"
},
{
"title": "3. Render",
"bounding": [860, 100, 380, 450],
"color": "#5e8f6e"
},
{
"title": "4. Onion Skinning",
"bounding": [1260, 100, 380, 500],
"color": "#e67e22"
},
{
"title": "5. Export",
"bounding": [1660, 100, 380, 450],
"color": "#9b59b6"
}
],
"config": {},
"extra": {
"info": {
"name": "Onion Skinning Demo (TRELLIS2)",
"description": "Demonstrate onion skinning effect on 3D mesh rendered video.",
"author": "ComfyUI-TRELLIS2_Motion",
"version": "1.0"
}
},
"version": 0.4
}
workflows/single_image_novel_view.json
+130
View File
@@ -0,0 +1,130 @@
{
"last_node_id": 5,
"last_link_id": 4,
"nodes": [
{
"id": 1,
"type": "LoadImage",
"pos": [100, 200],
"size": [300, 300],
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{"name": "IMAGE", "type": "IMAGE", "links": [1], "slot_index": 0},
{"name": "MASK", "type": "MASK", "links": null}
],
"properties": {"Node name for S&R": "LoadImage"},
"widgets_values": ["input_image.png", "image"]
},
{
"id": 2,
"type": "Trellis2Predict",
"pos": [500, 200],
"size": [280, 180],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{"name": "image", "type": "IMAGE", "link": 1}
],
"outputs": [
{"name": "mesh", "type": "TRELLIS_MESH", "links": [2], "slot_index": 0}
],
"properties": {"Node name for S&R": "Trellis2Predict"},
"widgets_values": ["1024", 0, true, false, 50000]
},
{
"id": 3,
"type": "CameraTrajectory_T2",
"pos": [500, 450],
"size": [280, 350],
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "links": [3], "slot_index": 0}
],
"properties": {"Node name for S&R": "CameraTrajectory_T2"},
"widgets_values": [
"orbit",
20.0,
-15.0,
0.0,
2.5,
60.0,
0.0,
0.0,
0.0,
false,
30.0,
"y"
]
},
{
"id": 4,
"type": "MeshRenderer",
"pos": [900, 280],
"size": [280, 220],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{"name": "mesh", "type": "TRELLIS_MESH", "link": 2},
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "link": 3}
],
"outputs": [
{"name": "rendered_image", "type": "IMAGE", "links": [4], "slot_index": 0}
],
"properties": {"Node name for S&R": "MeshRenderer"},
"widgets_values": [1920, 1080, "pbr", "#000000", "studio", 1]
},
{
"id": 5,
"type": "SaveImage",
"pos": [1300, 280],
"size": [300, 300],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 4}
],
"properties": {"Node name for S&R": "SaveImage"},
"widgets_values": ["trellis2_novel_view"]
}
],
"links": [
[1, 1, 0, 2, 0, "IMAGE"],
[2, 2, 0, 4, 0, "TRELLIS_MESH"],
[3, 3, 0, 4, 1, "CAMERA_TRAJECTORY"],
[4, 4, 0, 5, 0, "IMAGE"]
],
"groups": [
{
"title": "1. Load Image",
"bounding": [50, 100, 400, 450],
"color": "#3f789e"
},
{
"title": "2. Convert to 3D Mesh & Set Camera",
"bounding": [460, 100, 380, 750],
"color": "#8f5e8e"
},
{
"title": "3. Render & Save",
"bounding": [860, 180, 800, 450],
"color": "#5e8f6e"
}
],
"config": {},
"extra": {
"info": {
"name": "Single Image Novel View (TRELLIS2)",
"description": "Convert a single image to 3D mesh using TRELLIS2 and render from a new camera angle.",
"author": "ComfyUI-TRELLIS2_Motion",
"version": "1.0"
}
},
"version": 0.4
}
workflows/video_angle_shift.json
+134
View File
@@ -0,0 +1,134 @@
{
"last_node_id": 6,
"last_link_id": 5,
"nodes": [
{
"id": 1,
"type": "VHS_LoadVideo",
"pos": [100, 200],
"size": [300, 250],
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{"name": "IMAGE", "type": "IMAGE", "links": [1], "slot_index": 0},
{"name": "frame_count", "type": "INT", "links": null},
{"name": "audio", "type": "VHS_AUDIO", "links": null}
],
"properties": {"Node name for S&R": "VHS_LoadVideo"},
"widgets_values": ["input_video.mp4", "video", 0, false, true, 30]
},
{
"id": 2,
"type": "VideoToMeshes",
"pos": [500, 200],
"size": [280, 180],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 1}
],
"outputs": [
{"name": "mesh_sequence", "type": "TRELLIS_MESH_SEQUENCE", "links": [2], "slot_index": 0}
],
"properties": {"Node name for S&R": "VideoToMeshes"},
"widgets_values": [5, 0, "1024", true, 0]
},
{
"id": 3,
"type": "CameraTrajectory_T2",
"pos": [500, 450],
"size": [280, 350],
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "links": [3], "slot_index": 0}
],
"properties": {"Node name for S&R": "CameraTrajectory_T2"},
"widgets_values": [
"orbit",
15.0,
0.0,
0.0,
2.5,
60.0,
0.0,
0.0,
0.0,
true,
30.0,
"y"
]
},
{
"id": 4,
"type": "MeshSequenceRenderer",
"pos": [900, 280],
"size": [280, 260],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{"name": "mesh_sequence", "type": "TRELLIS_MESH_SEQUENCE", "link": 2},
{"name": "camera_trajectory", "type": "CAMERA_TRAJECTORY", "link": 3}
],
"outputs": [
{"name": "rendered_video", "type": "IMAGE", "links": [4], "slot_index": 0}
],
"properties": {"Node name for S&R": "MeshSequenceRenderer"},
"widgets_values": [1920, 1080, "pbr", "#000000", "studio", false, 3, 5]
},
{
"id": 5,
"type": "VHS_VideoCombine",
"pos": [1300, 280],
"size": [300, 300],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{"name": "images", "type": "IMAGE", "link": 4}
],
"outputs": [
{"name": "Filenames", "type": "VHS_FILENAMES", "links": null}
],
"properties": {"Node name for S&R": "VHS_VideoCombine"},
"widgets_values": ["trellis2_shifted", 30, "video/h264-mp4", true, true, ""]
}
],
"links": [
[1, 1, 0, 2, 0, "IMAGE"],
[2, 2, 0, 4, 0, "TRELLIS_MESH_SEQUENCE"],
[3, 3, 0, 4, 1, "CAMERA_TRAJECTORY"],
[4, 4, 0, 5, 0, "IMAGE"]
],
"groups": [
{
"title": "1. Load Video",
"bounding": [50, 100, 400, 400],
"color": "#3f789e"
},
{
"title": "2. Process & Configure Camera",
"bounding": [460, 100, 380, 750],
"color": "#8f5e8e"
},
{
"title": "3. Render & Export",
"bounding": [860, 180, 800, 450],
"color": "#5e8f6e"
}
],
"config": {},
"extra": {
"info": {
"name": "Video Angle Shift (TRELLIS2)",
"description": "Process video frames through TRELLIS2 and render from shifted camera angle.",
"author": "ComfyUI-TRELLIS2_Motion",
"version": "1.0"
}
},
"version": 0.4
}