Add AnyFlow Any-Step Video Diffusion Pipelines (Bidirectional + FAR Causal)

[Enderfga]

What does this PR do?

This PR adds pipelines for AnyFlow (paper, project page, official code, model weights), an any-step video diffusion framework built on flow maps. A single distilled checkpoint can be evaluated at 1, 2, 4, 8, 16, 32 NFE without retraining, and quality scales monotonically with steps — unlike consistency-based distillation, which often degrades as NFE grows.

Two new pipelines are added, both on top of a new FlowMapEulerDiscreteScheduler and reusing WanLoraLoaderMixin:

• AnyFlowPipeline → AnyFlowTransformer3DModel: bidirectional text-to-video built on the Wan2.1 backbone with an AnyFlowDualTimestepTextImageEmbedding conditioning on the source/target timestep pair (t, r).
• AnyFlowFARPipeline → AnyFlowFARTransformer3DModel: frame-level autoregressive variant (block-sparse causal flex_attention + KV cache + compressed-frame patch embedding) jointly handling T2V / I2V / V2V through one context_sequence argument.

Four checkpoints are released under the nvidia/anyflow collection (Wan2.1-T2V-{1.3B,14B} bidi + FAR-Wan2.1-{1.3B,14B} causal). All four have been validated bit-exact against the official NVlabs/AnyFlow reference on H200: forward L2 = 0.00e+00 for scheduler / transformer / bidi pipeline / FAR pipeline; backward grad delta is 4.88e-04, attributable to bf16 kernel non-determinism only (PR-vs-PR = PR-vs-reference, ratio 1.000); inference latency matches the reference at ±0.0% on both pipelines.

T2V inference example:

import torch
from diffusers import AnyFlowPipeline
from diffusers.utils import export_to_video

pipe = AnyFlowPipeline.from_pretrained(
    "nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers", torch_dtype=torch.bfloat16
).to("cuda")

prompt = "A red panda eating bamboo in a forest, cinematic lighting"
video = pipe(prompt, num_inference_steps=4, num_frames=33).frames[0]
export_to_video(video, "anyflow_t2v.mp4", fps=16)

I2V inference example with the FAR pipeline (single conditioning frame → autoregressive rollout):

import numpy as np
import torch
from diffusers import AnyFlowFARPipeline
from diffusers.utils import export_to_video, load_image

pipe = AnyFlowFARPipeline.from_pretrained(
    "nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers", torch_dtype=torch.bfloat16
).to("cuda")

first_frame = load_image("path/to/first_frame.png").resize((832, 480))
arr = np.asarray(first_frame).astype("float32") / 255.0
context = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).unsqueeze(2).to("cuda")

video = pipe(
    prompt="a cat walks across a sunlit lawn",
    context_sequence={"raw": context},
    num_inference_steps=4,
    num_frames=81,
).frames[0]
export_to_video(video, "anyflow_i2v.mp4", fps=16)

Documentation: EN tutorial at docs/source/en/using-diffusers/anyflow.md, ZH tutorial at docs/source/zh/using-diffusers/anyflow.md, and three API pages (pipelines + two transformer model pages). Tests: 22 fast tests (transformer + scheduler, CPU) plus four pipeline test files, with slow integration tests gated on RUN_SLOW=1 @require_torch_accelerator for the released checkpoints.

anyflow-pr-presentation.mp4

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the documentation guidelines, and here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@yiyixuxu @asomoza

[dg845]

Suggested change

	_attention_backend = None
	_parallel_config = None
	_SUPPORTED_BACKENDS = (None, "flex", "_native_flex")
	_attention_backend = "flex"
	_parallel_config = None
	_SUPPORTED_BACKENDS = ("flex", "_native_flex")

