Add Consistency Models Pipeline

docs/source/en/_toctree.yml

@@ -184,6 +184,8 @@
       title: Audio Diffusion
     - local: api/pipelines/audioldm
       title: AudioLDM
+    - local: api/pipelines/consistency_models
+      title: Consistency Models
     - local: api/pipelines/controlnet
       title: ControlNet
     - local: api/pipelines/cycle_diffusion
@@ -274,6 +276,8 @@
   - sections:
     - local: api/schedulers/overview
       title: Overview
+    - local: api/schedulers/cm_stochastic_iterative
+      title: Consistency Model Multistep Scheduler
     - local: api/schedulers/ddim
       title: DDIM
     - local: api/schedulers/ddim_inverse

docs/source/en/api/pipelines/consistency_models.mdx

@@ -0,0 +1,87 @@
+# Consistency Models
+
+Consistency Models were proposed in [Consistency Models](https://arxiv.org/abs/2303.01469) by Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever.
+
+The abstract of the [paper](https://arxiv.org/pdf/2303.01469.pdf) is as follows:
+
+*Diffusion models have significantly advanced the fields of image, audio, and video generation, but they depend on an iterative sampling process that causes slow generation. To overcome this limitation, we propose consistency models, a new family of models that generate high quality samples by directly mapping noise to data. They support fast one-step generation by design, while still allowing multistep sampling to trade compute for sample quality. They also support zero-shot data editing, such as image inpainting, colorization, and super-resolution, without requiring explicit training on these tasks. Consistency models can be trained either by distilling pre-trained diffusion models, or as standalone generative models altogether. Through extensive experiments, we demonstrate that they outperform existing distillation techniques for diffusion models in one- and few-step sampling, achieving the new state-of-the-art FID of 3.55 on CIFAR-10 and 6.20 on ImageNet 64x64 for one-step generation. When trained in isolation, consistency models become a new family of generative models that can outperform existing one-step, non-adversarial generative models on standard benchmarks such as CIFAR-10, ImageNet 64x64 and LSUN 256x256. *
+
+Resources:
+
+* [Paper](https://arxiv.org/abs/2303.01469)
+* [Original Code](https://github.com/openai/consistency_models)
+
+Available Checkpoints are:
+- *cd_imagenet64_l2 (64x64 resolution)* [openai/consistency-model-pipelines](https://huggingface.co/openai/consistency-model-pipelines)
+- *cd_imagenet64_lpips (64x64 resolution)* [openai/diffusers-cd_imagenet64_lpips](https://huggingface.co/openai/diffusers-cd_imagenet64_lpips)
+- *ct_imagenet64 (64x64 resolution)* [openai/diffusers-ct_imagenet64](https://huggingface.co/openai/diffusers-ct_imagenet64)
+- *cd_bedroom256_l2 (256x256 resolution)* [openai/diffusers-cd_bedroom256_l2](https://huggingface.co/openai/diffusers-cd_bedroom256_l2)
+- *cd_bedroom256_lpips (256x256 resolution)* [openai/diffusers-cd_bedroom256_lpips](https://huggingface.co/openai/diffusers-cd_bedroom256_lpips)
+- *ct_bedroom256 (256x256 resolution)* [openai/diffusers-ct_bedroom256](https://huggingface.co/openai/diffusers-ct_bedroom256)
+- *cd_cat256_l2 (256x256 resolution)* [openai/diffusers-cd_cat256_l2](https://huggingface.co/openai/diffusers-cd_cat256_l2)
+- *cd_cat256_lpips (256x256 resolution)* [openai/diffusers-cd_cat256_lpips](https://huggingface.co/openai/diffusers-cd_cat256_lpips)
+- *ct_cat256 (256x256 resolution)* [openai/diffusers-ct_cat256](https://huggingface.co/openai/diffusers-ct_cat256)
+
+## Available Pipelines
+
+| Pipeline | Tasks | Demo | Colab |
+|:---:|:---:|:---:|:---:|
+| [ConsistencyModelPipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pipeline_consistency_models.py) | *Unconditional Image Generation* | | |
+
+This pipeline was contributed by our community members [dg845](https://github.com/dg845) and [ayushtues](https://huggingface.co/ayushtues) :heart: 
+
+## Usage Example
+
+```python
+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("consistency_model_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("consistency_model_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(timesteps=[22, 0], class_labels=145).images[0]
+image.save("consistency_model_multistep_sample_penguin.png")
+```
+
+For an additional speed-up, one can also make use of `torch.compile`. Multiple images can be generated in <1 second as follows:
+
+```py
+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
+[[autodoc]] ConsistencyModelPipeline
+    - all
+    - __call__

docs/source/en/api/schedulers/cm_stochastic_iterative.mdx

@@ -0,0 +1,11 @@
+# Consistency Model Multistep Scheduler
+
+## Overview
+
+Multistep and onestep scheduler (Algorithm 1) introduced alongside consistency models in the paper [Consistency Models](https://arxiv.org/abs/2303.01469) by Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever.
+Based on the [original consistency models implementation](https://github.com/openai/consistency_models).
+Should generate good samples from [`ConsistencyModelPipeline`] in one or a small number of steps.
+
+## CMStochasticIterativeScheduler
+[[autodoc]] CMStochasticIterativeScheduler
+