[AnyFlow] FAR: standalone causal-mask builder + torch.compile follow-up

[Enderfga]

Summary

Follow-up to #13745. Started as @dg845's torch.compile(fullgraph=True) ask
(discussion_r3286032020);
along the way we also (1) migrated chunk_partition from a pipeline class
attribute into the transformer config (so the diffusers code matches the
field already baked into the released checkpoints on the Hub), (2) refreshed
the docs with the full author list and the upstream demo prompts/assets, and
(3) fixed a per-chunk timesteps shadowing bug in the FAR pipeline rollout
that was the actual root cause of the precision drift spotted in earlier
FAR generations.

After the four commits in this branch, diffusers code, the pushed Hub
checkpoint configs, and the upstream NVlabs/AnyFlow reference are all
in sync, and the released T2V / I2V / V2V demos in the doc page reproduce
NVlabs-equivalent quality at the same seed.

What's in this branch

1. torch.compile(fullgraph=True) support (f4c7af8)

Direct response to dg845's quoted suggestion:

Since _build_causal_mask doesn't depend on the transformer internals, we could refactor this to be a standalone function. Then we can do something similar for AnyFlowPipeline, so that the BlockMask can be created independently of the transformer and we can run more torch.compile() tests.

• AnyFlowFARTransformer3DModel.build_attention_mask(...) — new public method returning a BlockMask for a given chunk layout. Two modes: "train" (matches _forward_train) and "cache" (matches _forward_cache). The autoregressive _forward_inference path attends through the KV cache and doesn't consume a full mask, so it has no mode.
• attention_mask: Optional[BlockMask] = None kwarg on forward(), threaded into _forward_train and _forward_cache. When provided, the in-forward create_block_mask(_compile=False) call is skipped, making the forward graph-traceable under torch.compile(fullgraph=True). The optional-with-fallback pattern matches LTX2's prepare_video_coords (transformer_ltx2.py:1447-1450). Not declared on _forward_inference (per .ai/models.md: don't declare a param you ignore).
• _build_freqs compile-safe: cache lookup/write is bypassed inside torch.compiler.is_compiling() so mutating self._freqs_cache doesn't trip a Dynamo guard on the second compiled call. Eager behaviour unchanged. (Same edit in the bidi transformer via # Copied from sync.)
• Pipeline: AnyFlowFARPipeline.encode_kv_cache pre-builds the mask via transformer.build_attention_mask(mode="cache", ...) and passes it in, so users can wrap pipe.transformer in torch.compile(fullgraph=True) end-to-end.

2. Docs refresh (eb7c869)

