PromptFix：你提供提示，我们修复照片

NeurIPS 2024

   

---

本仓库提供了 PromptFix 的官方 PyTorch 实现，包括预训练权重、训练与推理代码，以及我们用于训练的精选数据集。

📢 PromptFix 旨在根据人类指令处理退化图像并移除不需要的元素。它支持广泛的任务，例如：

• 🎨 上色
• 🧹 物体移除
• 🌫️ 去雾
• 💨 去模糊
• 🖼️ 水印移除
• ❄️ 雪地清除
• 🌙 低光增强

PromptFix 基于扩散模型架构构建，在保留原始结构的同时，通过 20 步去噪过程实现了卓越的图像缺陷修复效果，并且能够有效泛化到不同宽高比的图像。

目录

• 环境搭建
• 推理
• 下载数据集
  • 数据集构成
• 🧑‍💻 训练
• 📝 引用 PromptFix
• 🙏 致谢
• ⚠️ 免责声明

环境搭建

请按照以下步骤克隆仓库、设置环境并安装依赖项。代码已在 Python 3.10 上测试通过。

git clone https://github.com/yeates/PromptFix.git
cd PromptFix
conda create -n promptfix python=3.10 -y
conda activate promptfix
pip install -r requirements.txt

推理

要处理默认的示例图像，请运行以下命令。预训练模型权重将自动从 Hugging Face 下载，并放置在 checkpoints/ 目录下：

bash scripts/inference.sh

下载数据集

我们精心整理了一个超过 100 万 个样本的 训练数据集。每个样本包含成对的图像以及指令和辅助文本提示。该数据集覆盖了多种低层图像处理任务。

要在项目根目录下下载数据集，请运行以下命令：

bash scripts/download_promptfix_dataset.sh

数据集构成

数据集包含以下任务：

任务	百分比
🎨 上色	29.3%
🌙 低光增强	20.7%
🖼️ 水印移除	12.4%
🧹 物体移除	11.9%
❄️ 雪地清除	9.7%
🌫️ 去雾	8.9%
💨 去模糊	7.1%
总计	100%

注意： 数据集被打包为 Parquet 文件，共分为 100 个部分。每个部分可以独立加载。如果您想在不下载整个数据集的情况下进行小规模实验，只需下载几个 Parquet 文件即可。

🧑‍💻 训练

要训练模型，请运行：

bash scripts/train.sh <GPU_NUMS>

将 <GPU_NUMS> 替换为您希望使用的 GPU 数量。

保存检查点后，您需要将 EMA（指数移动平均）格式的权重转换为可加载的检查点：

python scripts/convert_ckpt.py --ema-ckpt <EMA_CKPT_PATH> --out-ckpt <OUT_CKPT_PATH>

例如：

python scripts/convert_ckpt.py --ema-ckpt ./train_logs/promptfix/checkpoints/ckpt_epoch_0/state.pth --out-ckpt ./checkpoints/promptfix_epoch_1.ckpt

📝 引用 PromptFix

如果您使用了我们的数据集或代码，请为本仓库点赞 ⭐ 并引用我们的论文：

@inproceedings{yu2024promptfix,
  title={PromptFix: You Prompt and We Fix the Photo},
  author={Yu, Yongsheng and Zeng, Ziyun and Hua, Hang and Fu, Jianlong and Luo, Jiebo},
  booktitle={NeurIPS},
  year={2024}
}

🙏 致谢

我们感谢 InstructDiffusion、Stable Diffusion 和 InstructPix2Pix 的作者们分享他们的代码。

⚠️ 免责声明

本仓库是一个开源研究项目的一部分，仅用于学术和科研目的。我们尚未建立任何与该项目相关的官方商业服务、产品或 Web 应用程序。请自行承担使用本软件的风险；它可能无法满足您的所有期望或需求。

请注意，PromptFix 数据集来源于开源研究项目和公开可用的照片库。使用我们的数据集即表示您同意遵守所有与源数据相关的许可协议和使用条款。此外，您也承认并同意，无论是数据集本身，还是基于该数据集训练的任何模型，均不得用于任何商业用途。

PromptFix 快速上手指南

PromptFix 是一个基于扩散模型的 AI 图像修复工具，能够根据用户的文本指令处理退化图像并移除不需要的元素。支持上色、去水印、去雾、去模糊、低光照增强等多种任务。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python 版本: 3.10 (代码已在此版本测试通过)
• 硬件要求: 需要 NVIDIA GPU 以进行推理和训练（建议显存 8GB 以上）
• 依赖管理: 推荐使用 conda 管理虚拟环境
• 网络环境: 需要访问 Hugging Face 下载预训练权重和数据集（国内用户可能需要配置代理或使用镜像加速）

安装步骤

按照以下步骤克隆仓库并安装依赖：

# 1. 克隆项目仓库
git clone https://github.com/yeates/PromptFix.git
cd PromptFix

# 2. 创建并激活 Python 3.10 虚拟环境
conda create -n promptfix python=3.10 -y
conda activate promptfix

# 3. 安装项目依赖
# 国内用户若 pip 下载缓慢，可添加清华源：-i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -r requirements.txt

基本使用

安装完成后，您可以直接运行推理脚本。系统会自动从 Hugging Face 下载预训练模型权重并存放在 checkpoints/ 目录下。

运行默认示例：

bash scripts/inference.sh

运行结束后，处理后的图片将保存在输出目录中。您可以查看 examples/ 文件夹下的输入图片与对应的指令，了解如何构造自己的测试用例。

提示：如需处理自定义图片，请参考项目源码中的推理逻辑修改输入路径和提示词（Prompt）。