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]

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]

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]

It isn't the same implementation, unfortunately, so # Copied from would fail make fix-copies. Concretely, versus transformer_wan.WanRotaryPosEmbed:

• it calls get_1d_rotary_pos_embed(..., use_real=False, repeat_interleave_real=False) and stores a single complex freqs, whereas Wan uses use_real=True, repeat_interleave_real=True and registers two real freqs_cos / freqs_sin buffers;
• forward(fhw, device) -> Tensor returns that one complex tensor shaped (1, 1, ppf*pph*ppw, -1), while Wan's forward(hidden_states) -> (cos, sin) derives the patch counts from the input and returns a tuple shaped (1, N, 1, -1);
• it takes an extra fhw_dim argument with different dim-splitting logic.

Reconciling them would mean rewriting this class (and its CausalWanRotaryPosEmbed subclass plus the apply_rotary_emb call sites) onto Wan's cos/sin convention — a numerics-sensitive change I'd rather not fold into this PR. Happy to rename it to something less confusing (e.g. SanaWMRotaryPosEmbed) so it doesn't read as a copy of Wan's.

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

[lawrence-cj]

@dg845 gentle ping

[dg845]

@askserge can you do a review of the current Sana-WM pipeline implementation in this PR?

[sergereview]

🤗 Serge says:

SANA-WM Pipeline Review

This is a large addition (~8000+ lines across transformer, kernels, pipeline, refiner, camera utilities, docs, tests, and conversion script) for a camera-controlled image-to-video pipeline. The architecture is complex with a two-stage design (DiT + optional LTX-2 AR refiner). Overall the code is functional but has several issues that should be addressed before merge.

Correctness

• The transformer model file (transformer_sana_wm.py) is ~6500 lines and appears to be a near-verbatim port of the research repo rather than a clean diffusers-style implementation. It carries a large amount of dead code: unused activation/norm registries, training-time utilities (remove_bn, set_norm_eps), unused embedding paths (use_delta_actions, use_delta_translation, use_delta_pose_additive, cross_attn_image_embeds, pack_latents, etc.), and multiple attention variants that are never selected by the default config. Per the repo's coding style guide: "When porting from a research repo, delete training-time code paths, experimental flags, and ablation branches entirely — only keep the inference path you are actually integrating."

• _xformers_available is hardcoded to False (line 5461), yet the forward method has a branch (lines 6348-6352) that raises ValueError when mask is None and _xformers_available is False. While the pipeline always passes mask, this makes the transformer unusable standalone without a mask — a surprising API trap.

