Diffusers documentation
Модели согласованности
Модели согласованности
Модели согласованности были предложены в работе Consistency Models Янгом Сонгом, Прафуллой Дхаривалом, Марком Ченом и Ильей Сутскевером.
Аннотация из статьи:
Модели диффузии значительно продвинули области генерации изображений, аудио и видео, но они зависят от итеративного процесса выборки, что приводит к медленной генерации. Чтобы преодолеть это ограничение, мы предлагаем consistency models — новое семейство моделей, которые генерируют образцы высокого качества, напрямую отображая шум в данные. По своей конструкции они поддерживают быструю генерацию за один шаг, при этом позволяя многошаговую выборку для обмена вычислительных ресурсов на качество образцов. Они также поддерживают редактирование данных в режиме zero-shot, такое как inpainting изображений, колоризация и суперразрешение, без необходимости явного обучения на этих задачах. Consistency models можно обучать либо путём дистилляции предобученных диффузионных моделей, либо как самостоятельные генеративные модели. В ходе обширных экспериментов мы демонстрируем, что они превосходят существующие техники дистилляции для диффузионных моделей при одношаговой и малошаговой выборке, достигая нового state-of-the-art значения FID 3.55 на CIFAR-10 и 6.20 на ImageNet 64x64 для генерации за один шаг. При обучении в изоляции consistency models становятся новым семейством генеративных моделей, которые могут превосходить существующие одношаговые, неадверсариальные генеративные модели на стандартных бенчмарках, таких как CIFAR-10, ImageNet 64x64 и LSUN 256x256.
Исходный код можно найти в openai/consistency_models, а дополнительные контрольные точки доступны в openai.
Пайплайн был предложен dg845 и ayushtues. ❤️
Советы
Для дополнительного ускорения используйте torch.compile, чтобы генерировать несколько изображений менее чем за 1 секунду:
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
класс diffusers.ConsistencyModelPipeline
< source >( unet: UNet2DModelscheduler: CMStochasticIterativeScheduler )
Параметры
- unet (UNet2DModel) —
UNet2DModelдля удаления шума из закодированных латентов изображения. - scheduler (SchedulerMixin) —
планировщик, используемый в связке с
unetдля удаления шума из закодированных латентов изображения. В настоящее время совместим только с CMStochasticIterativeScheduler.
Пайплайн для безусловной или условной (по классу) генерации изображений.
Эта модель наследуется от DiffusionPipeline. Общую документацию по методам, реализованным для всех пайплайнов (загрузка, сохранение, запуск на конкретном устройстве и т. д.), см. в документации родительского класса.
__call__
< source >( batch_size: int = 1class_labels: typing.Union[torch.Tensor, list[int], int, NoneType] = Nonenum_inference_steps: int = 1timesteps: list[int] = Nonegenerator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Truecallback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = Nonecallback_steps: int = 1 ) → ImagePipelineOutput or tuple
Параметры
- batch_size (
int, необязательный, по умолчанию 1) — Количество изображений для генерации. - class_labels (
torch.Tensorилиlist[int]илиint, необязательный) — Необязательные метки классов для условного настраивания класс-условных моделей согласованности. Не используются, если модель не является класс-условной. - num_inference_steps (
int, необязательный, по умолчанию 1) — Количество шагов денойзинга. Как правило, большее число шагов денойзинга приводит к изображению более высокого качества, но за счёт более медленного вывода. - timesteps (
list[int], необязательный) — Пользовательские timesteps для процесса денойзинга. Если не заданы, используются равноотстоящиеnum_inference_stepstimesteps. Должны быть в порядке убывания. - 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. Если не указано, 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
класс diffusers.ImagePipelineOutput
< source >( images: list[PIL.Image.Image] | numpy.ndarray )
Класс вывода для пайплайнов обработки изображений.