Buckets:

mishig's picture
|
download
raw
15.5 kB

AuraFlow

AuraFlow se inspira en Stable Diffusion 3 y es, con diferencia, el modelo de generación de imágenes a partir de texto más grande que cuenta con una licencia Apache 2.0. Este modelo logra resultados de última generación en el benchmark GenEval.

Fue desarrollado por el equipo de Fal y más detalles al respecto se pueden encontrar en esta publicación de blog.

AuraFlow puede resultar bastante costoso de ejecutar en dispositivos de hardware de consumo. Sin embargo, puedes aplicar una serie de optimizaciones para ejecutarlo más rápido y de una manera más eficiente en memoria. Consulta esta sección para obtener más detalles.

Cuantización

La cuantización ayuda a reducir los requisitos de memoria de los modelos muy grandes almacenando los pesos del modelo en un tipo de datos de menor precisión. Sin embargo, la cuantización puede tener un impacto variable en la calidad del vídeo dependiendo del modelo de vídeo.

Consulte el resumen de Quantization para obtener más información sobre los backends de cuantización admitidos y la selección de un backend de cuantización que admita su caso de uso. El siguiente ejemplo demuestra cómo cargar un AuraFlowPipeline cuantizado para inferencia con 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")

La carga de checkpoints GGUF también está soportada:

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

Soporte para torch.compile()

AuraFlow se puede compilar con torch.compile() para reducir la latencia de inferencia incluso con diferentes resoluciones. Primero, instala PyTorch nightly siguiendo las instrucciones de aquí. El fragmento de código siguiente muestra los cambios necesarios para habilitar esto:

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

Especificar use_duck_shape como False indica al compilador si debe usar la misma variable simbólica para representar tamaños de entrada que son iguales. Para más detalles, consulta este comentario.

Esto permite mejoras de velocidad que van desde el 100% (en resoluciones bajas) hasta el 30% (en resolución 1536x1536).

Gracias a AstraliteHeart, quien nos ayudó a reescribir la clase AuraFlowTransformer2DModel para que lo anterior funcione con diferentes resoluciones (PR).

AuraFlowPipeline[[diffusers.AuraFlowPipeline]]

diffusers.AuraFlowPipeline[[diffusers.AuraFlowPipeline]]

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

Fuente

Parámetros:

tokenizer (T5TokenizerFast) : Tokenizador de la clase T5Tokenizer.

text_encoder (T5EncoderModel) : Codificador de texto congelado. AuraFlow utiliza T5, específicamente la variante EleutherAI/pile-t5-xl.

vae (AutoencoderKL) : Modelo de Variational Auto-Encoder (VAE) para codificar y decodificar imágenes hacia y desde representaciones latentes.

transformer (AuraFlowTransformer2DModel) : Arquitectura de Transformer condicional (MMDiT y DiT) para desruidar los latentes de imagen codificados.

scheduler (FlowMatchEulerDiscreteScheduler) : Un programador para usar en combinación con transformer para desruidar los latentes de imagen codificados.

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

Fuente

Parámetros:

prompt (str o list[str], opcional) : El prompt o los prompts para guiar la generación de la imagen. Si no se define, se debe pasar prompt_embeds en su lugar.

negative_prompt (str o list[str], opcional) : El prompt o prompts que no deben guiar la generación de la imagen. Si no se define, se debe pasar negative_prompt_embeds en su lugar. Se ignora cuando no se usa guía (es decir, se ignora si guidance_scale es menor que 1).

height (int, opcional, por defecto self.transformer.config.sample_size * self.vae_scale_factor) : La altura en píxeles de la imagen generada. Por defecto, se establece en 1024 para obtener los mejores resultados.

width (int, opcional, por defecto self.transformer.config.sample_size * self.vae_scale_factor) : El ancho en píxeles de la imagen generada. Por defecto, se establece en 1024 para obtener los mejores resultados.

num_inference_steps (int, opcional, valor predeterminado 50) : El número de pasos de desruido. Más pasos de desruido suelen conducir a una imagen de mayor calidad a costa de una inferencia más lenta.

sigmas (list[float], opcional) : Sigmas personalizados usados para reemplazar la estrategia de espaciado de pasos de tiempo del programador. Si se pasa sigmas, num_inference_steps y timesteps deben ser None.

guidance_scale (float, opcional, valor predeterminado 5.0) : Escala de orientación definida en Classifier-Free Diffusion Guidance. guidance_scale se define como w de la ecuación 2. de Imagen Paper. La escala de orientación se habilita estableciendo guidance_scale > 1. Una escala de orientación más alta fomenta la generación de imágenes estrechamente vinculadas al prompt de texto, generalmente a costa de una menor calidad de imagen.

