【AIGC】Diffusers:加载管道、模型和调度程序

前言

拥有一种使用扩散系统进行推理的简单方法对于扩散器至关重要。扩散系统通常由多个组件组成，例如参数化模型、分词器和调度器，它们以复杂的方式进行交互。这就是为什么我们设计了 DiffusionPipeline，将整个扩散系统的复杂性包装成一个易于使用的 API，同时保持足够的灵活性以适应其他用例，例如将每个组件单独加载为构建块以组装您自己的扩散系统。

推理或训练所需的一切都可以通过该 from_pretrained() 方法访问。

 本指南将向您展示如何加载：

• 来自中心和本地的管道
• 将不同的组件放入管道中
• 模型变体，例如不同的浮点类型或非指数平均平均 （EMA） 权重
• 模型和调度程序

扩散管线 

  如果您有兴趣更详细地了解 DiffusionPipeline 类的工作原理，请跳到DiffusionPipeline 说明部分。

DiffusionPipeline 类是从 Hub 加载最新趋势扩散模型的最简单、最通用的方法。DiffusionPipeline.from_pretrained（） 方法自动从检查点检测正确的管道类，下载并缓存所有必需的配置和权重文件，并返回准备进行推理的管道实例。

from diffusers import DiffusionPipeline

repo_id = "runwayml/stable-diffusion-v1-5"
pipe = DiffusionPipeline.from_pretrained(repo_id, use_safetensors=True)

 还可以加载具有特定管道类的检查点。上面的示例加载了一个稳定扩散模型;若要获得相同的结果，请使用 StableDiffusionPipeline 类：

from diffusers import StableDiffusionPipeline

repo_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(repo_id, use_safetensors=True)

 模型（如 CompVis/stable-diffusion-v1-4 或runwayml/stable-diffusion-v1-5 ）也可以用于多个任务，例如文本到图像或图像到图像。若要区分要将模型用于的任务，必须直接使用其相应的特定于任务的管道类加载它：

from diffusers import StableDiffusionImg2ImgPipeline

repo_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(repo_id)

 本地管道

 若要在本地加载扩散管道，请使用 git-lfs 手动将模型（在本例中为 runwayml/stable-diffusion-v1-5 ）下载到本地磁盘。这会在磁盘上创建一个本地文件夹 ./stable-diffusion-v1-5 ， ：

git-lfs install
git clone https://huggingface.co/runwayml/stable-diffusion-v1-5

 然后将本地路径传递给 from_pretrained（）：

from diffusers import DiffusionPipeline

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

 当 from_pretrained（） 方法检测到本地路径时，它不会从 Hub 下载任何文件，但这也意味着它不会下载和缓存检查点的最新更改。

 在管道中交换组件

 您可以使用另一个兼容组件自定义任何管道的默认组件。定制很重要，因为：

• 更改调度程序对于探索生成速度和质量之间的权衡非常重要。
• 模型的不同组件通常是独立训练的，您可以将组件换成性能更好的组件。
• 在微调过程中，通常只训练某些组件（如 UNet 或文本编码器）。

 若要找出哪些计划程序与自定义兼容，可以使用以下 compatibles 方法：

from diffusers import DiffusionPipeline

repo_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion = DiffusionPipeline.from_pretrained(repo_id, use_safetensors=True)
stable_diffusion.scheduler.compatibles

让我们使用 SchedulerMixin.from_pretrained（） 方法将默认的 PNDMScheduler 替换为性能更高的调度器 EulerDiscreteScheduler。从管道存储库的正确子文件夹加载调度程序配置需要该 subfolder="scheduler" 参数。

然后，您可以将新的 EulerDiscreteScheduler 实例传递给 DiffusionPipeline 中的 scheduler 参数：

from diffusers import DiffusionPipeline, EulerDiscreteScheduler

repo_id = "runwayml/stable-diffusion-v1-5"
scheduler = EulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
stable_diffusion = DiffusionPipeline.from_pretrained(repo_id, scheduler=scheduler, use_safetensors=True)

