Diffusers documentation

PixArt-α

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

PixArt-α

PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis — это работа Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu и Zhenguo Li.

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

Самые передовые модели генерации изображений по тексту (T2I) требуют значительных затрат на обучение (например, миллионы часов работы GPU), что серьёзно препятствует фундаментальным инновациям в сообществе AIGC и увеличивает выбросы CO2. В данной статье представлена PIXART-α — диффузионная T2I-модель на основе Transformer, качество генерации изображений которой конкурентоспособно с передовыми генераторами изображений (например, Imagen, SDXL и даже Midjourney), достигая стандартов, близких к коммерческому применению. Кроме того, модель поддерживает синтез изображений высокого разрешения вплоть до 1024px при низких затратах на обучение, как показано на рисунках 1 и 2. Для достижения этой цели предложены три ключевых решения: (1) Декомпозиция стратегии обучения: мы разработали три различных этапа обучения, которые независимо оптимизируют пиксельные зависимости, согласование текста и изображения, а также эстетическое качество изображения; (2) Эффективный T2I-Transformer: мы интегрировали модули кросс-внимания в Diffusion Transformer (DiT) для внедрения текстовых условий и упрощения вычислительно затратного ветвления с классовыми условиями; (3) Данные с высокой информативностью: мы подчёркиваем значимость плотности концепций в парах «текст-изображение» и используем крупную Vision-Language модель для автоматической разметки плотных псевдо-описаний, чтобы помочь в обучении согласованию текста и изображения. В результате скорость обучения PIXART-α значительно превосходит существующие крупномасштабные T2I-модели: например, обучение PIXART-α занимает лишь 10,8% времени обучения Stable Diffusion v1.5 (675 против 6250 GPU-дней A100), что экономит почти $300,000 ($26,000 против $320,000) и снижает выбросы CO2 на 90%. Более того, по сравнению с более крупной SOTA-моделью RAPHAEL, наши затраты на обучение составляют всего 1%. Многочисленные эксперименты демонстрируют, что PIXART-α превосходит другие модели по качеству изображений, художественности и семантическому контролю. Мы надеемся, что PIXART-α предоставит новые идеи сообществу AIGC и стартапам, чтобы ускорить создание собственных высококачественных, но недорогих генеративных моделей с нуля.

Исходный код можно найти в PixArt-alpha/PixArt-alpha, а все доступные контрольные точки — в PixArt-alpha.

Некоторые замечания по этому конвейеру:

  • Для шумоподавления используется базовая архитектура Transformer (вместо UNet). В связи с этим её архитектура похожа на архитектуру DiT.
  • Модель обучалась с использованием текстовых условий, вычисленных с помощью T5. Это позволяет пайплайну лучше следовать сложным текстовым промптам с детализированными описаниями.
  • Модель хорошо генерирует изображения высокого разрешения с различными соотношениями сторон. Для получения наилучших результатов авторы рекомендуют использовать определённые диапазоны размеров, которые можно найти здесь.
  • По качеству он соперничает с передовыми системами генерации изображений по тексту (на момент написания этой статьи), такими как Stable Diffusion XL, Imagen и DALL-E 2, при этом превосходя их по эффективности.

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

Вывод на GPU с менее чем 8 ГБ видеопамяти

Запустите PixArtAlphaPipeline на GPU с менее чем 8 ГБ видеопамяти, загрузив текстовый энкодер в 8-битной точности. Разберём полноценный пример.

Сначала установите библиотеку bitsandbytes:

pip install -U bitsandbytes

Затем загрузите текстовый энкодер в 8-битном формате:

from transformers import T5EncoderModel
from diffusers import PixArtAlphaPipeline
import torch

text_encoder = T5EncoderModel.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    subfolder="text_encoder",
    load_in_8bit=True,
    device_map="auto",

)
pipe = PixArtAlphaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    text_encoder=text_encoder,
    transformer=None,
    device_map="auto"
)

Теперь используйте pipe, чтобы закодировать промпт:

with torch.no_grad():
    prompt = "cute cat"
    prompt_embeds, prompt_attention_mask, negative_embeds, negative_prompt_attention_mask = pipe.encode_prompt(prompt)

Поскольку текстовые эмбеддинги уже вычислены, удалите text_encoder и pipe из памяти, чтобы освободить часть VRAM на GPU:

import gc

def flush():
    gc.collect()
    torch.cuda.empty_cache()

del text_encoder
del pipe
flush()

Затем вычислите латенты, используя эмбеддинги промпта в качестве входных данных:

pipe = PixArtAlphaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    text_encoder=None,
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"

latents = pipe(
    negative_prompt=None,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    prompt_attention_mask=prompt_attention_mask,
    negative_prompt_attention_mask=negative_prompt_attention_mask,
    num_images_per_prompt=1,
    output_type="latent",
).images

del pipe.transformer
flush()

Обратите внимание, что при инициализации pipe вы устанавливаете text_encoder в None, чтобы он не загружался.

После вычисления латентов передайте их в VAE для декодирования в реальное изображение:

with torch.no_grad():
    image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]
