TRELLIS.2 LoRA微调：提升3D生成细节精度

简介

在机器人、模拟仿真及媒体娱乐领域，3D 生成技术有着广泛的应用前景。然而，确保生成的对象符合严格的规则并实现模型与数据表示的标准化仍具挑战性。目前存在多种文件格式（如 GLB、PLY）和表示方法（如神经辐射场 NeRF、3D 高斯泼溅 3DGS），每种方法各有优劣。

TRELLIS 引入了一种称为“统一隐藏表示”的共同内部表示形式。其解码器可将该表示转换为所需输出格式的 3D 模型，从而避免了为每种表示格式使用不同模型的需求。随后发布的 TRELLIS.2 进一步优化了模型对 3D 结构的表示能力。尽管性能有所提升，但这两种模型缺乏适应特定任务的简便微调方法。

改进目标

通用模型（如 TRELLIS.2）在处理特定应用中的细节时可能存在偏差。例如，用户可能需要精确的家具模型几何结构，而未经微调的预训练模型往往难以满足这种精度要求。在针对 Google 扫描对象 (GSO) 数据集中 188 双鞋子的实验中，预训练模型生成的鞋子外观整体正确，但鞋底细节往往不准确。团队测试了 LoRA 技术是否能改善这一情况。

此次实验旨在利用 GSO 数据集提升 TRELLIS.2 在特定对象类别上的表现。重点在于明确哪些管道组件需要微调，哪些应保持冻结状态。由于管道将结构、详细几何形状和材料分离开来，因此可以仅针对任务中重要的组件进行精细调整。

技术选型

LoRA (低秩适配)

LoRA 通过训练小型适配器矩阵来微调模型，同时保持原始权重不变。这种方法显著减少了可训练参数的数量以及梯度计算和优化器状态所需的内存。该项目使用 Hugging Face 的 PEFT 库来添加和训练这些适配器。

Hugging Face Jobs

工作负载在云硬件上以容器化形式运行，用于执行预处理和训练任务。若需使用本地 GPU，可参考 Docker 命令。开始之前需确保具备以下条件：

• 有效的 HF_TOKEN。
• Hugging Face 账户拥有访问 DINOv3 模型的权限（DINOv3 为受限模型）。
• 若使用自定义数据集，需提前将其上传至 Hub。

TRELLIS.2 工作流程解析

在进行训练前，了解 TRELLIS.2 如何生成 3D 对象以及 LoRA 适配器的插入位置至关重要。TRELLIS.2 包含五个主要阶段：去除图像背景、编码图像、预测粗略结构、生成详细几何形状以及生成材料。最终，管道将结果输出为纹理网格。

输入图像编码

首先，BiRefNet 从输入图像中移除背景。接着，DINOv3-L 将图像编码为特征令牌（feature tokens）。生成模型利用这些令牌引导 3D 对象的构建。

稀疏结构生成 (Sparse Structure Flow)

Sparse Structure Flow DiT 负责预测物体的粗略形状。这是一个扩散 Transformer (DiT)，产生尺寸为 16×3×8 的小型潜在表示。随后，变分自编码器 (VAE) 解码器将该潜在表示转换为 64 × 64 × 64 的占用网格。网格中的每个体素（voxel）指示该位置是否属于物体的一部分。在此阶段，LoRA 适配器被添加到模型中，因为其输出直接引导后续的几何生成阶段。

O-Voxel 与结构化潜在表示

O-Voxel 是 TRELLIS.2 采用的 3D 格式，这是一种稀疏体素网格，仅在物体表面保留体素。每个体素存储局部几何信息（如网格顶点）和材料属性。由于格式紧凑，CPU 可在数秒内完成网格与 O-Voxel 之间的转换。结构化稀疏潜在表示来自 TRELLIS，它将 3D 数据存储在稀疏张量中以提高计算速度。模型仅处理网格中约 10% 的体素，避免了对空体素的无效计算。因此，Shape Flow DiT 和 Texture Flow DiT 采用稀疏 Transformer 架构，而 Sparse Structure Flow DiT 则采用密集 Transformer 架构。

几何生成 (Shape Flow)

Shape Flow DiT 结合占用网格和 DINOv3 特征令牌生成潜在表示。另一个 VAE 解码器随后将该潜在表示转换为分辨率为 512 × 512 × 512 的 O-Voxel 几何形状。在此阶段，LoRA 适配器同样被添加到 DiT 中，因为该模块控制着物体的具体形状。

材料生成与网格创建

在最后阶段，Texture Flow DiT 以隐式形状和 DINOv3 特征为输入，定义几何表面的材料和纹理。VAE 解码器将潜在表示转换为每个体素的基础颜色、金属度、粗糙度和 Alpha 值。此阶段未添加 LoRA 适配器，原因有二：一是数据集缺乏训练所需的独立材料地图（仅有混合了光照效应的漫反射色彩图）；二是问题主要集中在鞋子形状上，基础模型已具备较好的颜色和纹理处理能力。

最终，管道将几何形状与材料属性合并为带纹理的网格，并导出为 GLB 文件。

微调策略规划

组件选择