• The _tokens_per_frame calculation in _RefinerChunkRunner multiplies by patch_size_t instead of dividing. For the LTX-2 transformer with patch_size_t=1 this is a no-op, but the formula is semantically wrong: tokens per frame should be (H // patch_size) * (W // patch_size) / patch_size_t (or just the spatial product when patch_size_t=1).

Security

• torch.load(..., weights_only=False) at refiner.py:299 is the only instance in the entire diffusers codebase. The checkpoint payload contains only dicts, ints, tuples, and tensors — weights_only=True should work (or use safetensors). Loading arbitrary pickles from a user-supplied checkpoint_dir is a deserialization vulnerability.

• torch.load at transformer_sana_wm.py:6523 (null embed loading) also lacks weights_only=True, though null_embed_path defaults to None so it always hits the except branch in practice. Still should be cleaned up.

Style / Repo conventions

• The transformer file is enormous (~6500 lines) with many helper classes, registries, and code paths that only exist for training or ablation variants not used at inference. The diffusers guide explicitly says to inline small helpers and delete unused code paths. This file needs significant pruning.

• self.logger = print (line 5964) — the transformer uses print as its logger instead of the standard logging.get_logger(__name__) pattern used everywhere else in diffusers. Several self.logger(...) calls produce noisy stdout output during construction.

• Hard assert statements (e.g., lines 5981-5984, 6165, 6425) should be replaced with proper ValueError raises per Python best practices.

• The SanaWMPipelineOutput dataclass is missing a license header.

Tests

• The test file acknowledges that the transformer has hardcoded architecture dimensions and cannot be tested with a tiny dummy model via PipelineTesterMixin. This means there are no fast-path pipeline tests — only CPU unit tests for helpers and a @skip-decorated integration test. This is a significant coverage gap.

Dependencies

• The transformer requires fla (flash-linear-attention), timm, and optionally triton as runtime dependencies. These are not listed in diffusers' setup requirements. The fallback stubs raise at construction time, which is fine, but the dependency surface is unusually large for a diffusers model.

serge v0.1.0 · model: claude-opus-4-6 · 24 LLM turns · 29 tool calls · 204.6s · 2189757 in / 5733 out tokens

[dg845]

Hi @lawrence-cj, sorry for the delay and thanks for your patience! I have asked the Serge bot to review the current code and its comments all look reasonable to me. Hopefully this helps unblock you, will follow up with further review comments.

[lawrence-cj]

Hi @lawrence-cj, sorry for the delay and thanks for your patience! I have asked the Serge bot to review the current code and its comments all look reasonable to me. Hopefully this helps unblock you, will follow up with further review comments.

Congrads on the agent workflow for reviewing code. Hopefully it will speed up and perfect all PR processes.
I'll continue fix the problem and go through your comments.

[yiyixuxu]

Hey @lawrence-cj — sorry we've been slow on this one, and thanks for your patience.

The PR deviates quite a bit from our conventions, and I think a normal review round would take a lot of back-and-forth on both sides. But we have a self-review skill now: https://github.com/huggingface/diffusers/blob/main/.ai/skills/self-review/SKILL.md

Basically: point your agent at the skill and ask it for a self-review → you'll get back a report that separates blocking issues from everything else → you work through the blocking ones with the agent interactively. If something's ambiguous or you're not sure it's right, skip it — the agent will keep a note so we can weigh in during review.

It would speed things up a lot if you're able to do that. Otherwise we'll take over the PR and refactor it ourselves — we don't have the bandwidth right now, but we'll get to it as soon as we can. Just let us know which you'd prefer.

[lawrence-cj]

Thanks @yiyixuxu — I'd much rather do the work than hand it over. Running the self-review skill against .ai/review-rules.md now and working through the blocking findings; I'll post the report here once I've been through it.

[lawrence-cj]

Self-review report

Ran the self-review skill against .ai/review-rules.md + AGENTS.md / models.md / pipelines.md / testing.md. Thanks for pointing me at it — it caught a real bug and a lot of research-repo residue.

Net −2,559 lines (pushed in b9fd857e1). transformer_sana_wm.py 6555 → 4592, refiner.py 1286 → 1017.

🐛 Correctness bug it found

flow_shift was a dead knob. __call__ did self.scheduler.config.shift = flow_shift, but FlowMatchEulerDiscreteScheduler.set_timesteps reads self.shift → self._shift, not config.shift. So the documented kwarg did nothing and every run silently used the checkpoint's shift=9.8. It also half-mutated a FrozenDict (the freeze guard checks __frozen, stored name-mangled, so it never fires) — leaving config.shift and config["shift"] permanently disagreeing on a shared component.

Fixed by deleting the kwarg and letting the scheduler own the shift (pipelines.md gotcha 3). That also removes the gotcha-7 mutation and keeps the numerics every validation run was built on.

Fixed

Model (models.md)

• _no_split_modules = ["blocks"] → ["SanaVideoMSCamCtrlBlock"] — it was an attribute name, so accelerate matched nothing and device_map="auto" would have split a block across GPUs and crashed. Added _repeated_blocks and _skip_layerwise_casting_patterns. I deliberately left _keep_in_fp32_modules unset: adding the usual Wan-style list broke inference on GPU, because SANA-WM applies the timestep modulation inline (t2i_modulate), so holding scale_shift_table / t_embedder in fp32 upcasts the hidden states and feeds fp32 activations into bf16 weights. Supporting it needs explicit casts in the block forward — noted in the code.
• torch.float64 → torch.float32 on the three live RoPE sites (gotcha 5); torch.empty → torch.zeros for parameter init (gotcha 6); self.proj.weight.dtype → input dtype (gotcha 4); WanRotaryPosEmbed.freqs → non-persistent buffer instead of a plain attribute reassigned inside forward.
• Deleted ~1,960 lines of unreachable code: all weight-init/transfer-learning helpers (from_pretrained overwrites them), initialize_gemma_params (out-of-band from_pretrained, and broken — read an undefined attribute), token_drop + the training branch, the cam-debug apparatus, both unreachable recurrence variants, the xformers branches (_xformers_available was a literal False defined twice and xformers was never imported), 3 env-var escape hatches, 9 unused classes, 12 unused helpers, and the sincos family — which let import numpy as np go, satisfying the no-NumPy-in-forward rule.

Pipeline / refiner (pipelines.md)

• torch.randn → randn_tensor (gotcha 10). The signature advertises generator: Generator | list[Generator], but torch.randn raises TypeError on a list and RuntimeError on a CPU generator with a CUDA device — so that path was broken.
• Deleted the SLURM preemption/resume feature (checkpoint_dir, _atomic_save_state, _capture_state/_restore_state) and the dead single-shot path the docstring called a "debugging fallback" — ~270 lines of cluster-ops scaffolding AGENTS.md rules out.
• _empty_cuda_cache (CUDA-only) → empty_device_cache; @torch.inference_mode() → @torch.no_grad() and dropped from inner helpers (gotcha 2); removed dead _callback_tensor_inputs, _exclude_from_cpu_offload (a no-op — pipeline_utils already skips non-nn.Module components), _kv_max_frames, latents.detach().
• Missing license header; the + 1e-6 in a branch already guarded on norm > 0; ephemeral comments.

Verification

• State dict unchanged: 871/871 keys, exact against the released checkpoint.
• Differential test old-vs-new module: bit-exact (max|Δ| = 0.0) with float64 RoPE held constant; 4.98e-4 abs / 1.24e-4 rel with the prescribed float64→float32, entirely attributable to that change.
• End-to-end GPU smoke (stage 1 + refiner, public checkpoint): (24, 704, 1280, 3), no NaN, frames visually correct. Frame mean moved 0.5560 → 0.5554, which is exactly the float64→float32 RoPE change. (This run is what caught the _keep_in_fp32_modules problem above.)
• 15 CPU tests, ruff, doc-builder, docstring checks all clean.

Deliberately not changed — I'd like your call first

1. Nested-pipeline composition. SanaWMLTX2Refiner is a DiffusionPipeline registered as a component of SanaWMPipeline. Because it isn't an nn.Module, none of the base-class device/offload machinery reaches its ~87 GB, which is why both pipelines hand-roll .to(device) shuffling — a pipelines.md gotcha-7 violation I did not paper over, since the honest fix is structural. Options: flatten Kandinsky-combined-style (register refiner_* sub-modules so model_cpu_offload_seq covers them), or ship stage 2 as a separate top-level pipeline chained via output_type="latent". Related: pipelines.md opens with "prefer modular for new pipelines" — happy to go that route if you'd rather.

2. Attention pattern. The GDN linear-attention blocks don't use AttentionModuleMixin / a processor / dispatch_attention_fn, and cross-attention passes a dense additive -10000.0 mask that hard-raises on flash/FA3/sage. Converting cross-attention + the 5 softmax blocks is tractable; expressing the GDN recurrence as a processor is not obviously possible. Want me to convert the parts that fit and leave GDN as a documented exception?

3. Single-file policy / Triton. transformer_sana_wm_kernels.py has no precedent in-tree. The pure-PyTorch camera math and RoPE prep should move into the model file regardless. For the ~3 live Triton kernels the real choice is: ship pure-PyTorch only in this PR and land kernels as a follow-up, or publish to kernels-community and dispatch (as suggested for fla). Your call — I didn't want to pick unilaterally.

4. Tests. I removed the @slow integration stub (testing.md: none in the initial PR) and unblocked the real fix — num_layers / hidden_size / num_attention_heads / patch_size are now register_to_config params (defaults unchanged, so the released config.json still loads and the state dict is untouched), and a tiny 107-param variant now builds. What's still missing is the actual testing.md suite: test_pipeline_sana_wm.py with a SanaWMPipelineTesterConfig + PipelineTesterMixin/MemoryTesterMixin, and a generated model-level test. PipelineTesterMixin additionally needs the standard __call__ surface (batching, prompt_embeds, callback_on_step_end, num_videos_per_prompt), which this pipeline doesn't have. Since that surface is shaped by the decision in (1), I'd rather land the tests right after you weigh in than guess now.

Smaller items I left for review: seed alongside generator, check_inputs returning normalized values, _encode_first_frame/_build_camera_kwargs being private, TARGET_HEIGHT/WIDTH module constants, encode_prompt overlapping SanaPipeline's without # Copied from, and the estimate_intrinsics_with_pi3x helper that needs an undeclared pi3 dependency.