image = pipe.image_processor.postprocess(image, output_type="pil")[0]
image.save("cat.png")

Удалив неиспользуемые компоненты и очистив видеопамять GPU, вы сможете запустить PixArtAlphaPipeline с видеопамятью менее 8 ГБ.

Если вам нужен отчёт об использовании памяти, выполните этот скрипт.

Текстовые эмбеддинги, вычисленные в 8-битном формате, могут влиять на качество генерируемых изображений из-за потери информации в пространстве представлений, вызванной снижением точности. Рекомендуется сравнивать результаты с использованием 8-битного формата и без него.

При загрузке text_encoder вы установили load_in_8bit в True. Вы также можете указать load_in_4bit, чтобы ещё больше снизить требования к памяти — до менее чем 7 ГБ.

PixArtAlphaPipeline

класс diffusers.PixArtAlphaPipeline

< >

( tokenizer: T5Tokenizertext_encoder: T5EncoderModelvae: AutoencoderKLtransformer: PixArtTransformer2DModelscheduler: DPMSolverMultistepScheduler )

Параметры

  • vae (AutoencoderKL) — Модель вариационного автоэнкодера (VAE) для кодирования изображений в латентные представления и декодирования из них.
  • text_encoder (T5EncoderModel) — Замороженный текстовый энкодер. PixArt-Alpha использует T5, а именно вариант t5-v1_1-xxl.
  • tokenizer (T5Tokenizer) — токенизатор класса T5Tokenizer.
  • transformer (PixArtTransformer2DModel) — Модель PixArtTransformer2DModel, управляемая текстом, для удаления шума из закодированных латентов изображения. Изначально в конфигурации публиковалась как Transformer2DModel, но это несоответствие можно игнорировать.
  • scheduler (SchedulerMixin) — планировщик, используемый в связке с transformer для удаления шума из закодированных латентов изображения.

Pipeline для генерации изображений по тексту с использованием PixArt-Alpha.

Эта модель наследуется от DiffusionPipeline. См. документацию родительского класса для общих методов, которые библиотека реализует для всех пайплайнов (например, загрузка или сохранение, запуск на конкретном устройстве и т. д.)

__call__

< >

( prompt: str | list[str] = Nonenegative_prompt: str = ''num_inference_steps: int = 20timesteps: list[int] = Nonesigmas: list[float] = Noneguidance_scale: float = 4.5num_images_per_prompt: int | None = 1height: int | None = Nonewidth: int | None = Noneeta: float = 0.0generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Nonenegative_prompt_embeds: typing.Optional[torch.Tensor] = Nonenegative_prompt_attention_mask: 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 = 1clean_caption: bool = Trueuse_resolution_binning: bool = Truemax_sequence_length: int = 120**kwargs ) → ImagePipelineOutput or tuple

