VITRON：用于理解、生成、分割和编辑的统一像素级视觉大模型

NeurIPS 2024 论文

Hao Fei$^{1,2}$，Shengqiong Wu$^{1,2}$，Hanwang Zhang$^{1,3}$，Tat-Seng Chua$^{2}$，Shuicheng Yan$^{1}$

▶ $^{1}$ 新加坡 Skywork AI ▶ $^{2}$ 新加坡国立大学 ▶ $^{3}$ 南洋理工大学

📰 新闻

• [2024.09.26] 我们很高兴这项工作已被 NeurIPS 2024 接受。
• [2024.07.19] 我们发布了为 Text Invocation Instruction Tuning 构建的 数据集。
• [2024.06.28] 🤗 我们发布了检查点，更多详情请参阅 README。
• [2024.04.04] 👀👀👀 我们的 Vitron 现已上线！欢迎 关注 👀 此仓库以获取最新更新。

😮 亮点

现有的视觉大模型可能仍然面临一些挑战，例如对实例级别的理解较为表面化、未能统一支持图像和视频，以及在各类视觉任务上的覆盖不足。为了弥补这些不足，我们提出了 Vitron，一个通用的像素级视觉大模型，旨在对静态图像和动态视频内容进行全面的理解（感知与推理）、生成、分割（定位与跟踪）以及编辑（修复填充）。

🛠️ 系统要求与安装

• Python >= 3.8
• Pytorch == 2.1.0
• CUDA 版本 >= 11.8
• 安装所需包：

git clone https://github.com/SkyworkAI/Vitron
cd Vitron
conda create -n vitron python=3.10 -y
conda activate vitron
pip install --upgrade pip 
pip install -e .
pip install -e ".[train]"
pip install flash-attn --no-build-isolation
pip install decord opencv-python git+https://github.com/facebookresearch/pytorchvideo.git@28fe037d212663c6a24f373b94cc5d478c8c1a1d

🔥🔥🔥 安装或运行失败？ 🔥🔥🔥

1. 运行 ffmpeg 时出现 Unknown encoder 'x264' 错误：

   • 尝试重新安装 ffmpeg：

   conda uninstall ffmpeg
   conda install -c conda-forge ffmpeg   # `-c conda-forge` 不可省略
   

2. 安装 detectron2 失败时，尝试以下命令：

   python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'
   

   或参考此 网站。

3. gradio 出现错误。由于 gradio>=4.0.0 有较大更新，请确保安装与 requirements.txt 中相同版本的 gradio。

4. deepspeed 出现错误。如果您正在微调模型，可能会遇到以下错误：

   FAILED: cpu_adam.so
   /usr/bin/ld: cannot find -lcurand
   

   该错误是由安装 deepspeed 时软链接设置不当引起的。请尝试以下命令解决：

   cd ~/miniconda3/envs/vitron/lib
   ls -al libcurand*  # 检查链接
   rm libcurand.so   # 删除错误链接
   ln -s libcurand.so.10.3.5.119 libcurand.so  # 重建正确链接
   

   再次确认：

   python 
   from deepspeed.ops.op_builder import CPUAdamBuilder
   ds_opt_adam = CPUAdamBuilder().load()  # 如果加载成功，则说明 deepspeed 已成功安装。
   

代码结构

.
├── assets
├── checkpoints    # 保存预训练检查点
├── data          
├── examples
├── modules        # 项目中使用的各个模块
│   ├── GLIGEN
│   ├── i2vgen-xl
│   ├── SEEM
│   └── StableVideo
├── scripts
└── vitron
    ├── model
    │   ├── language_model
    │   ├── multimodal_encoder
    │   ├── multimodal_projector
    │   └── region_extractor
    └── train

👍 部署 Gradio 演示

• 首先，您需要准备好检查点，更多详情请参阅 README。
• 然后，您可以在本地运行演示：

python app.py

微调您的模型

• 首先，准备数据集。 我们发布了为 Invocation-oriented Instruction Tuning 构建的数据集。更多详情请参阅 README。
• 然后，在 finetune_lora.sh 中修改 image/video/data 路径。

