Add SANA-WM camera-controlled image-to-video pipeline

[lawrence-cj]

What does this PR do?

Hi @sayakpaul @dg845 , Long time no see. Hoping your are doing great. ♥️

Adds SANA-WM, the camera-controlled image-to-video world model from NVIDIA + MIT HAN Lab, as a first-class diffusers pipeline and transformer. Given a first-frame image, a text prompt, and a camera trajectory (explicit c2w poses or a WASD/IJKL action-DSL string), the pipeline generates a video whose motion follows the requested camera path. Trained natively for minute-scale generation at 704×1280.

The pipeline runs in two stages:

1. Stage 1 — SanaWMTransformer3DModel. A 1.6B-parameter bidirectional DiT with GDN-Triton linear attention and a UCPE camera-control branch; samples with an LTX-style flow-matching Euler scheduler at per-token timesteps. The first latent frame is the conditioning anchor.
2. Stage 2 — SanaWMLTX2Refiner (optional). A chunk-causal AR refiner that wraps diffusers' LTX2VideoTransformer3DModel + LTX2TextConnectors + Gemma-3 text encoder. Processes 3 latent frames at a time with a sliding window of [source_sink + recent_history + active_block] K/V, so per-block compute is bounded and total refinement cost is linear in video length.

Both stages decode through AutoencoderKLLTX2Video.

Layout

src/diffusers/
├── models/transformers/
│   ├── transformer_sana_wm.py          # SanaWMTransformer3DModel + blocks + helpers
│   └── transformer_sana_wm_kernels.py  # fused Triton kernels + camera math
└── pipelines/sana_wm/
    ├── __init__.py
    ├── pipeline_sana_wm.py             # SanaWMPipeline
    ├── pipeline_output.py              # SanaWMPipelineOutput
    ├── refiner.py                      # SanaWMLTX2Refiner + RefinerChunkRunner
    └── cam_utils.py                    # action DSL, intrinsics, resize+crop, Plücker/raymap

scripts/sana_wm/convert_sana_wm_to_diffusers.py
docs/source/en/api/{pipelines/sana_wm.md, models/sana_wm_transformer3d.md}

Usage

import torch
from PIL import Image
from diffusers import SanaWMPipeline
from diffusers.utils import export_to_video

pipe = SanaWMPipeline.from_pretrained(
    "Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
    torch_dtype=torch.bfloat16,
)
pipe.vae.to(torch.float32)
pipe.enable_model_cpu_offload()

out = pipe(
    image=Image.open("input.png").convert("RGB"),
    prompt="A car driving across a vast desert plain at golden hour.",
    action="w-80,jw-40,w-40",                    # WASD-style action DSL
    intrinsics=[800.0, 800.0, 845.0, 464.0],      # fx, fy, cx, cy in original-image pixels
    num_frames=161,
    num_inference_steps=60,
)
export_to_video(list(out.frames), "sana_wm.mp4", fps=16)

Demo

5-second sample (30 stage-1 steps + 3-step distilled AR refiner, official asset/sana_wm/demo_0 inputs, 704×1280 @ 16 fps) :

sana_wm_5s.mp4

Smoke tests

End-to-end on 1× H100 80GB with `enable_model_cpu_offload` and the official `asset/sana_wm/demo_0.{png,txt,_pose.npy,_intrinsics.npy}`:

Duration	Frames	Stage-1 (30 steps)	Refiner (AR, 3 blocks)	Output
5s	80	1:11	5:24 / step	525 KB
10s	160	1:11	28:55 (7 blocks)	1.4 MB
20s	320	1:57	≈ 4 min / block (14)	3.2 MB
50s	800	5:33	30:46 (34 blocks)	6.3 MB

Checkpoint conversion

scripts/sana_wm/convert_sana_wm_to_diffusers.py --src Efficient-Large-Model/SANA-WM_bidirectional --dst /local/path converts the public release into a `from_pretrained`-loadable directory (VAE, Gemma-2 tokenizer + text_encoder, transformer, scheduler, refiner subfolders, top-level `model_index.json`).

Related

Paper: https://arxiv.org/abs/2605.15178

