diffusers/main/ru/api/pipelines/aura_flow.md

AuraFlow

AuraFlow вдохновлён Stable Diffusion 3 и является крупнейшей моделью генерации изображений по текстовому описанию, распространяемой по лицензии Apache 2.0. Эта модель демонстрирует наилучшие результаты на бенчмарке GenEval.

Она была разработана командой Fal, дополнительные сведения можно найти в этой статье в блоге.

Запуск AuraFlow на потребительском оборудовании может быть довольно дорогим. Однако вы можете применить набор оптимизаций, чтобы ускорить его работу и снизить потребление памяти. Подробности — в этом разделе.

Квантование

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

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

import torch
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, AuraFlowTransformer2DModel, AuraFlowPipeline
from transformers import BitsAndBytesConfig as BitsAndBytesConfig, T5EncoderModel

quant_config = BitsAndBytesConfig(load_in_8bit=True)
text_encoder_8bit = T5EncoderModel.from_pretrained(
    "fal/AuraFlow",
    subfolder="text_encoder",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_8bit = AuraFlowTransformer2DModel.from_pretrained(
    "fal/AuraFlow",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

pipeline = AuraFlowPipeline.from_pretrained(
    "fal/AuraFlow",
    text_encoder=text_encoder_8bit,
    transformer=transformer_8bit,
    dtype=torch.float16,
    device_map="balanced",
)

prompt = "a tiny astronaut hatching from an egg on the moon"
image = pipeline(prompt).images[0]
image.save("auraflow.png")

Также поддерживается загрузка чекпоинтов GGUF:

import torch
from diffusers import (
    AuraFlowPipeline,
    GGUFQuantizationConfig,
    AuraFlowTransformer2DModel,
)

transformer = AuraFlowTransformer2DModel.from_single_file(
    "https://huggingface.co/city96/AuraFlow-v0.3-gguf/blob/main/aura_flow_0.3-Q2_K.gguf",
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
    dtype=torch.bfloat16,
)

pipeline = AuraFlowPipeline.from_pretrained(
    "fal/AuraFlow-v0.3",
    transformer=transformer,
    dtype=torch.bfloat16,
)

prompt = "a cute pony in a field of flowers"
image = pipeline(prompt).images[0]
image.save("auraflow.png")

Поддержка torch.compile()

AuraFlow можно скомпилировать с помощью torch.compile(), чтобы ускорить время вывода даже для разных разрешений. Сначала установите PyTorch nightly, следуя инструкциям здесь. Ниже приведён фрагмент кода, показывающий необходимые изменения для включения этой функции:

+ torch.fx.experimental._config.use_duck_shape = False
+ pipeline.transformer = torch.compile(
    pipeline.transformer, fullgraph=True, dynamic=True
)

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

Это обеспечивает ускорение от 100% (на низких разрешениях) до 30% (на разрешении 1536x1536).

Благодарим AstraliteHeart за помощь в переписывании класса AuraFlowTransformer2DModel, благодаря чему вышеописанное работает для разных разрешений (PR).

AuraFlowPipeline[[diffusers.AuraFlowPipeline]]

diffusers.AuraFlowPipeline[[diffusers.AuraFlowPipeline]]

diffusers.AuraFlowPipeline(tokenizer: T5Tokenizer, text_encoder: UMT5EncoderModel, vae: AutoencoderKL, transformer: AuraFlowTransformer2DModel, scheduler: FlowMatchEulerDiscreteScheduler)

Источник

Параметры:

tokenizer (T5TokenizerFast) : Токенизатор класса T5Tokenizer.

text_encoder (T5EncoderModel) : Замороженный текстовый энкодер. AuraFlow использует T5, а именно вариант EleutherAI/pile-t5-xl.

vae (AutoencoderKL) : Модель вариационного автоэнкодера (VAE) для кодирования изображений в латентные представления и декодирования из них.

transformer (AuraFlowTransformer2DModel) : архитектура условного Transformer (MMDiT и DiT) для удаления шума из закодированных латентов изображения.

scheduler (FlowMatchEulerDiscreteScheduler) : планировщик, используемый в связке с transformer для удаления шума из закодированных латентов изображения.

call[[diffusers.AuraFlowPipeline.call]]

__call__(prompt: str | list[str] = None, negative_prompt: str | list[str] = None, num_inference_steps: int = 50, sigmas: list[float] = None, guidance_scale: float = 3.5, num_images_per_prompt: int | None = 1, height: int | None = 1024, width: int | None = 1024, generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = None, latents: typing.Optional[torch.Tensor] = None, prompt_embeds: typing.Optional[torch.Tensor] = None, prompt_attention_mask: typing.Optional[torch.Tensor] = None, negative_prompt_embeds: typing.Optional[torch.Tensor] = None, negative_prompt_attention_mask: typing.Optional[torch.Tensor] = None, max_sequence_length: int = 256, output_type: str | None = 'pil', return_dict: bool = True, attention_kwargs: dict[str, typing.Any] | None = None, callback_on_step_end: typing.Union[typing.Callable[[int, int], NoneType], diffusers.callbacks.PipelineCallback, diffusers.callbacks.MultiPipelineCallbacks, NoneType] = None, callback_on_step_end_tensor_inputs: list[str] = ['latents'])

Источник

Параметры:

prompt (str или list[str], необязательный) : Промпт или промпты для управления генерацией изображения. Если не задан, необходимо передать prompt_embeds вместо него.

negative_prompt (str или list[str], необязательный) : Промпт или промпты, которые не должны направлять генерацию изображения. Если не задан, необходимо передать negative_prompt_embeds. Игнорируется, если не используется направляющее влияние (то есть игнорируется, если guidance_scale меньше 1).

height (int, optional, по умолчанию self.transformer.config.sample_size * self.vae_scale_factor) : Высота генерируемого изображения в пикселях. По умолчанию установлено значение 1024 для достижения наилучших результатов.

width (int, optional, по умолчанию self.transformer.config.sample_size * self.vae_scale_factor) : Ширина генерируемого изображения в пикселях. По умолчанию установлено значение 1024 для достижения наилучших результатов.

num_inference_steps (int, optional, по умолчанию 50) : Количество шагов денойзинга. Как правило, большее число шагов денойзинга приводит к изображению более высокого качества, но за счёт более медленного вывода.

sigmas (list[float], optional) : Пользовательские сигмы, используемые для переопределения стратегии интервала шагов планировщика. Если переданы sigmas, то num_inference_steps и timesteps должны быть None.

guidance_scale (float, optional, по умолчанию 5.0) : Масштаб руководства, определённый в Classifier-Free Diffusion Guidance. guidance_scale определяется как w из уравнения 2 в Imagen Paper. Масштаб руководства активируется при установке guidance_scale > 1. Более высокий масштаб руководства способствует генерации изображений, тесно связанных с текстовым prompt, обычно за счёт снижения качества изображения.

num_images_per_prompt (int, необязательный, по умолчанию 1) : Количество изображений, генерируемых на один запрос.

generator (torch.Generator или list[torch.Generator], необязательный) : Один или список генераторов torch для обеспечения детерминированности генерации.

latents (torch.FloatTensor, необязательный) : Предварительно сгенерированные шумовые латенты, выбранные из гауссова распределения, которые используются в качестве входных данных для генерации изображений. Позволяет корректировать одну и ту же генерацию с помощью разных промптов. Если не заданы, тензор латентов будет сгенерирован путём выборки с использованием переданного случайного generator.

prompt_embeds (torch.FloatTensor, необязательный) : Предварительно сгенерированные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, например для взвешивания промпта. Если не предоставлены, текстовые эмбеддинги будут сгенерированы из входного аргумента prompt.

prompt_attention_mask (torch.Tensor, необязательный) : Предварительно сгенерированная маска внимания для текстовых эмбеддингов.

negative_prompt_embeds (torch.FloatTensor, необязательный) : Предварительно сгенерированные отрицательные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, например для взвешивания промпта. Если не предоставлены, 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_xl.StableDiffusionXLPipelineOutput вместо обычного кортежа.

attention_kwargs (dict, необязательный) : Словарь kwargs, который при указании передается в AttentionProcessor, определенный в self.processor в diffusers.models.attention_processor.

callback_on_step_end (Callable, необязательный) : функция, вызываемая в конце каждого шага денойзинга во время инференса. Функция вызывается со следующими аргументами: callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict). callback_kwargs будет содержать список всех тензоров, как указано в callback_on_step_end_tensor_inputs.

callback_on_step_end_tensor_inputs (list, optional) : Список тензорных входных данных для функции callback_on_step_end. Тензоры, указанные в списке, будут переданы в качестве аргумента callback_kwargs. Вы сможете включать только переменные, перечисленные в атрибуте ._callback_tensor_inputs вашего класса pipeline.

max_sequence_length (int по умолчанию 256) : Максимальная длина последовательности, используемая с prompt.

Возвращает: ImagePipelineOutput или tuple

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

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

Примеры:

>>> import torch
>>> from diffusers import AuraFlowPipeline

>>> pipe = AuraFlowPipeline.from_pretrained("fal/AuraFlow", torch_dtype=torch.float16)
>>> pipe = pipe.to("cuda")
>>> prompt = "A cat holding a sign that says hello world"
>>> image = pipe(prompt).images[0]
>>> image.save("aura_flow.png")

encode_prompt[[diffusers.AuraFlowPipeline.encode_prompt]]

encode_prompt(prompt: str | list[str], negative_prompt: str | list[str] = None, do_classifier_free_guidance: bool = True, num_images_per_prompt: int = 1, device: typing.Optional[torch.device] = None, prompt_embeds: typing.Optional[torch.Tensor] = None, negative_prompt_embeds: typing.Optional[torch.Tensor] = None, prompt_attention_mask: typing.Optional[torch.Tensor] = None, negative_prompt_attention_mask: typing.Optional[torch.Tensor] = None, max_sequence_length: int = 256, lora_scale: float | None = None)

Источник

Параметры:

prompt (str или list[str], необязательный) : промпт для кодирования

negative_prompt (str или list[str], необязательный) : Промпт, который не должен направлять генерацию изображения. Если не задан, необходимо передать negative_prompt_embeds. Игнорируется, если не используется направляющий сигнал (то есть игнорируется, если guidance_scale меньше 1).

do_classifier_free_guidance (bool, необязательный, по умолчанию True) : использовать ли classifier free guidance

num_images_per_prompt (int, необязательный, по умолчанию 1) : количество изображений, которые должны быть сгенерированы для каждого промпта

device : (torch.device, необязательный): устройство torch, на котором будут размещены полученные эмбеддинги

prompt_embeds (torch.Tensor, необязательный) : Предварительно сгенерированные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, например для взвешивания промпта. Если не предоставлены, текстовые эмбеддинги будут сгенерированы из входного аргумента prompt.

prompt_attention_mask (torch.Tensor, необязательный) : Предварительно сгенерированная маска внимания для текстовых эмбеддингов.

negative_prompt_embeds (torch.Tensor, необязательный) : Предварительно сгенерированные отрицательные текстовые эмбеддинги.

negative_prompt_attention_mask (torch.Tensor, необязательный) : Предварительно сгенерированная маска внимания для отрицательных текстовых эмбеддингов.

max_sequence_length (int, по умолчанию 256) : Максимальная длина последовательности, используемая для промпта.

lora_scale (float, optional) : масштаб lora, который будет применён ко всем слоям LoRA текстового энкодера, если слои LoRA загружены.

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