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.

[dg845]

Hi @lawrence-cj, I am reviewing the code and will try to have a full review out soon. Thanks for your patience!

[dg845]

Thanks for the PR! I left an initial design review for the pipeline code, am still working on reviewing the modeling code.

[lawrence-cj]

Thanks for the PR! I left an initial design review for the pipeline code, am still working on reviewing the modeling code.

Thanks so much for the review. I'll keep reformatting the code as your requirements.

[lawrence-cj]

Pushed f259221 + 6317ce3 addressing the latest review:

• SanaWMPipeline cleanup: VAE strides read from the vae component; check_inputs + prepare_latents; stage-1 loop inlined into __call__ using self.scheduler and self.progress_bar; generator is now a __call__ arg; first-frame/intrinsics preprocessing moved into a SanaWMImageProcessor(VaeImageProcessor); output post-processed via VideoProcessor; README moved into the docs.
• Refiner → standalone pipeline: SanaWMLTX2Refiner now subclasses DiffusionPipeline (components via register_modules, standard load/save) and drives its Euler steps through a FlowMatchEulerDiscreteScheduler. Kept as an optional nested component of SanaWMPipeline; validated end-to-end on 1×H100.

[dg845]

Thanks for iterating! Here is an initial review of the modeling code. In general I think we should try to simplify it as much as possible since it is quite long and there still appears to be some unused code.

[dg845]

Could we use the standard diffusers gradient checkpoint idiom here? For example, this is what the Flux 2 transformer does:

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

Lines 1350 to 1351 in 59d4e5a

	if torch.is_grad_enabled() and self.gradient_checkpointing:
	hidden_states = self._gradient_checkpointing_func(

[lawrence-cj]

The custom auto_grad_checkpoint was a no-op here (it gated on a grad_checkpointing attribute that is never set, and _supports_gradient_checkpointing=False), so I removed it and call the blocks directly (878836a). Wiring the standard _gradient_checkpointing_func would need the block's dynamic kwargs (camera_conditions / prope_fns / chunk_index) threaded through its positional-only signature — since this is an inference-focused release I left full training-time GC as a follow-up. Happy to add it if you'd like.

[dg845]

In general, can we (re)use diffusers shared utilities when possible (as the modeling code is already very long)? For example, I think we should be able to reuse diffusers.models.normalization.FP32LayerNorm here:

diffusers/src/diffusers/models/normalization.py

Line 84 in 59d4e5a

class FP32LayerNorm(nn.LayerNorm):

We should also be able to reuse some other utilities like RMSNorm, embeddings like get_1d_rotary_pos_embed, etc.

[dg845]

Suggested change

	class WanRotaryPosEmbed(nn.Module):
	# Copied from diffusers.models.transformers.transformer_wan.WanRotaryPosEmbed
	class WanRotaryPosEmbed(nn.Module):

Is WanRotaryPosEmbed here intended to be the same as for Wan? If so we can use # Copied from to sync the implementations.

[lawrence-cj]

Pushed 6b19325 addressing most of these — reused diffusers' shared FP32LayerNorm + get_1d_rotary_pos_embed, removed the dead diagonal_mask / block-timing / qkv_store_buffer paths, dropped the unreached load_state_dict override and unused PatchEmbedMS, and collapsed forward_frame_aware into forward (−503 lines, GPU smoke output byte-identical). Three I left for a follow-up / with a note: the Sana+SanaMSVideoCamCtrl+SanaWMTransformer3DModel merge (worth its own pass since it drops the _inner. state-dict prefix and needs the checkpoint re-converted), keeping the local RMSNorm (diffusers' lacks the scale_factor/norm_dim this uses), and WanRotaryPosEmbed (it's a reimplementation, not a copy — complex-freqs vs cos/sin + extra fhw_dim — so # Copied from won't apply).