安全检查器

 像 Stable Diffusion 这样的扩散模型可能会产生有害内容，这就是为什么 Diffusers 有一个安全检查器来检查生成的输出是否符合已知的硬编码 NSFW 内容。如果您出于任何原因想要禁用安全检查器，请传递 None safety_checker 参数：

from diffusers import DiffusionPipeline

repo_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion = DiffusionPipeline.from_pretrained(repo_id, safety_checker=None, use_safetensors=True)
"""
You have disabled the safety checker for  by passing `safety_checker=None`. Ensure that you abide by the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend keeping the safety filter enabled in all public-facing circumstances, disabling it only for use cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .
"""

 跨管道重用组件

您还可以在多个管道中重复使用相同的组件，以避免两次将权重加载到 RAM 中。使用 components 方法保存组件： 

from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline

model_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained(model_id, use_safetensors=True)

components = stable_diffusion_txt2img.components

 然后，您可以将其 components 传递到另一个管道，而无需将权重重新加载到 RAM 中：

stable_diffusion_img2img = StableDiffusionImg2ImgPipeline(**components)

如果希望更灵活地重用或禁用哪些组件，还可以将组件单独传递到管道。例如，若要在图像到图像管道中重用文本到图像管道中的相同组件，但安全检查器和特征提取器除外：

from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline

model_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained(model_id, use_safetensors=True)
stable_diffusion_img2img = StableDiffusionImg2ImgPipeline(
    vae=stable_diffusion_txt2img.vae,
    text_encoder=stable_diffusion_txt2img.text_encoder,
    tokenizer=stable_diffusion_txt2img.tokenizer,
    unet=stable_diffusion_txt2img.unet,
    scheduler=stable_diffusion_txt2img.scheduler,
    safety_checker=None,
    feature_extractor=None,
    requires_safety_checker=False,
)

模型权重变体

 以不同的浮点类型存储，以实现较低的精度和较低的存储，例如 torch.float16 ，因为它只需要一半的带宽和存储即可下载。如果您正在继续训练或使用 CPU，则无法使用此变体。

非指数平均 （EMA） 权重，不应用于推理。您应该使用它们来继续微调模型。 

  当检查点具有相同的模型结构，但它们是在不同的数据集上训练的，并且使用不同的训练设置时，它们应该存储在单独的存储库中，而不是变体（例如， stable-diffusion-v1-4 和 stable-diffusion-v1-5 ）。

 否则，变体与原模型相同。它们具有完全相同的序列化格式（如 Safetensors）、模型结构和具有相同张量形状的权重。

checkpoint type 检查点类型	weight name 重量名称	argument for loading weights 加载砝码的参数
original 源语言	diffusion_pytorch_model.bin
floating point 浮点	diffusion_pytorch_model.fp16.bin	variant, torch_dtype
non-EMA 非 EMA	diffusion_pytorch_model.non_ema.bin	variant

 对于加载变体，有两个重要的参数需要了解：

 

• torch_dtype 定义加载的检查点的浮点精度。例如，如果要通过加载 fp16 变体来节省带宽，则应指定 torch_dtype=torch.float16 将权重转换为 fp16 。否则， fp16 权重将转换为默认 fp32 精度。您也可以在不定义 variant 参数的情况下加载原始检查点，并将其转换为 fp16 with torch_dtype=torch.float16 .在这种情况下，将首先下载默认 fp32 权重，然后在加载后将其 fp16 转换为默认权重。

• variant 定义应从存储库加载哪些文件。例如，如果要从存储库加载 non_ema 变体，则应指定 variant="non_ema" 下载 non_ema diffusers/stable-diffusion-variants 文件。

from diffusers import DiffusionPipeline
import torch

# load fp16 variant
stable_diffusion = DiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16, use_safetensors=True
)
# load non_ema variant
stable_diffusion = DiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", variant="non_ema", use_safetensors=True
)

若要将存储在不同浮点类型中的检查点保存为非 EMA 变体，请使用 DiffusionPipeline.save_pretrained（） 方法并指定 variant 参数。您应该尝试将变体保存到与原始检查点相同的文件夹中，以便您可以从同一文件夹加载两者：

from diffusers import DiffusionPipeline