JSON_FOLDER=None
IMAGE_FOLDER=None
VIDEO_FOLDER=None
DATA_PATH="./data/data.json"

• 接着，准备 检查点。
• 最后，运行代码：

bash scripts/fine_lora.sh

🙌 相关项目

您可以参考作为我们框架和代码库基础的相关工作， Vicuna， SEEM， i2vgenxl， StableVideo，以及 Zeroscope。 我们还部分借鉴了 Video-LLaVA, 和 LanguageBind 的灵感。 感谢他们出色的工作。

🔒 许可证

• 本项目的大部分内容根据 LICENSE 文件中的 Apache 2.0 许可协议发布。
• 本服务为研究预览版，仅供非商业用途使用，同时需遵守 LLaMA 模型的 许可证，OpenAI 生成数据的 使用条款，以及 ShareGPT 的 隐私政策。如发现任何潜在违规行为，请联系我们。

✏️ 引用

如果您在研究中认为我们的论文和代码有用，请考虑给予一颗星 :star: 和引用 :pencil:。

@inproceedings{fei2024vitron,
  title={VITRON: A Unified Pixel-level Vision LLM for Understanding, Generating, Segmenting, Editing},
  author={Fei, Hao and Wu, Shengqiong and Zhang, Hanwang and Chua, Tat-Seng and Yan, Shuicheng},
  year={2024},
  journal={Proceedings of the Advances in neural information processing systems},
}

✨ 星标历史

Vitron 快速上手指南

Vitron 是一个统一的像素级视觉大语言模型（Vision LLM），支持图像和视频的理解、生成、分割及编辑任务。本指南将帮助开发者快速完成环境配置并运行基础功能。

环境准备

在开始之前，请确保您的系统满足以下硬件和软件要求：

• 操作系统: Linux (推荐 Ubuntu)
• Python: >= 3.8 (推荐 3.10)
• PyTorch: == 2.1.0
• CUDA: >= 11.8
• 显存: 建议 24GB 及以上以运行完整功能

安装步骤

1. 克隆项目与创建环境

首先克隆代码仓库并创建独立的 Conda 环境：

git clone https://github.com/SkyworkAI/Vitron
cd Vitron
conda create -n vitron python=3.10 -y
conda activate vitron

2. 安装依赖包

升级 pip 并安装核心依赖。为了提升国内下载速度，建议临时使用清华或阿里镜像源（如下所示）：

pip install --upgrade pip 
# 使用国内镜像加速安装
pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -e ".[train]" -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install flash-attn --no-build-isolation -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install decord opencv-python git+https://github.com/facebookresearch/pytorchvideo.git@28fe037d212663c6a24f373b94cc5d478c8c1a1d -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：如果 flash-attn 编译失败，请确保已正确安装 CUDA Toolkit 且版本匹配。若遇到 ffmpeg 编码错误，可尝试执行 conda uninstall ffmpeg 后运行 conda install -c conda-forge ffmpeg。

3. 准备预训练权重

运行模型前需下载预训练检查点。请访问 checkpoints/README.md 获取权重文件，并将其放置在项目指定的 checkpoints 目录下。

基本使用

启动本地演示界面 (Gradio Demo)

这是体验 Vitron 功能最直观的方式。确保已完成上述权重准备工作后，在项目根目录运行：

python app.py

运行成功后，终端会显示本地访问地址（通常为 http://127.0.0.1:7860）。在浏览器中打开该地址，即可上传图像或视频，通过自然语言指令进行理解、分割、生成或编辑操作。

微调模型 (可选)

如果您希望使用自定义数据集进行微调，请按以下步骤操作：

1. 准备数据集（参考 data/README.md）。
2. 修改脚本 scripts/finetune_lora.sh 中的数据路径：

   JSON_FOLDER=None
   IMAGE_FOLDER=None
   VIDEO_FOLDER=None
   DATA_PATH="./data/data.json"
   

3. 执行微调命令：

   bash scripts/fine_lora.sh