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。
• 当你第一次运行推理脚本时，权重会自动下载。
• 在使用这些权重之前，请注意许可证。

可以在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}},
}

许可证

• 整个OpenLRM项目根据Apache许可证，版本2.0授权，而某些组件则由NVIDIA的专有许可证涵盖。用户有责任遵守每个组件的相应许可条款。
• 模型权重根据创作共用署名-非商业性使用4.0国际许可协议授权。它们仅提供用于研究目的，并且不能用于商业用途。