diffusers/pr_14371/en/api/models/anyflow_far_transformer3d.md

AnyFlowFARTransformer3DModel

The causal (FAR) 3D Transformer used by AnyFlowFARPipeline — the FAR variant of AnyFlow. See the AnyFlowFARPipeline page for paper, authors, and released checkpoints. It extends the v0.35.1 Wan2.1 backbone with three additions:

1. FAR causal block-mask via torch.nn.attention.flex_attention, supporting chunk-wise autoregressive generation as introduced in FAR.
2. Compressed-frame patch embedding (far_patch_embedding) for context (already-generated) frames, warm-started from the full-resolution patch_embedding at construction time via trilinear interpolation.
3. Dual-timestep flow-map embedding (same as AnyFlowTransformer3DModel) — every forward call conditions on both the source timestep t and the target timestep r.

The default chunk schedule (chunk_partition) is stored in the model config; the released NVIDIA AnyFlow-FAR checkpoints use [1, 3, 3, 3, 3, 3, 3, 2] for the canonical 81-frame setting. forward accepts a per-call chunk_partition override, so the same checkpoint also handles other num_frames configurations without retraining.

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

Source

Parameters:

patch_size (Tuple[int], defaults to (1, 2, 2)) : 3D patch dimensions for full-resolution chunks.

compressed_patch_size (Tuple[int], defaults to (1, 4, 4)) : Larger patch dimensions for the FAR-compressed (context) chunks.

full_chunk_limit (int, defaults to 3) : Maximum number of full-resolution chunks before earlier chunks are demoted to compressed FAR context. The released checkpoints use 3.

num_attention_heads (int, defaults to 40) : Number of attention heads.

attention_head_dim (int, defaults to 128) : The number of channels in each head.

in_channels (int, defaults to 16) : The number of channels in the input latent.

out_channels (int, defaults to 16) : The number of channels in the output latent.

text_dim (int, defaults to 4096) : Input dimension for text embeddings (UMT5).

freq_dim (int, defaults to 256) : Dimension for sinusoidal time embeddings.

ffn_dim (int, defaults to 13824) : Intermediate dimension in feed-forward network.

num_layers (int, defaults to 40) : Number of transformer blocks.

cross_attn_norm (bool, defaults to True) : Enable cross-attention normalization.

eps (float, defaults to 1e-6) : Epsilon for normalization layers.

image_dim (Optional[int], optional, defaults to None) : Image embedding dimension for I2V conditioning.

rope_max_seq_len (int, defaults to 1024) : Maximum sequence length used to precompute rotary position frequencies.

gate_value (float, defaults to 0.25) : Mixing gate between source-timestep and delta-timestep embeddings.

deltatime_type (str, defaults to 'r') : Either "r" (delta is the target timestep) or "t-r" (delta is the absolute interval).

chunk_partition (Tuple[int, ...], defaults to (1, 3, 3, 3, 3, 3, 3, 2)) : Default per-chunk frame counts used by the pipeline. The released NVIDIA AnyFlow-FAR checkpoints target num_frames=81 (21 latent frames at VAE temporal stride 4) split as 1 + 3*6 + 2. A different num_frames requires a matching chunk_partition override passed to AnyFlowFARPipeline.call() (and likewise to forward).

Causal (FAR) 3D Transformer for AnyFlow flow-map sampling with chunk-wise autoregressive generation.

Extends the v0.35.1 Wan2.1 backbone with:

• FAR causal block-mask via torch.nn.attention.flex_attention, supporting chunk-wise autoregressive generation (FAR).
• Compressed-frame patch embedding far_patch_embedding for context (already-generated) frames, initialized from patch_embedding via trilinear interpolation so a freshly constructed model is already at a reasonable starting point even before LoRA fine-tuning.
• Dual-timestep flow-map embedding for any-step sampling (same as AnyFlowTransformer3DModel).

Use AnyFlowTransformer3DModel instead for plain bidirectional T2V — that variant skips the FAR causal masking and far_patch_embedding and is ~5–10% smaller.

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

