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 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.

[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.

[yiyixuxu]

I agree with @dg845 here - let's make the refiner a transformer model.

Whether to implement the refiner transformer as a new model: