ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

AnyFlowFARTransformer3DModel 深度解析:基于 FAR 因果掩码的分块自回归视频扩散 Transformer

AnyFlowFARTransformer3DModel 深度解析:基于 FAR 因果掩码的分块自回归视频扩散 Transformer AnyFlowFARTransformer3DModel 深度解析基于 FAR 因果掩码的分块自回归视频扩散 Transformer【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers导读AnyFlowFARTransformer3DModel是 diffusers 中驱动AnyFlowFARPipeline的因果FAR3D Transformer它在 Wan2.1 v0.35.1 骨干网络上叠加了三项结构增强基于torch.nn.attention.flex_attention的 FAR 因果块掩码支持分块自回归生成、面向已生成上下文帧的压缩 patch 嵌入以及与双向变体相同的双时间步 flow-map 嵌入。本文以该模型的 API 文档为主线结合 transformer_anyflow_far.py 源码与测试用例讲解其架构原理、配置参数、三种前向路径训练 / KV 缓存预填充 / 自回归推理以及与之配套的流水线用法帮助读者理解任意步数视频扩散 长视频因果生成这一组合在 diffusers 中的具体实现。背景AnyFlow 与 FAR 因果变体AnyFlow 是由 NVIDIA、新加坡国立大学与 MIT 合作提出的视频扩散蒸馏框架论文编号 2605.13724其核心思想是学习任意时间区间上的flow map$\Phi_{r\leftarrow t}: \mathbf{z}_t \to \mathbf{z}_r$而非 consistency models 的固定端点映射 $\mathbf{z}_t \to \mathbf{z}_0$。这使得同一个模型无需重新调度即可在任意步数1、2、4、8、16……下推理质量随步数增加单调提升。AnyFlowFARTransformer3DModel是 AnyFlow 的causal (FAR) 变体的骨干模型被AnyFlowFARPipeline使用。所谓 FAR来自论文 2503.19325指分块自回归生成视频被切成多个 chunk 逐块去噪每个 chunk 只关注过去已生成的 chunk而不是像双向模型那样一次性对整个视频张量做全局自注意力。这一设计让流水线在所有 latent 帧分配完成之前就能开始采样对超长序列如 81 帧甚至更长尤其重要同时也天然支持首帧图生视频I2V与视频续写V2V。官方发布的相关检查点包括nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers与nvidia/AnyFlow-FAR-Wan2.1-14B-Diffusers因果 T2V / I2V / V2V其对应的双向版本nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers等则由AnyFlowTransformer3DModel驱动。相对双向骨干的三项结构增强按照 API 文档的归纳AnyFlowFARTransformer3DModel在 v0.35.1 Wan2.1 骨干上做了三处关键改动源码中的实现与之一一对应1. FAR 因果块掩码BlockMask通过torch.nn.attention.flex_attention的BlockMask机制实现分块因果注意力支持 chunk 级的自回归生成。源码中掩码的构建由模块级函数_build_anyflow_far_causal_block_mask完成mask_mod定义了frame_idx[q_idx] frame_idx[kv_idx]的分块因果比较序列长度按 128 的倍数补齐以满足 BlockMask 的块大小要求模型暴露了build_attention_mask方法允许调用方在forward之外预构建掩码并经由attention_mask参数传入从而让整个 transformer 兼容torch.compile(fullgraphTrue)——因为flex_attention.create_block_mask(_compileFalse)在编译作用域内会打断计算图消费该掩码的是专用的AnyFlowCausalAttnProcessor它把_attention_backend固定为flex支持flex与_native_flex若通过dispatch_attention_fn配置了非 flex 后端会直接抛出ValueError同时在__init__中要求 PyTorch 2.0 及以上依赖F.scaled_dot_product_attention。2. 压缩帧 patch 嵌入far_patch_embedding对已经生成的上下文帧使用更大的 patch 尺寸进行压缩嵌入降低历史帧的 token 数量。模型配置中patch_size(1, 2, 2)对应全分辨率 chunkcompressed_patch_size(1, 4, 4)对应压缩上下文 chunk。关键实现细节是热启动warm-start# 源码片段transformer_anyflow_far.py original_weight self.patch_embedding.weight.data.view(-1, 1, *patch_size) new_weight F.interpolate(original_weight, sizecompressed_patch_size, modetrilinear, align_cornersFalse)在构造时通过三线性插值把全分辨率patch_embedding的权重复制到far_patch_embedding偏置也一并复制匹配 FAR-Dev 的setup_far_model()初始化策略——这意味着即使不做 LoRA 微调新构造的模型也已经处于一个合理的起点加载已训练检查点时state-dict 会覆盖这些权重。3. 双时间步 flow-map 嵌入dual-timestep与AnyFlowTransformer3DModel相同时间步嵌入被替换为AnyFlowDualTimestepTextImageEmbeddingCausal每一次前向都以源时间步t和目标时间步r为条件从而学习 flow map $\Phi_{r\leftarrow t}$。该模块支持任意步采样any-step并支持deltatime_type两种取值rdelta 时间步即目标时间步rt-rdelta 时间步为绝对区间t - r。此外gate_value0.25控制源时间步嵌入与 delta 时间步嵌入的混合门控计算公式为(1 - gate) * temb gate * delta_emb。快速上手加载 FAR 因果 TransformerAPI 文档给出了最直接的用法——用from_pretrained从因果版 AnyFlow 检查点单独加载 transformer 子模块from diffusers import AnyFlowFARTransformer3DModel # Causal AnyFlow checkpoint (FAR): transformer AnyFlowFARTransformer3DModel.from_pretrained( nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers, subfoldertransformer )在 diffusers 的顶层导出中AnyFlowFARTransformer3DModel可直接从diffusers命名空间导入见 src/diffusers/init.py 中的导出列表。更常见的做法是连同流水线一起加载import torch from diffusers import AnyFlowFARPipeline from diffusers.utils import export_to_video pipe AnyFlowFARPipeline.from_pretrained( nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers, torch_dtypetorch.bfloat16 ).to(cuda) prompt An astronaut runs smoothly and appears almost weightless on the lunar surface video pipe(prompt, num_inference_steps4, num_frames81).frames[0] export_to_video(video, anyflow_far.mp4, fps16)模型内部还支持WanLoraLoaderMixin与上游 Wan 流水线相同的 LoRA 加载机制、PeftAdapterMixin与FromOriginalModelMixin并启用了_supports_gradient_checkpointing_skip_layerwise_casting_patterns覆盖patch_embedding、far_patch_embedding、condition_embedder与各 norm 层。配置参数详解含默认值与含义AnyFlowFARTransformer3DModel.__init__的全部参数均通过register_to_config注册进模型配置以下是源码中的默认值与作用参数默认值说明patch_size(1, 2, 2)全分辨率 chunk 的 3D patch 尺寸compressed_patch_size(1, 4, 4)FAR 压缩上下文chunk 的更大 patch 尺寸full_chunk_limit3在更早 chunk 被降级为压缩 FAR 上下文前最多保留的全分辨率 chunk 数量已发布检查点使用3num_attention_heads40注意力头数attention_head_dim128每个注意力头的通道数in_channels/out_channels16/16输入 / 输出 latent 通道数text_dim4096文本嵌入UMT5输入维度freq_dim256正弦时间嵌入维度ffn_dim13824前馈网络中间维度num_layers40Transformer 块数量cross_attn_normTrue是否启用交叉注意力归一化eps1e-6归一化层 epsilonimage_dimNoneI2V 条件图像嵌入维度rope_max_seq_len1024预计算旋转位置频率的最大序列长度gate_value0.25源时间步嵌入与 delta 时间步嵌入的混合门控deltatime_typerrdelta 为目标时间步或t-rdelta 为绝对区间chunk_partition(1, 3, 3, 3, 3, 3, 3, 2)默认的逐 chunk 帧数见下文chunk_partition默认分块调度默认分块(1, 3, 3, 3, 3, 3, 3, 2)共 8 个 chunk总和为 21正好对应 81 帧原始视频在 VAE 时间步压缩率为 4 时得到的(81 - 1) // 4 1 21个 latent 帧。其中第一个 chunk 只有 1 帧最后的 chunk 只有 2 帧中间 6 个 chunk 各 3 帧写成1 3*6 2。forward接受每次调用级别的chunk_partition覆盖因此同一个检查点无需重新训练即可适配其他num_frames配置——前提是传入的chunk_partition总和必须等于 latent 帧数(num_frames - 1) // vae_scale_factor_temporal 1否则流水线会抛出ValueError。这也是AnyFlowFARPipeline中chunk_partition参数的语义。三条前向路径训练 / 缓存预填充 / 自回归推理forward根据kv_cache与kv_cache_flag的取值把计算路由到三个内部路径见 forward 方法理解这三条路径是掌握 FAR 因果生成的关键条件路径返回kv_cache is None_forward_train因果训练 rolloutTransformer2DModelOutput含samplekv_cache非空且kv_cache_flag[is_cache_step]为真_forward_cache缓存预填充只写入 KV cacheAnyFlowFARTransformerOutput(sampleNone, kv_cache...)其他_forward_inference自回归推理步经 KV cache 读取历史AnyFlowFARTransformerOutput(sample..., kv_cache...)各路径的差异点训练路径使用_forward_far_patchify同时处理全分辨率帧patch_embedding与压缩帧far_patch_embedding通过clean_hidden_states/clean_timestep支持训练 rolloutV2V / I2V输出前会裁剪掉 far 上下文与 clean 副本缓存预填充路径AnyFlowCausalAttnProcessor在is_cache_step分支把 key/value 分别写入compressed_cache与full_cachesample为None不产生可用的去噪结果自回归推理路径processor 从缓存中拼接历史 token 再与当前 chunk 的 token 拼接rotary_emb只取当前 chunk 对应的尾部位置输出只反 patchify 当前 chunkchunk_partition[-1]的帧。与 KV 缓存相关的内部数据结构流水线在 pipeline_anyflow_far.py 中按num_layers为每层分配两个缓存张量full_cache形状为(2, batch, heads, full_chunk_limit * max(chunk) * full_token_per_frame, head_dim)保存全分辨率历史 chunkcompressed_cache形状为(2, batch, heads, (len(chunk) - full_chunk_limit 1) * max(chunk) * compressed_token_per_frame, head_dim)保存已压缩的历史 chunk。kv_cache_flag中维护num_cached_chunks、is_cache_step以及各 token 计数配合AnyFlowCausalRotaryPosEmbed对压缩帧频率表做复数平均池化下采样avg_pool_complex保证压缩网格与全分辨率网格的位置编码保持对齐。输出数据类AnyFlowFARTransformerOutput源码位于 transformer_anyflow_far.py包含两个字段sample自回归 chunk 的预测去噪目标缓存预填充路径中为Nonekv_cache供后续自回归步使用的逐块 KV 缓存状态。在 AnyFlowFARPipeline 中的角色与调用链AnyFlowFARPipeline以tokenizer UMT5 文本编码器 AnyFlowFARTransformer3DModel AutoencoderKLWan FlowMapEulerDiscreteScheduler为组件其model_cpu_offload_seq text_encoder-transformer-vae。流水线通过video像素空间(B, T, C, H, W)T 4n 1或video_latents已编码的 FAR 布局 latent二选一决定任务模式两者都不传 → 纯文生视频T2V单帧video→ 图生视频I2V多帧video→ 视频续写V2V。去噪循环是外层遍历 chunk、内层遍历时间步的双层结构每个非上下文 chunk 都运行num_inference_steps步FlowMapEulerDiscreteScheduler.stepchunk 完成后调用encode_kv_cache把该 chunk 的 KV 写入缓存供后续 chunk 使用。与推理相关的关键参数num_inference_steps每个 chunk 的去噪步数默认50蒸馏检查点支持任意步数1、2、4、8……也可通过sigmas/timesteps自定义调度guidance_scale默认1.0——已发布检查点把 CFG 融合进了权重推理时不会再跑一次无条件前向use_mean_velocityTrue以源时间步t和目标时间步r为条件预测均值速度关闭则退化为原始 Euler 步进use_kv_cacheTrue跨因果 chunk 复用 FAR 注意力 KV 缓存仅调试时建议关闭chunk_partition覆盖默认分块改变num_frames时必须同步提供匹配的覆盖值。调度器方面FlowMapEulerDiscreteScheduler源码见 scheduling_flow_map_euler_discrete.py默认shift1.0可通过apply_shift施加乘法时间步偏移shift5与 AnyFlow 训练配置一致并通过set_timesteps结合自定义sigmas/timesteps支持任意步采样。与双向变体 AnyFlowTransformer3DModel 的对比维度AnyFlowTransformer3DModel双向AnyFlowFARTransformer3DModel因果 FAR驱动流水线AnyFlowPipeline双向 T2VAnyFlowFARPipelineT2V / I2V / V2V自注意力全局双向AnyFlowAttnProcessorSDPA 兼容后端分块因果AnyFlowCausalAttnProcessor仅 flex 后端掩码无 BlockMaskflex_attention.BlockMask序列按 128 补齐额外模块无far_patch_embedding压缩上下文帧三线性插值热启动KV 缓存无逐层full_cachecompressed_cache跨 chunk 复用模型体量约小 5–10%略大含压缩分支适用场景纯 prompt、固定 NFE 求最大质量I2V、V2V、流式 / 渐进式生成从源码结构看AnyFlowTransformerBlock通过is_causal标志在构造时选择自注意力处理器双向变体传FalseFAR 变体传True其余前向逻辑完全一致——所有因果专属机制BlockMask、KV 缓存都被封装在 processor 内部。两个文件之间的共享子模块通过# Copied from注释同步make fix-copies可保持两者一致。测试与验证仓库为 FAR 变体提供了完整的测试覆盖见 test_models_transformer_anyflow_far.py可以作为理解与验证模型行为的参考核心模型测试TestAnyFlowFARTransformer3D使用微型配置2 头、head_dim12、2 层、chunk_partition[2, 2]覆盖训练 rollout 路径训练测试test_gradient_checkpointing_is_applied等注释指出 flex attention 的 backward kernel 仅支持 GPUFAR 训练正确性在 H200 上通过 pipeline replay与 NVlabs/AnyFlow 对比 L20端到端验证编译测试TestAnyFlowFARTransformer3DCompile预构建 BlockMask 注入attention_mask保证forward不在编译作用域内调用create_block_masktorch.export因 BlockMask 含 Python 可调用mask_mod而无法作为 pytree 输入故 AOT 导出测试被跳过处理器测试TestAnyFlowCausalAttnProcessor验证默认后端为flex、非 flex 后端抛ValueError、输出数据类属性存在。延伸阅读流水线完整用法T2V / I2V / V2V 示例与检查点列表docs/source/en/api/pipelines/anyflow.md双向变体模型文档docs/source/en/api/models/anyflow_transformer3d.md模型源码FAR 变体transformer_anyflow_far.py流水线源码pipeline_anyflow_far.py任意步调度器源码scheduling_flow_map_euler_discrete.py检查点转换脚本convert_anyflow_to_diffusers.py【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表