num_images_per_prompt (int, opcional, valor predeterminado 1) : El número de imágenes a generar por indicación.

generator (torch.Generator o list[torch.Generator], opcional) : Uno o una lista de generadores de torch para hacer la generación determinista.

latents (torch.FloatTensor, opcional) : Latentes ruidosos pregenerados, muestreados de una distribución gaussiana, que se usarán como entradas para la generación de imágenes. Se pueden utilizar para ajustar la misma generación con diferentes prompts. Si no se proporcionan, se generará un tensor de latentes mediante muestreo usando el generator aleatorio suministrado.

prompt_embeds (torch.FloatTensor, opcional) : Embeddings de texto pregenerados. Se pueden usar para ajustar fácilmente las entradas de texto, p. ej. ponderación de prompts. Si no se proporcionan, los embeddings de texto se generarán a partir del argumento de entrada prompt.

prompt_attention_mask (torch.Tensor, opcional) : Máscara de atención pregenerada para los embeddings de texto.

negative_prompt_embeds (torch.FloatTensor, opcional) : Embeddings de texto negativos pregenerados. Se pueden utilizar para ajustar fácilmente las entradas de texto, p. ej. el ponderado de prompts. Si no se proporcionan, negative_prompt_embeds se generará a partir del argumento de entrada negative_prompt.

negative_prompt_attention_mask (torch.Tensor, opcional) : Máscara de atención pregenerada para los embeddings de texto negativos.

output_type (str, opcional, por defecto "pil") : El formato de salida de la imagen generada. Elija entre PIL: PIL.Image.Image o np.array.

return_dict (bool, opcional, valor predeterminado True) : Indica si se debe devolver un ~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput en lugar de una tupla simple.

attention_kwargs (dict, opcional) : Un diccionario de kwargs que, si se especifica, se pasa al AttentionProcessor tal como se define bajo self.processor en diffusers.models.attention_processor.

callback_on_step_end (Callable, opcional) : Una función que se llama al final de cada paso de desruido durante la inferencia. La función se llama con los siguientes argumentos: callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict). callback_kwargs incluirá una lista de todos los tensores especificados por callback_on_step_end_tensor_inputs.

callback_on_step_end_tensor_inputs (list, opcional) : La lista de entradas de tensor para la función callback_on_step_end. Los tensores especificados en la lista se pasarán como argumento callback_kwargs. Solo podrá incluir variables enumeradas en el atributo ._callback_tensor_inputs de su clase de pipeline.

max_sequence_length (int por defecto 256) : Longitud máxima de la secuencia a utilizar con el prompt.

Devuelve: ImagePipelineOutput o tuple

Si return_dict es True, se devuelve ImagePipelineOutput; en caso contrario, se devuelve una tuple cuyo primer elemento es una lista con las imágenes generadas.

Función invocada al llamar a la pipeline para la generación.

Ejemplos:

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

Fuente

Parámetros:

prompt (str o list[str], opcional) : prompt a codificar

negative_prompt (str o list[str], opcional) : El prompt que no debe guiar la generación de la imagen. Si no se define, se debe pasar negative_prompt_embeds en su lugar. Se ignora cuando no se usa guía (es decir, se ignora si guidance_scale es menor que 1).

do_classifier_free_guidance (bool, opcional, por defecto True) : si se debe usar classifier free guidance o no

num_images_per_prompt (int, opcional, valor predeterminado 1) : número de imágenes que deben generarse por indicación

device : (torch.device, opcional): dispositivo de torch en el que colocar las incrustaciones resultantes

prompt_embeds (torch.Tensor, opcional) : Embeddings de texto pregenerados. Se pueden usar para ajustar fácilmente las entradas de texto, p. ej. ponderación de prompts. Si no se proporcionan, los embeddings de texto se generarán a partir del argumento de entrada prompt.

prompt_attention_mask (torch.Tensor, opcional) : Máscara de atención pregenerada para los embeddings de texto.

negative_prompt_embeds (torch.Tensor, opcional) : Embeddings de texto negativo pregenerados.

negative_prompt_attention_mask (torch.Tensor, opcional) : Máscara de atención pregenerada para los embeddings de texto negativos.

max_sequence_length (int, por defecto 256) : Longitud máxima de la secuencia a utilizar para el prompt.

lora_scale (float, opcional) : Una escala lora que se aplicará a todas las capas LoRA del codificador de texto si se cargan capas LoRA.

Codifica el prompt en estados ocultos del codificador de texto.

Xet Storage Details

Size:
15.5 kB
·
Xet hash:
5a609c850d0bbd7538b48869532ee6610e13c83e95570fee55d320fcb4c925db

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.