I think setting the default _attention_backend to "flex" rather than None and removing None from the _SUPPORTED_BACKENDS is cleaner, as only Flex Attention backends are compatible with AnyFlowCausalAttnProcessor. (Using None would generally default to the "native" backend, which isn't compatible.)

[Enderfga]

Done in ffdc969 — _attention_backend = "flex" default and _SUPPORTED_BACKENDS = ("flex", "_native_flex"). Caught a real bug while verifying: the previous None default would silently fall through to SDPA on backends that ignore BlockMask (visible on mps locally — now raises loudly instead of returning wrong outputs).

[dg845]

Suggested change

	backend="flex",
	backend=self._attention_backend,

Follow up to #13745 (comment): using self._attention_backend instead of hardcoding flex here allows us to use other supported backends such as _native_flex.

[Enderfga]

Done in ffdc969 — backend=self._attention_backend so _native_flex can be selected explicitly.

[dg845]

Suggested change

	def _build_freqs(self, device: torch.device) -> torch.Tensor:
	# Copied from diffusers.models.transformers.transformer_anyflow.AnyFlowRotaryPosEmbed._build_freqs
	def _build_freqs(self, device: torch.device) -> torch.Tensor:

I think _build_freqs should be the same for both the causal and non-causal RoPE embedding modules, so sync their implementations.

[Enderfga]

Done in ffdc969 — added # Copied from diffusers.models.transformers.transformer_anyflow.AnyFlowRotaryPosEmbed._build_freqs. make fix-copies runs clean.

[dg845]

Suggested change

	def _forward_full_frame(self, num_frames, height, width, device) -> torch.Tensor:
	# Copied from diffusers.models.transformers.transformer_anyflow.AnyFlowRotaryPosEmbed._forward_full_frame
	def _forward_full_frame(self, num_frames, height, width, device) -> torch.Tensor:

Similarly, I think _forward_full_frame should be the same between the causal and non-causal RoPE modules.

[Enderfga]

Done in ffdc969 — same pattern (# Copied from for _forward_full_frame).

[dg845]

Suggested change

	Pre-VAE conditioning frames of shape `(B, C, T, H, W)` in `[0, 1]`. When provided, the pipeline
	VAE-encodes them and keeps the corresponding latent prefix fixed during sampling. Mutually exclusive
	with `video_latents`.
	Pre-VAE conditioning frames of shape `(B, T, C, H, W)` in `[0, 1]`. When provided, the pipeline
	VAE-encodes them and keeps the corresponding latent prefix fixed during sampling. Mutually exclusive
	with `video_latents`.

I think this needs to be updated as VideoProcessor.preprocess_video expects 5D torch.Tensor inputs to have shape [B, T, C, H, W] instead of [B, C, T, H, W].

[Enderfga]

Done in ffdc969 — docstring + EXAMPLE_DOC_STRING flipped to (B, T, C, H, W) everywhere in both pipelines. Good catch — video_processor.preprocess_video's 5D contract is (B, T, C, H, W), so the previous (B, C, T, H, W) doc would have silently broken users who followed it literally.

[dg845]

Can we move this check to check_inputs so that we fail earlier?

[Enderfga]

Done in 7a6643b — both bidi and FAR pipelines now do the video / video_latents mutual-exclusion check inside check_inputs. The FAR-specific (num_frames - 1) % 4 == 0 constraint moved there too, so both fail before any work runs.

[dg845]

Suggested change

	@torch.no_grad()
	@torch.no_grad()
	def encode_video(self, video: torch.Tensor, height: int, width: int) -> torch.Tensor:
	@torch.no_grad()
	def encode_video(self, video: torch.Tensor, height: int, width: int) -> torch.Tensor:

nit: remove extra @torch.no_grad() decorator.

[Enderfga]

Done in ffdc969 — duplicate @torch.no_grad() removed. (Per the bot follow-up, the non-duplicate @torch.no_grad() was also dropped from bidi encode_video and FAR encode_kv_cache since __call__ already wraps the no-grad scope.)

[dg845]

Suggested change

	>>> # Single-frame I2V: wrap the conditioning image as a (1, 3, 1, H, W) tensor in [0, 1].
	>>> first_frame = load_image("path/to/first_frame.png").resize((832, 480))
	>>> arr = np.asarray(first_frame).astype("float32") / 255.0
	>>> context = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).unsqueeze(2).to("cuda")
	>>> # Single-frame I2V: wrap the conditioning image as a (1, 1, 3, H, W) tensor in [0, 1].
	>>> first_frame = load_image("path/to/first_frame.png").resize((832, 480))
	>>> arr = np.asarray(first_frame).astype("float32") / 255.0
	>>> context = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).unsqueeze(1).to("cuda")

