Hugging Face's logo

Diffusers documentation

加载调度器与模型

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.34.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

加载调度器与模型

Diffusion管道是由可互换的调度器(schedulers)和模型(models)组成的集合，可通过混合搭配来定制特定用例的流程。调度器封装了整个去噪过程（如去噪步数和寻找去噪样本的算法），其本身不包含可训练参数，因此内存占用极低。模型则主要负责从含噪输入到较纯净样本的前向传播过程。

本指南将展示如何加载调度器和模型来自定义流程。我们将全程使用stable-diffusion-v1-5/stable-diffusion-v1-5检查点，首先加载基础管道：

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")

通过pipeline.scheduler属性可查看当前管道使用的调度器：

pipeline.scheduler
PNDMScheduler {
  "_class_name": "PNDMScheduler",
  "_diffusers_version": "0.21.4",
  "beta_end": 0.012,
  "beta_schedule": "scaled_linear",
  "beta_start": 0.00085,
  "clip_sample": false,
  "num_train_timesteps": 1000,
  "set_alpha_to_one": false,
  "skip_prk_steps": true,
  "steps_offset": 1,
  "timestep_spacing": "leading",
  "trained_betas": null
}

加载调度器

调度器通过配置文件定义，同一配置文件可被多种调度器共享。使用SchedulerMixin.from_pretrained()方法加载时，需指定subfolder参数以定位配置文件在仓库中的正确子目录。

例如加载DDIMScheduler：

from diffusers import DDIMScheduler, DiffusionPipeline

ddim = DDIMScheduler.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="scheduler")

然后将新调度器传入管道：

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", scheduler=ddim, torch_dtype=torch.float16, use_safetensors=True
).to("cuda")

调度器对比

不同调度器各有优劣，难以定量评估哪个最适合您的流程。通常需要在去噪速度与质量之间权衡。我们建议尝试多种调度器以找到最佳方案。通过pipeline.scheduler.compatibles属性可查看兼容当前管道的所有调度器。

下面我们使用相同提示词和随机种子，对比LMSDiscreteScheduler、EulerDiscreteScheduler、EulerAncestralDiscreteScheduler和DPMSolverMultistepScheduler的表现：

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")

prompt = "A photograph of an astronaut riding a horse on Mars, high resolution, high definition."
generator = torch.Generator(device="cuda").manual_seed(8)

使用from_config()方法加载不同调度器的配置来切换管道调度器：

LMSDiscreteScheduler
EulerDiscreteScheduler
EulerAncestralDiscreteScheduler
DPMSolverMultistepScheduler

LMSDiscreteScheduler通常能生成比默认调度器更高质量的图像。

from diffusers import LMSDiscreteScheduler

pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[0]
image
LMSDiscreteScheduler
EulerDiscreteScheduler
EulerAncestralDiscreteScheduler
DPMSolverMultistepScheduler

多数生成图像质量相近，实际选择需根据具体场景测试多种调度器进行比较。

Flax调度器

对比Flax调度器时，需额外将调度器状态加载到模型参数中。例如将FlaxStableDiffusionPipeline的默认调度器切换为超高效的FlaxDPMSolverMultistepScheduler：

[!警告] FlaxLMSDiscreteScheduler和FlaxDDPMScheduler目前暂不兼容FlaxStableDiffusionPipeline。

import jax
import numpy as np
from flax.jax_utils import replicate
from flax.training.common_utils import shard
from diffusers import FlaxStableDiffusionPipeline, FlaxDPMSolverMultistepScheduler

scheduler, scheduler_state = FlaxDPMSolverMultistepScheduler.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    subfolder="scheduler"
)
pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    scheduler=scheduler,
    variant="bf16",
    dtype=jax.numpy.bfloat16,
)
params["scheduler"] = scheduler_state

利用Flax对TPU的兼容性实现并行图像生成。需为每个设备复制模型参数，并分配输入数据：

# 每个并行设备生成1张图像（TPUv2-8/TPUv3-8支持8设备并行）
prompt = "一张宇航员在火星上骑马的高清照片，高分辨率，高画质。"
num_samples = jax.device_count()
prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)

prng_seed = jax.random.PRNGKey(0)
num_inference_steps = 25

# 分配输入和随机种子
params = replicate(params)
prng_seed = jax.random.split(prng_seed, jax.device_count())
prompt_ids = shard(prompt_ids)

images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images
images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))

模型加载

通过ModelMixin.from_pretrained()方法加载模型，该方法会下载并缓存模型权重和配置的最新版本。若本地缓存已存在最新文件，则直接复用缓存而非重复下载。

通过subfolder参数可从子目录加载模型。例如stable-diffusion-v1-5/stable-diffusion-v1-5的模型权重存储在unet子目录中：

from diffusers import UNet2DConditionModel

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

也可直接从仓库加载：

from diffusers import UNet2DModel

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

加载和保存模型变体时，需在ModelMixin.from_pretrained()和ModelMixin.save_pretrained()中指定variant参数：

from diffusers import UNet2DConditionModel

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

使用from_pretrained()的torch_dtype参数指定模型加载精度：

from diffusers import AutoModel

unet = AutoModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", torch_dtype=torch.float16
)

也可使用torch.Tensor.to方法即时转换精度，但会转换所有权重（不同于torch_dtype参数会保留_keep_in_fp32_modules中的层）。这对某些必须保持fp32精度的层尤为重要（参见示例）。

< > Update on GitHub

←有效和高效的扩散 Accelerate inference→