Hugging Face's logo

Diffusers documentation

AnyFlowFARTransformer3DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

AnyFlowFARTransformer3DModel

AnyFlowFARPipelineで使用される因果的（FAR）3D Transformer — AnyFlowのFARバリアントです。論文、著者、公開されたチェックポイントについては、AnyFlowFARPipelineのページを参照してください。v0.35.1のWan2.1バックボーンを拡張し、3つの追加要素を含みます：

  1. FARで導入されたチャンク単位の自己回帰生成をサポートする、torch.nn.attention.flex_attentionによるFAR causal block-mask。
  2. コンテキスト（すでに生成された）フレーム用の圧縮フレームパッチ埋め込み（far_patch_embedding）。構築時に、フル解像度のpatch_embeddingから三線形補間によってウォームスタートされます。
  3. デュアルタイムステップフローマップ埋め込み（AnyFlowTransformer3DModelと同様）— すべてのフォワード呼び出しは、ソースタイムステップtとターゲットタイムステップrの両方を条件として使用します。

デフォルトのチャンクスケジュール(chunk_partition)はモデルの設定に保存されます。公開されているNVIDIA AnyFlow-FARのチェックポイントは、標準的な81フレームの設定に対して[1, 3, 3, 3, 3, 3, 3, 2]を使用しています。forwardは呼び出しごとのchunk_partitionの上書きを受け入れるため、同じチェックポイントで再学習なしに他のnum_frames設定にも対応できます。

from diffusers import AnyFlowFARTransformer3DModel

# Causal AnyFlow checkpoint (FAR):
transformer = AnyFlowFARTransformer3DModel.from_pretrained(
    "nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers", subfolder="transformer"
)

AnyFlowFARTransformer3DModel

クラス diffusers.AnyFlowFARTransformer3DModel

< >

( patch_size: typing.Tuple[int] = (1, 2, 2)compressed_patch_size: typing.Tuple[int] = (1, 4, 4)full_chunk_limit: int = 3num_attention_heads: int = 40attention_head_dim: int = 128in_channels: int = 16out_channels: int = 16text_dim: int = 4096freq_dim: int = 256ffn_dim: int = 13824num_layers: int = 40cross_attn_norm: bool = Trueeps: float = 1e-06image_dim: typing.Optional[int] = Nonerope_max_seq_len: int = 1024gate_value: float = 0.25deltatime_type: str = 'r'chunk_partition: typing.Tuple[int, ...] = (1, 3, 3, 3, 3, 3, 3, 2) )

パラメータ

  • patch_size (Tuple[int]、デフォルトは (1, 2, 2)) — フルレゾリューションのチャンクに対する3Dパッチの次元です。
  • compressed_patch_size (Tuple[int]、デフォルトは (1, 4, 4)) — FAR圧縮（コンテキスト）チャンク用の、より大きなパッチ次元。
  • full_chunk_limit (int、デフォルトは 3) — 以前のチャンクが圧縮 FAR コンテキストに降格されるまでの、フル解像度チャンクの最大数です。リリースされたチェックポイントでは 3 が使用されています。
  • num_attention_heads (int、デフォルトは 40) — アテンションヘッドの数。
  • attention_head_dim (int、デフォルトは 128) — 各ヘッドのチャネル数です。
  • in_channels (int、デフォルトは 16) — 入力潜在変数のチャネル数です。
  • out_channels (int、デフォルトは 16) — 出力潜在変数のチャネル数。
  • text_dim (int、デフォルトは 4096) — テキスト埋め込み (UMT5) の入力次元数。
  • freq_dim (int、デフォルトは 256) — 正弦波時間埋め込みの次元数。
  • ffn_dim (int、デフォルトは 13824) — フィードフォワードネットワークの中間次元。
  • num_layers (int、デフォルトは 40) — Transformerブロックの数。
  • cross_attn_norm (bool、デフォルトは True) — クロスアテンションの正規化を有効にします。
  • eps (float、デフォルトは 1e-6) — 正規化レイヤー用のエプシロン。
  • image_dim (Optional[int]、任意、デフォルトは None) — I2V 条件付け用の画像埋め込み次元数。
  • rope_max_seq_len (int、デフォルトは 1024) — 回転位置周波数を事前計算するために使用される最大シーケンス長。
  • gate_value (float、デフォルトは 0.25) — ソースタイムステップ埋め込みとデルタタイムステップ埋め込みの間の混合ゲート。
  • deltatime_type (str、デフォルトは ‘r’) — "r"（delta がターゲットのタイムステップ）または "t-r"（delta が絶対区間）のいずれかです。
  • chunk_partition (Tuple[int, …]、デフォルトは (1, 3, 3, 3, 3, 3, 3, 2)) — パイプラインで使用されるチャンクごとのフレーム数のデフォルト値です。公開されている NVIDIA AnyFlow-FAR チェックポイントは num_frames=81（VAE の時間方向ストライド 4 における潜在フレーム 21 相当）をターゲットとしており、1 + 3*6 + 2 に分割されます。異なる num_frames を使用する場合は、AnyFlowFARPipeline.call()（および forward）に渡す chunk_partition のオーバーライドが必要です。