For the same reason as #13745 (comment), we should input a [B, T, C, H, W] rather than a [B, C, T, H, W] video tensor.

[Enderfga]

Done in ffdc969 — unsqueeze(0).unsqueeze(1) to produce the (1, 1, 3, H, W) shape per the suggestion.

[dg845]

Suggested change

	video (`torch.Tensor`, *optional*):
	Pre-VAE conditioning frames of shape `(B, C, T, H, W)` in `[0, 1]` (`T = 4n + 1`). When provided, the
	pipeline VAE-encodes them and keeps the corresponding latent prefix fixed during sampling. Mutually
	exclusive with `video_latents`.
	video (`torch.Tensor`, *optional*):
	Pre-VAE conditioning frames of shape `(B, T, C, H, W)` in `[0, 1]` (`T = 4n + 1`). When provided, the
	pipeline VAE-encodes them and keeps the corresponding latent prefix fixed during sampling. Mutually
	exclusive with `video_latents`.

Analogous suggestion to #13745 (comment) for the FAR causal pipeline.

[Enderfga]

Done in ffdc969 — same fix as the bidi side.

[dg845]

Instead of using a show_progress argument here, I think we should use nested progress bars like LLaDA-2 does. We can define an outer progress bar:

diffusers/src/diffusers/pipelines/llada2/pipeline_llada2.py

Lines 426 to 428 in f502538

	outer_progress_bar_config = getattr(self, "_progress_bar_config", {}).copy()
	block_progress_bar_config = {**outer_progress_bar_config, "position": 0, "desc": "Blocks"}
	for num_block in tqdm(range(prefill_blocks, num_blocks), **block_progress_bar_config):

and an inner progress bar:

diffusers/src/diffusers/pipelines/llada2/pipeline_llada2.py

Lines 444 to 450 in f502538

	inner_progress_bar_config = {
	**outer_progress_bar_config,
	"position": 1,
	"leave": False,
	"desc": f"Block {num_block} Inference Steps",
	}
	progress_bar = tqdm(total=num_inference_steps, **inner_progress_bar_config)

using the pipeline's _progress_bar_config and appropriate arguments to make sure that the inner progress bars don't stack up. This should respect any configuration set through DiffusionPipeline.set_progress_bar_config better (for example, using pipe.set_progress_bar_config(disable=None) to disable the progress bars).

[Enderfga]

Done in 7a6643b — show_progress argument removed; replaced with nested tqdm bars in the LLaDA-2 pattern (outer Chunks at position=0, inner Inference Steps per chunk at position=1, leave=False). Both pick up DiffusionPipeline._progress_bar_config, so set_progress_bar_config(disable=None) etc. now work as expected.

[dg845]

nit: I think getting the underlying t_sigma and r_sigma corresponding to timestep and r_timestep via something like the logic in _resolve_next_timestep or an internal step_idx like FlowMatchEulerDiscreteScheduler uses:

diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

Lines 501 to 503 in f502538

	sigma_idx = self.step_index
	sigma = self.sigmas[sigma_idx]
	sigma_next = self.sigmas[sigma_idx + 1]

would be slightly better here to cover cases where the timesteps and sigmas aren't related through scaling by self.config.num_train_timesteps.

[Enderfga]

Done across 7a6643b + 84605d5 + 89128cf:

• 7a6643b: introduced index_for_timestep() and rewrote step() to resolve both t_sigma and r_sigma via self.sigmas[idx] lookups. For the shipped linspace + shift schedule this is bit-identical to the previous t / num_train_timesteps formulation (max abs diff = 0.0 on an 8-step replay), but it stays correct under any future schedule whose timestep / sigma mapping isn't strictly linear.
• 84605d5: added the full FlowMatchEulerDiscreteScheduler-style state machine — _step_index, _begin_index, step_index / begin_index properties, set_begin_index, _init_step_index. step() lazily initializes and advances the counter on every call so downstream callbacks / composable schedulers observe it. Sigma resolution stays a pure function of the passed-in (timestep, r_timestep) so step() is idempotent (calling it twice with the same args returns identical prev_sample).
• 89128cf: audit fix — the earlier _init_step_index raised when the first timestep was off-schedule, which contradicted step()'s documented any-step support. _init_step_index now falls back to 0 for off-schedule starts (still a valid observable counter); _resolve_next_timestep was removed since its callers were all inlined.

