Diffusers documentation
AuraFlow
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
< source >( tokenizer: T5Tokenizertext_encoder: UMT5EncoderModelvae: AutoencoderKLtransformer: AuraFlowTransformer2DModelscheduler: 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__
< source >( prompt: str | list[str] = Nonenegative_prompt: str | list[str] = Nonenum_inference_steps: int = 50sigmas: list = Noneguidance_scale: float = 3.5num_images_per_prompt: int | None = 1height: int | None = 1024width: int | None = 1024generator: 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] = Nonemax_sequence_length: int = 256output_type: str | None = 'pil'return_dict: bool = Trueattention_kwargs: dict[str, typing.Any] | None = Nonecallback_on_step_end: typing.Union[typing.Callable[[int, int], NoneType], diffusers.callbacks.PipelineCallback, diffusers.callbacks.MultiPipelineCallbacks, NoneType] = Nonecallback_on_step_end_tensor_inputs: list = ['latents'] ) → ImagePipelineOutput or tuple
Параметры
- prompt (
strилиlist[str], необязательный) — Промпт или промпты для управления генерацией изображения. Если не задан, необходимо передатьprompt_embedsвместо него. - negative_prompt (
strorlist[str], необязательный) — Промпт или промпты, которые не должны направлять генерацию изображения. Если не задан, необходимо передатьnegative_prompt_embeds. Игнорируется, если не используется направляющее поле (то есть игнорируется, еслиguidance_scaleменьше1). - height (
int, необязательный, по умолчанию self.transformer.config.sample_size * self.vae_scale_factor) — Высота генерируемого изображения в пикселях. По умолчанию установлено значение 1024 для достижения наилучших результатов. - width (
int, необязательный, по умолчанию self.transformer.config.sample_size * self.vae_scale_factor) — Ширина генерируемого изображения в пикселях. По умолчанию установлено значение 1024 для достижения наилучших результатов. - num_inference_steps (
int, необязательный, по умолчанию 50) — Количество шагов денойзинга. Как правило, большее число шагов денойзинга приводит к изображению более высокого качества, но за счёт более медленного вывода. - sigmas (
list[float], необязательный) — Пользовательские сигмы, используемые для переопределения стратегии интервала шагов по времени планировщика. Если переданыsigmas,num_inference_stepsиtimestepsдолжны бытьNone. - guidance_scale (
float, необязательный, по умолчанию 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, необязательный) — Список тензорных входных данных для функцииcallback_on_step_end. Тензоры, указанные в списке, будут переданы в качестве аргументаcallback_kwargs. Вы сможете включать только переменные, перечисленные в атрибуте._callback_tensor_inputsвашего класса пайплайна. - 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
< source >( prompt: str | list[str]negative_prompt: str | list[str] = Nonedo_classifier_free_guidance: bool = Truenum_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] = Nonemax_sequence_length: int = 256lora_scale: float | None = None )
Параметры
- prompt (
strилиlist[str], необязательный) — промпт для кодирования - negative_prompt (
strorlist[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, необязательный) — Масштаб LoRA, который будет применён ко всем слоям LoRA текстового кодировщика, если слои LoRA загружены.
Кодирует промпт в скрытые состояния текстового энкодера.