Diffusers 文档

高效且高效的扩散

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

高效且高效的扩散

DiffusionPipeline 生成特定风格的图像或包含您想要的内容可能很棘手。通常,您必须多次运行 DiffusionPipeline 才能得到您满意的图像。但是,无中生有地生成内容是一个计算密集型过程,特别是当您一遍又一遍地运行推理时。

这就是为什么从 pipeline 中获得最大的计算(速度)和内存(GPU vRAM)效率非常重要,以减少推理周期之间的时间,以便您可以更快地迭代。

本教程将引导您了解如何使用 DiffusionPipeline 更快、更好地生成图像。

首先加载 stable-diffusion-v1-5/stable-diffusion-v1-5 模型

from diffusers import DiffusionPipeline

model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
pipeline = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True)

您将使用的示例提示是一位老战士酋长的肖像,但您可以随意使用自己的提示

prompt = "portrait photo of a old warrior chief"

速度

💡 如果您没有 GPU,您可以从像 Colab 这样的 GPU 提供商那里免费使用一个!

加速推理的最简单方法之一是将 pipeline 放在 GPU 上,就像您对任何 PyTorch 模块所做的那样

pipeline = pipeline.to("cuda")

为了确保您可以使用相同的图像并对其进行改进,请使用 Generator 并为 可重复性 设置种子

import torch

generator = torch.Generator("cuda").manual_seed(0)

现在您可以生成图像了

image = pipeline(prompt, generator=generator).images[0]
image

此过程在 T4 GPU 上花费了约 30 秒(如果您的分配的 GPU 比 T4 更好,可能会更快)。默认情况下,DiffusionPipeline 以完整的 float32 精度运行推理 50 个推理步骤。您可以通过切换到较低的精度(如 float16)或运行更少的推理步骤来加速此过程。

让我们首先以 float16 加载模型并生成图像

import torch

pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_safetensors=True)
pipeline = pipeline.to("cuda")
generator = torch.Generator("cuda").manual_seed(0)
image = pipeline(prompt, generator=generator).images[0]
image

这次,生成图像仅花费了约 11 秒,这几乎比以前快了 3 倍!

💡 我们强烈建议始终以 float16 运行您的 pipelines,到目前为止,我们很少看到输出质量有任何下降。

另一种选择是减少推理步骤的数量。选择更高效的 scheduler 可以帮助减少步骤数量,而不会牺牲输出质量。您可以通过调用 compatibles 方法在 DiffusionPipeline 中找到哪些 schedulers 与当前模型兼容

pipeline.scheduler.compatibles
[
    diffusers.schedulers.scheduling_lms_discrete.LMSDiscreteScheduler,
    diffusers.schedulers.scheduling_unipc_multistep.UniPCMultistepScheduler,
    diffusers.schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteScheduler,
    diffusers.schedulers.scheduling_deis_multistep.DEISMultistepScheduler,
    diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler,
    diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler,
    diffusers.schedulers.scheduling_ddpm.DDPMScheduler,
    diffusers.schedulers.scheduling_dpmsolver_singlestep.DPMSolverSinglestepScheduler,
    diffusers.schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteScheduler,
    diffusers.utils.dummy_torch_and_torchsde_objects.DPMSolverSDEScheduler,
    diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler,
    diffusers.schedulers.scheduling_pndm.PNDMScheduler,
    diffusers.schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteScheduler,
    diffusers.schedulers.scheduling_ddim.DDIMScheduler,
]

Stable Diffusion 模型默认使用 PNDMScheduler,通常需要约 50 个推理步骤,但性能更高的 schedulers,如 DPMSolverMultistepScheduler,仅需要约 20 或 25 个推理步骤。使用 from_config() 方法加载新的 scheduler

from diffusers import DPMSolverMultistepScheduler

pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)

现在将 num_inference_steps 设置为 20

generator = torch.Generator("cuda").manual_seed(0)
image = pipeline(prompt, generator=generator, num_inference_steps=20).images[0]
image

太棒了,您已成功将推理时间缩短至仅 4 秒!⚡️

内存

提高 pipeline 性能的另一个关键是减少内存消耗,这间接意味着更快的速度,因为您通常试图最大化每秒生成的图像数量。查看您可以一次生成多少图像的最简单方法是尝试不同的批处理大小,直到出现 OutOfMemoryError (OOM)。

创建一个函数,该函数将从提示列表和 Generators 生成一批图像。确保为每个 Generator 分配一个种子,以便在生成良好结果时可以重复使用它。

