Mô hình nhất quán

Consistency Models được đề xuất trong Consistency Models bởi Yang Song, Prafulla Dhariwal, Mark Chen và Ilya Sutskever.

Tóm tắt từ bài báo là:

Các mô hình khuếch tán đã thúc đẩy đáng kể các lĩnh vực tạo hình ảnh, âm thanh và video, nhưng chúng phụ thuộc vào quá trình lấy mẫu lặp lại, dẫn đến tốc độ tạo chậm. Để khắc phục hạn chế này, chúng tôi đề xuất các mô hình nhất quán (consistency models), một họ mô hình mới tạo ra các mẫu chất lượng cao bằng cách ánh xạ trực tiếp từ nhiễu sang dữ liệu. Chúng hỗ trợ tạo mẫu một bước nhanh chóng theo thiết kế, đồng thời vẫn cho phép lấy mẫu nhiều bước để đánh đổi giữa chi phí tính toán và chất lượng mẫu. Chúng cũng hỗ trợ chỉnh sửa dữ liệu zero-shot, chẳng hạn như inpainting hình ảnh, tô màu và siêu phân giải, mà không cần huấn luyện rõ ràng trên các tác vụ này. Mô hình nhất quán có thể được huấn luyện bằng cách chưng cất (distill) các mô hình khuếch tán đã huấn luyện trước, hoặc như các mô hình tạo độc lập hoàn toàn. Thông qua các thí nghiệm rộng rãi, chúng tôi chứng minh rằng chúng vượt trội hơn các kỹ thuật chưng cất hiện có cho mô hình khuếch tán trong lấy mẫu một bước và ít bước, đạt được FID mới là 3.55 trên CIFAR-10 và 6.20 trên ImageNet 64x64 cho tạo mẫu một bước. Khi được huấn luyện độc lập, mô hình nhất quán trở thành một họ mô hình tạo mới có thể vượt trội hơn các mô hình tạo một bước, không đối kháng hiện có trên các benchmark tiêu chuẩn như CIFAR-10, ImageNet 64x64 và LSUN 256x256.

Mã nguồn gốc có thể được tìm thấy tại openai/consistency_models, và các checkpoint bổ sung có sẵn tại openai.

Pipeline này được đóng góp bởi dg845 và ayushtues. ❤️

Mẹo

Để tăng tốc độ thêm, hãy sử dụng torch.compile để tạo nhiều hình ảnh trong <1 giây:

  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

>>> 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