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

AuraFlow

AuraFlow é inspirado no Stable Diffusion 3 e é, de longe, o maior modelo de geração de imagem a partir de texto que vem com a licença Apache 2.0. Este modelo alcança resultados de estado da arte no benchmark GenEval.

Foi desenvolvido pela equipe da Fal e mais detalhes sobre ele podem ser encontrados em este post de blog.

O AuraFlow pode ser bastante caro de executar em dispositivos de hardware de consumo. No entanto, você pode aplicar uma série de otimizações para executá-lo mais rápido e de uma forma mais amigável à memória. Consulte esta seção para mais detalhes.

Quantização

A quantização ajuda a reduzir os requisitos de memória de modelos muito grandes armazenando os pesos do modelo em um tipo de dados de menor precisão. No entanto, a quantização pode ter um impacto variável na qualidade do vídeo, dependendo do modelo de vídeo.

Consulte a visão geral de Quantization para saber mais sobre os backends de quantização suportados e sobre como selecionar um backend de quantização que atenda ao seu caso de uso. O exemplo abaixo demonstra como carregar um AuraFlowPipeline quantizado para inferência com 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")

O carregamento de checkpoints GGUF também é suportado:

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")

Suporte para torch.compile()

O AuraFlow pode ser compilado com torch.compile() para acelerar a latência de inferência, mesmo para diferentes resoluções. Primeiro, instale o PyTorch nightly seguindo as instruções em aqui. O trecho abaixo mostra as alterações necessárias para habilitar isso:

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

Especificar use_duck_shape como False instrui o compilador se ele deve usar a mesma variável simbólica para representar tamanhos de entrada que são iguais. Para mais detalhes, confira este comentário.

Isso permite melhorias de velocidade de 100% (em resoluções baixas) a 30% (em resolução 1536x1536).

Agradecemos a AstraliteHeart, que nos ajudou a reescrever a classe AuraFlowTransformer2DModel para que o acima funcione para diferentes resoluções (PR).

AuraFlowPipeline[[diffusers.AuraFlowPipeline]]

diffusers.AuraFlowPipeline[[diffusers.AuraFlowPipeline]]

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

Fonte

Parâmetros:

tokenizer (T5TokenizerFast) : Tokenizer da classe T5Tokenizer.

text_encoder (T5EncoderModel) : Codificador de texto congelado. O AuraFlow usa T5, especificamente a variante EleutherAI/pile-t5-xl.

vae (AutoencoderKL) : Modelo de Variational Auto-Encoder (VAE) para codificar e decodificar imagens para e a partir de representações latentes.

transformer (AuraFlowTransformer2DModel) : Arquitetura de Transformer condicional (MMDiT e DiT) para desruidar os latentes de imagem codificados.

scheduler (FlowMatchEulerDiscreteScheduler) : Um scheduler a ser usado em combinação com transformer para desruidar os latentes da imagem codificada.

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'])

Fonte

Parâmetros:

prompt (str ou list[str], opcional) : O prompt ou os prompts para orientar a geração da imagem. Se não definido, é necessário passar prompt_embeds em vez disso.

negative_prompt (str ou list[str], opcional) : O prompt ou os prompts que não devem orientar a geração da imagem. Se não definido, é necessário passar negative_prompt_embeds em vez disso. Ignorado quando não se usa orientação (ou seja, ignorado se guidance_scale for menor que 1).

height (int, opcional, padrão é self.transformer.config.sample_size * self.vae_scale_factor) : A altura em pixels da imagem gerada. Isso é definido como 1024 por padrão para os melhores resultados.

width (int, opcional, padrão é self.transformer.config.sample_size * self.vae_scale_factor) : A largura em pixels da imagem gerada. Isso é definido como 1024 por padrão para os melhores resultados.

num_inference_steps (int, opcional, padrão 50) : O número de etapas de desruidamento. Mais etapas de desruidamento geralmente resultam em uma imagem de maior qualidade, em detrimento de uma inferência mais lenta.

sigmas (list[float], opcional) : Sigmas personalizados usados para substituir a estratégia de espaçamento de passos de tempo do agendador. Se sigmas for passado, num_inference_steps e timesteps devem ser None.

guidance_scale (float, opcional, padrão 5.0) : Escala de orientação definida em Classifier-Free Diffusion Guidance. guidance_scale é definido como w da equação 2. do Imagen Paper. A escala de orientação é ativada ao definir guidance_scale > 1. Uma escala de orientação maior incentiva a geração de imagens que estão intimamente ligadas ao prompt de texto, geralmente em detrimento de uma menor qualidade de imagem.

