[CVPR 2024] MACE：扩散模型中的大规模概念擦除

MACE：扩散模型中的大规模概念擦除的官方实现。

MACE：扩散模型中的大规模概念擦除<br> Shilin Lu, Zilan Wang, Leyang Li, Yanzhu Liu, Adams Wai-Kin Kong <br> CVPR 2024

摘要：<br> 大规模文本到图像扩散模型的快速发展引发了人们对其可能被滥用于创建有害或误导性内容的日益关注。在本文中，我们提出了MACE，一个用于大规模概念擦除任务的微调框架。该任务旨在防止模型在接收提示时生成包含不需要的概念的图像。现有的概念擦除方法通常限于同时处理不到五个概念，并且难以在擦除概念同义词（普遍性）和维持无关概念（特异性）之间取得平衡。相比之下，MACE的不同之处在于成功地将擦除范围扩大到100个概念，并在普遍性和特异性之间取得了有效平衡。这是通过利用闭式交叉注意力细化和LoRA微调来共同消除不需要的概念信息实现的。此外，MACE集成了多个LoRA，而不会相互干扰。我们对MACE与先前方法在四个不同任务上进行了广泛评估：物体擦除、名人擦除、露骨内容擦除和艺术风格擦除。我们的结果表明，MACE在所有评估任务中都超越了先前的方法。

teaser

framework

(a) 我们的框架专注于调整交叉注意力（CA）块内与提示相关的投影矩阵。(b) 使用闭式解对预训练U-Net的CA块进行细化，阻止模型将目标短语的残余信息嵌入到周围的词中。(c) 对于每个要移除的概念，学习一个独特的LoRA模块以消除其内在信息。(d) 引入闭式解以集成多个LoRA模块，避免相互干扰，同时防止灾难性遗忘。

设置

创建Conda环境

git clone https://github.com/Shilin-LU/MACE.git
conda create -n mace python=3.10
conda activate mace
conda install pytorch==2.0.1 torchvision==0.15.2 pytorch-cuda=11.7 -c pytorch -c nvidia

安装Grounded-SAM以准备用于LoRA调优的掩码

您可以选择使用其他分割模型并跳过此部分；但请注意，如果掩码不精确或未使用，性能可能会受到影响。

export AM_I_DOCKER=False
export BUILD_WITH_CUDA=True
# export CUDA_HOME=/path/to/cuda-11.7/

cd MACE
git clone https://github.com/IDEA-Research/Grounded-Segment-Anything.git
cd Grounded-Segment-Anything

# 安装Segment Anything：
python -m pip install -e segment_anything

# 安装Grounding DINO：
pip install --no-build-isolation -e GroundingDINO

# 安装osx：
git submodule update --init --recursive
cd grounded-sam-osx && bash install.sh

# 安装RAM和Tag2Text：
git clone https://github.com/xinyu1205/recognize-anything.git
pip install -r ./recognize-anything/requirements.txt
pip install -e ./recognize-anything/

下载Grounded-SAM的预训练权重。

cd ..    # cd Grounded-Segment-Anything

# 下载预训练的groundingdino-swin-tiny模型：
wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth

# 下载预训练的SAM模型：
wget https://huggingface.co/lkeab/hq-sam/resolve/main/sam_hq_vit_h.pth

安装其他依赖项

pip install diffusers==0.22.0 transformers==4.38.1
pip install accelerate openai omegaconf

准备MACE训练数据

要擦除概念，应为每个概念生成8张图像及其相应的分割掩码。要为您的目标概念准备数据，请在configs/object/ship.yaml中配置您的设置，并执行以下命令：

CUDA_VISIBLE_DEVICES=0 python data_preparation.py configs/object/ship.yaml

在开始大规模概念擦除过程之前，确保您已预先缓存了您希望保留的先验知识（例如MSCOCO）和特定领域知识（例如某些名人、艺术风格或物体）。

您可以从这个OneDrive文件夹下载我们预先缓存的文件。下载后，将这些文件放在./cache/目录中使用。
或者，要保留您选择的额外知识，您可以通过修改脚本src/cache_coco.py来缓存信息。

训练MACE以擦除概念

准备好数据后，您可以在同一配置文件configs/object/ship.yaml中指定训练参数，并运行以下命令：

CUDA_VISIBLE_DEVICES=0 python training.py configs/object/ship.yaml

从微调模型采样

可以通过运行以下命令简单地测试微调模型，生成几张图像：

CUDA_VISIBLE_DEVICES=0 python inference.py \
          --num_images 3 \
          --prompt '你的提示词' \
          --model_path /path/to/saved_model/LoRA_fusion_model \
          --save_path /path/to/save/folder

要基于预设种子的提示词列表（例如，来自CSV文件./prompts_csv/celebrity_100_concepts.csv）生成大量图像，请执行以下命令（超参数step应设置为与num_processes相同的值）：

CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
          --multi_gpu --num_processes=4 --main_process_port 31372 \
          src/sample_images_from_csv.py \
          --prompts_path ./prompts_csv/celebrity_100_concepts.csv \
          --save_path /path/to/save/folder \
          --model_name /path/to/saved_model/LoRA_fusion_model \
          --step 4

MACE微调模型权重

我们提供了几个使用MACE微调的Stable Diffusion v1.4模型。

要擦除的概念类型	微调模型
物体擦除	OneDrive链接
名人擦除	OneDrive链接
艺术风格擦除	OneDrive链接
露骨内容擦除	OneDrive链接

指标评估

在我们的评估中，我们使用了多种指标，包括FID、CLIP分数、CLIP分类准确率、GCD准确率和NudeNet检测结果。

评估FID：

CUDA_VISIBLE_DEVICES=0 python metrics/evaluate_fid.py --dir1 'path/to/generated/image/folder' --dir2 'path/to/coco/GT/folder'

评估CLIP分数：

CUDA_VISIBLE_DEVICES=0 python metrics/evaluate_clip_score.py --image_dir 'path/to/generated/image/folder' --prompts_path './prompts_csv/coco_30k.csv'

评估GCD准确率。使用此脚本进行检测时，请确保输入目录中只包含图像，无需进入子目录。这样可以防止在处理过程中出现错误。（请参阅GCD安装指南）：

conda activate GCD
CUDA_VISIBLE_DEVICES=0 python metrics/evaluate_by_GCD.py --image_folder 'path/to/generated/image/folder'

评估NudeNet检测结果（请参阅NudeNet安装指南）：

CUDA_VISIBLE_DEVICES=0 python metrics/evaluate_by_nudenet.py --folder 'path/to/generated/image/folder'

评估CLIP分类准确率：

CUDA_VISIBLE_DEVICES=0 python metrics/evaluate_clip_accuracy.py --base_folder '生成图像文件夹的路径'

致谢

我们感谢以下贡献者，我们的代码基于他们的工作：Diffusers、Concept-Ablation、Forget-Me-Not、UCE。

引用

如果您觉得本仓库有用，请考虑引用。

@article{lu2024mace,
 title={MACE: Mass Concept Erasure in Diffusion Models}, 
 author={Lu, Shilin and Wang, Zilan and Li, Leyang and Liu, Yanzhu and Kong, Adams Wai-Kin},
 journal={arXiv preprint arXiv:2403.06135},
 year={2024}
}