diffusers/v0.43.0/zh/using-diffusers/write_own_pipeline.md

理解流水线、模型和调度器

🧨 Diffusers 旨在成为一个用户友好且灵活的工具箱，用于构建针对特定用例定制的扩散系统。该工具箱的核心是模型和调度器。虽然 DiffusionPipeline 为了方便起见将这些组件打包在一起，但您也可以解包该流水线，并单独使用模型和调度器来创建新的扩散系统。

在本教程中，你将学习如何使用模型和调度器来组装一个用于推理的扩散系统，从基础流水线开始，逐步过渡到 Stable Diffusion 流水线。

拆解一个基础流水线

Pipeline 是一种快速便捷的模型推理方式，只需四行代码即可生成图像：

>>> from diffusers import DDPMPipeline

>>> ddpm = DDPMPipeline.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
>>> image = ddpm(num_inference_steps=25).images[0]
>>> image

这非常简单，但 pipeline 是如何做到的呢？让我们拆解一下 pipeline，看看其内部究竟发生了什么。

在上面的示例中，pipeline 包含一个 UNet2DModel 模型和一个 DDPMScheduler。该 pipeline 通过获取与期望输出尺寸相同的随机噪声，并将其多次传入模型来对图像进行去噪。在每个时间步，模型预测 noise residual，调度器则利用它来预测噪声更少的图像。pipeline 会重复此过程，直到达到指定的推理步数。

若要分别使用模型和调度器来重建流水线，让我们编写自己的去噪过程。

1. 加载模型和调度器：

>>> from diffusers import DDPMScheduler, UNet2DModel

>>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
>>> model = UNet2DModel.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

2. 设置去噪过程运行的时间步数：

>>> scheduler.set_timesteps(50)

3. 设置调度器的时间步会创建一个包含均匀间隔元素的张量，本例中为 50 个。每个元素对应模型对图像进行去噪的时间步。之后创建去噪循环时，你将遍历此张量以对图像进行去噪：

>>> scheduler.timesteps
tensor([980, 960, 940, 920, 900, 880, 860, 840, 820, 800, 780, 760, 740, 720,
    700, 680, 660, 640, 620, 600, 580, 560, 540, 520, 500, 480, 460, 440,
    420, 400, 380, 360, 340, 320, 300, 280, 260, 240, 220, 200, 180, 160,
    140, 120, 100,  80,  60,  40,  20,   0])

4. 创建与期望输出形状相同的随机噪声：

>>> import torch

>>> sample_size = model.config.sample_size
>>> noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")  # or "mps", "xpu", "cpu"

5. 现在编写一个循环来遍历各个时间步。在每个时间步，模型执行一次 UNet2DModel.forward() 前向传播，并返回带噪残差。调度器的 step() 方法接收带噪残差、时间步和输入，并预测上一个时间步的图像。该输出将成为去噪循环中模型的下一个输入，如此反复，直到遍历完 timesteps 数组的末尾。

>>> input = noise

>>> for t in scheduler.timesteps:
...     with torch.no_grad():
...         noisy_residual = model(input, t).sample
...     previous_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
...     input = previous_noisy_sample

这就是完整的去噪过程，你可以使用相同的模式来编写任何扩散系统。

6. 最后一步是将去噪后的输出转换为图像：

>>> from PIL import Image
>>> import numpy as np

>>> image = (input / 2 + 0.5).clamp(0, 1).squeeze()
>>> image = (image.permute(1, 2, 0) * 255).round().to(torch.uint8).cpu().numpy()
>>> image = Image.fromarray(image)
>>> image

在下一节中，你将检验自己的技能，并深入剖析更为复杂的 Stable Diffusion 流水线。这些步骤大体相同。你将初始化必要的组件，并设置时间步数以创建一个 timestep 数组。timestep 数组用于去噪循环，对于该数组中的每个元素，模型都会预测一张噪声更少的图像。去噪循环会遍历 timestep，在每个时间步，它会输出一个带噪声的残差，调度器（scheduler）则利用该残差来预测前一个时间步下噪声更少的图像。此过程会不断重复，直到遍历完 timestep 数组的末尾。

让我们试一试！

拆解 Stable Diffusion 流水线

Stable Diffusion 是一个文本到图像的潜在扩散模型。之所以称为潜在扩散模型，是因为它处理的是图像的降维表示，而非实际的像素空间，从而具有更高的内存效率。编码器将图像压缩为更小的表示，解码器再将压缩后的表示还原为图像。对于文本到图像模型，你需要一个分词器和一个编码器来生成文本嵌入。根据前面的示例，你已经知道需要一个 UNet 模型和一个调度器。

如你所见，这已经比仅包含 UNet 模型的 DDPM 流水线更为复杂。Stable Diffusion 模型包含三个独立的预训练模型。

💡 阅读 Stable Diffusion 是如何工作的？ 博客，了解更多关于 VAE、UNet 和文本编码器模型工作原理的详细信息。

现在您已经了解了 Stable Diffusion 流水线所需的组件，请使用 from_pretrained() 方法加载所有这些组件。您可以在预训练的 stable-diffusion-v1-5/stable-diffusion-v1-5 检查点中找到它们，每个组件都存储在一个独立的子文件夹中：

