Diffusers documentation

AuraFlow

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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

classe diffusers.AuraFlowPipeline

< >

( tokenizer: T5Tokenizertext_encoder: UMT5EncoderModelvae: AutoencoderKLtransformer: AuraFlowTransformer2DModelscheduler: FlowMatchEulerDiscreteScheduler )

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 da imagem codificada.
  • scheduler (FlowMatchEulerDiscreteScheduler) — Um scheduler a ser usado em combinação com transformer para desruidar os latentes da imagem codificada.

__call__

< >

( prompt: str | list[str] = Nonenegative_prompt: str | list[str] = Nonenum_inference_steps: int = 50sigmas: list[float] = 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[str] = ['latents'] ) → ImagePipelineOutput or tuple

Parâmetros

  • prompt (str ou list[str], opcional) — O prompt ou os prompts para orientar a geração da imagem. Se não for 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 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).
  • height (int, opcional, padrão é self.transformer.config.sample_size * self.vae_scale_factor) — A altura em pixels da imagem gerada. O padrão é 1024 para obter os melhores resultados.
  • width (int, opcional, padrão é self.transformer.config.sample_size * self.vae_scale_factor) — A largura em pixels da imagem gerada. O padrão é 1024 para obter 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, à custa 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) — Incorporações de texto pré-geradas. Podem ser usadas para ajustar facilmente as entradas de texto, por exemplo, ponderação de prompts. Se não forem fornecidas, as incorporações de texto serão geradas 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) — Incorporações de texto negativas pré-geradas. 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 negativos.
  • 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 ao 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, onde 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

< >

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

Parâmetros

  • prompt (str ou list[str], opcional) — prompt a ser codificado
  • negative_prompt (str ou list[str], opcional) — O prompt que não deve 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 em que os embeddings resultantes serão colocados
  • prompt_embeds (torch.Tensor, opcional) — Incorporações de texto pré-geradas. Podem ser usadas para ajustar facilmente as entradas de texto, por exemplo, ponderação de prompts. Se não forem fornecidas, as incorporações de texto serão geradas 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) — Incorporações de texto negativo pré-geradas.
  • negative_prompt_attention_mask (torch.Tensor, opcional) — Máscara de atenção pré-gerada para embeddings de texto negativos.
  • 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 codificador de texto, caso as camadas LoRA sejam carregadas.

Codifica o prompt em estados ocultos do codificador de texto.

Update on GitHub