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はStable Diffusion 3に着想を得ており、Apache 2.0ライセンスが付与されているテキストから画像を生成するモデルの中で、現時点で最大規模のものです。このモデルは、GenEvalベンチマークにおいて最先端の結果を達成しています。

Falチームによって開発され、詳細についてはこのブログ記事を参照してください。

AuraFlowは、コンシューマー向けハードウェアで実行するとコストが高くなることがあります。ただし、一連の最適化を行うことで、より高速かつメモリ効率の良い実行が可能です。詳細については、このセクションを参照してください。

量子化

Quantizationは、モデルの重みを低精度のデータ型で保存することで、非常に大規模なモデルのメモリ要件を軽減するのに役立ちます。ただし、動画モデルによって、Quantizationが動画の品質に与える影響は異なります。

サポートされている量子化バックエンドと、ユースケースに対応する量子化バックエンドの選択方法については、Quantizationの概要を参照してください。以下の例では、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,
    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%、1536x1536解像度では30%の高速化が実現されます。

上記が異なる解像度でも動作するように、AuraFlowTransformer2DModelクラスの書き直しを手伝ってくれたAstraliteHeartに感謝します(PR)。

AuraFlowPipeline

クラス diffusers.AuraFlowPipeline

< >

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

< >

( 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

パラメータ

  • prompt (str または list[str]、省略可) — 画像生成をガイドするためのプロンプトです。未定義の場合、代わりに prompt_embeds を渡す必要があります。
  • negative_prompt (str または list[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 は、Imagen Paper の数式 2 の w として定義されています。guidance_scale > 1 を設定することで、ガイダンススケールが有効になります。ガイダンススケールを高くすると、prompt のテキストと密接に関連する画像の生成が促されますが、通常は画像の品質が低下する傾向があります。
  • num_images_per_prompt (int、任意、デフォルトは 1) — プロンプトごとに生成する画像の数。
  • generator (torch.Generator または list[torch.Generator]、省略可) — 生成を決定論的にするためのtorch ジェネレーター1つまたはリスト。
  • 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.Image または np.array から選択します。
  • return_dict (bool、省略可能、デフォルトは True) — ~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput を返すかどうか。プレーンなタプルを返すこともできます。
  • attention_kwargs (dict, 省略可) — 指定された場合、diffusers.models.attention_processor で定義されている self.processor 内の AttentionProcessor に渡される kwargs 辞書です。
  • 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

< >

( 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 (str または list[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 スケールです。

プロンプトをテキストエンコーダーの隠れ状態にエンコードします。

Update on GitHub