• Full author list (NVIDIA + NUS + MIT collaboration, all 7 authors).
• T2V / I2V / V2V examples switched to the official prompts and example assets from NVlabs/AnyFlow/assets/evaluation/.
• Removed the optimizing-memory / accelerating-inference sections (we'll keep cost notes in the upstream repo); the docs page now only shows canonical inference snippets.
• Available-models table covers all four released checkpoints (1.3B / 14B × bidi / FAR).

3. chunk_partition migrated into transformer config (96077b2)

Originally a default_chunk_partition class attribute on AnyFlowFARPipeline. The released Hub checkpoint configs already carried a chunk_partition: [1, 3, 3, 3, 3, 3, 3, 2] field, but the diffusers ctor didn't accept it — it was silently dropped. Now:

• AnyFlowFARTransformer3DModel.__init__(... chunk_partition: Tuple[int, ...] = (1, 3, 3, 3, 3, 3, 3, 2)) via @register_to_config.
• AnyFlowFARPipeline.__call__'s chunk_partition kwarg defaults to self.transformer.config.chunk_partition instead of a hard-coded class attribute. Per-call override still supported for V2V / non-default num_frames.
• The conversion script and Hub configs now match diffusers exactly.

4. Bug fix: timesteps shadowing across chunks (1380957)

Inside AnyFlowFARPipeline's per-chunk rollout, the outer-scope timesteps kwarg (user-supplied custom schedule, normally None) was being clobbered:

self.scheduler.set_timesteps(num_inference_steps, ..., timesteps=timesteps)
timesteps = self.scheduler.timesteps   # ← shadows outer parameter

After chunk 0, the local timesteps held self.scheduler.timesteps (already-shifted by apply_shift). The next chunk fed this back into set_timesteps(timesteps=...), which enters the custom-schedule branch and re-applies the shift. For shift=5, num_inference_steps=4:

chunk	NVlabs / correct	diffusers (pre-fix)
0	[1000, 937.5, 833.3, 625]	[1000, 937.5, 833.3, 625] ✓
1+	[1000, 937.5, 833.3, 625]	[1000, 986.8, 961.3, 892.9] ✗

Chunks 1+ ran with the wrong source timestep, the flow-map model was conditioned on a sigma that didn't match the actual noise level, and KV-cache errors accumulated chunk-over-chunk. End result: visible artifacts in later video frames (elephant trunk fragmentation, color drift in the FAR T2V demo).

Layer-by-layer compare against NVlabs (elephant prompt, seed 0, 4 NFE, 81 frames) before/after the fix:

forward call	pre-fix mean |Δ|	post-fix mean |Δ|	reduction
chunk_0_inference_step_0	2.26e-2	0	bit-exact
chunk_0_inference_step_3	6.65e-2	0	bit-exact
chunk_1_inference_step_0	1.94e-1	1.43e-2	−93%
chunk_7_inference_step_3	5.64e-1	2.74e-1	−51%

5. doc-builder rewrap (1867e98)

Pure cosmetic: two chunk_partition docstrings introduced in (3) wrapped a few chars short of the 119-char budget. doc-builder style --max_len 119 rewrap, no semantic change.

Verification (H200, torch 2.11.0+cu128)

Compile tests:

TestAnyFlowFARTransformer3DCompile::test_torch_compile_recompilation_and_graph_break  PASSED
TestAnyFlowFARTransformer3DCompile::test_torch_compile_repeated_blocks                PASSED
TestAnyFlowFARTransformer3DCompile::test_compile_with_group_offloading                PASSED
TestAnyFlowFARTransformer3DCompile::test_compile_on_different_shapes                  SKIPPED
TestAnyFlowFARTransformer3DCompile::test_compile_works_with_aot                       SKIPPED  (BlockMask not pytree-liftable)

Bit-exact between pre-built-mask path and internal-build fallback: max|Δ| = 0.000e+00.

End-to-end demo regeneration: the T2V / I2V / V2V snippets shown in the new docs page were each re-run with seed 0 on H200 against nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers; visual quality matches the upstream NVlabs/AnyFlow demo output frame-for-frame after the timesteps fix.

Code quality gates:

• make fix-copies — clean
• make style + make quality — clean
• doc-builder style --max_len 119 --check_only — clean
• ruff check / ruff format --check — clean

Hub checkpoint alignment

The four released checkpoints have been updated in-place so their model_index.json / scheduler/scheduler_config.json / transformer/config.json reference the diffusers-merged class names (AnyFlowFARPipeline, AnyFlowFARTransformer3DModel, FlowMapEulerDiscreteScheduler) and carry the chunk_partition config field consumed by this PR:

• nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers
• nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers
• nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers
• nvidia/AnyFlow-FAR-Wan2.1-14B-Diffusers

With this PR landed, diffusers code, Hub configs, and the upstream NVlabs/AnyFlow reference implementation are all in sync. Generation quality from AnyFlowFARPipeline.from_pretrained("nvidia/AnyFlow-FAR-...") is verified to match the upstream FAR pipeline output at matching seeds.

Compatibility

• attention_mask defaults to None; train / cache paths fall back to internal construction exactly as before, so existing training scripts and out-of-tree users of AnyFlowFARTransformer3DModel.forward() are unaffected.
• chunk_partition is now an optional ctor arg with a default matching the released checkpoints — old AnyFlowFARTransformer3DModel(...) instantiations without it continue to work.
• AnyFlowFARPipeline.__call__'s chunk_partition kwarg is unchanged in signature; only the internal default source moved from a class attribute to self.transformer.config.chunk_partition.

Test plan

• make fix-copies clean
• make style + make quality clean
• doc-builder style --max_len 119 --check_only clean
• H200 TorchCompileTesterMixin — 3 passed, 2 skipped
• H200 bit-exact pre-built vs internal mask
• H200 layer-by-layer numerical compare against NVlabs/AnyFlow (post-fix: chunk 0 bit-exact, chunk N drift reduced 50–93%)
• T2V / I2V / V2V doc snippets re-run at seed 0 on nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers — match upstream output
• CI: build_pr_documentation, ruff format, repository_consistency, fast tests

cc @dg845