Source

Parameters:

chunk_partition : per-chunk frame counts (must sum to the number of latent frames).

height, width : latent spatial dimensions.

has_clean_context : True when clean_hidden_states will be threaded through forward (training V2V/I2V); only this presence flag affects the mask layout.

device : device for the resulting BlockMask. The mask is not auto-moved by device_map="auto"; build it on the same device the transformer's inputs will live on.

mode : "train" (matches _forward_train) or "cache" (matches _forward_cache). The autoregressive _forward_inference path attends through the KV cache and has no mode here.

Returns: ~torch.nn.attention.flex_attention.BlockMask

causal mask spanning the FAR layout, padded to a multiple of 128 along the sequence dimension (the BlockMask block-size requirement).

Raises: ValueError

• ValueError -- if mode is neither "train" nor "cache".

Pre-build the causal ~torch.nn.attention.flex_attention.BlockMask outside forward.

Pass the result via forward's attention_mask kwarg to make the whole transformer compatible with torch.compile(fullgraph=True). Without a pre-built mask, forward falls back to constructing it internally — that path uses flex_attention.create_block_mask(_compile=False) and breaks the compile graph.

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)

Source

Parameters:

hidden_states (torch.Tensor) : Latent input of shape (B, F, C, H, W).

timestep (torch.Tensor) : Source (noisier) flow-map timestep t.

r_timestep (torch.Tensor) : Target (cleaner) flow-map timestep r.

encoder_hidden_states (torch.Tensor) : UMT5 text embeddings.

chunk_partition (List[int]) : Per-chunk frame counts; total must match the number of latent frames in hidden_states.

encoder_hidden_states_image (torch.Tensor, optional) : I2V image embedding; concatenated before text tokens when provided.

clean_hidden_states (torch.Tensor, optional) : Clean (noise-free) conditioning frames used by the training rollout.

clean_timestep (torch.Tensor, optional) : Timesteps for the clean conditioning frames in the training rollout.

kv_cache (List[Dict[str, torch.Tensor]], optional) : Per-block KV cache for autoregressive inference. None selects the training path.

kv_cache_flag (Dict[str, Any], optional) : KV-cache metadata (e.g. is_cache_step flag and token counts).

attention_mask (BlockMask, optional) : Pre-built causal mask, typically constructed via build_attention_mask. Consumed by the train and KV-cache prefill paths; the autoregressive inference path attends through the KV cache and does not use a full mask. When None, the train / cache paths build the mask internally; that fallback is not compile-safe (the underlying flex_attention.create_block_mask breaks the graph under fullgraph=True), so pass a pre-built mask whenever wrapping forward in torch.compile.

attention_kwargs (dict, optional) : Forwarded to the attention processors.

return_dict (bool, optional, defaults to True) : If False, returns positional tuples instead of an output dataclass.

Returns: [~models.transformer_2d.Transformer2DModelOutput], [AnyFlowFARTransformerOutput] or tuple

When return_dict is False, a plain tuple is returned. Otherwise, the causal training rollout (kv_cache is None) returns a [~models.transformer_2d.Transformer2DModelOutput], while the cache-prefill and autoregressive inference paths return an [AnyFlowFARTransformerOutput].

FAR causal forward pass. Dispatches to one of three internal paths:

• kv_cache is None → causal training rollout (returns Transformer2DModelOutput).
• kv_cache is not None and kv_cache_flag["is_cache_step"] → cache-prefill (returns AnyFlowFARTransformerOutput with sample=None).
• Otherwise → autoregressive inference step (returns 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)

Source

Parameters:

sample (torch.Tensor or None) : Predicted denoising target for the autoregressive chunk. None for the cache-prefill path, which only writes the KV cache and produces no usable sample.

kv_cache (list[dict[str, torch.Tensor]], optional) : Per-block KV cache state used by subsequent autoregressive steps.

Output dataclass for AnyFlowFARTransformer3DModel's causal forward paths.