[dg845]

As a preliminary comment, would it be possible to use PyTorch ops instead of custom Triton kernels (or add pure PyTorch fallback paths) for now? We will work on supporting the custom kernels through kernels. CC @sayakpaul

[lawrence-cj]

As a preliminary comment, would it be possible to use PyTorch ops instead of custom Triton kernels (or add pure PyTorch fallback paths) for now? We will work on supporting the custom kernels through kernels. CC @sayakpaul

Yes, love to do that.

[lawrence-cj]

Done in c0712d3f8 — Triton is now optional, with an automatic pure-PyTorch fallback at dispatch time. Mapping when Triton isn't usable:

Requested	Falls back to
BidirectionalGDNTriton	BidirectionalGDN
BidirectionalGDNUCPESinglePathLiteLATriton	BidirectionalGDNUCPESinglePathLiteLA
BidirectionalGDNUCPESinglePathLiteLABothTriton	BidirectionalGDNUCPESinglePathLiteLA

Triton remains the default on CUDA + Triton ≥ 3. CPU tests added under tests/pipelines/sana_wm/.

[lawrence-cj]

@dg845 @yiyixuxu Gentle ping here.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

@lawrence-cj
for these 3 questions

Refiner composition — keep it nested, flatten Kandinsky-style, or split into two top-level pipelines?

have not reviewed in this PR

Attention — GDN linear attention doesn't fit the processor / dispatch_attention_fn pattern. Convert cross-attn + the softmax blocks only, and document GDN as an exception?

let's try to remove the timm dependency first in this round, don't worry about dispatch_attention_fn for now

Triton kernels — pure-PyTorch only in this PR (kernels as a follow-up), or publish to kernels-community and dispatch?

either one

[lawrence-cj]

Thanks @yiyixuxu — timm is gone in 64f24508c. All three imports turned out to be trivially replaceable:

• Mlp — both local subclasses already overrode forward entirely and only inherited fc1/act/drop1/fc2/drop2, so it's now self-contained (and the redundant class Mlp(Mlp) wrapper is deleted).
• Attention — GDN inherited only qkv and proj from it (it defines its own q_norm/k_norm), so those two layers are declared inline and it subclasses nn.Module.
• DropPath — drop_path was hardcoded 0.0, so it was always nn.Identity(); removed with the rest of the plumbing.

No dummy_timm_objects.py needed since the dependency is removed rather than made optional, and setup.py is untouched (other models still use it). State dict is unchanged — 871/871 keys match the released checkpoint — and the end-to-end GPU smoke gives byte-identical output.

Leaving dispatch_attention_fn and the Triton question for later rounds as you suggested, and the refiner composition until you've had a chance to look.

One small thing whenever you get a moment: the workflow runs are still sitting on action_required for this fork, so nothing beyond the labeler has actually executed on the recent commits.

[yiyixuxu]

@lawrence-cj thanks, can you address all the inline review comments as well?

[lawrence-cj]

Thanks for the thorough pass @yiyixuxu — worked through all of them. transformer_sana_wm.py is 6555 → 4123 lines, and transformer_sana_wm_kernels.py (3234 lines) is gone.

