扩散模型Diffusers Pipeline API使用介绍

1 关于Diffusers Pipeline

1.1 简介

大部分扩散模型包含多个独立训练的子模型和组件模块组合而成，例如StableDiffusion 有：

• 3个独立训练的子模型：Autoencoder、 Conditional Unet、CLIP text encoder
• 调度器组件scheduler,
• CLIPImageProcessor,
• safety checker.

为了让开发者以最简单的方式使用最新最先进的扩散模型，diffusers开发了pipeline管理和使用这些类，使得开发者可以以端对端方式使用扩散模型。

注意：pipeline本身没有提供任何训练相关功能，如果想要实现训练，可以参考官方的训练样例

1.2 官方Pipeline

以下表格是diffusers官方实现的Pipeline，每个Pipeline有对应的论文。

2 Pipeline API接口

扩散模型包含多个独立的模型和组件，不同任务中模型独立训练，并且可以用其他模型替换。不同的Pipeline可能包含专有的函数接口，但所有Pipeline都有的共同函数如下：

• from_pretrained(cls, pretrained_model_name_or_path, **kwargs)： 参数pretrained_model_name_or_path可以是 Hugging Face Hub repository的 id, 例如： runwayml/stable-diffusion-v1-5 或本地路径："./stable-diffusion". 为了确保所有模型和组件能被正确加载，需要提供一个 model_index.json 文件, 例如： runwayml/stable-diffusion-v1-5/model_index.json, 这个文件定义了所有要被加载的组件。其格式如下： : ["", ""]，其中 是类实例的名称。此类可以在库""中加载到。
• save_pretrained(self, save_directory) : 参数save_directory为本地目录路径，例如： ./stable-diffusion ，所有的模型和组件会被保存。每个模型和组件创建一个对应的子目录，子目录名称为模型或组件的属性名称如./stable_diffusion/unet. 此外，还会再根目录创建 model_index.json 文件如：./stable_diffusion/model_index.json
• to(self, torch_device: Optional[Union[str, torch.device]] = None) 参数torch_device为 string 或 torch.device 类型，将所有torch.nn.Module 类型的对象转移到指定的device上，此函数与pytorch的to函数功能一致。
• __call__函数执行推理，此函数定义了pipeline的推理逻辑，不同的Pipeline对应的推理输入差别很大，例如文生图Pipeline StableDiffusionPipeline 的输入应该是文本prompt，输出是生成的图。 而DDPMPipeline 则无需提供任何输入。因此读者需要根据实际的Pipeline功能以及查看相应的官方文档使用。

注意: 所有的Pipeline的__call__函数会自动调用torch.no_grad函数禁用梯度，因为Pipeline不是用于训练。如果你在前向推理后有保存梯度的需求，可以自定义Pipeline，参考官方示例

3 使用示例

1 扩散模型：文生图

# make sure you're logged in with `huggingface-cli login`
from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]  
    
image.save("astronaut_rides_horse.png")
1
2
3
4
5
6
7
8
9
10

2 扩散模型：图生图

类StableDiffusionImg2ImgPipeline 接受一个文本prompt和初始图片作为条件，指导生成新图。

import requests
from PIL import Image
from io import BytesIO

from diffusers import StableDiffusionImg2ImgPipeline

# load the pipeline
device = "cuda"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to(device)

# let's download an initial image
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"

response = requests.get(url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
init_image = init_image.resize((768, 512))

prompt = "A fantasy landscape, trending on artstation"

images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images

images[0].save("fantasy_landscape.png")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

可以在colab中直接运行colab

3 扩充模型：In-painting

类 StableDiffusionInpaintPipeline 接受文本prompt和mask，用于编辑图像指定区域。

import PIL
import requests
import torch
from io import BytesIO

from diffusers import StableDiffusionInpaintPipeline

def download_image(url):
    response = requests.get(url)
    return PIL.Image.open(BytesIO(response.content)).convert("RGB")

img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = download_image(img_url).resize((512, 512))
mask_image = download_image(mask_url).resize((512, 512))

pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

可以在colab中直接运行 colab