DiffEdit
图像编辑通常需要提供要编辑区域的蒙版。DiffEdit 会根据文本查询自动为您生成蒙版,从而更容易在没有图像编辑软件的情况下创建蒙版。DiffEdit 算法分三个步骤进行
- 扩散模型根据某些查询文本和参考文本对图像进行去噪,从而为图像的不同区域生成不同的噪声估计;差异用于推断蒙版,以识别需要更改哪些图像区域以匹配查询文本
- 输入图像使用 DDIM 编码到潜在空间
- 使用扩散模型根据文本查询对潜在变量进行解码,使用蒙版作为引导,以便蒙版之外的像素与输入图像中的像素保持相同
本指南将向您展示如何使用 DiffEdit 编辑图像,而无需手动创建蒙版。
在开始之前,请确保已安装以下库
# uncomment to install the necessary libraries in Colab
#!pip install -q diffusers transformers accelerate
该 StableDiffusionDiffEditPipeline 需要图像蒙版和一组部分反转的潜在变量。图像蒙版是从 generate_mask() 函数生成的,并包含两个参数,source_prompt
和 target_prompt
。这些参数确定要编辑图像的哪些内容。例如,如果要将一碗水果更改为一碗梨,则
source_prompt = "a bowl of fruits"
target_prompt = "a bowl of pears"
部分反转的潜在变量是从 invert() 函数生成的,通常最好包含一个prompt
或标题来描述图像,以帮助指导反向潜在变量采样过程。标题通常可以是您的source_prompt
,但可以随意尝试其他文本描述!
让我们加载管道、调度器、逆调度器,并启用一些优化以减少内存使用量
import torch
from diffusers import DDIMScheduler, DDIMInverseScheduler, StableDiffusionDiffEditPipeline
pipeline = StableDiffusionDiffEditPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1",
torch_dtype=torch.float16,
safety_checker=None,
use_safetensors=True,
)
pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
pipeline.inverse_scheduler = DDIMInverseScheduler.from_config(pipeline.scheduler.config)
pipeline.enable_model_cpu_offload()
pipeline.enable_vae_slicing()
加载要编辑的图像
from diffusers.utils import load_image, make_image_grid
img_url = "https://github.com/Xiang-cd/DiffEdit-stable-diffusion/raw/main/assets/origin.png"
raw_image = load_image(img_url).resize((768, 768))
raw_image
使用 generate_mask() 函数生成图像蒙版。您需要将source_prompt
和 target_prompt
传递给它,以指定要编辑图像的哪些内容
from PIL import Image
source_prompt = "a bowl of fruits"
target_prompt = "a basket of pears"
mask_image = pipeline.generate_mask(
image=raw_image,
source_prompt=source_prompt,
target_prompt=target_prompt,
)
Image.fromarray((mask_image.squeeze()*255).astype("uint8"), "L").resize((768, 768))
接下来,创建反转的潜在变量并将其传递给描述图像的标题
inv_latents = pipeline.invert(prompt=source_prompt, image=raw_image).latents
最后,将图像蒙版和反转的潜在变量传递给管道。target_prompt
现在变为prompt
,source_prompt
用作negative_prompt
output_image = pipeline(
prompt=target_prompt,
mask_image=mask_image,
image_latents=inv_latents,
negative_prompt=source_prompt,
).images[0]
mask_image = Image.fromarray((mask_image.squeeze()*255).astype("uint8"), "L").resize((768, 768))
make_image_grid([raw_image, mask_image, output_image], rows=1, cols=3)
生成源和目标嵌入
可以使用 Flan-T5 模型自动生成源和目标嵌入,而不是手动创建它们。
从 🤗 Transformers 库加载 Flan-T5 模型和分词器
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large", device_map="auto", torch_dtype=torch.float16)
提供一些初始文本以提示模型生成源和目标提示。
source_concept = "bowl"
target_concept = "basket"
source_text = f"Provide a caption for images containing a {source_concept}. "
"The captions should be in English and should be no longer than 150 characters."
target_text = f"Provide a caption for images containing a {target_concept}. "
"The captions should be in English and should be no longer than 150 characters."
接下来,创建一个实用函数来生成提示
@torch.no_grad()
def generate_prompts(input_prompt):
input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids.to("cuda")
outputs = model.generate(
input_ids, temperature=0.8, num_return_sequences=16, do_sample=True, max_new_tokens=128, top_k=10
)
return tokenizer.batch_decode(outputs, skip_special_tokens=True)
source_prompts = generate_prompts(source_text)
target_prompts = generate_prompts(target_text)
print(source_prompts)
print(target_prompts)
如果您有兴趣详细了解生成不同质量文本的策略,请查看 生成策略 指南。
加载 StableDiffusionDiffEditPipeline 使用的文本编码器模型来编码文本。您将使用文本编码器来计算文本嵌入
import torch
from diffusers import StableDiffusionDiffEditPipeline
pipeline = StableDiffusionDiffEditPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16, use_safetensors=True
)
pipeline.enable_model_cpu_offload()
pipeline.enable_vae_slicing()
@torch.no_grad()
def embed_prompts(sentences, tokenizer, text_encoder, device="cuda"):
embeddings = []
for sent in sentences:
text_inputs = tokenizer(
sent,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
prompt_embeds = text_encoder(text_input_ids.to(device), attention_mask=None)[0]
embeddings.append(prompt_embeds)
return torch.concatenate(embeddings, dim=0).mean(dim=0).unsqueeze(0)
source_embeds = embed_prompts(source_prompts, pipeline.tokenizer, pipeline.text_encoder)
target_embeds = embed_prompts(target_prompts, pipeline.tokenizer, pipeline.text_encoder)
最后,将嵌入传递给 generate_mask() 和 invert() 函数以及管道以生成图像
from diffusers import DDIMInverseScheduler, DDIMScheduler
from diffusers.utils import load_image, make_image_grid
from PIL import Image
pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
pipeline.inverse_scheduler = DDIMInverseScheduler.from_config(pipeline.scheduler.config)
img_url = "https://github.com/Xiang-cd/DiffEdit-stable-diffusion/raw/main/assets/origin.png"
raw_image = load_image(img_url).resize((768, 768))
mask_image = pipeline.generate_mask(
image=raw_image,
- source_prompt=source_prompt,
- target_prompt=target_prompt,
+ source_prompt_embeds=source_embeds,
+ target_prompt_embeds=target_embeds,
)
inv_latents = pipeline.invert(
- prompt=source_prompt,
+ prompt_embeds=source_embeds,
image=raw_image,
).latents
output_image = pipeline(
mask_image=mask_image,
image_latents=inv_latents,
- prompt=target_prompt,
- negative_prompt=source_prompt,
+ prompt_embeds=target_embeds,
+ negative_prompt_embeds=source_embeds,
).images[0]
mask_image = Image.fromarray((mask_image.squeeze()*255).astype("uint8"), "L")
make_image_grid([raw_image, mask_image, output_image], rows=1, cols=3)
为反转生成标题
虽然您可以使用source_prompt
作为标题来帮助生成部分反转的潜在变量,但您也可以使用 BLIP 模型自动生成标题。
从 🤗 Transformers 库加载 BLIP 模型和处理器
import torch
from transformers import BlipForConditionalGeneration, BlipProcessor
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base", torch_dtype=torch.float16, low_cpu_mem_usage=True)
创建一个实用函数,根据输入图像生成标题
@torch.no_grad()
def generate_caption(images, caption_generator, caption_processor):
text = "a photograph of"
inputs = caption_processor(images, text, return_tensors="pt").to(device="cuda", dtype=caption_generator.dtype)
caption_generator.to("cuda")
outputs = caption_generator.generate(**inputs, max_new_tokens=128)
# offload caption generator
caption_generator.to("cpu")
caption = caption_processor.batch_decode(outputs, skip_special_tokens=True)[0]
return caption
加载输入图像并使用generate_caption
函数为其生成标题
from diffusers.utils import load_image
img_url = "https://github.com/Xiang-cd/DiffEdit-stable-diffusion/raw/main/assets/origin.png"
raw_image = load_image(img_url).resize((768, 768))
caption = generate_caption(raw_image, model, processor)
现在,您可以将标题放入 invert() 函数中以生成部分反转的潜在变量!
< > 在 GitHub 上更新