Modèles de cohérence

Les Consistency Models ont été proposés dans Consistency Models par Yang Song, Prafulla Dhariwal, Mark Chen et Ilya Sutskever.

L'abstract de l'article est :

Les modèles de diffusion ont considérablement fait progresser les domaines de la génération d’images, d’audio et de vidéo, mais ils reposent sur un processus d’échantillonnage itératif qui entraîne une génération lente. Pour surmonter cette limitation, nous proposons les modèles de cohérence, une nouvelle famille de modèles qui génèrent des échantillons de haute qualité en cartographiant directement le bruit vers les données. Ils prennent en charge la génération rapide en une étape par conception, tout en permettant un échantillonnage multi-étapes pour échanger le calcul contre la qualité des échantillons. Ils prennent également en charge l’édition de données en zéro tirage, comme l’interpolation d’images, la colorisation et la super-résolution, sans nécessiter d’entraînement explicite sur ces tâches. Les modèles de cohérence peuvent être entraînés soit en distillant des modèles de diffusion pré-entraînés, soit comme des modèles génératifs autonomes. À travers des expériences exhaustives, nous démontrons qu’ils surpassent les techniques de distillation existantes pour les modèles de diffusion en échantillonnage en une ou quelques étapes, atteignant le nouvel état de l’art avec un FID de 3,55 sur CIFAR-10 et 6,20 sur ImageNet 64x64 pour la génération en une étape. Lorsqu’ils sont entraînés de manière isolée, les modèles de cohérence deviennent une nouvelle famille de modèles génératifs qui peuvent surpasser les modèles génératifs existants en une étape et non adversariaux sur des benchmarks standard tels que CIFAR-10, ImageNet 64x64 et LSUN 256x256.

Le code source original peut être trouvé sur openai/consistency_models, et des points de contrôle supplémentaires sont disponibles sur openai.

Le pipeline a été contribué par dg845 et ayushtues. ❤️

Astuces

Pour un gain de vitesse supplémentaire, utilisez torch.compile pour générer plusieurs images en moins d'une seconde :

  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