AnyFlowのフローマップサンプリングにおけるチャンク単位の自己回帰生成を行う、因果的(FAR)3D Transformer。

v0.35.1のWan2.1バックボーンを以下のように拡張します：

  • torch.nn.attention.flex_attention による FAR causal block-mask、チャンク単位の自己回帰生成（FAR）をサポートします。
  • コンテキスト（すでに生成された）フレーム用の Compressed-frame patch embedding far_patch_embedding は、patch_embedding から三線形補間によって初期化されており、LoRAファインチューニングを行う前でも、新しく構築されたモデルがすでに適切な初期点にあることを保証します。
  • 任意ステップサンプリング用のデュアルタイムステップフローマップ埋め込み（AnyFlowTransformer3DModelと同じ）。

通常の双方向 T2V には AnyFlowTransformer3DModel を使用してください。このバリアントは FAR の因果マスクと far_patch_embedding をスキップするため、サイズが約 5〜10% 小さいです。

build_attention_mask

< >

( chunk_partition: typing.List[int]height: intwidth: inthas_clean_context: bool = Falsedevice: typing.Optional[torch.device] = Nonemode: str = 'train' ) → ~torch.nn.attention.flex_attention.BlockMask

パラメータ

  • chunk_partition — チャンクごとのフレーム数（潜在フレームの総数と一致する必要があります）。
  • height, width — latent の空間次元。
  • has_clean_context — clean_hidden_statesがforwardにスレッドされる場合（V2V/I2Vの学習時）にTrueになります。マスクのレイアウトに影響するのは、この存在フラグのみです。
  • device — 結果として得られるBlockMaskのデバイスです。マスクはdevice_map="auto"によって自動的に移動されません。トランスフォーマーの入力が配置されるのと同じデバイス上で構築してください。
  • mode — "train"（_forward_trainに一致）または"cache"（_forward_cacheに一致）。 自己回帰的な_forward_inferenceパスはKVキャッシュを介してアテンションを行い、ここではモードを持ちません。

戻り値

~torch.nn.attention.flex_attention.BlockMask

FARレイアウトにまたがるcausal mask。シーケンス次元に沿って128の倍数にパディングされています（BlockMaskのブロックサイズ要件）。

例外

ValueError

  • ValueError — mode が"train"でも"cache"でもない場合。

forward の外で、causal な ~torch.nn.attention.flex_attention.BlockMask を事前構築してください。

forward の attention_mask キーワード引数に結果を渡すことで、トランスフォーマー全体を torch.compile(fullgraph=True) と互換性のある状態にします。事前構築されたマスクがない場合、forward は内部でマスクを構築するフォールバック処理に切り替わります。この処理は flex_attention.create_block_mask(_compile=False) を使用するため、コンパイルグラフが破綻します。

forward

< >