def get_inputs(batch_size=1):
    generator = [torch.Generator("cuda").manual_seed(i) for i in range(batch_size)]
    prompts = batch_size * [prompt]
    num_inference_steps = 20

    return {"prompt": prompts, "generator": generator, "num_inference_steps": num_inference_steps}

batch_size=4 开始,看看您消耗了多少内存

from diffusers.utils import make_image_grid

images = pipeline(**get_inputs(batch_size=4)).images
make_image_grid(images, 2, 2)

除非您的 GPU 具有更多 vRAM,否则上面的代码可能返回了 OOM 错误!大多数内存都被交叉注意力层占用。您可以顺序运行此操作而不是批量运行,以节省大量内存。您只需配置 pipeline 以使用 enable_attention_slicing() 函数即可

pipeline.enable_attention_slicing()

现在尝试将 batch_size 增加到 8!

images = pipeline(**get_inputs(batch_size=8)).images
make_image_grid(images, rows=2, cols=4)

之前您甚至无法生成 4 张图像的批次,而现在您可以生成 8 张图像的批次,每张图像约 3.5 秒!这可能是您在 T4 GPU 上不牺牲质量的情况下能够达到的最快速度。

质量

在最后两节中,您学习了如何通过使用 fp16、使用性能更高的 scheduler 减少推理步骤数量以及启用注意力切片来减少内存消耗来优化 pipeline 的速度。现在,您将专注于如何提高生成图像的质量。

更好的 checkpoints

最明显的步骤是使用更好的 checkpoints。Stable Diffusion 模型是一个很好的起点,自其正式发布以来,还发布了几个改进版本。但是,使用较新版本并不意味着您会自动获得更好的结果。您仍然需要自己尝试不同的 checkpoints,并做一些研究(例如使用 负面提示)以获得最佳结果。

随着该领域的不断发展,越来越多的高质量 checkpoints 被微调以产生特定风格。尝试浏览 HubDiffusers Gallery,找到您感兴趣的那个!

更好的 pipeline 组件

您还可以尝试用更新的版本替换当前的 pipeline 组件。让我们尝试将 Stability AI 最新的 autoencoder 加载到 pipeline 中,并生成一些图像

from diffusers import AutoencoderKL

vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16).to("cuda")
pipeline.vae = vae
images = pipeline(**get_inputs(batch_size=8)).images
make_image_grid(images, rows=2, cols=4)

更好的 prompt 工程

您用于生成图像的文本 prompt 非常重要,以至于它被称为prompt 工程。在 prompt 工程期间要记住的一些注意事项是

  • 我想要生成的图像或类似图像在互联网上是如何存储的?
  • 我可以提供哪些额外的细节来引导模型朝着我想要的风格发展?

考虑到这一点,让我们改进 prompt 以包含颜色和更高质量的细节

prompt += ", tribal panther make up, blue on red, side profile, looking away, serious eyes"
prompt += " 50mm portrait photography, hard rim lighting photography--beta --ar 2:3  --beta --upbeta"

使用新的 prompt 生成一批图像

images = pipeline(**get_inputs(batch_size=8)).images
make_image_grid(images, rows=2, cols=4)

非常令人印象深刻!让我们稍微调整第二张图像 - 对应于种子为 1Generator - 通过添加一些关于主体年龄的文本

prompts = [
    "portrait photo of the oldest warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3  --beta --upbeta",
    "portrait photo of an old warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3  --beta --upbeta",
    "portrait photo of a warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3  --beta --upbeta",
    "portrait photo of a young warrior chief, tribal panther make up, blue on red, side profile, looking away, serious eyes 50mm portrait photography, hard rim lighting photography--beta --ar 2:3  --beta --upbeta",
]

generator = [torch.Generator("cuda").manual_seed(1) for _ in range(len(prompts))]
images = pipeline(prompt=prompts, generator=generator, num_inference_steps=25).images
make_image_grid(images, 2, 2)

下一步

在本教程中,您学习了如何优化 DiffusionPipeline 以提高计算和内存效率,以及提高生成输出的质量。如果您有兴趣让您的 pipeline 变得更快,请查看以下资源

  • 了解 PyTorch 2.0torch.compile 如何使推理速度提高 5 - 300%。在 A100 GPU 上,推理速度最多可提高 50%!
  • 如果您不能使用 PyTorch 2,我们建议您安装 xFormers。其内存高效的注意力机制与 PyTorch 1.13.1 配合使用效果极佳,可提高速度并减少内存消耗。
  • 其他优化技术,例如模型卸载,在本指南中进行了介绍。
< > 更新 在 GitHub 上