microsoft/TRELLIS.2

GitHub: microsoft/TRELLIS.2

TRELLIS.2 是一个基于扩散变换器的大规模图像到 3D 生成模型,能高效创建复杂拓扑和 PBR 材质的 3D 资产。

Stars: 8181 | Forks: 989

![](https://static.pigsec.cn/wp-content/uploads/repos/2026/05/9bc2d0af9b202428.webp) # 用于3D生成的原生紧凑结构化潜码 Paper Hugging Face Project Page License https://github.com/user-attachments/assets/63b43a7e-acc7-4c81-a900-6da450527d8f *(由于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. 丰富的纹理建模 ### 4. 极简处理 数据处理流程经过简化,可实现即时转换,完全**无需渲染**和**无需优化**。 * **< 10秒** (单CPU): 带纹理网格 → O-Voxel * **< 100毫秒** (CUDA): O-Voxel → 带纹理网格 ## 🗺️ 路线图 - [x] 论文发布 - [x] 发布图像到3D推理代码 - [x] 发布预训练检查点(40亿参数) - [x] Hugging Face Spaces 演示 - [x] 发布形状条件化纹理生成推理代码 - [x] 发布训练代码 ## 🛠️ 安装说明 ### 前置条件 - **系统**:该代码目前仅在 **Linux** 上测试。 - **硬件**:需要一块至少24GB显存的NVIDIA GPU。代码已在NVIDIA A100和H100 GPU上验证。 - **软件**: - 编译某些包需要 [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit-archive)。推荐版本为12.4。 - 推荐使用 [Conda](https://docs.anaconda.com/miniconda/install/#quick-command-line-install) 管理依赖。 - 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](https://pytorch.org/get-started/previous-versions/) 的安装命令。 - 如果您安装了多个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`。详见[最小示例](#minimal-example)。 - 由于依赖项众多,安装过程可能需要一些时间,请耐心等待。如果遇到任何问题,可以尝试逐个安装依赖项,每次指定一个标志。 - 如果在安装过程中遇到任何问题,请随时提交 issue 或联系我们。 创建一个名为 `trellis2` 的新 conda 环境并安装依赖项: . ./setup.sh --new-env --basic --flash-attn --nvdiffrast --nvdiffrec --cumesh --o-voxel --flexgemm `setup.sh` 的详细用法可通过运行 `. ./setup.sh --help` 查看。 用法: 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](https://huggingface.co/microsoft/TRELLIS.2-4B) | ## 🚀 用法 ### 1. 图像到3D生成 #### 最小示例 这里是一个[示例](example.py),展示如何使用预训练模型进行3D资产生成。 ``` import os os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # Can save GPU memory 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 # 设置 Environment Map envmap = EnvMap(torch.tensor( cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda' )) # 加载 Pipeline pipeline = Trellis2ImageTo3DPipeline.from_pretrained("microsoft/TRELLIS.2-4B") pipeline.cuda() # 加载图像并运行 image = Image.open("assets/example_image/T.png") mesh = pipeline.run(image)[0] mesh.simplify(16777216) # nvdiffrast limit # 渲染视频 video = render_utils.make_pbr_vis_frames(render_utils.render_video(mesh, envmap=envmap)) imageio.mimsave("sample.mp4", video, fps=15) # 导出为 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`:一个视频,可视化展示生成的带有PBR材质和环境光照的3D资产。 - `sample.glb`:提取的、可直接用于PBR的3D资产,格式为GLB。 **注意:** `.glb` 文件默认以 `OPAQUE` 模式导出。虽然Alpha通道保留在纹理贴图中,但最初并未激活。要启用透明度,请将资产导入您的3D软件,并手动将纹理的Alpha通道连接到材质的不透明度或Alpha输入。 #### 网络演示 [app.py](app.py) 提供了一个用于图像到3D资产生成的简单网络演示。您可以使用以下命令运行演示: ``` python app.py ``` 然后,您可以在终端所示地址访问演示。 ### 2. PBR纹理生成 请参考 [example_texturing.py](example_texturing.py) 了解如何为给定3D形状生成PBR纹理的示例。此外,您可以使用 [app_texturing.py](app_texturing.py) 运行PBR纹理生成的网络演示。 ## 🏋️ 训练 我们提供完整的训练代码库,使用户能够从头开始训练 **TRELLIS.2** 或在自定义数据集上进行微调。 ### 1. 数据准备 在训练之前,原始3D资产必须转换为 **O-Voxel** 表示。此过程包括网格转换、紧凑结构化潜码生成和元数据准备。 ### 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`: 包含条件渲染图像的目录。 用于形状和纹理生成的第二阶段和第三阶段流模型可以使用以下配置进行训练: * 形状流:`slat_flow_img2shape_dit_1_3B_512_bf16.json` * 纹理流:`slat_flow_imgshape2tex_dit_1_3B_512_bf16.json` 示例命令: ``` # Shape flow 模型 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\"}}" # Texture flow 模型 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):** 处理带纹理网格与O-Voxel表示之间转换逻辑的核心库,确保即时双向转换。 * **[FlexGEMM](https://github.com/JeffreyXiang/FlexGEMM):** 基于Triton的高效稀疏卷积实现,能够快速处理稀疏体素结构。 * **[CuMesh](https://github.com/JeffreyXiang/CuMesh):** CUDA加速的网格工具,用于高速后处理、重网格化、减面和UV展开。 ## ⚖️ 许可证 该模型和代码在 **[MIT 许可证](LICENSE)** 下发布。 请注意,某些依赖项在单独的许可条款下运行: - [**nvdiffrast**](https://github.com/NVlabs/nvdiffrast):用于渲染生成的3D资产。该软件包受其自身[许可证](https://github.com/NVlabs/nvdiffrast/blob/main/LICENSE.txt)约束。 - [**nvdiffrec**](https://github.com/NVlabs/nvdiffrec):实现了用于PBR材质的分裂求和渲染器。该软件包受其自身[许可证](https://github.com/NVlabs/nvdiffrec/blob/main/LICENSE.txt)约束。 ## 📚 引用 如果您觉得这个模型对您的研究有用,请引用我们的工作: ``` @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} } ```
标签:3D建模, 3D生成, 3D资产生成, 3D重建, Diffusion Transformers, O-Voxel, PBR材质, Vectored Exception Handling, 人工智能, 任意拓扑处理, 保真度高, 凭据扫描, 图像到3D, 增强现实, 大型模型, 快速生成, 深度学习, 游戏开发, 潜空间学习, 生成式AI, 生成模型, 用户模式Hook绕过, 神经网络, 稀疏3D VAE, 稀疏体素表示, 索引, 紧凑潜空间, 结构化潜变量, 虚拟现实, 视觉生成, 计算机图形学, 计算机视觉, 逆向工具, 高分辨率生成, 高效生成