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

importtorchfromPILimportImagefromdiffusersimportSanaWMPipelinefromdiffusers.utilsimportexport_to_videopipe=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 DSLintrinsics=[800.0, 800.0, 845.0, 464.0], # fx, fy, cx, cy in original-image pixelsnum_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]

@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)

[yiyixuxu]

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

Whether to implement the refiner transformer as a new model:

[lawrence-cj]

Pushed 8f9af9ae2. 93 of 96 threads resolved.

The refiner is now its own model and its own pipeline, per @yiyixuxu's call:

latents=SanaWMPipeline(...) # stage 1video=SanaWMLTX2Refiner(...) # stage 2, pass vae=pipe.vae to share weights

SanaWMLTX2RefinerTransformer3DModel subclasses LTX2VideoTransformerBlock and overrides only forward, so the state dict is unchanged at 3510 keys; the KV cache follows the Flux2KVCache pattern and is threaded through forward(..., kv_cache=, kv_cache_mode=) instead of being set as attributes on LTX-2's submodules. refiner.py drops 982 → 638 lines and the manual offload is gone with the nesting that forced it.

Two real bugs surfaced on the way, both worth recording:

1. torch_dtype never reached the nested sub-pipeline.SanaWMPipeline.from_pretrained(..., torch_dtype=torch.bfloat16) was silently running all of stage 2 in float32 — transformer, text encoder and connectors alike. That is where the "~87 GB of refiner weights" came from, and half the reason the manual offload existed. Loading the refiner on its own halves its footprint. This is why the GPU smoke moves 0.5561 → 0.5548: I chased that delta specifically because it was bigger than anything else in this PR, and it's the dtype fix, not a regression. Decoded video is visually identical.

2. DiffusionPipeline.from_pretrained has no subfolder argument. My first attempt put the refiner in a refiner/ subfolder; the loader silently picked up the base pipeline's identically-named transformer/, tokenizer/, text_encoder/ and scheduler/ folders and only failed because connectors/ doesn't exist at the top level. Had it existed, it would have run to completion on the wrong weights. The refiner now ships as its own repo (--dst-refiner, default <dst>-refiner), matching how SDXL splits base and refiner.

Also in this round: MultiHeadCrossAttention moved to AttentionModuleMixin + a processor using dispatch_attention_fn with a bool mask, and the shared RMSNorm replaced the local copy (measured rather than argued — 0.5560 → 0.5561, inside bf16 noise).

Still open (3):

• Timesteps/TimestepEmbedding — would rename t_embedder.mlp.0/2 to linear_1/linear_2 and break the released checkpoint.
• y_embedder.y_embedding — unused in forward but one of the 871 checkpoint keys.
• Flattening the GDN inheritance chain — GDN is registered standalone in ATTENTION_BLOCKS, so I asked whether dropping those registry entries is acceptable first.

All three are one-line changes gated on whether you want the checkpoint re-exported; happy to do them together if so.

The heavy CI jobs are sitting on action_required for this fork again whenever one of you gets a chance.

[dg845]

@bot /style

[github-actions]

Style bot fixed some files and pushed the changes.

[yiyixuxu]

thanks, i did another round of review on the transformer model and left more comments

[yiyixuxu]

Suggested change

	The state-dict layout matches the public SANA-WM release one-to-one — the diffusers wrapper places the inner DiT
	under a `_inner.` prefix. See [`SanaWMTransformer3DModel.add_inner_prefix`] for the helper used by the conversion
	script.

diffusers has its own checkpoint, no?

[yiyixuxu]

Suggested change

	# String-keyed registry for the GDN/softmax attention block variants used by the SANA-WM DiT.
	# `SanaWMTransformer3DModel` looks classes up here by its `attn_type` / `camctrl_type` config strings.
	# Populated after the class definitions below.
	ATTENTION_BLOCKS: dict[str, type] = {}
	def_resolve_attention_block(name: str, *, role: str) ->type:
	"""Look up a registered attention class by its config string."""
	cls=ATTENTION_BLOCKS.get(name)
	ifclsisNone:
	raiseValueError(f"Unknown {role}: {name!r}. Available: {sorted(ATTENTION_BLOCKS)}")
	returncls

can we just use the map inline?

[yiyixuxu]

Suggested change

	# Name used by the `camctrl_type` config string and the block-name mappings below.
	BidirectionalSoftmaxUCPESinglePathLiteLA=_SoftmaxUCPESinglePathLiteLA
	# The released `config.json` names the fused-Triton variants (`attn_type="BidirectionalGDNTriton"`,
	# `camctrl_type="BidirectionalGDNUCPESinglePathLiteLABothTriton"`). The Triton kernels now live outside
	# `diffusers`, so those names resolve to the equivalent pure-PyTorch implementations.
	ATTENTION_BLOCKS.update(
	{
	"GDN": GDN,
	"BidirectionalGDN": BidirectionalGDN,
	"BidirectionalGDNTriton": BidirectionalGDN,
	"BidirectionalGDNUCPESinglePathLiteLA": BidirectionalGDNUCPESinglePathLiteLA,
	"BidirectionalGDNUCPESinglePathLiteLATriton": BidirectionalGDNUCPESinglePathLiteLA,
	"BidirectionalGDNUCPESinglePathLiteLABothTriton": BidirectionalGDNUCPESinglePathLiteLA,
	}
	)

let;s just support the attn_type and camctrl_type of released checkpoint

[yiyixuxu]

Suggested change

nn.init.zeros_(self.t_conv.weight)

[yiyixuxu]

Suggested change

	defrestore_shape(tensor, target_d):
	returntensor.permute(0, 1, 3, 2, 4).reshape(B, H, target_d, N)

can you inline this? just one line

[yiyixuxu]

Suggested change

attn_type=attn_type_list[i],

[yiyixuxu]

Suggested change

	camctrl_type=camctrl_type_list[i],
	attn_cls=attn_cls_list[i],

[yiyixuxu]

Suggested change

	attn_type_list, camctrl_type_list=_inject_softmax_layers(
	attn_type_list,
	camctrl_type_list,
	softmax_every_n,
	)
	foriinrange(depth):
	ifsoftmax_every_n>0and (i+1) %softmax_every_n==0:
	attn_cls=_SoftmaxUCPESinglePathLiteLA
	else:
	attn_cls=xx
	attn_cls_list.append(attn_cls)

[yiyixuxu]

can we remove unuused branch here?

[yiyixuxu]

is it possible for us to only support what the released checkopint actually uses? if so, can we clean up all the class inherit from this accordingly? i.e remove all the classes not used by the released checkopint, try to flatten the ones that we have to keep?