# save as fp16 variant
stable_diffusion.save_pretrained("runwayml/stable-diffusion-v1-5", variant="fp16")
# save as non-ema variant
stable_diffusion.save_pretrained("runwayml/stable-diffusion-v1-5", variant="non_ema")

如果不将变体保存到现有文件夹，则必须指定参数， variant 否则它将抛出一个 Exception ，因为它找不到原始检查点：

#  this won't work
stable_diffusion = DiffusionPipeline.from_pretrained(
    "./stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
#  this works
stable_diffusion = DiffusionPipeline.from_pretrained(
    "./stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16, use_safetensors=True
)

模型

 模型从 ModelMixin.from_pretrained（） 方法加载，该方法下载并缓存最新版本的模型权重和配置。如果本地缓存中有最新的文件，则 from_pretrained（） 会重用缓存中的文件，而不是重新下载它们。

 可以从带有参数的 subfolder 子文件夹加载模型。例如，模型权重 runwayml/stable-diffusion-v1-5 存储在 unet 子文件夹中：

from diffusers import UNet2DConditionModel

repo_id = "runwayml/stable-diffusion-v1-5"
model = UNet2DConditionModel.from_pretrained(repo_id, subfolder="unet", use_safetensors=True)

或者直接从存储库的目录：

from diffusers import UNet2DModel

repo_id = "google/ddpm-cifar10-32"
model = UNet2DModel.from_pretrained(repo_id, use_safetensors=True)

您还可以通过在 ModelMixin.from_pretrained（） 和 ModelMixin.save_pretrained（） 中指定 variant 参数来加载和保存模型变体：

from diffusers import UNet2DConditionModel

model = UNet2DConditionModel.from_pretrained(
    "runwayml/stable-diffusion-v1-5", subfolder="unet", variant="non_ema", use_safetensors=True
)
model.save_pretrained("./local-unet", variant="non_ema")

 调度程序

调度器是从 SchedulerMixin.from_pretrained（） 方法加载的，与模型不同，调度器没有参数化或训练;它们由配置文件定义。 

加载调度程序不会消耗任何大量内存，并且相同的配置文件可用于各种不同的调度程序。例如，以下调度程序与 StableDiffusionPipeline 兼容，这意味着您可以在以下任何类中加载相同的调度程序配置文件： 

from diffusers import StableDiffusionPipeline
from diffusers import (
    DDPMScheduler,
    DDIMScheduler,
    PNDMScheduler,
    LMSDiscreteScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    DPMSolverMultistepScheduler,
)

repo_id = "runwayml/stable-diffusion-v1-5"

ddpm = DDPMScheduler.from_pretrained(repo_id, subfolder="scheduler")
ddim = DDIMScheduler.from_pretrained(repo_id, subfolder="scheduler")
pndm = PNDMScheduler.from_pretrained(repo_id, subfolder="scheduler")
lms = LMSDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
euler = EulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
dpm = DPMSolverMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")

# replace `dpm` with any of `ddpm`, `ddim`, `pndm`, `lms`, `euler_anc`, `euler`
pipeline = StableDiffusionPipeline.from_pretrained(repo_id, scheduler=dpm, use_safetensors=True)

DiffusionPipeline 解释

 作为类方法，DiffusionPipeline.from_pretrained（） 负责两件事：

• 下载推理所需的最新版本的文件夹结构并对其进行缓存。如果本地缓存中提供了最新的文件夹结构，则 DiffusionPipeline.from_pretrained（） 会重用缓存，并且不会重新下载文件。
• 将缓存的权重加载到正确的管道类（从文件中 model_index.json 检索），并返回该类的实例。

 管道的基础文件夹结构与其类实例直接对应。例如，StableDiffusionPipeline 对应于 runwayml/stable-diffusion-v1-5 中的文件夹结构。

from diffusers import DiffusionPipeline

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

 你将看到 pipeline 是 StableDiffusionPipeline 的一个实例，它由七个组件组成：

 

• "feature_extractor" ：来自 Transformer 的 CLIPImageProcessor。
• "safety_checker" ：用于筛选有害内容的组件。
• "scheduler" ：PNDMScheduler 的实例。
• "text_encoder" ：来自 Transformer 的 CLIPTextModel。
• "tokenizer" ：来自Transformer的 CLIPTokenizer。
• "unet" ：UNet2DConditionModel 的实例。
• "vae" ：AutoencoderKL 的实例。

StableDiffusionPipeline {
  "feature_extractor": [
    "transformers",
    "CLIPImageProcessor"
  ],
  "safety_checker": [
    "stable_diffusion",
    "StableDiffusionSafetyChecker"
  ],
  "scheduler": [
    "diffusers",
    "PNDMScheduler"
  ],
  "text_encoder": [
    "transformers",
    "CLIPTextModel"
  ],
  "tokenizer": [
    "transformers",
    "CLIPTokenizer"
  ],
  "unet": [
    "diffusers",
    "UNet2DConditionModel"
  ],
  "vae": [
    "diffusers",
    "AutoencoderKL"
  ]
}

将管道实例的组件与文件夹结构进行比较，你将看到存储库中的每个组件都有一个单独的 runwayml/stable-diffusion-v1-5 文件夹：

.
├── feature_extractor
│   └── preprocessor_config.json
├── model_index.json
├── safety_checker
│   ├── config.json
|   ├── model.fp16.safetensors
│   ├── model.safetensors
│   ├── pytorch_model.bin
|   └── pytorch_model.fp16.bin
├── scheduler
│   └── scheduler_config.json
├── text_encoder
│   ├── config.json
|   ├── model.fp16.safetensors
│   ├── model.safetensors
│   |── pytorch_model.bin
|   └── pytorch_model.fp16.bin
├── tokenizer
│   ├── merges.txt
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   └── vocab.json
├── unet
│   ├── config.json
│   ├── diffusion_pytorch_model.bin
|   |── diffusion_pytorch_model.fp16.bin
│   |── diffusion_pytorch_model.f16.safetensors
│   |── diffusion_pytorch_model.non_ema.bin
│   |── diffusion_pytorch_model.non_ema.safetensors
│   └── diffusion_pytorch_model.safetensors
|── vae
.   ├── config.json
.   ├── diffusion_pytorch_model.bin
    ├── diffusion_pytorch_model.fp16.bin
    ├── diffusion_pytorch_model.fp16.safetensors
    └── diffusion_pytorch_model.safetensors

您可以将管道的每个组件作为属性进行访问，以查看其配置：

pipeline.tokenizer
CLIPTokenizer(
    name_or_path="/root/.cache/huggingface/hub/models--runwayml--stable-diffusion-v1-5/snapshots/39593d5650112b4cc580433f6b0435385882d819/tokenizer",
    vocab_size=49408,
    model_max_length=77,
    is_fast=False,
    padding_side="right",
    truncation_side="right",
    special_tokens={
        "bos_token": AddedToken("<|startoftext|>", rstrip=False, lstrip=False, single_word=False, normalized=True),
        "eos_token": AddedToken("<|endoftext|>", rstrip=False, lstrip=False, single_word=False, normalized=True),
        "unk_token": AddedToken("<|endoftext|>", rstrip=False, lstrip=False, single_word=False, normalized=True),
        "pad_token": "<|endoftext|>",
    },
    clean_up_tokenization_spaces=True
)

每个管道都需要一个 model_index.json 文件，该文件告诉 DiffusionPipeline：

• 要从 _class_name 哪个管道类加载
• 哪个版本的扩散器用于创建模型 _diffusers_version

{
  "_class_name": "StableDiffusionPipeline",
  "_diffusers_version": "0.6.0",
  "feature_extractor": [
    "transformers",
    "CLIPImageProcessor"
  ],
  "safety_checker": [
    "stable_diffusion",
    "StableDiffusionSafetyChecker"
  ],
  "scheduler": [
    "diffusers",
    "PNDMScheduler"
  ],
  "text_encoder": [
    "transformers",
    "CLIPTextModel"
  ],
  "tokenizer": [
    "transformers",
    "CLIPTokenizer"
  ],
  "unet": [
    "diffusers",
    "UNet2DConditionModel"
  ],
  "vae": [
    "diffusers",
    "AutoencoderKL"
  ]
}

参考链接：

 https://huggingface.co/docs/diffusers/using-diffusers/loading