40亿参数3D生成模型TRELLIS.2:高效多能,多平台资源与训练代码全公开!
用于 3D 生成的原生紧凑结构化潜空间 trellis2.mp4
因 GitHub 大小限制,此为压缩版。请在我们的项目页面查看高质量视频!
TRELLIS.2 是款先进的大型 3D 生成模型,拥有 40 亿参数,专为高保真图像到 3D 生成设计。它采用新颖“无场”稀疏体素结构 O-Voxel,能重建和生成具复杂拓扑结构、尖锐特征和完整 PBR 材质的任意 3D 资产。
特性
1.高质量、高分辨率与高效率:有 40 亿参数的模型借助普通的 DiTs,能高效且高质量生成高分辨率全纹理资产。它利用空间下采样率为 16 倍的稀疏 3D VAE,将资产编码到紧凑潜空间中。不同分辨率下的总时间*分解(形状 + 材质)如下:
- 512³:约 3 秒(2 秒 + 1 秒)
- 1024³:约 17 秒(10 秒 + 7 秒)
- 1536³:约 60 秒(35 秒 + 25 秒)
*在 NVIDIA H100 GPU 上测试所得。
2.任意拓扑结构处理:O-Voxel 表示法突破等值面场限制,能稳健处理复杂结构,无需有损转换,可处理的结构有:
- 开放表面(如衣物、树叶)
- 非流形几何结构
- 内部封闭结构
3.丰富的纹理建模:除基本颜色外,TRELLIS.2 还能对包括基础颜色、粗糙度、金属度和不透明度在内的任意表面属性建模,支持逼真渲染和透明效果。
4.极简处理流程:数据处理流程简化,可即时转换,无需渲染和优化。
- 纹理网格转换为 O-Voxel:< 10 秒(单 CPU)
- O-Voxel 转换为纹理网格:< 100 毫秒(CUDA)
路线图
- 发布论文
- 发布图像到 3D 推理代码
- 发布预训练检查点(40 亿参数)
- 在 Hugging Face Spaces 上进行演示
- 发布形状条件纹理生成推理代码
- 发布训练代码
安装
前提条件
- 系统:目前代码仅在 Linux 系统上测试过。
- 硬件:需要至少 24GB 内存的 NVIDIA GPU,代码已在 NVIDIA A100 和 H100 GPU 上验证。
- 软件:需要 CUDA Toolkit 编译某些包,推荐版本为 12.4。建议用 Conda 管理依赖项,Python 版本需 3.8 或更高。
安装步骤
1. 克隆仓库:
git clone -b main https://github.com/microsoft/TRELLIS.2.git --recursive
cd TRELLIS.2
2. 安装依赖项:运行以下命令前需注意:
- 添加 `--new-env` 会创建名为 trellis2 的新 Conda 环境。若用现有 Conda 环境,移除该标志。
- 默认情况下,trellis2 环境将使用 CUDA 12.4 的 PyTorch 2.6.0。若用不同版本的 CUDA,移除 `--new-env` 标志并手动安装所需依赖项,安装命令参考 PyTorch 官方文档。
- 若安装多个版本的 CUDA Toolkit,运行命令前将 `CUDA_HOME` 设置为正确版本。例如,若同时安装 CUDA Toolkit 12.4 和 13.0,运行命令前执行 `export CUDA_HOME=/usr/local/cuda-12.4`。
- 默认情况下,代码使用 flash-attn 后端进行注意力计算。对于不支持 flash-attn 的 GPU(如 NVIDIA V100),手动安装 xformers 并在运行代码前将 `ATTN_BACKEND` 环境变量设置为 xformers,更多详情参考最小示例。
安装可能需些时间,因依赖项多,请耐心等待。若遇问题,可尝试逐个安装依赖项,每次指定一个标志。若安装中遇任何问题,可随时创建 issue 或联系我们。
创建名为 trellis2 的新 Conda 环境并安装依赖项:
./setup.sh --new-env --basic --flash-attn --nvdiffrast --nvdiffrec --cumesh --o-voxel --flexgemm
运行 `./setup.sh --help` 可查看 `setup.sh` 详细用法。
`setup.sh` 选项:
- -h, --help:显示此帮助信息
- --new-env:创建新的 Conda 环境
- --basic:安装基本依赖项
- --flash-attn:安装 flash-attention
- --cumesh:安装 cumesh
- --o-voxel:安装 o-voxel
- --flexgemm:安装 flexgemm
- --nvdiffrast:安装 nvdiffrast
- --nvdiffrec:安装 nvdiffrec
预训练权重
预训练模型 TRELLIS.2-4B 可在 Hugging Face 上获取,更多详情参考该平台上的模型卡片。
| 模型 | 参数 | 分辨率 | 链接 |
|---|---|---|---|
| TRELLIS.2-4B | 40 亿 | 512³ - 1536³ | Hugging Face |
使用方法
1. 图像到 3D 生成最小示例
以下是用预训练模型进行 3D 资产生成的示例代码:
import os
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # 可节省 GPU 内存
import cv2
import imageio
from PIL import Image
import torch
from trellis2.pipelines import Trellis2ImageTo3DPipeline
from trellis2.utils import render_utils
from trellis2.renderers import EnvMap
import o_voxel
# 1. 设置环境映射
envmap = EnvMap(torch.tensor(
cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED),
cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda'
))
# 2. 加载管道
pipeline = Trellis2ImageTo3DPipeline.from_pretrained("microsoft/TRELLIS.2-4B")
pipeline.cuda()
# 3. 加载图像并运行
image = Image.open("assets/example_image/T.png")
mesh = pipeline.run(image)[0]
mesh.simplify(16777216) # nvdiffrast 限制
# 4. 渲染视频
video = render_utils.make_pbr_vis_frames(render_utils.render_video(mesh, envmap=envmap))
imageio.mimsave("sample.mp4", video, fps=15)
# 5. 导出为 GLB
glb = o_voxel.postprocess.to_glb(
vertices=mesh.vertices,
faces=mesh.faces,
attr_volume=mesh.attrs,
coords=mesh.coords,
attr_layout=mesh.layout,
voxel_size=mesh.voxel_size,
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
decimation_target=1000000,
texture_size=4096,
remesh=True,
remesh_band=1,
remesh_project=0,
verbose=True
)
glb.export("sample.glb", extension_webp=True)
执行脚本后,将生成以下文件:
- `sample.mp4`:可视化生成 3D 资产的视频,含 PBR 材质和环境光照效果。
- `sample.glb`:提取的 PBR 就绪的 3D 资产,格式为 GLB。
注意:`.glb` 文件默认以不透明模式导出,尽管纹理图中保留了 alpha 通道,但初始时未激活。若启用透明度,需将资产导入 3D 软件,并手动将纹理的 alpha 通道连接到材质的不透明度或 alpha 输入。
`app.py` 提供简单的 Web 演示,用于图像到 3D 资产的生成,用以下命令运行演示:
python app.py
然后在终端显示的地址访问演示页面。
2. PBR 纹理生成
若了解如何为给定 3D 形状生成 PBR 纹理,参考 `example_texturing.py`。此外,可用 `app_texturing.py` 运行 PBR 纹理生成的 Web 演示。
训练
我们提供完整的训练代码库,用户可从头开始训练 TRELLIS.2 或在自定义数据集上微调。
1. 数据准备
训练前,需将原始 3D 资产转换为 O-Voxel 表示,此过程包括网格转换、紧凑结构化潜空间生成和元数据准备。详细的数据预处理和数据集组织说明参考 `data_toolkit/README.md`。
2. 运行训练
训练通过 `train.py` 脚本管理,该脚本接受多个命令行参数配置实验:
- --config:实验配置文件的路径。
- --output_dir:训练输出的目录。
- --load_dir:加载检查点的目录(默认为 `output_dir`)。
- --ckpt:恢复训练的检查点步骤(默认为最新检查点)。
- --data_dir:数据集路径或指定数据集位置的 JSON 字符串。
- --auto_retry:失败时自动重试的次数。
- --tryrun:进行试运行,不进行实际训练。
- --profile:启用训练性能分析。
- --num_nodes:分布式训练的节点数。
- --node_rank:当前节点的排名。
- --num_gpus:每个节点的 GPU 数量(默认为所有可用 GPU)。
- --master_addr:分布式训练的主节点地址。
- --master_port:分布式训练通信的端口。
SC-VAE 训练
-形状 SC-VAE 训练:
python train.py \ --config configs/scvae/shape_vae_next_dc_f16c32_fp16.json \ --output_dir results/shape_vae_next_dc_f16c32_fp16 \ --data_dir "{\"ObjaverseXL_sketchfab\": {\"base\": \"datasets/ObjaverseXL_sketchfab\", \"mesh_dump\": \"datasets/ObjaverseXL_sketchfab/mesh_dumps\", \"dual_grid\": \"datasets/ObjaverseXL_sketchfab/dual_grid_256\", \"asset_stats\": \"datasets/ObjaverseXL_sketchfab/asset_stats\"}}"此命令用 `shape_vae_next_dc_f16c32_fp16.json` 配置在 Objaverse-XL 数据集上训练形状 SC-VAE,训练输出保存到 `results/shape_vae_next_dc_f16c32_fp16`。数据集以 JSON 字符串形式指定,每个数据集条目包含:
- base:数据集的根目录。
- mesh_dump:包含预处理网格转储的目录。
- dual_grid:包含预计算双网格表示的目录。
- asset_stats:包含预计算资产统计信息的目录。
若在更高分辨率下微调模型,可用 `shape_vae_next_dc_f16c32_fp16_ft_512.json` 配置,并相应更新 `finetune_ckpt` 字段和调整数据集路径。
-纹理 SC-VAE 训练:
python train.py \ --config configs/scvae/tex_vae_next_dc_f16c32_fp16.json \ --output_dir results/tex_vae_next_dc_f16c32_fp16 \ --data_dir "{\"ObjaverseXL_sketchfab\": {\"base\": \"datasets/ObjaverseXL_sketchfab\", \"pbr_dump\": \"datasets/ObjaverseXL_sketchfab/pbr_dumps\", \"pbr_voxel\": \"datasets/ObjaverseXL_sketchfab/pbr_voxels_256\", \"asset_stats\": \"datasets/ObjaverseXL_sketchfab/asset_stats\"}}"流模型训练
-稀疏结构流模型训练:
python train.py \ --config configs/gen/ss_flow_img_dit_1_3B_64_bf16.json \ --output_dir results/ss_flow_img_dit_1_3B_64_bf16 \ --data_dir "{\"ObjaverseXL_sketchfab\": {\"base\": \"datasets/ObjaverseXL_sketchfab\", \"ss_latent\": \"datasets/ObjaverseXL_sketchfab/ss_latents/ss_enc_conv3d_16l8_fp16_64\", \"render_cond\": \"datasets/ObjaverseXL_sketchfab/renders_cond\"}}"此命令用指定的配置文件在 Objaverse-XL 数据集上训练稀疏结构流模型,输出保存到 `results/ss_flow_img_dit_1_3B_64_bf16`。数据集配置包括:
- base:数据集根目录。
- ss_latent:包含预计算稀疏结构潜空间的目录。
- render_cond:包含条件渲染图像的目录。
形状和纹理生成的第二阶段和第三阶段流模型可用以下配置训练:
-形状流模型:
python train.py \ --config configs/gen/slat_flow_img2shape_dit_1_3B_512_bf16.json \ --output_dir results/slat_flow_img2shape_dit_1_3B_512_bf16 \ --data_dir "{\"ObjaverseXL_sketchfab\": {\"base\": \"datasets/ObjaverseXL_sketchfab\", \"shape_latent\": \"datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_512\", \"render_cond\": \"datasets/ObjaverseXL_sketchfab/renders_cond\"}}"-纹理流模型:
python train.py \ --config configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16.json \ --output_dir results/slat_flow_imgshape2tex_dit_1_3B_512_bf16 \ --data_dir "{\"ObjaverseXL_sketchfab\": {\"base\": \"datasets/ObjaverseXL_sketchfab\", \"shape_latent\": \"datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_512\", \"pbr_latent\": \"datasets/ObjaverseXL_sketchfab/pbr_latents/tex_enc_next_dc_f16c32_fp16_512\", \"render_cond\": \"datasets/ObjaverseXL_sketchfab/renders_cond\"}}"通过更新以下配置文件中的 `finetune_ckpt` 字段并相应调整数据集路径,可进行更高分辨率的微调:
- slat_flow_img2shape_dit_1_3B_512_bf16_ft1024.json
- slat_flow_imgshape2tex_dit_1_3B_512_bf16_ft1024.json
相关包
TRELLIS.2 基于我们团队开发的几个专业高性能包构建:
- O-Voxel:核心库,处理纹理网格和 O-Voxel 表示之间的转换逻辑,确保即时双向转换。
- FlexGEMM:基于 Triton 的高效稀疏卷积实现,可快速处理稀疏体素结构。
- CuMesh:CUDA 加速的网格实用工具,用于高速后处理、重新网格化、简化和 UV 展开。
许可证
此模型和代码遵循 MIT 许可证发布。注意,某些依赖项遵循单独的许可条款:
- nvdiffrast:用于渲染生成的 3D 资产,该包遵循其自身的许可证。
- nvdiffrec:实现 PBR 材质的分裂求和渲染器,该包遵循其自身的许可证。
引用
若发现此模型对您的研究有用,请引用我们的工作:
@article{xiang2025trellis2,
title={Native and Compact Structured Latents for 3D Generation},
author={Xiang, Jianfeng and Chen, Xiaoxue and Xu, Sicheng and Wang, Ruicheng and Lv, Zelong and Deng, Yu and Zhu, Hongyuan and Dong, Yue and Zhao, Hao and Yuan, Nicholas Jing and Yang, Jiaolong},
journal={Tech report},
year={2025}
}