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

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

diffusers.AuraFlowPipeline[[diffusers.AuraFlowPipeline]]

diffusers.AuraFlowPipeline(tokenizer: T5Tokenizer, text_encoder: UMT5EncoderModel, vae: AutoencoderKL, transformer: AuraFlowTransformer2DModel, scheduler: 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[[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'])

ソース

パラメータ:

prompt (str または list[str]、省略可) : 画像生成をガイドするためのプロンプトです。未定義の場合、代わりに prompt_embeds を渡す必要があります。

negative_prompt (str または list[str]、省略可能) : 画像生成をガイドしないプロンプトです。定義されていない場合、代わりに negative_prompt_embeds を渡す必要があります。ガイドを使用していない場合は無視されます（guidance_scale が 1 未満の場合は無視されます）。

height (int、optional、デフォルトは self.transformer.config.sample_size * self.vae_scale_factor) : 生成される画像のピクセル単位の縦幅。最適な結果を得るため、デフォルトでは 1024 に設定されています。

width (int、optional、デフォルトは 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, 任意) : 画像生成の入力として使用する、ガウス分布からサンプリングされた事前生成されたノイズ付き潜在変数（latents）。異なるプロンプトで同じ生成を調整するために使用できます。指定されない場合、指定されたランダム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[[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)

ソース

パラメータ:

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、optional、デフォルトは 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スケールです。

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