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.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

一致性模型

一致性模型(Consistency Models)由 Yang Song、Prafulla Dhariwal、Mark Chen 和 Ilya Sutskever 在 Consistency Models 中提出。

该论文的摘要如下:

扩散模型在图像、音频和视频生成领域取得了显著进展,但它们依赖于迭代采样过程,导致生成速度缓慢。为了克服这一局限性,我们提出了一致性模型(consistency models),这是一种新型模型,通过直接将噪声映射到数据来生成高质量样本。它们在设计上支持快速的一步生成,同时仍允许多步采样,以在计算量和样本质量之间进行权衡。它们还支持零样本数据编辑,例如图像修复、上色和超分辨率,而无需针对这些任务进行显式训练。一致性模型可以通过蒸馏预训练的扩散模型来训练,也可以作为独立的生成模型进行训练。通过大量实验,我们证明它们在单步和少步采样中优于现有的扩散模型蒸馏技术,在 CIFAR-10 上实现了 3.55 的新 SOTA FID,在 ImageNet 64x64 上实现了 6.20 的新 SOTA FID(均为单步生成)。当独立训练时,一致性模型成为一种新的生成模型家族,在 CIFAR-10、ImageNet 64x64 和 LSUN 256x256 等标准基准上,能够超越现有的单步、非对抗性生成模型。

原始代码库可在 openai/consistency_models 中找到,更多检查点可在 openai 上获取。

该流水线由 dg845 和 ayushtues 贡献。❤️

提示

如需进一步提速,请使用 torch.compile 在 1 秒内生成多张图片:

  import torch
  from diffusers import ConsistencyModelPipeline

  device = "cuda"  # or "mps", "xpu", "cpu"
  # Load the cd_bedroom256_lpips checkpoint.
  model_id_or_path = "openai/diffusers-cd_bedroom256_lpips"
  pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, dtype=torch.float16)
  pipe.to(device)

+ pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

  # Multistep sampling
  # Timesteps can be explicitly specified; the particular timesteps below are from the original GitHub repo:
  # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L83
  for _ in range(10):
      image = pipe(timesteps=[17, 0]).images[0]
      image.show()

ConsistencyModelPipeline

类 diffusers.ConsistencyModelPipeline

< >

( unet: UNet2DModelscheduler: CMStochasticIterativeScheduler )

参数

用于无条件或类别条件图像生成的 Pipeline。

该模型继承自 DiffusionPipeline。有关为所有 pipeline 实现的通用方法(下载、保存、在特定设备上运行等),请参阅父类文档。

__call__

< >

( batch_size: int = 1class_labels: typing.Union[torch.Tensor, list[int], int, NoneType] = Nonenum_inference_steps: int = 1timesteps: list = Nonegenerator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Truecallback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = Nonecallback_steps: int = 1 ) → ImagePipelineOutput or tuple

参数

  • batch_size (int,可选,默认为 1) — 要生成的图像数量。
  • class_labels (torch.Tensor 或 list[int] 或 int,可选) — 用于条件化类条件一致性模型的可选类别标签。如果模型不是类条件的,则不会使用此参数。
  • num_inference_steps (int, 可选,默认为 1) — 去噪步数。更多的去噪步数通常能生成更高质量的图像,但推理速度会变慢。
  • timesteps (list[int], 可选) — 用于去噪过程的自定义时间步。如果未定义,则使用等间距的 num_inference_steps 个时间步。必须按降序排列。
  • generator (torch.Generator,可选) — 用于使生成过程确定性的 torch.Generator。
  • latents (torch.Tensor, 可选) — 预先生成的、从高斯分布中采样的含噪潜变量,用作图像生成的输入。可用于以不同的提示词微调同一生成过程。如果未提供,则使用提供的随机generator进行采样,生成一个潜变量张量。
  • output_type (str,可选,默认为 "pil") — 生成图像的输出格式。可选 PIL.Image 或 np.array。
  • return_dict (bool,可选,默认为 True) — 是否返回 ImagePipelineOutput 而不是普通元组。
  • callback (Callable, 可选) — 一个在推理过程中每隔 callback_steps 步调用一次的函数。该函数以以下参数被调用:callback(step: int, timestep: int, latents: torch.Tensor)。
  • callback_steps (int,可选,默认为 1) — 调用 callback 函数的频率。如果未指定,则会在每一步调用该回调。

返回值

ImagePipelineOutput 或 tuple

如果 return_dict 为 True,则返回 ImagePipelineOutput;否则返回一个 tuple,其中第一个元素是包含生成图像的列表。

示例:

>>> import torch

>>> from diffusers import ConsistencyModelPipeline

>>> device = "cuda"
>>> # Load the cd_imagenet64_l2 checkpoint.
>>> model_id_or_path = "openai/diffusers-cd_imagenet64_l2"
>>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
>>> pipe.to(device)

>>> # Onestep Sampling
>>> image = pipe(num_inference_steps=1).images[0]
>>> image.save("cd_imagenet64_l2_onestep_sample.png")

>>> # Onestep sampling, class-conditional image generation
>>> # ImageNet-64 class label 145 corresponds to king penguins
>>> image = pipe(num_inference_steps=1, class_labels=145).images[0]
>>> image.save("cd_imagenet64_l2_onestep_sample_penguin.png")

>>> # Multistep sampling, class-conditional image generation
>>> # Timesteps can be explicitly specified; the particular timesteps below are from the original GitHub repo:
>>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77
>>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]
>>> image.save("cd_imagenet64_l2_multistep_sample_penguin.png")

ImagePipelineOutput

类 diffusers.ImagePipelineOutput

< >

( images: list[PIL.Image.Image] | numpy.ndarray )

参数

  • images (List[PIL.Image.Image] 或 np.ndarray) — 长度为 batch_size 的去噪 PIL 图像列表,或形状为 (batch_size, height, width, num_channels) 的 NumPy 数组。

图像流水线输出类。

Update on GitHub