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 主干网络的基础上增加了三项内容：

  1. 通过 torch.nn.attention.flex_attention 实现的 FAR 因果块掩码，支持 FAR 中介绍的按块自回归生成。
  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 压缩（上下文）块的更大 patch 维度。
  • 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) — 归一化层的 epsilon 值。
  • 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 不同，则需要传入匹配的 chunk_partition 覆盖值到 AnyFlowFARPipeline.call()（同样也适用于 forward）。

用于 AnyFlow 流图采样的因果（FAR）3D Transformer，支持分块自回归生成。

在 v0.35.1 Wan2.1 主干网络的基础上扩展了以下功能：

  • 通过 torch.nn.attention.flex_attention 实现的 FAR 因果块掩码，支持分块自回归生成（FAR）。
  • 压缩帧补丁嵌入 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 — 潜在空间维度。
  • 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 缓存进行注意力计算，此处没有 mode。

返回值

~torch.nn.attention.flex_attention.BlockMask

跨越 FAR 布局的因果掩码，沿序列维度填充为 128 的倍数（BlockMask 的块大小要求）。

引发

ValueError

  • ValueError — 如果 mode 既不是 "train" 也不是 "cache"。

在 forward 外部预先构建因果 ~torch.nn.attention.flex_attention.BlockMask。

通过 forward 的 attention_mask 关键字参数传入结果，可使整个 transformer 兼容 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，可选) — I2V 图像嵌入；如果提供，则拼接在文本 token 之前。
  • clean_hidden_states (torch.Tensor, optional) — 训练 rollout 中使用的干净（无噪声）条件帧。
  • clean_timestep (torch.Tensor, optional) — 训练 rollout 中干净条件帧的时间步。
  • kv_cache (List[Dict[str, torch.Tensor]], 可选) — 用于自回归推理的按块 KV 缓存。None 表示选择训练路径。
  • kv_cache_flag (Dict[str, Any]，可选) — KV 缓存元数据（例如 is_cache_step 标志和 token 数量）。
  • attention_mask (BlockMask, 可选) — 预构建的因果掩码，通常通过 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。否则，因果训练 rollout（kv_cache is None）返回 [~models.transformer_2d.Transformer2DModelOutput]，而缓存预填充和自回归推理路径返回 [AnyFlowFARTransformerOutput]。

FAR 因果前向传播。分发到三个内部路径之一：

  • kv_cache is None → 因果训练 rollout（返回 Transformer2DModelOutput）。
  • kv_cache is not None 且 kv_cache_flag["is_cache_step"] → 缓存预填充（返回 AnyFlowFARTransformerOutput，其中 sample=None）。
  • 否则 → 自回归推理步骤（返回 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) — 自回归块的预测去噪目标。对于仅写入 KV 缓存且不生成可用样本的缓存预填充路径，值为 None。
  • kv_cache (list[dict[str, torch.Tensor]], optional) — 用于后续自回归步骤的按块划分的 KV 缓存状态。

AnyFlowFARTransformer3DModel 因果前向路径的输出数据类。

Update on GitHub

←AllegroTransformer3DModel AnyFlowTransformer3DModel→