Buckets:
AnyFlowFARTransformer3DModel
AnyFlowFARPipeline が使用する因果的(FAR)3D Transformer — AnyFlow の FAR バリアントです。論文、著者、公開されたチェックポイントについては、AnyFlowFARPipeline のページを参照してください。v0.35.1 の Wan2.1 バックボーンを拡張し、3 つの追加要素を含みます。
- FARで導入されたチャンク単位の自己回帰生成をサポートする、
torch.nn.attention.flex_attentionによるFAR causal block-mask。 - コンテキスト(すでに生成された)フレーム用の圧縮フレームパッチ埋め込み(
far_patch_embedding)。構築時に、フル解像度のpatch_embeddingから三線形補間によってウォームスタートされます。 - デュアルタイムステップフローマップ埋め込み(
AnyFlowTransformer3DModelと同じ)— すべての forward 呼び出しは、ソースタイムステップ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]]
diffusers.AnyFlowFARTransformer3DModel[[diffusers.AnyFlowFARTransformer3DModel]]
diffusers.AnyFlowFARTransformer3DModel(patch_size: typing.Tuple[int] = (1, 2, 2), compressed_patch_size: typing.Tuple[int] = (1, 4, 4), full_chunk_limit: int = 3, num_attention_heads: int = 40, attention_head_dim: int = 128, in_channels: int = 16, out_channels: int = 16, text_dim: int = 4096, freq_dim: int = 256, ffn_dim: int = 13824, num_layers: int = 40, cross_attn_norm: bool = True, eps: float = 1e-06, image_dim: typing.Optional[int] = None, rope_max_seq_len: int = 1024, gate_value: float = 0.25, deltatime_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 のオーバーライドを対応する値に設定する必要があります。
Causal (FAR) AnyFlow のフローマップサンプリングに、チャンク単位の自己回帰生成を用いた 3D Transformer。
v0.35.1のWan2.1バックボーンを以下のように拡張します:
torch.nn.attention.flex_attentionによる FAR causal block-mask、チャンク単位の自己回帰生成(FAR)をサポートします。- コンテキスト(すでに生成済みの)フレーム用の圧縮フレームパッチ埋め込み
far_patch_embeddingは、patch_embeddingから三線形補間によって初期化されます。これにより、LoRAファインチューニングを行う前でも、新しく構築したモデルがすでに妥当な初期点にあることを保証します。 - 任意ステップサンプリング用のデュアルタイムステップフローマップ埋め込み(
AnyFlowTransformer3DModelと同じ)。
通常の双方向 T2V には AnyFlowTransformer3DModel を使用してください。このバリアントは FAR の因果マスクと far_patch_embedding をスキップするため、サイズが約 5〜10% 小さくなります。
build_attention_mask[[diffusers.AnyFlowFARTransformer3DModel.build_attention_mask]]
build_attention_mask(chunk_partition: typing.List[int], height: int, width: int, has_clean_context: bool = False, device: typing.Optional[torch.device] = None, mode: str = 'train')
パラメータ:
chunk_partition : チャンクごとのフレーム数(潜在フレームの総数と一致する必要があります)。
height, width : 潜在空間の空間次元。
has_clean_context : clean_hidden_states が forward に渡される場合(V2V/I2V の学習時)に True になります。マスクのレイアウトに影響するのは、この存在フラグのみです。
device : 生成されるBlockMaskのデバイス。device_map="auto"ではマスクは自動的に移動されません。Transformerの入力が配置されるデバイス上で構築してください。
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[[diffusers.AnyFlowFARTransformer3DModel.forward]]
forward(hidden_states: Tensor, timestep: Tensor, r_timestep: Tensor, encoder_hidden_states: Tensor, chunk_partition: typing.List[int], encoder_hidden_states_image: typing.Optional[torch.Tensor] = None, clean_hidden_states: typing.Optional[torch.Tensor] = None, clean_timestep: typing.Optional[torch.Tensor] = None, kv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = None, kv_cache_flag: typing.Optional[typing.Dict[str, typing.Any]] = None, attention_mask: typing.Optional[torch.nn.attention.flex_attention.BlockMask] = None, attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, return_dict: bool = True)
パラメータ:
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]、optional) : KV-cache のメタデータ(例: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、optional) : attention processors に転送されます。
return_dict (bool、optional、デフォルトは 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 および autoregressive inference のパスでは [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]]
diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput[[diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput]]
diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput(sample: typing.Optional[torch.Tensor] = None, kv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = None)
パラメータ:
sample (torch.Tensor または None) : オートリグレッシブチャンクに対する予測されたデノイジングターゲット。KVキャッシュのみを書き込み、使用可能なサンプルを生成しないキャッシュプリフィルパスでは None になります。
kv_cache (list[dict[str, torch.Tensor]]、optional) : 後続の自己回帰ステップで使用されるブロックごとのKVキャッシュ状態。
AnyFlowFARTransformer3DModelの因果的前方パス用の出力データクラスです。
Xet Storage Details
- Size:
- 14.6 kB
- Xet hash:
- 42ade6f0334219c48a719b5c4c3a89d896dc1c5025bb2a3bb689e1d5616ca8f7
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.