仅向 Sparse Structure Flow DiT 和 Shape Flow DiT 添加 LoRA 适配器，因为这两个模块控制物体的结构和几何形状。Texture Flow DiT 保持冻结状态以保留其现有能力。由于数据集中的色彩图混合了环境光照效应，若在此类地图上训练 Texture Flow DiT，模型可能错误地将光照学习为表面属性，甚至导致灾难性遗忘。

为何选用 LoRA

实验子集包含 1,030 个 GSO 对象中的 188 双鞋。对于如此小规模的数据集，LoRA 比全量微调更为合适。全量微调不仅成本高昂，还可能覆盖模型已有的先验知识优势。配置方案为两个目标 DiT 共训练约 930 万个适配器参数。在 Sparse Structure Flow DiT 的注意力和前馈（MLP）层中使用 rank 8 适配器；在 Shape Flow DiT 的注意力层中使用 rank 16 适配器。Rank 决定了适配器矩阵的大小，较高的 rank 允许适配器学习更多信息。

执行流程

3D 模型的设置通常涉及复杂的依赖关系。项目提供了两个 Docker 镜像（可在 Hugging Face Jobs 或本地机器运行）：一个用于预处理，一个用于微调。

1. 准备数据：将 GLB 或 OBJ 网格转换为 TRELLIS.2 训练所需的隐藏表示，并对输入图像进行渲染。
2. 启动适配器：使用准备好的数据集和 LoRA 配置为每个 DiT 运行任务。
3. 代码解析：检查容器如何加载数据、添加适配器及训练模型。

预处理后的数据集会被保存，以便所有训练运行重复使用。

操作指南

在 Hugging Face Jobs 上运行

预处理

预处理容器使用 Blender 渲染每个对象的 16 个视图作为条件图像，并通过 xvfb 运行实用程序。容器还通过预处理管道处理每一个物体。对于 188 个物体，在 A100-large 实例上大约需要 1 小时。

export HF_TOKEN=hf_xxx  

hf jobs run -d --flavor a100-large --timeout 6h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-prep:latest \
  trellis2-prep \
    --meshes your-username/my-meshes \
    --push-to your-username/my-meshes-trellis2 \
    --private --holdout 0.05 \
    --views 16 --workers 8

关键参数说明：

LoRA 微调

训练容器加载预处理数据集和配置预设，并使用 Trackio 记录训练指标。ss 模式仅微调 Sparse Structure Flow DiT，shape 模式微调 Shape Flow DiT。

export HF_TOKEN=hf_xxx  

hf jobs run -d --flavor a100-large --timeout 12h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train --config ss \
    --latents your-username/my-meshes-trellis2 \
    --push-to your-username/my-trellis2-lora \
    --run-name ss_lora \
    --push-path ss_lora \
    --trackio-space your-username/trellis2-trackio

hf jobs run -d --flavor a100-large --timeout 12h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train --config shape \
    --latents your-username/my-meshes-trellis2 \
    --push-to your-username/my-trellis2-lora \
    --run-name shape_lora \
    --push-path shape_lora \
    --trackio-space your-username/trellis2-trackio \
    --set max_steps=3000 i_save=500

若设置 --push-to，容器会将适配器检查点上传至 Hub。也可直接使用已处理的示例数据集：

export HF_TOKEN=hf_xxx  

hf jobs run -d --flavor a100-large --timeout 12h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train \
    --config ss \
    --latents suvadityamuk/gso-footwear-trellis2-latents \
    --push-to your-username/trellis2-footwear-lora \
    --run-name ss_lora \
    --push-path ss_lora \
    --trackio-space your-username/trellis2-trackio \
    --set max_steps=3000 i_save=500


hf jobs run -d --flavor a100-large --timeout 12h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train \
    --config shape \
    --latents suvadityamuk/gso-footwear-trellis2-latents \
    --push-to your-username/trellis2-footwear-lora \
    --run-name shape_lora \
    --push-path shape_lora \
    --trackio-space your-username/trellis2-trackio \
    --set max_steps=3000 i_save=500

本地 Docker 运行

预处理

命令与 Hugging Face Jobs 类似，仅需调整 Docker 选项：

docker run --gpus all --rm --shm-size 16g -e HF_TOKEN=hf_xxx \
  -v hf-cache:/root/.cache/huggingface -v "$PWD/work":/workspace/work \
  cr.hf.co/suvadityamuk/trellis2-finetuning-prep:latest \
  trellis2-prep --meshes your-username/my-meshes --push-to your-username/my-meshes-trellis2

LoRA 微调

docker run --gpus all --rm --shm-size 16g \
  -v hf-cache:/root/.cache/huggingface -v "$PWD/work":/workspace/work \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train --config ss --latents your-username/my-meshes-trellis2 \
    --push-to your-username/my-trellis2-lora --trackio-space your-username/trellis2-trackio


docker run --gpus all --rm --shm-size 16g \
  -v hf-cache:/root/.cache/huggingface -v "$PWD/work":/workspace/work \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train --config shape --latents your-username/my-meshes-trellis2 \
    --push-to your-username/my-trellis2-lora --trackio-space your-username/trellis2-trackio