Modelos de consistencia

Los Consistency Models fueron propuestos en Consistency Models por Yang Song, Prafulla Dhariwal, Mark Chen e Ilya Sutskever.

El resumen del artículo es:

Los modelos de difusión han avanzado significativamente en los campos de generación de imágenes, audio y vídeo, pero dependen de un proceso de muestreo iterativo que provoca una generación lenta. Para superar esta limitación, proponemos los modelos de consistencia, una nueva familia de modelos que generan muestras de alta calidad mapeando directamente el ruido a los datos. Por diseño, admiten una generación rápida en un solo paso, al tiempo que permiten el muestreo multipaso para intercambiar cómputo por calidad de muestra. También admiten la edición de datos sin ejemplos (zero-shot), como el relleno de imágenes (inpainting), la colorización y la superresolución, sin necesidad de entrenamiento explícito en estas tareas. Los modelos de consistencia se pueden entrenar ya sea destilando modelos de difusión preentrenados, o como modelos generativos independientes. A través de experimentos extensos, demostramos que superan a las técnicas de destilación existentes para modelos de difusión en el muestreo de uno y pocos pasos, logrando el nuevo estado del arte en FID de 3.55 en CIFAR-10 y 6.20 en ImageNet 64x64 para la generación en un solo paso. Cuando se entrenan de forma aislada, los modelos de consistencia se convierten en una nueva familia de modelos generativos que pueden superar a los modelos generativos existentes de un solo paso y no adversariales en benchmarks estándar como CIFAR-10, ImageNet 64x64 y LSUN 256x256.

El código original se puede encontrar en openai/consistency_models, y hay puntos de control adicionales disponibles en openai.

El pipeline fue aportado por dg845 y ayushtues. ❤️

Consejos

Para una aceleración adicional, use torch.compile para generar varias imágenes en <1 segundo:

  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