Bit-exact replay on H200 (random-init bidi + FAR forward, fp32, comparing d0181ea baseline to 84605d5): state_dict missing=0 / unexpected=0, L2 = 0.0e+00, max|Δ| = 0.0e+00.

[dg845]

Suggested change

	@unittest.skip("AnyFlow has no optional components.")
	def test_save_load_optional_components(self):
	pass

I think we don't need to skip test_save_load_optional_components here as this test actually passes.

[Enderfga]

Done in ffdc969 — skip removed; test passes on CPU.

[dg845]

Suggested change

	@unittest.skip("AnyFlow has no optional components.")
	def test_save_load_optional_components(self):
	pass

Analogous comment as #13745 (comment).

[Enderfga]

Done in ffdc969 — same as the bidi side.

[dg845]

Hi @claude, can you do a review to see if there is anything I missed?

[github-actions]

Claude finished @dg845's task in 5m 22s —— View job

---

Review Complete

Hi @dg845 — here's my review of the remaining items. Your review was thorough; I'll focus on things I found that you didn't already flag, and then include a summary of your open items for tracking.

---

New findings (beyond your existing review comments)

1. Unused attention_kwargs parameter in both transformers

Both AnyFlowTransformer3DModel.forward() (transformer_anyflow.py:633) and AnyFlowFARTransformer3DModel.forward() (transformer_anyflow_far.py:923) declare attention_kwargs: Optional[Dict[str, Any]] = None but never use or forward it. In the FAR transformer, it's propagated through the common dict to _forward_cache/_forward_inference/_forward_train, where it's accepted as a parameter but never referenced.

Per models.md: "Don't declare attention_mask (or encoder_hidden_states_mask) in the forward signature if you ignore it." The same principle applies to attention_kwargs. Either use it (pass to the attention processors) or remove it.

2. Misleading far_cfg variable name in bidirectional transformer

The bidirectional transformer (transformer_anyflow.py:667) constructs a far_cfg dict, but "FAR" (Frame-level AutoRegressive) is the causal variant specifically. This name is confusing in the bidi context. Consider renaming to something like layout_cfg or frame_cfg in the bidi transformer. The causal transformer's far_cfg is fine since it actually is FAR-specific.

3. FlowMapEulerDiscreteScheduler is missing _step_index / _begin_index infrastructure

Beyond the _resolve_next_timestep improvement you flagged, the FlowMatchEulerDiscreteScheduler (which this scheduler is modeled after) implements a full step-tracking API: _step_index, _begin_index, step_index property, set_begin_index(), and _init_step_index(). The new scheduler has none of this. While not strictly needed for the current pipelines (since they resolve r from sigmas[i + 1] directly), downstream code may expect schedulers to expose step_index — e.g., IP-Adapter scheduled timesteps, certain callbacks, or composable schedulers.

4. Redundant @torch.no_grad() on encode_video (bidi pipeline) and encode_kv_cache (FAR pipeline)

You flagged the duplicate @torch.no_grad() on the FAR encode_video (line 378). I'd extend this: the non-duplicate @torch.no_grad() on the bidi pipeline's encode_video (pipeline_anyflow.py:354) and the FAR's encode_kv_cache (pipeline_anyflow_far.py:396) are also redundant per pipelines.md: "Redundant inside helpers that __call__ already covers. The decorator puts every descendant in no-grad, so an inner with torch.no_grad(): is noise — and worse, it forecloses callers who want to invoke pipe.encode_prompt(...) with grads enabled."

5. Dead temb.ndim == 2 branch in bidirectional transformer output norm

In transformer_anyflow.py:702-707, the else branch (for temb.ndim != 3) is unreachable. temb comes from condition_embedder.forward() → forward_timestep() which always returns a 3D tensor (B, total_tokens, dim) because of the unflatten + repeat_interleave at line 279. The else branch handling a 2D temb is dead code.

