Hugging Face's logo

Diffusers 文档

一致性模型

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

一致性模型

一致性模型在 Consistency Models 中被 Yang Song, Prafulla Dhariwal, Mark Chen, 和 Ilya Sutskever 提出。

该论文的摘要是

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

原始代码库可以在 openai/consistency_models 找到，其他检查点可以在 openai 找到。

该 pipeline 由 dg845 和 ayushtues 贡献。 ❤️

提示

为了进一步加速，可以使用 torch.compile 在 <1 秒的时间内生成多张图像

  import torch
  from diffusers import ConsistencyModelPipeline

  device = "cuda"
  # Load the cd_bedroom256_lpips checkpoint.
  model_id_or_path = "openai/diffusers-cd_bedroom256_lpips"
  pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_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

class diffusers.ConsistencyModelPipeline

< >

( unet: UNet2DModel scheduler: CMStochasticIterativeScheduler )

参数

用于无条件或类条件图像生成的管线。

此模型继承自 DiffusionPipeline。请查看超类文档以了解为所有管线实现的通用方法（下载、保存、在特定设备上运行等）。

__call__

< >

( batch_size: int = 1 class_labels: typing.Union[torch.Tensor, typing.List[int], int, NoneType] = None num_inference_steps: int = 1 timesteps: typing.List[int] = None generator: typing.Union[torch._C.Generator, typing.List[torch._C.Generator], NoneType] = None latents: typing.Optional[torch.Tensor] = None output_type: typing.Optional[str] = 'pil' return_dict: bool = True callback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = None callback_steps: int = 1 ) → ImagePipelineOutput 或 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

class diffusers.ImagePipelineOutput

< >

( images: typing.Union[typing.List[PIL.Image.Image], numpy.ndarray] )

参数

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

图像管线的输出类。

< > 在 GitHub 上更新

←CogView3 ControlNet→