num_images_per_prompt (int, opcional, padrão 1) : O número de imagens a serem geradas por prompt.

generator (torch.Generator ou list[torch.Generator], opcional) : Um ou uma lista de geradores do torch para tornar a geração determinística.

latents (torch.FloatTensor, opcional) : Latentes ruidosos pré-gerados, amostrados de uma distribuição gaussiana, a serem usados como entradas para a geração de imagens. Pode ser usado para ajustar a mesma geração com diferentes prompts. Se não for fornecido, um tensor de latentes será gerado por amostragem usando o generator aleatório fornecido.

prompt_embeds (torch.FloatTensor, opcional) : Emojis de texto pré-gerados. Pode ser usado para ajustar facilmente as entradas de texto, por exemplo, ponderação de prompts. Se não for fornecido, os emojis de texto serão gerados a partir do argumento de entrada prompt.

prompt_attention_mask (torch.Tensor, opcional) : Máscara de atenção pré-gerada para embeddings de texto.

negative_prompt_embeds (torch.FloatTensor, opcional) : Embeddings de texto negativos pré-gerados. Pode ser usado para ajustar facilmente as entradas de texto, por exemplo, ponderação de prompts. Se não for fornecido, negative_prompt_embeds será gerado a partir do argumento de entrada negative_prompt.

negative_prompt_attention_mask (torch.Tensor, opcional) : Máscara de atenção pré-gerada para embeddings de texto negativo.

output_type (str, opcional, padrão "pil") : O formato de saída da imagem gerada. Escolha entre PIL: PIL.Image.Image ou np.array.

return_dict (bool, opcional, padrão True) : Se deve ou não retornar um ~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput em vez de uma tupla simples.

attention_kwargs (dict, opcional) : Um dicionário de kwargs que, se especificado, é passado para o AttentionProcessor conforme definido em self.processor em diffusers.models.attention_processor.

callback_on_step_end (Callable, opcional) : Uma função que é chamada no final de cada etapa de desruidamento durante a inferência. A função é chamada com os seguintes argumentos: callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict). callback_kwargs incluirá uma lista de todos os tensores, conforme especificado por callback_on_step_end_tensor_inputs.

callback_on_step_end_tensor_inputs (list, opcional) : A lista de entradas de tensor para a função callback_on_step_end. Os tensores especificados na lista serão passados como argumento callback_kwargs. Você só poderá incluir variáveis listadas no atributo ._callback_tensor_inputs da sua classe de pipeline.

max_sequence_length (int padrão é 256) : Comprimento máximo da sequência a ser usado com o prompt.

Retorna: ImagePipelineOutput ou tuple

Se return_dict for True, ImagePipelineOutput é retornado; caso contrário, uma tuple é retornada, em que o primeiro elemento é uma lista com as imagens geradas.

Função invocada ao chamar o pipeline para geração.

Exemplos:

>>> 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)

Fonte

Parâmetros:

prompt (str ou list[str], opcional) : prompt a ser codificado

negative_prompt (str ou list[str], opcional) : O prompt para não orientar a geração da imagem. Se não for definido, é necessário passar negative_prompt_embeds em vez disso. Ignorado quando não se usa orientação (ou seja, ignorado se guidance_scale for menor que 1).

do_classifier_free_guidance (bool, opcional, padrão True) : se deve ou não usar classifier free guidance

num_images_per_prompt (int, opcional, padrão 1) : número de imagens que devem ser geradas por prompt

device : (torch.device, opcional): dispositivo torch onde os embeddings resultantes serão colocados

prompt_embeds (torch.Tensor, opcional) : Embeddings de texto pré-gerados. Podem ser usados para ajustar facilmente as entradas de texto, por exemplo, ponderação de prompts. Se não forem fornecidos, os embeddings de texto serão gerados a partir do argumento de entrada prompt.

prompt_attention_mask (torch.Tensor, opcional) : Máscara de atenção pré-gerada para embeddings de texto.

negative_prompt_embeds (torch.Tensor, opcional) : Embeddings de texto negativo pré-gerados.

negative_prompt_attention_mask (torch.Tensor, opcional) : Máscara de atenção pré-gerada para embeddings de texto negativo.

max_sequence_length (int, padrão 256) : Comprimento máximo da sequência a ser usado para o prompt.

lora_scale (float, opcional) : Uma escala lora que será aplicada a todas as camadas LoRA do text encoder, caso as camadas LoRA sejam carregadas.

Codifica o prompt em estados ocultos do codificador de texto.