OpenLRM：开源大型重建模型

新闻

[2024.03.13] 更新训练代码并发布OpenLRM v1.1.1。
[2024.03.08] 我们发布了用于渲染Objaverse图像的核心blender脚本。
[2024.03.05] Huggingface演示现在默认使用openlrm-mix-base-1.1模型。请参考模型卡以了解更新的模型架构和训练设置的详细信息。
[2024.03.04] 版本更新v1.1。发布在Objaverse和MVImgNet上训练的模型权重。代码库进行了主要的重构，使其在可用性和扩展性方面得到提升。详细信息请参考v1.1.0。
[2024.01.09] 更新了所有在Objaverse上训练的v1.0模型。请参考HF模型并覆盖之前的模型权重。
[2023.12.21] Hugging Face演示上线了。赶紧试试吧！
[2023.12.20] 发布在Objaverse上训练的基础和大型模型的权重。
[2023.12.20] 我们发布了项目OpenLRM，这是LRM论文的开源实现。

安装

git clone https://github.com/3DTopia/OpenLRM.git
cd OpenLRM

环境

首先安装OpenLRM的依赖。
```
pip install -r requirements.txt
```
然后请按照xFormers安装指南开启DINOv2 encoder内部的内存高效注意力机制。

快速开始

预训练模型

模型权重发布在Hugging Face。
当你第一次运行推理脚本时，权重会自动下载。
在使用这些权重之前，请注意许可证。

模型	训练数据	层数	特征维度	三重维度	输入分辨率	链接
openlrm-obj-small-1.1	Objaverse	12	512	32	224	HF
openlrm-obj-base-1.1	Objaverse	12	768	48	336	HF
openlrm-obj-large-1.1	Objaverse	16	1024	80	448	HF
openlrm-mix-small-1.1	Objaverse + MVImgNet	12	512	32	224	HF
openlrm-mix-base-1.1	Objaverse + MVImgNet	12	768	48	336	HF
openlrm-mix-large-1.1	Objaverse + MVImgNet	16	1024	80	448	HF

可以在model_card.md文档中找到包含更多详细信息的模型卡。

准备图像

我们在assets/sample_input中放置了一些示例输入，你可以快速尝试。
准备RGBA图像或带有白色背景的RGB图像（使用一些背景移除工具，例如Rembg，Clipdrop）。

推理

运行推理脚本以获取3D资产。
你可以通过设置标志EXPORT_VIDEO=true和EXPORT_MESH=true来指定生成哪种形式的输出。
请根据你想要使用的模型设置默认INFER_CONFIG。例如，基础模型使用infer-b.yaml，小型模型使用infer-s.yaml。

一个示例用法如下：

# 示例用法
EXPORT_VIDEO=true
EXPORT_MESH=true
INFER_CONFIG="./configs/infer-b.yaml"
MODEL_NAME="zxhezexin/openlrm-mix-base-1.1"
IMAGE_INPUT="./assets/sample_input/owl.png"

python -m openlrm.launch infer.lrm --infer $INFER_CONFIG model_name=$MODEL_NAME image_input=$IMAGE_INPUT export_video=$EXPORT_VIDEO export_mesh=$EXPORT_MESH

提示

建议使用的PyTorch版本为>=2.1。代码在PyTorch 2.1.2版本下开发和测试。
如果你遇到CUDA OOM问题，请尝试减少推理配置中的frame_size。
如果xFormers确实在工作，你应该能看到UserWarning: xFormers is available。

训练

配置

我们在configs/accelerate-train.yaml下提供了一个示例accelerate配置文件，默认使用8个GPU并进行bf16混合精度训练。
你可以修改配置文件以适应你自己的环境。

数据准备

我们提供了用于渲染Objaverse图像的核心Blender脚本。
请参考Objaverse渲染以获取包括分布式渲染在内的其他脚本。

运行训练

在configs/train-sample.yaml下提供了一个示例训练配置文件。
请用你自己的路径替换配置文件中的数据相关路径，并自定义训练设置。

一个示例训练用法如下：

# 示例用法
ACC_CONFIG="./configs/accelerate-train.yaml"
TRAIN_CONFIG="./configs/train-sample.yaml"

accelerate launch --config_file $ACC_CONFIG -m openlrm.launch train.lrm --config $TRAIN_CONFIG

推理训练过的模型

推理管道与huggingface工具兼容，以提高便利性。

你需要通过运行以下脚本将训练检查点转换为推理模型。

python scripts/convert_hf.py --config <YOUR_EXACT_TRAINING_CONFIG> convert.global_step=null

转换后的模型将默认保存在exps/releases下，可以按照推理指南进行推理。

致谢

感谢原论文的作者们的精彩工作！特别感谢张凯和洪一聪在重现过程中提供的帮助。
本项目得到了上海人工智能实验室提供的计算资源支持。
本项目得到了刘子蔚和贾佳亚的指导。

引用

如果你发现这项工作对你的研究有所帮助，请考虑引用：

@article{hong2023lrm,
  title={Lrm: Large reconstruction model for single image to 3d},
  author={Hong, Yicong and Zhang, Kai and Gu, Jiuxiang and Bi, Sai and Zhou, Yang and Liu, Difan and Liu, Feng and Sunkavalli, Kalyan and Bui, Trung and Tan, Hao},
  journal={arXiv preprint arXiv:2311.04400},
  year={2023}
}

@misc{openlrm,
  title = {OpenLRM: Open-Source Large Reconstruction Models},
  author = {Zexin He and Tengfei Wang},
  year = {2023},
  howpublished = {\url{https://github.com/3DTopia/OpenLRM}},
}