( hidden_states: Tensortimestep: Tensorr_timestep: Tensorencoder_hidden_states: Tensorchunk_partition: typing.List[int]encoder_hidden_states_image: typing.Optional[torch.Tensor] = Noneclean_hidden_states: typing.Optional[torch.Tensor] = Noneclean_timestep: typing.Optional[torch.Tensor] = Nonekv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = Nonekv_cache_flag: typing.Optional[typing.Dict[str, typing.Any]] = Noneattention_mask: typing.Optional[torch.nn.attention.flex_attention.BlockMask] = Noneattention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = Nonereturn_dict: bool = True ) → [~models.transformer_2d.Transformer2DModelOutput], [AnyFlowFARTransformerOutput] or tuple

パラメータ

  • hidden_states (torch.Tensor) — 形状が(B, F, C, H, W)の潜在入力です。
  • timestep (torch.Tensor) — ソース（ノイズの多い）フローマップのタイムステップ t。
  • r_timestep (torch.Tensor) — ターゲット（よりクリーンな）フローマップのタイムステップ r。
  • encoder_hidden_states (torch.Tensor) — UMT5のテキスト埋め込み。
  • chunk_partition (List[int]) — チャンクごとのフレーム数。hidden_states内の潜在フレームの総数と一致する必要があります。
  • encoder_hidden_states_image (torch.Tensor、optional) — I2V 画像埋め込み。指定された場合、テキストトークンの前に連結されます。
  • clean_hidden_states (torch.Tensor、optional) — 学習ロールアウトで使用される、ノイズのない（クリーンな）条件フレーム。
  • clean_timestep (torch.Tensor, optional) — 学習ロールアウトにおけるクリーンな条件付けフレームのタイムステップです。
  • kv_cache (List[Dict[str, torch.Tensor]], optional) — 自己回帰推論用のブロックごとの KV キャッシュです。None を指定すると、学習パスが選択されます。
  • kv_cache_flag (Dict[str, Any]、任意) — KVキャッシュのメタデータ（例：is_cache_stepフラグやトークン数）。
  • attention_mask (BlockMask, optional) — 事前構築済みの因果マスク。通常は build_attention_mask を通じて構築されます。学習パスと KV キャッシュのプリフィルパスで使用されます。自動回帰推論パスは KV キャッシュを介してアテンションを行うため、完全なマスクは使用しません。None の場合、学習 / キャッシュパスは内部でマスクを構築しますが、このフォールバックはコンパイル安全ではありません（基盤となる flex_attention.create_block_mask は fullgraph=True 下でグラフを破断するため）、forward を torch.compile でラップする場合は、事前構築済みのマスクを渡してください。
  • attention_kwargs (dict、任意) — アテンションプロセッサに転送されます。
  • return_dict (bool、任意、デフォルトは True) — False の場合、出力データクラスではなく位置付きタプルを返します。

戻り値

[~models.transformer_2d.Transformer2DModelOutput]、[AnyFlowFARTransformerOutput] または tuple

return_dict が False の場合、プレーンな tuple が返されます。それ以外の場合、causal training rollout（kv_cache is None）は [~models.transformer_2d.Transformer2DModelOutput] を返し、cache-prefill および自己回帰推論のパスは [AnyFlowFARTransformerOutput] を返します。

FARの因果的前方パス。3つの内部パスのいずれかにディスパッチします:

  • kv_cache is None → causal training rollout（Transformer2DModelOutputを返します）。
  • kv_cache is not None かつ kv_cache_flag["is_cache_step"] → cache-prefill（sample=None を持つ AnyFlowFARTransformerOutput を返します）。
  • それ以外 → 自己回帰推論ステップ（AnyFlowFARTransformerOutputを返します）。

AnyFlowFARTransformerOutput

クラス diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput

< >

( sample: typing.Optional[torch.Tensor] = Nonekv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = None )

パラメータ

  • sample (torch.Tensor または None) — 自己回帰チャンクに対する予測されたデノイジングターゲット。None はキャッシュプリフィルパスであり、KVキャッシュへの書き込みのみを行い、使用可能なサンプルは生成しません。
  • kv_cache (list[dict[str, torch.Tensor]]、optional) — 後続の自己回帰ステップで使用されるブロックごとのKVキャッシュ状態。

AnyFlowFARTransformer3DModelの因果的前方パス用の出力データクラスです。

Update on GitHub

←AllegroTransformer3DModel AnyFlowTransformer3DModel→