Diffusers 文件

DiffEdit

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

DiffEdit

影像編輯通常需要提供要編輯區域的掩碼。DiffEdit 會根據文字查詢為您自動生成掩碼,從而使建立掩碼變得更容易,無需使用影像編輯軟體。DiffEdit 演算法分三步工作:

  1. 擴散模型根據某個查詢文字和參考文字對影像進行去噪,這會對影像的不同區域產生不同的噪聲估計;利用此差異來推斷一個掩碼,以識別影像中需要更改以匹配查詢文字的區域。
  2. 使用 DDIM 將輸入影像編碼到潛在空間中。
  3. 在文字查詢的條件下,使用擴散模型對潛在變數進行解碼,並以掩碼為指導,使掩碼外的畫素與輸入影像保持一致。

本指南將向您展示如何使用 DiffEdit 在不手動建立掩碼的情況下編輯影像。

開始之前,請確保已安裝以下庫:

# uncomment to install the necessary libraries in Colab
#!pip install -q diffusers transformers accelerate

StableDiffusionDiffEditPipeline 需要一個影像掩碼和一組部分反轉的潛在變數。影像掩碼由 generate_mask() 函式生成,幷包含兩個引數:source_prompttarget_prompt。這些引數決定了要在影像中編輯的內容。例如,如果您想將一碗 *水果* 更改為一碗 *梨*,則:

source_prompt = "a bowl of fruits"
target_prompt = "a bowl of pears"

部分反轉的潛在變數由 invert() 函式生成,通常最好包含一個 prompt 或 *標題* 來描述影像,以幫助指導反向潛在取樣過程。標題通常可以是您的 source_prompt,但請隨時嘗試其他文字描述!

讓我們載入 pipeline、排程器、逆向排程器,並啟用一些最佳化以減少記憶體使用:

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_prompttarget_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

最後,將影像掩碼和反轉的潛在變數傳遞給 pipeline。現在 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() 函式,以及 pipeline 來生成影像。

  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 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.