6. Public methods forward_far_patchify and forward_far_patchify_inference should be private

In transformer_anyflow_far.py, forward_far_patchify (line 1000) and forward_far_patchify_inference (line 1024) are public methods only called internally by _forward_train/_forward_cache and _forward_inference respectively. Per models.md, module calls should be visible in forward directly, and methods that are only called internally should be private. These should be _forward_far_patchify / _forward_far_patchify_inference.

7. Minor: comment numbering skips step 4 in both pipelines

Both pipeline_anyflow.py and pipeline_anyflow_far.py jump from comment # 3. Encode input prompt to # 5. Prepare latent variables — step 4 is missing. (Very minor readability nit.)

8. Minor: _build_causal_mask comment numbering

In transformer_anyflow_far.py:1076-1082, the mask_mod function comments skip from # 1) to # 3) (no # 2)).

---

Suggestions / additional info: dead code trace

I traced the call paths from both pipelines into the transformers:

Bidirectional pipeline → AnyFlowTransformer3DModel:

• The pipeline never passes encoder_hidden_states_image (always None). The image embedding path at line 689 is therefore unreachable via the current pipeline — this is fine as advisory since I2V may be added later or users may call the model directly.
• attention_kwargs is never consumed (see finding Add glide modeling files #1 above).
• The temb.ndim == 2 branch is unreachable (see finding Add UNet for Latent Diffusion #5 above).

FAR pipeline → AnyFlowFARTransformer3DModel:

• The pipeline only uses _forward_cache (via encode_kv_cache) and _forward_inference (via the denoising loop). _forward_train is never called by the pipeline. This is expected — it's for training use.
• clean_hidden_states and clean_timestep are only used in _forward_train and _forward_cache. The pipeline only passes them as None from the denoising loop (and through encode_kv_cache where they default to None), so the clean_hidden_states is not None branches in _forward_cache are unreachable during pipeline inference. Again expected for a training code path.
• attention_kwargs is propagated but never consumed (see finding Add glide modeling files #1).

---

Summary of dg845's open review items (for tracking)

These are your existing review comments that appear to still be open based on the latest code:

#	File	Issue	Status
1	transformer_anyflow_far.py:87	_attention_backend should default to "flex", remove None from _SUPPORTED_BACKENDS	Open
2	transformer_anyflow_far.py:197	Use self._attention_backend instead of hardcoded "flex"	Open
3	transformer_anyflow_far.py:644	_build_freqs needs # Copied from annotation	Open
4	transformer_anyflow_far.py:720	_forward_full_frame needs # Copied from annotation	Open
5	pipeline_anyflow.py:412	Video shape docstring: (B, T, C, H, W) not (B, C, T, H, W)	Open
6	pipeline_anyflow.py:540	Move video/video_latents mutual-exclusion check to check_inputs	Open
7	pipeline_anyflow_far.py:379	Remove duplicate @torch.no_grad()	Open
8	pipeline_anyflow_far.py:58	Example: unsqueeze(1) not unsqueeze(2) for (B, T, C, H, W)	Open
9	pipeline_anyflow_far.py:475	FAR docstring video shape correction	Open
10	pipeline_anyflow_far.py:692	Use nested progress bars instead of show_progress/tqdm	Open
11	scheduling_flow_map_euler_discrete.py:189	Use _step_index or _resolve_next_timestep for sigma lookup	Open
12	test_anyflow.py:143	Remove unnecessary test_save_load_optional_components skip	Open
13	test_anyflow_far.py:154	Remove unnecessary test_save_load_optional_components skip	Open

---

Overall the PR is well-structured — the transformer split is clean, the scheduler API is sensible, and the # Copied from annotations are correctly used. The main remaining work is the items above and dg845's open comments.

[Enderfga]

Thanks for the second pass @dg845 — really appreciate the careful read. Working through all 13 now. A quick map of how I'm going to handle them:

• One-click suggestions (FAR processor backend defaults + # Copied from syncs, the video shape fix in both pipelines + EXAMPLE_DOC_STRING, the duplicate @torch.no_grad(), lifting the two test_save_load_optional_components skips) — applying as-is.
• Architectural moves — also taking all three: hoist the video/video_latents mutual-exclusion check into check_inputs, switch the FAR pipeline to nested progress bars in the LLaDA-2 pattern (outer chunk / inner timestep, respecting _progress_bar_config), and rework _resolve_next_timestep to index via step_idx like FlowMatchEulerDiscreteScheduler does (no timestep/num_train reverse-derivation).

I'll re-run the bit-exact replay against NVlabs/AnyFlow on H200 after the scheduler rework lands, since that one touches the numeric path. Pushing in a few hours; will reply per-thread as each is done.

[Enderfga]

Hi @dg845 @sayakpaul — second-pass review fully addressed. Per-thread replies are inline; this is the high-level summary.

dg845's 13 review threads — all applied

Area	Outcome
FAR processor backend (1, 2)	Default _attention_backend = "flex", _SUPPORTED_BACKENDS = ("flex", "_native_flex"), backend=self._attention_backend
FAR rope # Copied from (3, 4)	_build_freqs + _forward_full_frame synced to bidi; make fix-copies clean
video shape docs (5, 8, 9)	Pipelines + EXAMPLE_DOC_STRING + user guides flipped to (B, T, C, H, W) to match video_processor.preprocess_video's 5D contract
Pre-flight validation (6)	video / video_latents mutual-exclusion + (num_frames - 1) % 4 == 0 moved into check_inputs
FAR pipeline cleanup (7, 10)	Duplicate @torch.no_grad() removed; show_progress arg dropped in favour of nested tqdm bars (LLaDA-2 pattern, respects set_progress_bar_config)
Scheduler (11)	New index_for_timestep; step() resolves both t_sigma / r_sigma via sigma lookup; full _step_index / _begin_index state machine; off-schedule any-step path repaired (89128cf audit fix)
Test skip cleanup (12, 13)	test_save_load_optional_components un-skipped on both pipelines

Claude bot follow-up — also applied

1. attention_kwargs: both transformer forward() methods now decorate with @apply_lora_scale("attention_kwargs") (matching WanTransformer3DModel); pipelines plumb the kwarg through to transformer(...) and encode_kv_cache(...); unused parameter is removed from the inner FAR _forward_* signatures.
2. far_cfg → layout_cfg in the bidirectional transformer (the bidi path isn't FAR; FAR file keeps far_cfg).
3. Full _step_index / _begin_index state machine on FlowMapEulerDiscreteScheduler — same as item 11 above.
4. Redundant @torch.no_grad() removed from bidi encode_video + FAR encode_kv_cache (already inside __call__'s no-grad scope).
5. Dead temb.ndim == 2 branch removed from the bidi transformer's output norm.
6. forward_far_patchify[_inference] → _forward_far_patchify[_inference] (only called internally).
7. Pipeline step comments renumbered so # 4. is no longer skipped.
8. _build_causal_mask mask-mod comments renumbered (1, 2, 3, …).

Bit-exact validation (H200, fp32, random-init bidi + FAR)

Replay comparing d0181ea (pre-second-pass) to current HEAD:

• state_dict key match: missing = 0 / unexpected = 0 for both models
• bidi forward: L2 = 0.0e+00, max|Δ| = 0.0e+00
• FAR forward: L2 = 0.0e+00, max|Δ| = 0.0e+00

So the entire second-pass refactor is provably numerically equivalent on the shipped linspace + shift schedule. The earlier L2 = 0 guarantee against NVlabs/AnyFlow is preserved.

Commits this round

• ffdc969 — A (1-click suggestions, 10 of dg845's)
• 7a6643b — B (refactors: check_inputs hoist, nested progress bars, scheduler sigma-lookup)
• 84605d5 — Claude bot's 8 findings
• 89128cf — audit fix: off-schedule any-step regression in _init_step_index
• be2aefe — user-guide doc updates for the shape + kwarg fixes

Ready when you are. Happy to iterate further if anything's still off.

[dg845]

@bot /style

[github-actions]

Style bot fixed some files and pushed the changes.