Forgedit: 基于学习和遗忘的文本引导图像编辑

摘要

仅以原始图像本身和目标文本提示作为输入的文本引导图像编辑,无论是对真实还是合成图像,都是一项非常通用且具有挑战性的任务。它要求编辑模型自行估计应该编辑图像的哪一部分,然后执行刚性或非刚性编辑,同时保持原始图像的特征。在本文中,我们设计了一种新颖的文本引导图像编辑方法,称为Forgedit。首先,我们提出了一个视觉-语言联合优化框架,能够在30秒内重建原始图像,比之前的最先进方法更快,过拟合也更少。然后我们在扩散模型的文本嵌入空间中提出了一种新颖的向量投影机制,能够分别控制身份相似性和编辑强度。最后,我们发现了扩散模型中UNet的一个普遍特性,即UNet编码器学习空间和结构,UNet解码器学习外观和身份。基于这一特性,我们设计了遗忘机制,成功解决了在一张图像上微调扩散模型时不可避免的致命过拟合问题,从而显著提升了扩散模型的编辑能力。我们的方法Forgedit基于Stable Diffusion构建,在具有挑战性的文本引导图像编辑基准TEdBench上取得了新的最先进结果,在CLIP分数和LPIPS分数方面都超越了之前的最先进方法,如基于Imagen的Imagic。

替代文本

致谢

本代码基于Diffusers实现的Imagic

安装

代码在requirements.txt环境下测试通过,在一张A100 GPU上训练512x512分辨率的图像需要30秒以上,40秒以下。然而,这并不是最佳设置,因为之前我在阿里云的一个默认docker环境中使用一张A800 GPU测试这段代码时更快,正好需要30秒。

TEdBench

谷歌研究院在Imagic中提出的TEdBench可以在TEdBench找到。我们的原始Forgedit在TEdBench上的完整编辑结果可以在原始Forgedit tedbench仓库中找到。我们没有在TEdBench上使用DreamBooth+Forgedit进行定量比较,因此没有提供DreamBooth+Forgedit的结果。

基于Stable Diffusion的Forgedit

要重现我们在TEdBench上的结果,请使用Stable Diffusion 1.4。要重现我们在视觉讲故事上的结果,请使用SG161222/Realistic_Vision_V6.0_B1_noVAE。BLIP模型是Salesforce/blip-image-captioning-base。请注意,其他BLIP变体可能会导致性能下降。在此代码发布中,实现了原始Forgedit和DreamBoothForgedit。

原始Forgedit

例如,要重现视觉讲故事的结果,我们使用SG161222/Realistic_Vision_V6模型通过原始Forgedit进行训练、保存和编辑,文本嵌入插值是向量减法,遗忘策略是encoderkv,分辨率为768x768,gamma范围从0.5到0.7。请注意,不同提示的gamma范围可能不同。由于分辨率为768x768,这个训练过程需要一分钟以上。如果分辨率设置为512x512,在A100上的训练时间应该是30到40秒。

accelerate launch src/sample_forgedit_batch_textencoder.py --train=True --edit=True  --save=True --forget='encoderkv'  --interpolation=vs  --targeth=768 --targetw=768 --gammastart=5 --gammaend=8

模型将保存到./vanillaforgedit

要使用保存的模型编辑图像,我们使用参数--loadfrom

accelerate launch src/sample_forgedit_batch_textencoder.py --train=False --edit=True --save=False --forget='encoderkv' --loadfrom='/mnt/bn/editdiffusion/Forgedit/vanillaforgedit/img=test.jpg_textsteps=400_bsz=10_unetlr=6e-05_textlr=0.001' --interpolation=vs  --targeth=768 --targetw=768 --gammastart=5 --gammaend=8

需要在src/sample_forgedit_batch_textencoder.py中更改img_url和prompt。

在src/forgedit_stable_diffusion/pipelineattentionparallel_bsz=1.py中实现了7种典型的遗忘策略,可以通过--forget参数设置。请注意,默认值是--forget='donotforget',不使用遗忘策略。

在EditEval v1上重现示例的完整超参数可以在此说明中找到

DreamBoothForgedit

要使用向量投影进行DreamBoothForgedit的微调、保存和编辑,

accelerate launch src/sample_dreambooth_batch_textencoder.py --save=True --interpolation=vp

要使用保存的编辑模型进行编辑,

accelerate launch src/sample_dreambooth_batch_textencoder.py --train=False --interpolation=vp

遗忘策略在src/forgedit_stable_diffusion/pipelinedreamboothparallel_bsz=1_textencoder.py中实现, 可以在sample_dreambooth_batch_textencoder.py的freeze_list中使用

替代文本

引用

我们的论文可以在Forgedit: 基于学习和遗忘的文本引导图像编辑中找到:

@article{zhang2023forgedit,
  title={Forgedit: Text Guided Image Editing via Learning and Forgetting},
  author={Zhang, Shiwen and Xiao, Shuai and Huang, Weilin},
  journal={arXiv preprint arXiv:2309.10556},
  year={2023}
}