Applied: Triton removed entirely (the *Triton config strings now map onto the pure-PyTorch classes, so existing checkpoints still load — I'll publish the kernels to kernels-community in a follow-up); timm dropped; get_activation(); GLUMBConv/DWMlp and the local registries deleted; all closure-returning helpers, @torch.compile decorators, partial, and the training-only inits removed.

Not applied — 4 hard blockers, not preferences:

	Why
RMSNorm	Shared class has no scale_factor; config sets y_norm_scale_factor=0.01. Also changes bf16 numerics. Can upstream scale_factor instead.
Timesteps/TimestepEmbedding	Renames state-dict keys (t_embedder.mlp.0/2 vs linear_1/linear_2).
uniform-only chunking	Config is first_chunk_plus_one and that path is live.
drop y_embedding	Unused in forward, but it's one of the 871 checkpoint keys.

Each is 1 line to change if you'd rather take the checkpoint re-export — just say which.

Verified: state dict unchanged at 871/871 keys; GPU smoke on the public checkpoint gives correct video.

Workflow runs are still on action_required for this fork whenever you get a chance.

[dg845]

@lawrence-cj thanks for your patience! I have reviewed the refactored code. It would be helpful if you could run another self-review as described in #13881 (comment) after addressing the comments, as this will help speed up the review process.

[dg845]

Since MultiHeadCrossAttention uses the torch native F.scaled_dot_product_attention, I think we should follow the standard diffusers attention pattern and use dispatch_attention_fn here instead, as this will allow us to support different attention backends such as Flash Attention.

[dg845]

Can you migrate MultiHeadCrossAttention to use the standard diffusers AttentionModuleMixin + attention processor design? For example, here is how Flux 2 implements the pattern:

diffusers/src/diffusers/models/transformers/transformer_flux2.py

Lines 497 to 499 in bda3386

	class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
	_default_processor_cls = Flux2AttnProcessor
	_available_processors = [Flux2AttnProcessor, Flux2KVAttnProcessor]

diffusers/src/diffusers/models/transformers/transformer_flux2.py

Lines 327 to 329 in bda3386

	class Flux2AttnProcessor:
	_attention_backend = None
	_parallel_config = None

[dg845]

If I understand correctly, _forward_video_block effectively implements a different forward method for LTX2VideoTransformerBlock. Given this, I think it would be better to implement the Sana-WM refiner transformer as a new model. We can use the # Copied from mechanism to sync the parts of the implementation that are shared with LTX2VideoTransformer3DModel, and add a new transformer block module whose forward follows the current _forward_video_block logic. CC @yiyixuxu

[lawrence-cj]

Thanks @dg845 — pushed in 99d51ebd6, 82c902cf7, d44abcb7a. 15 of 21 threads resolved.

Highlights: forward no longer stashes the patch grid on self (a real hazard under concurrency/torch.compile, not just style); asserts → ValueError; the pipeline no longer mutates the components it's handed (VAE settings are documented instead); WanRotaryPosEmbed → SanaWMRotaryPosEmbed, since it genuinely can't be # Copied from Wan; model tests added and pipeline tests migrated to pytest — 47 passed, 27 skipped.

Two flagged rather than done, both with a prerequisite:

• Gradient checkpointing needs the block **kwargs bag (12 keys, 41 pass-through sites) flattened first, since diffusers' default checkpoint function is positional-only. Worth its own PR.
• Manual offloading exists only because the refiner is nested; un-nesting it as you suggest makes the block disappear, so I'd rather do both together.

One correction: class SanaWMCamUtilsTests: would have silently disabled those tests — no pytest config, so python_classes is the default Test* and the unittest.TestCase base is the only reason they're collected today. Used TestSanaWMCamUtils.

Still yours + @yiyixuxu's call, since they reshape the public API: the dispatch_attention_fn migration, un-nesting the pipelines, and making the refiner transformer its own model. Happy to take all three.

Verified throughout: 871/871 state-dict keys, GPU smoke unchanged at frame mean 0.5560.

[dg845]

Thanks for the changes! I have left some follow up comments; can you also run make style and make quality to fix any code style errors?

@yiyixuxu could you take a look at the following comments?

• Whether to implement the refiner transformer as a new model: #13881 (comment)
  • I think #13881 (comment) (about refiner KV caching) is also relevant here
• Whether the refiner pipeline should be nested or separate: #13881 (comment)

[dg845]

As with SanaWMPipeline, we should use the standard model offloading utilities rather than manually implementing it here.

[dg845]

I don't think we should be modifying the LTX2VideoTransformer3DModel transformer here by setting attributes such as _tf_kv_prefix on its submodules. For KV caching, we can follow the Flux2KVCache pattern instead:

diffusers/src/diffusers/models/transformers/transformer_flux2.py

Lines 88 to 92 in 2ff5e58

	class Flux2KVCache:
	"""Container for all layers' reference-token KV caches.
	Holds separate cache lists for double-stream and single-stream transformer blocks.
	"""

CC @yiyixuxu