>>> from PIL import Image
>>> import torch
>>> from transformers import CLIPTextModel, CLIPTokenizer
>>> from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler

>>> vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_safetensors=True)
>>> tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="tokenizer")
>>> text_encoder = CLIPTextModel.from_pretrained(
...     "CompVis/stable-diffusion-v1-4", subfolder="text_encoder", use_safetensors=True
... )
>>> unet = UNet2DConditionModel.from_pretrained(
...     "CompVis/stable-diffusion-v1-4", subfolder="unet", use_safetensors=True
... )

将默认的 PNDMScheduler 替换为 UniPCMultistepScheduler，看看插入不同的调度器有多简单：

>>> from diffusers import UniPCMultistepScheduler

>>> scheduler = UniPCMultistepScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

为了加速推理，请将模型移至 GPU，因为与调度器不同，它们具有可训练的权重：

>>> torch_device = "cuda"  # or "mps", "xpu", "cpu"
>>> vae.to(torch_device)
>>> text_encoder.to(torch_device)
>>> unet.to(torch_device)

创建文本嵌入

下一步是对文本进行分词以生成嵌入。文本用于为 UNet 模型提供条件，并引导扩散过程生成与输入提示相似的内容。

💡 guidance_scale 参数决定了在生成图像时，提示词应被赋予多大的权重。

如果您想生成其他内容，请随意选择您喜欢的提示词！

>>> prompt = ["a photograph of an astronaut riding a horse"]
>>> height = 512  # default height of Stable Diffusion
>>> width = 512  # default width of Stable Diffusion
>>> num_inference_steps = 25  # Number of denoising steps
>>> guidance_scale = 7.5  # Scale for classifier-free guidance
>>> generator = torch.manual_seed(0)  # Seed generator to create the initial latent noise
>>> batch_size = len(prompt)

对文本进行分词，并根据提示生成嵌入：

>>> text_input = tokenizer(
...     prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt"
... )

>>> with torch.no_grad():
...     text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

您还需要生成无条件文本嵌入，即填充标记的嵌入。这些嵌入需要与条件text_embeddings具有相同的形状（batch_size和seq_length）：

>>> max_length = text_input.input_ids.shape[-1]
>>> uncond_input = tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")
>>> uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]

让我们将条件嵌入和无条件嵌入拼接成一个批次，以避免进行两次前向传播：

>>> text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

创建随机噪声

接下来，生成一些初始随机噪声作为扩散过程的起点。这是图像的潜在表示，随后会被逐步去噪。此时，latent 图像的尺寸小于最终图像尺寸，但这没关系，因为模型稍后会将其转换为最终的 512x512 图像尺寸。

💡 由于 vae 模型包含 3 个下采样层，因此高度和宽度会除以 8。你可以通过运行以下代码进行验证：

2 ** (len(vae.config.block_out_channels) - 1) == 8

>>> latents = torch.randn(
...     (batch_size, unet.config.in_channels, height // 8, width // 8),
...     generator=generator,
...     device=torch_device,
... )

对图像进行去噪

首先，使用初始噪声分布 sigma（即噪声尺度值）对输入进行缩放，这对于 UniPCMultistepScheduler 等改进型调度器是必需的：

>>> latents = latents * scheduler.init_noise_sigma

最后一步是创建去噪循环，该循环将逐步把 latents 中的纯噪声转换为由你的提示词所描述的图片。请记住，去噪循环需要完成三件事：

1. 设置去噪过程中调度器所使用的 timesteps。
2. 遍历时间步。
3. 在每个时间步，调用 UNet 模型预测噪声残差，并将其传递给调度器以计算前一个含噪样本。

>>> from tqdm.auto import tqdm

>>> scheduler.set_timesteps(num_inference_steps)

>>> for t in tqdm(scheduler.timesteps):
...     # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
...     latent_model_input = torch.cat([latents] * 2)

...     latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)

...     # predict the noise residual
...     with torch.no_grad():
...         noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

...     # perform guidance
...     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
...     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

...     # compute the previous noisy sample x_t -> x_t-1
...     latents = scheduler.step(noise_pred, t, latents).prev_sample

解码图像

最后一步是使用 vae 将潜在表示解码为图像，并通过 sample 获取解码后的输出：

# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

最后，将图像转换为 PIL.Image，即可查看生成的图像！

>>> image = (image / 2 + 0.5).clamp(0, 1).squeeze()
>>> image = (image.permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
>>> image = Image.fromarray(image)
>>> image

后续步骤

从基础到复杂的流水线，你已经看到，编写自己的扩散系统真正需要的只是一个去噪循环。该循环应设置调度器的时间步，遍历这些时间步，并在调用 UNet 模型预测噪声残差与将其传递给调度器以计算前一个含噪样本之间交替进行。

这正是 🧨 Diffusers 的设计初衷：让你能够直观、轻松地使用模型和调度器来编写自己的扩散系统。

接下来，您可以：

• 了解如何构建并贡献一个 pipeline 到 🧨 Diffusers。我们迫不及待想看看你会带来什么！
• 探索库中的现有流水线，看看你是否能够拆解流水线，并使用模型和调度器分别从头构建一个流水线。