Параметры

  • prompt (str или list[str], необязательный) — Промпт или промпты для управления генерацией изображения. Если не задан, необходимо передать prompt_embeds вместо него.
  • negative_prompt (str or list[str], необязательный) — Промпт или промпты, которые не должны направлять генерацию изображения. Если не задан, необходимо передать negative_prompt_embeds. Игнорируется, если не используется направляющее поле (то есть игнорируется, если guidance_scale меньше 1).
  • num_inference_steps (int, необязательный, по умолчанию 100) — Количество шагов денойзинга. Как правило, большее число шагов денойзинга приводит к изображению более высокого качества, но за счет более медленного вывода.
  • timesteps (list[int], необязательный) — Пользовательские шаги для процесса удаления шума в планировщиках, поддерживающих аргумент timesteps в методе set_timesteps. Если не задан, будет использоваться поведение по умолчанию при передаче num_inference_steps. Значения должны быть в порядке убывания.
  • sigmas (list[float], необязательный) — Пользовательские сигмы для процесса удаления шума в планировщиках, поддерживающих аргумент sigmas в их методе set_timesteps. Если не задано, будет использоваться поведение по умолчанию при передаче num_inference_steps.
  • guidance_scale (float, необязательный, по умолчанию 4.5) — Масштаб руководства, как определено в Classifier-Free Diffusion Guidance. guidance_scale определяется как w из уравнения 2 статьи Imagen Paper. Масштаб руководства включается при установке guidance_scale > 1. Более высокий масштаб руководства способствует генерации изображений, тесно связанных с текстом prompt, обычно за счёт снижения качества изображения.
  • num_images_per_prompt (int, необязательный, по умолчанию 1) — Количество изображений, генерируемых для каждого промпта.
  • height (int, необязательный, по умолчанию self.unet.config.sample_size) — Высота генерируемого изображения в пикселях.
  • width (int, необязательный, по умолчанию self.unet.config.sample_size) — Ширина генерируемого изображения в пикселях.
  • eta (float, необязательный, по умолчанию 0.0) — Соответствует параметру eta (η) в статье DDIM: https://huggingface.co/papers/2010.02502. Применяется только к schedulers.DDIMScheduler, для остальных игнорируется.
  • generator (torch.Generator или list[torch.Generator], необязательный) — Один или список генераторов torch для обеспечения детерминированности генерации.
  • latents (torch.Tensor, необязательный) — Предварительно сгенерированные зашумлённые латенты, выбранные из гауссова распределения, которые используются в качестве входных данных для генерации изображений. Позволяет корректировать один и тот же процесс генерации с помощью различных промптов. Если не заданы, тензор латентов будет сгенерирован путём выборки с использованием переданного случайного generator.
  • prompt_embeds (torch.Tensor, необязательный) — Предварительно сгенерированные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, например, для взвешивания промпта. Если не передано, текстовые эмбеддинги будут сгенерированы из входного аргумента prompt.
  • prompt_attention_mask (torch.Tensor, необязательный) — Предварительно сгенерированная маска внимания для текстовых эмбеддингов.
  • negative_prompt_embeds (torch.Tensor, необязательный) — Предварительно сгенерированные эмбеддинги отрицательного текста. Для PixArt-Alpha этот отрицательный промпт должен быть "". Если не предоставлен, negative_prompt_embeds будет сгенерирован из входного аргумента negative_prompt.
  • negative_prompt_attention_mask (torch.Tensor, необязательный) — Предварительно сгенерированная маска внимания для отрицательных текстовых эмбеддингов.
  • output_type (str, необязательный, по умолчанию "pil") — Формат вывода сгенерированного изображения. Выберите один из вариантов: PIL: PIL.Image.Image или np.array.
  • return_dict (bool, необязательный, по умолчанию True) — Определяет, следует ли возвращать ~pipelines.stable_diffusion.IFPipelineOutput вместо обычного кортежа.
  • callback (Callable, необязательный) — Функция, которая будет вызываться каждые callback_steps шагов во время инференса. Функция будет вызываться со следующими аргументами: callback(step: int, timestep: int, latents: torch.Tensor).
  • callback_steps (int, необязательный, по умолчанию 1) — Частота вызова функции callback. Если не указано, callback будет вызываться на каждом шаге.
  • clean_caption (bool, необязательный, по умолчанию True) — Определяет, нужно ли очищать подпись перед созданием эмбеддингов. Требует установки beautifulsoup4 и ftfy. Если зависимости не установлены, эмбеддинги будут созданы из исходного промпта.
  • use_resolution_binning (bool, по умолчанию True) — Если установлено в True, запрашиваемые высота и ширина сначала сопоставляются с ближайшими разрешениями с помощью ASPECT_RATIO_1024_BIN. После декодирования полученных латентов в изображения они масштабируются обратно до запрошенного разрешения. Полезно для генерации изображений, не являющихся квадратными.
  • max_sequence_length (int, по умолчанию 120) — максимальная длина последовательности, используемая с prompt.

Возвращает

ImagePipelineOutput или tuple

Если return_dict равно True, возвращается ImagePipelineOutput, иначе возвращается tuple, где первый элемент — это список с сгенерированными изображениями

Функция, вызываемая при обращении к конвейеру для генерации.

Примеры:

>>> import torch
>>> from diffusers import PixArtAlphaPipeline

>>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-XL-2-512x512" too.
>>> pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16)
>>> # Enable memory optimizations.
>>> pipe.enable_model_cpu_offload()

>>> prompt = "A small cactus with a happy face in the Sahara desert."
>>> image = pipe(prompt).images[0]

encode_prompt

< >

( prompt: str | list[str]do_classifier_free_guidance: bool = Truenegative_prompt: str = ''num_images_per_prompt: int = 1device: typing.Optional[torch.device] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Nonenegative_prompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.Tensor] = Noneclean_caption: bool = Falsemax_sequence_length: int = 120**kwargs )

Параметры

  • prompt (str или list[str], необязательный) — промпт для кодирования
  • negative_prompt (str or list[str], необязательный) — Промпт, который не должен направлять генерацию изображения. Если не задан, необходимо передать negative_prompt_embeds вместо него. Игнорируется, если не используется направляющий сигнал (т. е. игнорируется, если guidance_scale меньше 1). Для PixArt-Alpha это должно быть "".
  • do_classifier_free_guidance (bool, необязательный, по умолчанию True) — использовать ли classifier free guidance
  • num_images_per_prompt (int, необязательный, по умолчанию 1) — количество изображений, которые должны быть сгенерированы для каждого промпта
  • device — (torch.device, необязательный): устройство torch, на котором будут размещены полученные эмбеддинги
  • prompt_embeds (torch.Tensor, необязательный) — Предварительно сгенерированные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, например, для взвешивания промпта. Если не передано, текстовые эмбеддинги будут сгенерированы из входного аргумента prompt.
  • negative_prompt_embeds (torch.Tensor, необязательный) — Предварительно сгенерированные эмбеддинги отрицательного текста. Для PixArt-Alpha это должны быть эмбеддинги строки "".
  • clean_caption (bool, по умолчанию False) — Если True, функция выполнит предобработку и очистку переданного описания перед кодированием.
  • max_sequence_length (int, по умолчанию 120) — максимальная длина последовательности, используемая для промпта.

Кодирует промпт в скрытые состояния текстового энкодера.

Update on GitHub