Diffusers 文档

AuraFlow

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

AuraFlow

AuraFlow 的灵感来自 Stable Diffusion 3,是迄今为止最大的文本到图像生成模型,采用 Apache 2.0 许可证。该模型在 GenEval 基准测试中取得了最先进的结果。

它由 Fal 团队开发,更多详细信息可在 这篇博客文章 中找到。

AuraFlow 在消费级硬件设备上运行可能成本很高。但是,您可以执行一系列优化,使其运行更快,并更节省内存。有关更多详细信息,请查看 本节

量化

量化有助于通过以较低精度数据类型存储模型权重来减少大型模型的内存需求。但是,量化对视频质量的影响可能因视频模型而异。

请参阅 量化 概述,了解有关支持的量化后端以及选择支持您用例的量化后端的更多信息。以下示例演示了如何使用 bitsandbytes 加载量化 AuraFlowPipeline 进行推理。

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,
    torch_dtype=torch.float16,
)

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

pipeline = AuraFlowPipeline.from_pretrained(
    "fal/AuraFlow",
    text_encoder=text_encoder_8bit,
    transformer=transformer_8bit,
    torch_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),
    torch_dtype=torch.bfloat16,
)

pipeline = AuraFlowPipeline.from_pretrained(
    "fal/AuraFlow-v0.3",
    transformer=transformer,
    torch_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

class diffusers.AuraFlowPipeline

< >

( tokenizer: T5Tokenizer text_encoder: UMT5EncoderModel vae: AutoencoderKL transformer: AuraFlowTransformer2DModel scheduler: FlowMatchEulerDiscreteScheduler )

参数

  • 分词器 (T5TokenizerFast) — T5Tokenizer 类的分词器。
  • text_encoder (T5EncoderModel) — 冻结的文本编码器。AuraFlow 使用 T5,特别是 EleutherAI/pile-t5-xl 变体。
  • vae (AutoencoderKL) — 变分自编码器 (VAE) 模型,用于将图像编码和解码为潜在表示。
  • transformer (AuraFlowTransformer2DModel) — 用于对编码图像潜在表示进行去噪的条件 Transformer (MMDiT 和 DiT) 架构。
  • 调度器 (FlowMatchEulerDiscreteScheduler) — 与 `transformer` 结合使用的调度器,用于对编码图像潜在表示进行去噪。

__call__

< >

( prompt: typing.Union[str, typing.List[str]] = None negative_prompt: typing.Union[str, typing.List[str]] = None num_inference_steps: int = 50 sigmas: typing.List[float] = None guidance_scale: float = 3.5 num_images_per_prompt: typing.Optional[int] = 1 height: typing.Optional[int] = 1024 width: typing.Optional[int] = 1024 generator: typing.Union[torch._C.Generator, typing.List[torch._C.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: typing.Optional[str] = 'pil' return_dict: bool = True attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None callback_on_step_end: typing.Union[typing.Callable[[int, int, typing.Dict], NoneType], diffusers.callbacks.PipelineCallback, diffusers.callbacks.MultiPipelineCallbacks, NoneType] = None callback_on_step_end_tensor_inputs: typing.List[str] = ['latents'] )

参数

  • prompt (strList[str], 可选) — 用于引导图像生成的提示词。如果未定义,则必须传入 `prompt_embeds`。
  • negative_prompt (strList[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], 可选) — 用于覆盖调度器时间步长间距策略的自定义 sigma 值。如果传入 `sigmas`,则 `num_inference_steps` 和 `timesteps` 必须为 `None`。
  • guidance_scale (float, 可选, 默认为 5.0) — Classifier-Free Diffusion Guidance 中定义的引导比例。`guidance_scale` 被定义为 Imagen Paper 方程 2 中的 `w`。通过设置 `guidance_scale > 1` 来启用引导比例。更高的引导比例鼓励生成与文本 `prompt` 紧密相关的图像,通常以牺牲图像质量为代价。
  • num_images_per_prompt (int, 可选, 默认为 1) — 每个提示词生成的图像数量。
  • generator (torch.GeneratorList[torch.Generator], 可选) — 一个或多个 torch 生成器,用于使生成具有确定性。
  • latents (torch.FloatTensor, 可选) — 预先生成的噪声潜在表示,从高斯分布中采样,用作图像生成的输入。可用于使用不同的提示词调整相同的生成。如果未提供,将通过使用提供的随机 `generator` 进行采样来生成潜在张量。
  • prompt_embeds (torch.FloatTensor, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入,例如提示词加权。如果未提供,将从 `prompt` 输入参数生成文本嵌入。
  • prompt_attention_mask (torch.Tensor, 可选) — 预生成的文本嵌入注意力掩码。
  • negative_prompt_embeds (torch.FloatTensor, 可选) — 预生成的负面文本嵌入。可用于轻松调整文本输入,例如提示词加权。如果未提供,将从 `negative_prompt` 输入参数生成 negative_prompt_embeds。
  • negative_prompt_attention_mask (torch.Tensor, 可选) — 负面文本嵌入的预生成注意力掩码。
  • output_type (str, 可选, 默认为 "pil") — 生成图像的输出格式。在 PIL: PIL.Image.Imagenp.array 之间选择。
  • return_dict (bool, 可选, 默认为 True) — 是否返回 ~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput 而不是普通的元组。
  • attention_kwargs (dict, 可选) — 一个 kwargs 字典,如果指定,则会传递给 diffusers.models.attention_processor 中定义的 `self.processor` 下的 `AttentionProcessor`。
  • 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` 一起使用的最大序列长度。

调用管道进行生成时调用的函数。

示例

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

返回:ImagePipelineOutputtuple:如果 `return_dict` 为 `True`,则返回 ImagePipelineOutput,否则返回一个 `tuple`,其中第一个元素是包含生成图像的列表。

encode_prompt

< >

( prompt: typing.Union[str, typing.List[str]] negative_prompt: typing.Union[str, typing.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: typing.Optional[float] = None )

参数

  • prompt (strList[str], 可选) — 要编码的提示词。
  • negative_prompt (strList[str], 可选) — 不用于引导图像生成的提示词。如果未定义,则必须传递 negative_prompt_embeds。当不使用引导时(即,如果 guidance_scale 小于 1 时),此参数将被忽略。
  • do_classifier_free_guidance (bool, 可选, 默认为 True) — 是否使用分类器自由引导。
  • 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 比例。

将提示编码为文本编码器隐藏状态。

< > 在 GitHub 上更新