DiT

Масштабируемые диффузионные модели на основе трансформеров (DiT) — работа Уильяма Пиблса и Сэйнинга Ши.

Аннотация из статьи:

Мы исследуем новый класс диффузионных моделей, основанных на архитектуре трансформера. Мы обучаем латентные диффузионные модели изображений, заменяя обычно используемый U-Net-бэкбон трансформером, работающим с латентными патчами. Мы анализируем масштабируемость наших Diffusion Transformers (DiTs) с точки зрения сложности прямого прохода, измеряемой в Gflops. Мы обнаруживаем, что DiTs с более высоким значением Gflops — за счёт увеличения глубины/ширины трансформера или увеличения количества входных токенов — последовательно демонстрируют более низкий FID. Помимо хороших свойств масштабируемости, наши самые большие модели DiT-XL/2 превосходят все предыдущие диффузионные модели на бенчмарках ImageNet 512x512 и 256x256 с классовым условием, достигая передового значения FID 2.27 на последнем из них.

Исходный код можно найти в facebookresearch/dit.

Обязательно ознакомьтесь с руководством по планировщикам, чтобы узнать, как исследовать баланс между скоростью и качеством планировщика, а также с разделом повторное использование компонентов в разных конвейерах, чтобы узнать, как эффективно загружать одни и те же компоненты в несколько конвейеров.

DiTPipeline

>>> from diffusers import DiTPipeline, DPMSolverMultistepScheduler
>>> import torch

>>> pipe = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256", torch_dtype=torch.float16)
>>> pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
>>> pipe = pipe.to("cuda")

>>> # pick words from Imagenet class labels
>>> pipe.labels  # to print all available words

>>> # pick words that exist in ImageNet
>>> words = ["white shark", "umbrella"]

>>> class_ids = pipe.get_label_ids(words)

>>> generator = torch.manual_seed(33)
>>> output = pipe(class_labels=class_ids, num_inference_steps=25, generator=generator)

>>> image = output.images[0]  # label 'white shark'

ImagePipelineOutput