diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 525734a2e4cb..985dd7315440 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -413,6 +413,10 @@ title: SanaTransformer2DModel - local: api/models/sana_video_transformer3d title: SanaVideoTransformer3DModel + - local: api/models/sana_wm_refiner_transformer3d + title: SanaWMLTX2RefinerTransformer3DModel + - local: api/models/sana_wm_transformer3d + title: SanaWMTransformer3DModel - local: api/models/sd3_transformer2d title: SD3Transformer2DModel - local: api/models/skyreels_v2_transformer_3d @@ -651,6 +655,8 @@ title: Sana Sprint - local: api/pipelines/sana_video title: Sana Video + - local: api/pipelines/sana_wm + title: SANA-WM - local: api/pipelines/shap_e title: Shap-E - local: api/pipelines/stable_cascade diff --git a/docs/source/en/api/models/sana_wm_refiner_transformer3d.md b/docs/source/en/api/models/sana_wm_refiner_transformer3d.md new file mode 100644 index 000000000000..4e64ae6841e1 --- /dev/null +++ b/docs/source/en/api/models/sana_wm_refiner_transformer3d.md @@ -0,0 +1,55 @@ + + +# SanaWMLTX2RefinerTransformer3DModel + +The chunk-causal autoregressive refiner DiT used as stage 2 of [`SanaWMPipeline`], driven by +[`SanaWMLTX2Refiner`]. + +It is architecturally identical to [`LTX2VideoTransformer3DModel`] — same config arguments, same submodules, same +parameter names — so a released LTX-2 checkpoint loads into it unchanged. The forward pass differs: + +* only the video stream is run (the audio and audio/video cross-attention branches are skipped), +* self-attention runs against an explicit sliding-window KV cache ([`SanaWMRefinerKVCache`]) holding the attention + sink plus the recent refined history, so per-block compute is bounded and total refinement cost scales linearly + with video length, +* the caller supplies the video RoPE, which lets each autoregressive window keep every frame's absolute index in the + source video (see + [`SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions`]). + +The model can be loaded with: + +```python +import torch +from diffusers import SanaWMLTX2RefinerTransformer3DModel + +transformer = SanaWMLTX2RefinerTransformer3DModel.from_pretrained( + "Efficient-Large-Model/SANA-WM_bidirectional-diffusers", + subfolder="refiner/transformer", + torch_dtype=torch.bfloat16, +) +``` + +## SanaWMLTX2RefinerTransformer3DModel + +[[autodoc]] SanaWMLTX2RefinerTransformer3DModel + +## SanaWMRefinerKVCache + +[[autodoc]] models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache + +## SanaWMRefinerKVLayerCache + +[[autodoc]] models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache + +## Transformer2DModelOutput + +[[autodoc]] models.modeling_outputs.Transformer2DModelOutput diff --git a/docs/source/en/api/models/sana_wm_transformer3d.md b/docs/source/en/api/models/sana_wm_transformer3d.md new file mode 100644 index 000000000000..a4ca954caf68 --- /dev/null +++ b/docs/source/en/api/models/sana_wm_transformer3d.md @@ -0,0 +1,42 @@ + + +# SanaWMTransformer3DModel + +A 3D Diffusion Transformer (1.6B parameters) for camera-controlled image-to-video generation, used as the stage-1 +sampler of [`SanaWMPipeline`]. The transformer combines: + +* a bidirectional GDN-Triton linear-attention main branch (depth 20, hidden 2240, 20 heads), +* a UCPE (Unified Camera Pose Embedding) camera-control branch that consumes a raymap + Plücker representation of + the requested trajectory, and +* a Wan-style 3D rotary position embedding plus periodic softmax-attention blocks injected every `softmax_every_n` + layers. + +The model can be loaded with: + +```python +import torch +from diffusers import SanaWMTransformer3DModel + +transformer = SanaWMTransformer3DModel.from_pretrained( + "Efficient-Large-Model/SANA-WM_bidirectional-diffusers", + subfolder="transformer", + torch_dtype=torch.bfloat16, +) +``` + +## SanaWMTransformer3DModel + +[[autodoc]] SanaWMTransformer3DModel + +## Transformer2DModelOutput + +[[autodoc]] models.modeling_outputs.Transformer2DModelOutput diff --git a/docs/source/en/api/pipelines/sana_wm.md b/docs/source/en/api/pipelines/sana_wm.md new file mode 100644 index 000000000000..561a809dc224 --- /dev/null +++ b/docs/source/en/api/pipelines/sana_wm.md @@ -0,0 +1,170 @@ + + +# SANA-WM + +SANA-WM is a camera-controlled image-to-video world model built on top of SANA. Given a first-frame image, a text +prompt, and a camera trajectory (either explicit `c2w` poses or a WASD/IJKL action string), it generates a video +whose motion follows the requested camera path. + +Inference runs in two stages: + +1. **Stage 1 — SANA-WM DiT.** A 1.6B-parameter bidirectional DiT with GDN-Triton linear attention and a UCPE + camera-control branch. Sampling uses an LTX-style flow-matching Euler scheduler with per-token timesteps; the + first latent frame is the conditioning anchor. +2. **Stage 2 — LTX-2 refiner (optional).** A separate sink-bidirectional Euler refiner pipeline + ([`SanaWMLTX2Refiner`]) that wraps + [`SanaWMLTX2RefinerTransformer3DModel`] + `LTX2TextConnectors` and a Gemma-3 text encoder, run for 3 + distilled sigma steps. + +Both stages decode through the [`AutoencoderKLLTX2Video`] VAE. + +Available models: + +| Model | Recommended dtype | +|:-----:|:-----------------:| +| [`Efficient-Large-Model/SANA-WM_bidirectional-diffusers`](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional-diffusers) | `torch.bfloat16` | +| [`Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refiner`](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refiner) | `torch.bfloat16` | + +> [!TIP] +> SANA-WM is trained at a fixed 704×1280 resolution. The recommended dtype is for the transformer weights — keep +> the text encoder in `torch.bfloat16` and the VAE in `torch.float32` for best numerics. The pipeline expects +> camera intrinsics `[fx, fy, cx, cy]` in *original-image* pixel coordinates; the resize-and-center-crop transform +> is applied internally. + +## Inference + +```python +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.enable_model_cpu_offload() # ~45 GB of weights — offload between stages + +# SANA-WM was trained on the LTX-2 VAE in framewise mode with tiling enabled. Without these +# settings the VAE encodes the whole (B, C, T, H, W) clip in one shot, which gives subtly +# different numerics from the released checkpoint. +pipe.vae.enable_tiling() +pipe.vae.use_framewise_encoding = True +pipe.vae.use_framewise_decoding = True +pipe.vae.tile_sample_stride_num_frames = 64 +pipe.vae.tile_sample_min_num_frames = 96 + +prompt = "A car driving across a vast desert plain at golden hour." +output = pipe( + image=Image.open("input.png").convert("RGB"), + prompt=prompt, + action="w-80,jw-40,w-40", # WASD-style action DSL: forward 80f, jump+forward 40f, forward 40f + intrinsics=[800.0, 800.0, 845.0, 464.0], # fx, fy, cx, cy in original-image pixels + num_frames=161, + num_inference_steps=60, + guidance_scale=5.0, + generator=torch.Generator(device="cuda").manual_seed(42), + output_type="latent", # hand the latents to the refiner below +) +``` + +Pass `action=None` and supply your own `c2w` poses (`(F, 4, 4)` numpy array) to drive the camera trajectory +explicitly. Drop `output_type="latent"` to get video straight out of stage 1 and skip the refiner. + +### Stage 2 — the LTX-2 refiner + +[`SanaWMLTX2Refiner`] ships as its own repository, the way SDXL splits base and refiner. Pass the base pipeline's +VAE so the weights are shared rather than loaded twice: + +```python +from diffusers import SanaWMLTX2Refiner + +refiner = SanaWMLTX2Refiner.from_pretrained( + "Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refiner", + vae=pipe.vae, + torch_dtype=torch.bfloat16, +) +refiner.enable_model_cpu_offload() + +frames = refiner(output.latent, prompt, fps=16) +export_to_video(list(frames), "sana_wm.mp4", fps=16) +``` + +Without a `vae` the refiner returns refined latents instead of video, which is useful if you want to decode +yourself. + +> [!TIP] +> Enable offloading on **both** pipelines, or free stage 1 before stage 2 (`pipe.transformer.to("cpu")`). The two +> stages together are around 45 GB in `torch.bfloat16`, and keeping both resident on one 80 GB card leaves too +> little room for activations. Note also that stage 2 honours `torch_dtype`: loading the refiner in +> `torch.float32` roughly doubles its memory and changes the output slightly. + +If you don't have camera intrinsics, a hosted [modular block](../../modular_diffusers/overview) can estimate them +from a single frame. It lives outside `diffusers` because it pulls in Pi3X — an extra dependency and a second +checkpoint — so nothing is downloaded until you ask for it: + +```python +from diffusers import ModularPipeline + +# One-time per image. Requires `pip install pi3-vision`. +estimator = ModularPipeline.from_pretrained( + "Efficient-Large-Model/pi3x-intrinsics-estimator", trust_remote_code=True +) +estimator.load_components(dtype=torch.bfloat16) +intrinsics = estimator(image=Image.open("input.png").convert("RGB"), output="intrinsics") +``` + +## Converting the released checkpoint + +If you have the source SANA-WM release (not the pre-converted diffusers snapshot), run the conversion script once: + +```bash +python scripts/convert_sana_wm_to_diffusers.py \ + --src Efficient-Large-Model/SANA-WM_bidirectional \ + --dst ./SANA-WM_bidirectional-diffusers +``` + +This writes two directories: the base pipeline at `--dst`, and the stage-2 refiner alongside it at +`./SANA-WM_bidirectional-diffusers-refiner` (override with `--dst-refiner`). Then load each from its local path +as usual. + +## Components + +- `tokenizer` — [`GemmaTokenizerFast`] +- `text_encoder` — Gemma-2 (returns decoder hidden states) +- `vae` — [`AutoencoderKLLTX2Video`] (LTX-2, spatial ×32 / temporal ×8) +- `transformer` — [`SanaWMTransformer3DModel`], 1.6B-parameter bidirectional DiT +- `scheduler` — [`FlowMatchEulerDiscreteScheduler`] + +The stage-2 refiner is a separate repository with its own `transformer` +([`SanaWMLTX2RefinerTransformer3DModel`]), `connectors` (`LTX2TextConnectors`), `tokenizer`, Gemma-3 +`text_encoder` and `scheduler`. It has no `vae` of its own — pass the base pipeline's. + +## SanaWMPipeline + +[[autodoc]] SanaWMPipeline + - all + - __call__ + +## SanaWMLTX2Refiner + +The LTX-2 stage-2 refiner is a standalone [`DiffusionPipeline`] that takes stage-1 latents. Give it a `vae` (the +base pipeline's, so the weights are shared) to have it decode to video; without one it returns refined latents. + +[[autodoc]] SanaWMLTX2Refiner + - all + - __call__ + +## SanaWMPipelineOutput + +[[autodoc]] pipelines.sana_wm.pipeline_output.SanaWMPipelineOutput diff --git a/scripts/convert_sana_wm_to_diffusers.py b/scripts/convert_sana_wm_to_diffusers.py new file mode 100644 index 000000000000..5d266168a745 --- /dev/null +++ b/scripts/convert_sana_wm_to_diffusers.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert the public SANA-WM release into a diffusers-loadable directory. + +Reads the ``Efficient-Large-Model/SANA-WM_bidirectional`` HF repo (or a local +mirror) and writes a directory ready for ``SanaWMPipeline.from_pretrained(path)``: + + / + ├── model_index.json + ├── tokenizer/ + ├── text_encoder/ + ├── vae/ + ├── transformer/ + ├── scheduler/ + └── refiner/ + ├── transformer/ + ├── connectors/ + ├── text_encoder/ + └── tokenizer/ + +Usage: + python scripts/convert_sana_wm_to_diffusers.py \\ + --src Efficient-Large-Model/SANA-WM_bidirectional \\ + --dst /path/to/SANA-WM_bidirectional-diffusers \\ + [--no-refiner] + +The output is local-only; no upload to the Hub. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +from pathlib import Path + +import torch +from huggingface_hub import snapshot_download +from safetensors import safe_open +from safetensors.torch import save_file + + +def _copy_subdir(src: Path, dst: Path) -> None: + if dst.exists(): + shutil.rmtree(dst) + shutil.copytree(src, dst, symlinks=False) + + +def _cast_to_bfloat16(component: Path) -> None: + """Rewrite a component's safetensors shards in bfloat16. + + SANA-WM is served in `torch.bfloat16`, so `from_pretrained(torch_dtype=torch.bfloat16)` casts these weights + anyway; storing them cast halves the download without changing a single value. Non-float32 tensors are left + alone, and the VAE is deliberately not passed through here since it runs in `torch.float32`. + """ + shards = sorted(component.glob("*.safetensors")) + if not shards: + return + total = 0 + for shard in shards: + with safe_open(shard, framework="pt") as f: + metadata = f.metadata() + tensors = { + k: (lambda t: t.to(torch.bfloat16) if t.dtype == torch.float32 else t)(f.get_tensor(k)) + for k in f.keys() + } + tmp = shard.with_suffix(shard.suffix + ".tmp") + save_file(tensors, tmp, metadata=metadata) + tmp.replace(shard) + total += shard.stat().st_size + del tensors + index = component / "model.safetensors.index.json" + if index.exists(): + payload = json.loads(index.read_text()) + payload["metadata"]["total_size"] = total + index.write_text(json.dumps(payload, indent=2)) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--src", default="Efficient-Large-Model/SANA-WM_bidirectional", help="HF repo or local dir") + parser.add_argument("--dst", required=True, type=Path, help="Output directory") + parser.add_argument("--no-refiner", action="store_true", help="Skip refiner export") + parser.add_argument( + "--dst-refiner", + type=Path, + default=None, + help="Output directory for the stage-2 refiner pipeline (default: `-refiner`)", + ) + parser.add_argument( + "--torch-dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"], help="Weight dtype" + ) + args = parser.parse_args() + + torch_dtype = getattr(torch, args.torch_dtype) + dst: Path = args.dst.absolute() + dst.mkdir(parents=True, exist_ok=True) + + # Resolve the source on disk (snapshot_download for HF repos, otherwise use as-is). + src_path = Path(args.src) + if not src_path.is_dir(): + print(f"[convert] snapshot_download({args.src}) …") + src_path = Path(snapshot_download(args.src)) + print(f"[convert] source: {src_path}") + + # 1. VAE (already diffusers format under /vae). + print("[convert] vae …") + _copy_subdir(src_path / "vae", dst / "vae") + + # 2. Tokenizer + text encoder — fetch via the configured Gemma-2 repo. + # We save the full ``Gemma2ForCausalLM``; the pipeline grabs the decoder + # at runtime via ``self.text_encoder.model(...)``. This matches the sana + # inference recipe of ``AutoModelForCausalLM.from_pretrained(...).get_decoder()`` + # and avoids subtle state-dict prefix differences when saving just the + # decoder submodule. + print("[convert] tokenizer + text_encoder (gemma-2-2b-it) …") + from transformers import AutoModelForCausalLM, AutoTokenizer + + gemma_repo = "Efficient-Large-Model/gemma-2-2b-it" + tokenizer = AutoTokenizer.from_pretrained(gemma_repo) + tokenizer.padding_side = "right" + tokenizer.save_pretrained(dst / "tokenizer") + text_encoder = AutoModelForCausalLM.from_pretrained(gemma_repo, torch_dtype=torch_dtype) + text_encoder.save_pretrained(dst / "text_encoder") + del text_encoder + + # 3. Transformer (SanaWMTransformer3DModel) — load the public DiT, save in diffusers format. + print("[convert] transformer (SanaWMTransformer3DModel) …") + from diffusers import SanaWMTransformer3DModel + + transformer = SanaWMTransformer3DModel().to(torch_dtype).eval() + dit_ckpt = src_path / "dit" / "sana_wm_1600m_720p.safetensors" + if not dit_ckpt.is_file(): + raise FileNotFoundError(f"DiT checkpoint not found at {dit_ckpt}") + from safetensors.torch import load_file + + sd = load_file(str(dit_ckpt)) + sd.pop("pos_embed", None) # unused at inference (wan_rope is computed on-the-fly) + # The public release keys (``blocks.0...``) load directly into the merged + # SanaWMTransformer3DModel — no ``_inner.`` prefix anymore. + # `.pos_embed` entries are non-persistent buffers rebuilt at construction, so they are + # expected to be absent from the converted state dict; anything else means the mapping + # is wrong and would silently produce a broken transformer. + missing, unexpected = transformer.load_state_dict(sd, strict=False) + missing = [k for k in missing if not k.endswith(".pos_embed")] + if missing or unexpected: + raise RuntimeError( + "State dict does not match `SanaWMTransformer3DModel`.\n" + f" missing keys ({len(missing)}): {missing[:10]}{' …' if len(missing) > 10 else ''}\n" + f" unexpected keys ({len(unexpected)}): {unexpected[:10]}{' …' if len(unexpected) > 10 else ''}" + ) + transformer.save_pretrained(dst / "transformer") + del transformer, sd + + # 4. Scheduler — FlowMatchEulerDiscreteScheduler config. + print("[convert] scheduler …") + from diffusers import FlowMatchEulerDiscreteScheduler + + FlowMatchEulerDiscreteScheduler(shift=9.8).save_pretrained(dst / "scheduler") + + # 5. Refiner (LTX-2): a standalone DiffusionPipeline written to its own output + # directory, the way SDXL ships base and refiner as separate repos. + # `DiffusionPipeline.from_pretrained` has no `subfolder` argument, so a nested + # folder would silently load the *base* pipeline's components instead. Copy the LTX-2 sub-model folders as-is, split out a ``tokenizer/`` + # folder, add a ``scheduler/`` (FlowMatchEulerDiscreteScheduler), and write the + # manifest. The transformer weights are LTX-2's, but the refiner drives them + # through its own model class, so the manifest names that class. + if not args.no_refiner: + print("[convert] refiner …") + from transformers import AutoTokenizer + + refiner_src = src_path / "refiner" + refiner_dst = args.dst_refiner or dst.with_name(dst.name + "-refiner") + refiner_dst = Path(refiner_dst) + refiner_dst.mkdir(parents=True, exist_ok=True) + for sub in ("transformer", "connectors", "text_encoder"): + if (refiner_src / sub).is_dir(): + _copy_subdir(refiner_src / sub, refiner_dst / sub) + _cast_to_bfloat16(refiner_dst / sub) + + # The weights are LTX-2's, but they are driven by `SanaWMLTX2RefinerTransformer3DModel` + # (same submodule layout, streaming-attention forward), so point the config at that class. + transformer_config_path = refiner_dst / "transformer" / "config.json" + transformer_config = json.loads(transformer_config_path.read_text()) + transformer_config["_class_name"] = "SanaWMLTX2RefinerTransformer3DModel" + transformer_config_path.write_text(json.dumps(transformer_config, indent=2)) + + # Tokenizer lives co-located with the Gemma-3 text encoder in the release; + # re-save it into its own subfolder so it registers as a pipeline component. + refiner_tokenizer = AutoTokenizer.from_pretrained(refiner_src / "text_encoder") + refiner_tokenizer.save_pretrained(refiner_dst / "tokenizer") + + # Scheduler carries the distilled sigma schedule; shift=1.0 leaves the + # explicit sigmas passed at inference time unmodified. + FlowMatchEulerDiscreteScheduler(shift=1.0).save_pretrained(refiner_dst / "scheduler") + + refiner_index = { + "_class_name": "SanaWMLTX2Refiner", + "_diffusers_version": "0.38.0", + "transformer": ["diffusers", "SanaWMLTX2RefinerTransformer3DModel"], + # LTX2TextConnectors lives in diffusers.pipelines.ltx2 (not top-level), + # so the loader resolves it via the pipeline-module path ("ltx2", ...). + "connectors": ["ltx2", "LTX2TextConnectors"], + "tokenizer": ["transformers", type(refiner_tokenizer).__name__], + "text_encoder": ["transformers", "Gemma3ForConditionalGeneration"], + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + "text_max_sequence_length": 1024, + } + (refiner_dst / "model_index.json").write_text(json.dumps(refiner_index, indent=2)) + + # 6. model_index.json — the top-level diffusers manifest. + print("[convert] model_index.json …") + index = { + "_class_name": "SanaWMPipeline", + "_diffusers_version": "0.38.0", + "tokenizer": ["transformers", "GemmaTokenizerFast"], + "text_encoder": ["transformers", "Gemma2ForCausalLM"], + "vae": ["diffusers", "AutoencoderKLLTX2Video"], + "transformer": ["diffusers", "SanaWMTransformer3DModel"], + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + } + (dst / "model_index.json").write_text(json.dumps(index, indent=2)) + + print(f"[convert] done — wrote {dst}") + if not args.no_refiner: + print(f"[convert] done — wrote {refiner_dst}") + + +if __name__ == "__main__": + main() diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2825e9888c98..bed507e9a096 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -342,6 +342,8 @@ "SanaControlNetModel", "SanaTransformer2DModel", "SanaVideoTransformer3DModel", + "SanaWMLTX2RefinerTransformer3DModel", + "SanaWMTransformer3DModel", "SD3ControlNetModel", "SD3MultiControlNetModel", "SD3Transformer2DModel", @@ -787,6 +789,9 @@ "SanaSprintPipeline", "SanaVideoPipeline", "SanaVideoPipeline", + "SanaWMLTX2Refiner", + "SanaWMPipeline", + "SanaWMPipelineOutput", "SemanticStableDiffusionPipeline", "ShapEImg2ImgPipeline", "ShapEPipeline", @@ -1224,6 +1229,8 @@ SanaControlNetModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, + SanaWMLTX2RefinerTransformer3DModel, + SanaWMTransformer3DModel, SD3ControlNetModel, SD3MultiControlNetModel, SD3Transformer2DModel, @@ -1643,6 +1650,9 @@ SanaSprintImg2ImgPipeline, SanaSprintPipeline, SanaVideoPipeline, + SanaWMLTX2Refiner, + SanaWMPipeline, + SanaWMPipelineOutput, SemanticStableDiffusionPipeline, ShapEImg2ImgPipeline, ShapEPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 1a396b312441..9c73f8f86d31 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -147,6 +147,8 @@ _import_structure["transformers.transformer_qwenimage"] = ["QwenImageTransformer2DModel"] _import_structure["transformers.transformer_qwenimage21"] = ["QwenImage21Transformer2DModel"] _import_structure["transformers.transformer_sana_video"] = ["SanaVideoTransformer3DModel"] + _import_structure["transformers.transformer_sana_wm"] = ["SanaWMTransformer3DModel"] + _import_structure["transformers.transformer_sana_wm_refiner"] = ["SanaWMLTX2RefinerTransformer3DModel"] _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"] _import_structure["transformers.transformer_skyreels_v2"] = ["SkyReelsV2Transformer3DModel"] _import_structure["transformers.transformer_stable_audio3"] = ["StableAudio3DiTModel"] @@ -294,6 +296,8 @@ QwenImageTransformer2DModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, + SanaWMLTX2RefinerTransformer3DModel, + SanaWMTransformer3DModel, SD3Transformer2DModel, SkyReelsV2Transformer3DModel, StableAudio3DiTModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index ffb0cbc0318b..087f4be68674 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -62,6 +62,8 @@ from .transformer_qwenimage import QwenImageTransformer2DModel from .transformer_qwenimage21 import QwenImage21Transformer2DModel from .transformer_sana_video import SanaVideoTransformer3DModel + from .transformer_sana_wm import SanaWMTransformer3DModel + from .transformer_sana_wm_refiner import SanaWMLTX2RefinerTransformer3DModel from .transformer_sd3 import SD3Transformer2DModel from .transformer_skyreels_v2 import SkyReelsV2Transformer3DModel from .transformer_stable_audio3 import StableAudio3DiTModel diff --git a/src/diffusers/models/transformers/transformer_sana_wm.py b/src/diffusers/models/transformers/transformer_sana_wm.py new file mode 100644 index 000000000000..e422833c67a8 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_sana_wm.py @@ -0,0 +1,3168 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# This file is modified from https://github.com/PixArt-alpha/PixArt-sigma + +from __future__ import annotations + +import math +from copy import deepcopy +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import logging +from ..activations import get_activation +from ..attention import AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..embeddings import get_1d_rotary_pos_embed +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin, get_parameter_dtype +from ..normalization import RMSNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class Mlp(nn.Module): + """Two-layer feed-forward block (`fc1` -> activation -> `fc2`).""" + + def __init__( + self, + in_features: int, + hidden_features: int | None = None, + out_features: int | None = None, + act_layer: type[nn.Module] = nn.GELU, + bias: bool = True, + drop: float = 0.0, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.drop1 = nn.Dropout(drop) + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop2 = nn.Dropout(drop) + + def forward(self, hidden_states: torch.Tensor, HW: tuple[int, int] | None = None) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.drop1(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = self.drop2(hidden_states) + return hidden_states + + +class SanaWMTemporalShortConvolution(nn.Module): + """Depthwise short convolution over the temporal axis, run in both directions. + + SANA-WM's GDN attention applies a short depthwise conv to Q/K/V before the linear-attention kernel. The reference + implementation used the *causal* `fla.modules.ShortConvolution` layer (with `activation=None`); this is a + self-contained PyTorch implementation -- so the model needs no `fla-core` dependency and can be built on any device + -- that runs the causal kernel forwards and backwards to obtain the non-causal filter the bidirectional model + needs. + + A causal depthwise Conv1d with kernel ``[w_0, w_1, ..., w_{k-1}]`` computes at time *t*: + + ``y_fwd[t] = w_0 * x[t-k+1] + ... + w_{k-1} * x[t]`` + + Running the same kernel on the time-flipped input and flipping back gives: + + ``y_bwd[t] = w_{k-1} * x[t] + ... + w_0 * x[t+k-1]`` + + Both passes include the current timestep ``x[t]`` with the center weight ``w_{k-1}``. To avoid double-counting one + copy of the center contribution is subtracted: + + ``y = y_fwd + y_bwd - w_{k-1} * x`` + + The result is a symmetric temporal filter where every position in the window ``[t-k+1, t+k-1]`` is counted exactly + once. + + Args: + hidden_size (`int`): Number of channels (the conv is depthwise, one group per channel). + kernel_size (`int`): Temporal kernel width. + bias (`bool`, defaults to `False`): Whether to add a per-channel bias. + """ + + def __init__(self, hidden_size: int, kernel_size: int, bias: bool = False) -> None: + super().__init__() + self.hidden_size = hidden_size + self.kernel_size = kernel_size + # Same parameter layout as the reference implementation: (C, 1, K). + self.weight = nn.Parameter(torch.zeros(hidden_size, 1, kernel_size)) + self.bias = nn.Parameter(torch.zeros(hidden_size)) if bias else None + + def _causal_conv(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Depthwise causal conv over `(batch, seq_len, hidden_size)` inputs.""" + seq_len = hidden_states.shape[1] + # Left-pad by (K - 1) and drop the tail so output[t] only sees inputs <= t. + hidden_states = F.conv1d( + hidden_states.transpose(1, 2), + self.weight.to(hidden_states.dtype), + None if self.bias is None else self.bias.to(hidden_states.dtype), + groups=self.hidden_size, + padding=self.kernel_size - 1, + )[..., :seq_len] + return hidden_states.transpose(1, 2) + + def forward(self, hidden_states: torch.Tensor, num_frames: int) -> torch.Tensor: + """Apply the bidirectional conv along the temporal axis, with the spatial axis merged into the batch. + + Args: + hidden_states (`torch.Tensor`): Input of shape `(batch, num_frames * spatial_size, hidden_size)`. + num_frames (`int`): Number of frames the sequence axis is split into. + + Returns: + `torch.Tensor`: Tensor of the same shape and dtype as `hidden_states`. + """ + batch_size, seq_len, channels = hidden_states.shape + spatial_size = seq_len // num_frames + dtype_in = hidden_states.dtype + + # (B, T * S, C) -> (B * S, T, C). The causal conv backward is not reliable on the non-contiguous + # strided layout this permutation produces, hence the explicit `contiguous()`. + hidden_states = ( + hidden_states.reshape(batch_size, num_frames, spatial_size, channels) + .permute(0, 2, 1, 3) + .contiguous() + .reshape(batch_size * spatial_size, num_frames, channels) + ) + + causal_forward = self._causal_conv(hidden_states) + causal_backward = self._causal_conv(hidden_states.flip(1)).flip(1) + + # Subtract the shared center tap (last weight of the causal kernel). Weight shape: (channels, 1, kernel_size), + # so the last element along dim=-1 is the weight applied to x[t]. + center_term = hidden_states * self.weight[:, 0, -1].unsqueeze(0).unsqueeze(0) + + hidden_states = causal_forward + causal_backward - center_term + if hidden_states.dtype != dtype_in: + hidden_states = hidden_states.to(dtype_in) + + return ( + hidden_states.reshape(batch_size, spatial_size, num_frames, channels) + .permute(0, 2, 1, 3) + .reshape(batch_size, seq_len, channels) + ) + + +# Safe element-count threshold for a single conv call: PyTorch's 2D conv kernels (both cuDNN and the ATEN fallback) +# use 32-bit indexing internally, so very large ``(batch * frames, channels, height, width)`` inputs (e.g. minute-scale +# video at default CFG) can overflow. Empirically a single call up to ~1B elements is safe; above that we split along +# the leading dim. Set so short videos stay on the original fused path (no chunking, no overhead). +_INT32_SAFE_CONV_ELEMENTS = 1 << 30 # 1,073,741,824 + + +class SanaWMConvLayer(nn.Module): + """2D convolution with an optional activation. + + Wraps the convolution in a ``conv`` submodule to keep the checkpoint's parameter names + (``mlp.inverted_conv.conv.weight``, ...) unchanged. + + Args: + in_dim (`int`): Input channels. + out_dim (`int`): Output channels. + kernel_size (`int`, defaults to 3): Spatial kernel size (odd, so ``same`` padding is exact). + groups (`int`, defaults to 1): Convolution groups. + use_bias (`bool`, defaults to `False`): Whether the convolution has a bias. + act (`str`, *optional*): Activation name resolved through + [`~models.activations.get_activation`], or `None` for no activation. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + kernel_size: int = 3, + groups: int = 1, + use_bias: bool = False, + act: Optional[str] = None, + ) -> None: + super().__init__() + self.conv = nn.Conv2d( + in_dim, + out_dim, + kernel_size=(kernel_size, kernel_size), + padding=kernel_size // 2, + groups=groups, + bias=use_bias, + ) + self.act = get_activation(act) if act is not None else None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv(hidden_states) + if self.act is not None: + hidden_states = self.act(hidden_states) + return hidden_states + + +class GLUMBConvTemp(nn.Module): + """SANA-WM feed-forward block: a gated inverted-bottleneck conv over space plus a residual temporal conv. + + Args: + in_features (`int`): Input channels. + hidden_features (`int`): Width of the inverted bottleneck (doubled internally for the GLU gate). + out_feature (`int`, *optional*): Output channels, defaults to `in_features`. + kernel_size (`int`, defaults to 3): Spatial kernel size of the depthwise convolution. + use_bias (`tuple[bool, bool, bool]`, defaults to `(False, False, False)`): Bias flag per convolution. + act (`tuple`, defaults to `("silu", "silu", None)`): Activation for the inverted conv, the GLU gate and the + point conv respectively; `None` means no activation. + t_kernel_size (`int`, defaults to 3): Temporal kernel size of the residual temporal convolution. + """ + + def __init__( + self, + in_features: int, + hidden_features: int, + out_feature: Optional[int] = None, + kernel_size: int = 3, + use_bias: Tuple[bool, bool, bool] = (False, False, False), + act: Tuple[Optional[str], Optional[str], Optional[str]] = ("silu", "silu", None), + t_kernel_size: int = 3, + ) -> None: + super().__init__() + out_feature = out_feature or in_features + + self.glu_act = get_activation(act[1]) + self.inverted_conv = SanaWMConvLayer( + in_features, hidden_features * 2, kernel_size=1, use_bias=use_bias[0], act=act[0] + ) + self.depth_conv = SanaWMConvLayer( + hidden_features * 2, + hidden_features * 2, + kernel_size=kernel_size, + groups=hidden_features * 2, + use_bias=use_bias[1], + act=None, + ) + self.point_conv = SanaWMConvLayer( + hidden_features, out_feature, kernel_size=1, use_bias=use_bias[2], act=act[2] + ) + self.t_conv = nn.Conv2d( + out_feature, + out_feature, + kernel_size=(t_kernel_size, 1), + padding=(t_kernel_size // 2, 0), + bias=False, + ) + + def forward(self, hidden_states: torch.Tensor, HW: Tuple[int, int, int]) -> torch.Tensor: + batch_size, seq_len, channels = hidden_states.shape + num_frames, height, width = HW + hidden_states = hidden_states.reshape(batch_size * num_frames, height, width, channels).permute(0, 3, 1, 2) + + # Split the leading dim so each conv launch stays under PyTorch's 32-bit indexing limit (no-op for short clips). + rows_per_call = max(1, _INT32_SAFE_CONV_ELEMENTS // (self.inverted_conv.conv.out_channels * height * width)) + spatial_chunks = [] + for start in range(0, hidden_states.shape[0], rows_per_call): + chunk = self.inverted_conv(hidden_states[start : start + rows_per_call]) + chunk = self.depth_conv(chunk) + value, gate = torch.chunk(chunk, 2, dim=1) + spatial_chunks.append(self.point_conv(value * self.glu_act(gate))) + hidden_states = spatial_chunks[0] if len(spatial_chunks) == 1 else torch.cat(spatial_chunks, dim=0) + + # Residual temporal aggregation over the frame axis. + hidden_states = hidden_states.view(batch_size, num_frames, channels, height * width).permute(0, 2, 1, 3) + hidden_states = hidden_states + self.t_conv(hidden_states) + return hidden_states.permute(0, 2, 3, 1).reshape(batch_size, seq_len, channels) + + +class SanaWMCrossAttnProcessor: + """Cross-attention from image queries to the text condition.""" + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "MultiHeadCrossAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size, _, channels = hidden_states.shape + + query = attn.q_linear(hidden_states) + key, value = attn.kv_linear(encoder_hidden_states).view(batch_size, -1, 2, channels).unbind(2) + + query = attn.q_norm(query).view(batch_size, -1, attn.heads, attn.head_dim) + key = attn.k_norm(key).view(batch_size, -1, attn.heads, attn.head_dim) + value = value.view(batch_size, -1, attn.heads, attn.head_dim) + + # A boolean mask (rather than an additive float one) keeps the varlen backends usable. + if attention_mask is not None and attention_mask.ndim == 2: + attention_mask = attention_mask.bool()[:, None, None, :] + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.reshape(batch_size, -1, channels).type_as(query) + return attn.proj(hidden_states) + + +class MultiHeadCrossAttention(torch.nn.Module, AttentionModuleMixin): + _default_processor_cls = SanaWMCrossAttnProcessor + _available_processors = [SanaWMCrossAttnProcessor] + + def __init__(self, d_model, num_heads, qk_norm=False, processor=None, **block_kwargs): + super().__init__() + if not (d_model % num_heads == 0): + raise ValueError("d_model must be divisible by num_heads") + + self.d_model = d_model + self.heads = num_heads + self.head_dim = d_model // num_heads + self.inner_dim = d_model + + self.q_linear = nn.Linear(d_model, d_model) + self.kv_linear = nn.Linear(d_model, d_model * 2) + self.proj = nn.Linear(d_model, d_model) + if qk_norm: + self.q_norm = RMSNorm(d_model, eps=1e-6) + self.k_norm = RMSNorm(d_model, eps=1e-6) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + self.set_processor(processor if processor is not None else self._default_processor_cls()) + + def forward(self, x, cond, mask=None): + return self.processor(self, x, cond, attention_mask=mask) + + +class T2IFinalLayer(nn.Module): + """ + The final layer of Sana. + """ + + def __init__(self, hidden_size, patch_size, out_channels): + super().__init__() + if isinstance(patch_size, int): + patch_size = [patch_size, patch_size] + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, math.prod(patch_size) * out_channels, bias=True) + self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5) + self.out_channels = out_channels + + def forward_frame_aware(self, x, t): + # t: B,1,F,D + B, N, C = x.shape + num_frames = t.shape[2] + # shift, scale: 2, hidden_size -> 1,1,2,hidden_size -> B,F,2,hidden_size + shift, scale = (self.scale_shift_table[None, None, :, :] + t.transpose(1, 2)).chunk( + 2, dim=-2 + ) # each chunk: B,F,1,D + x = (self.norm_final(x).reshape(B, num_frames, -1, C) * (1 + scale) + shift).reshape(B, N, C) + x = self.linear(x) + return x + + def forward(self, x, t): + if len(t.shape) > 2: + return self.forward_frame_aware(x, t) + shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1) + x = self.norm_final(x) * (1 + scale) + shift + x = self.linear(x) + return x + + +################################################################################# +# Embedding Layers for Timesteps and Class Labels # +################################################################################# +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(self.dtype) + t_emb = self.mlp(t_freq) + return t_emb + + @property + def dtype(self): + # `get_parameter_dtype` is layerwise-casting aware: under layerwise casting the storage dtype + # (e.g. FP8) differs from the compute dtype, and `next(self.parameters()).dtype` returns the former. + return get_parameter_dtype(self) + + +class CaptionEmbedder(nn.Module): + """ + Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance. + """ + + def __init__( + self, + in_channels, + hidden_size, + act_layer=nn.GELU(approximate="tanh"), + token_num=120, + ): + super().__init__() + self.y_proj = Mlp( + in_features=in_channels, hidden_features=hidden_size, out_features=hidden_size, act_layer=act_layer, drop=0 + ) + self.register_buffer("y_embedding", nn.Parameter(torch.randn(token_num, in_channels) / in_channels**0.5)) + + def forward(self, caption): + return self.y_proj(caption) + + +class PatchEmbedMS3D(nn.Module): + """3D Image to Patch Embedding""" + + def __init__( + self, + patch_size=(1, 2, 2), + in_chans=3, + embed_dim=768, + kernel_size=None, + padding=0, + norm_layer=None, + flatten=True, + bias=True, + ): + super().__init__() + kernel_size = tuple(kernel_size or patch_size) + patch_size = tuple(patch_size) + self.kernel_size = kernel_size + self.patch_size = patch_size + self.flatten = flatten + if patch_size[0] != 1: + raise ValueError(f"Patch size for 3D embedding must be (1, *, *), got {patch_size}.") + if not padding and kernel_size[-1] % 2 > 0: + padding = tuple(k // 2 for k in kernel_size) + self.proj = nn.Conv3d( + in_chans, embed_dim, kernel_size=kernel_size, stride=patch_size, padding=padding, bias=bias + ) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCTHW -> BNC + x = self.norm(x) + return x + + +class SanaWMRotaryPosEmbed(nn.Module): + """Rotary position embedding for SANA-WM. + + Deliberately not shared with Wan's rotary embedding: the per-axis split is configurable through `fhw_dim`, and the + frequencies stay complex in a single `freqs` buffer rather than being split into real cos/sin buffers. + """ + + def __init__( + self, + attention_head_dim: int, + patch_size: Tuple[int, int, int], + max_seq_len: int, + theta: float = 10000.0, + fhw_dim: Optional[Tuple[int, int, int]] = None, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + + if fhw_dim is not None: + if not (attention_head_dim == sum(fhw_dim)): + raise ValueError(f"attention_head_dim {attention_head_dim} must match sum(fhw_dim) {sum(fhw_dim)}") + t_dim, h_dim, w_dim = fhw_dim + else: + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + + freqs = [] + for dim in [t_dim, h_dim, w_dim]: + freq = get_1d_rotary_pos_embed( + dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=torch.float32 + ) + freqs.append(freq) + self.register_buffer("freqs", torch.cat(freqs, dim=1), persistent=False) + + def forward(self, fhw: Tuple[int, int, int]) -> torch.Tensor: + ppf, pph, ppw = fhw + + freqs = self.freqs.split_with_sizes( + [ + self.attention_head_dim // 2 - 2 * (self.attention_head_dim // 6), + self.attention_head_dim // 6, + self.attention_head_dim // 6, + ], + dim=1, + ) + + freqs_f = freqs[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_h = freqs[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_w = freqs[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + freqs = torch.cat([freqs_f, freqs_h, freqs_w], dim=-1).reshape(1, 1, ppf * pph * ppw, -1) + return freqs + + +# --------------------------------------------------------------------------- +# UCM (Unified Camera Model) projection / unprojection and per-pixel ray +# transformation (world <-> ray) used by UCPE camera conditioning. +# --------------------------------------------------------------------------- + + +def compute_fov_from_fx_xi( + fx: Union[torch.Tensor, float], + xi: Union[torch.Tensor, float], + width: int, + device="cpu", + dtype=torch.float32, +): + """Inverse of :func:`compute_fx_from_fov_xi`.""" + + def to_tensor_1d(x): + if torch.is_tensor(x): + return x.to(device=device, dtype=dtype) + return torch.tensor([x], dtype=dtype, device=device) + + fx = to_tensor_1d(fx).reshape(-1) + xi = to_tensor_1d(xi).reshape(-1) + B = max(fx.shape[0], xi.shape[0]) + fx = fx.expand(B) + xi = xi.expand(B) + A = 2.0 * fx / width + phi = torch.atan(1.0 / A) + denom = torch.sqrt(A * A + 1.0) + ratio = (xi / denom).clamp(-1.0, 1.0) + theta = torch.asin(ratio) + phi + x_fov = torch.rad2deg(2.0 * theta) + return x_fov + + +def ucm_unproject_grid_fov( + x_fov: Union[float, torch.Tensor], + y_fov: Union[float, torch.Tensor], + xi: Union[float, torch.Tensor], + height: int, + width: int, + cx: Union[float, torch.Tensor], + cy: Union[float, torch.Tensor], + device: Union[torch.device, str] = "cpu", + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Unproject grid with intrinsics expressed as FoV (degrees) + xi.""" + is_batched = any(torch.is_tensor(p) and p.numel() > 1 for p in [x_fov, y_fov, xi, cx, cy]) + fx = compute_fx_from_fov_xi(x_fov, xi, width, device, dtype) + fy = compute_fx_from_fov_xi(y_fov, xi, height, device, dtype) + d_cam = ucm_unproject_grid( + height=height, + width=width, + fx=fx, + fy=fy, + cx=cx, + cy=cy, + xi=xi if torch.is_tensor(xi) else torch.tensor([xi], dtype=dtype, device=device), + dtype=dtype, + device=device, + y_down=True, + ) + if not is_batched: + d_cam = d_cam[0] + return d_cam + + +def world_to_ray_mats( + d_cam: torch.Tensor, # [H, W, 3], [B, H, W, 3], or [B, T, H, W, 3] + c2w: torch.Tensor, # [B, T, 4, 4] +) -> torch.Tensor: + """Build per-pixel ``ray<-world`` transforms from camera unit rays + C2W poses.""" + if d_cam.ndim == 3: + d_cam = d_cam.unsqueeze(0) + if d_cam.ndim == 4: + B, H, W, _ = d_cam.shape + T = c2w.shape[1] + d_cam = d_cam.unsqueeze(1).expand(-1, T, -1, -1, -1) + elif d_cam.ndim == 5: + B, T, H, W, _ = d_cam.shape + else: + raise ValueError(f"Unsupported d_cam shape: {d_cam.shape}") + + device = d_cam.device + dtype = d_cam.dtype + R_cam = c2w[..., :3, :3] + t_cam = c2w[..., :3, 3] + d_world = torch.einsum("btij,bthwj->bthwi", R_cam, d_cam) + cam_y = R_cam[..., :, 1] + # (B, T, 3) -> (B, T, H, W, 3) + cam_y = cam_y[:, :, None, None, :].expand(-1, -1, H, W, -1) + z_ray = F.normalize(d_world, dim=-1, eps=1e-6) + x_ray = torch.cross(cam_y, z_ray, dim=-1) + x_ray = F.normalize(x_ray, dim=-1, eps=1e-6) + y_ray = torch.cross(z_ray, x_ray, dim=-1) + y_ray = F.normalize(y_ray, dim=-1, eps=1e-6) + R_l2w = torch.stack([x_ray, y_ray, z_ray], dim=-1) + # (B, T, H, W, 3, 3) — transpose last two dims for the world->local rotation. + R_w2l = R_l2w.transpose(-1, -2) + # (B, T, 3) -> (B, T, H, W, 3) + t_world = t_cam[:, :, None, None, :].expand(-1, -1, H, W, -1) + t_w2l = -torch.einsum("bthwij,bthwj->bthwi", R_w2l, t_world) + raymats = torch.zeros(B, T, H, W, 4, 4, device=device, dtype=dtype) + raymats[..., :3, :3] = R_w2l + raymats[..., :3, 3] = t_w2l + raymats[..., 3, 3] = 1.0 + mask = torch.isnan(d_world).any(-1) + raymats[mask] = torch.eye(4, device=device, dtype=dtype) + return raymats + + +def create_grid( + height: int, + width: int, + batch: Optional[int] = None, + dtype: torch.dtype = torch.float32, + device: torch.device = torch.device("cpu"), +) -> torch.Tensor: + """Create a pixel coordinate grid of shape ``(H, W, 3)`` or ``(B, H, W, 3)``.""" + if device.type == "cpu": + if dtype not in (torch.float32, torch.float64): + raise ValueError( + f"ERR: {dtype} is not supported by {device.type}\nIf device is `cpu`, use float32 or float64" + ) + _xs = torch.linspace(0, width - 1, width, dtype=dtype, device=device) + _ys = torch.linspace(0, height - 1, height, dtype=dtype, device=device) + ys, xs = torch.meshgrid([_ys, _xs], indexing="ij") + zs = torch.ones_like(xs, dtype=dtype, device=device) + grid = torch.stack((xs, ys, zs), dim=2) + if batch is not None: + # Prepend a batch dim and broadcast. + grid = grid.unsqueeze(0).expand(batch, *grid.shape) + return grid + + +def ucm_unproject_grid( + height: int, + width: int, + fx: Union[float, torch.Tensor], + fy: Union[float, torch.Tensor], + cx: Union[float, torch.Tensor], + cy: Union[float, torch.Tensor], + xi: Union[float, torch.Tensor], + dtype: torch.dtype = torch.float32, + device: torch.device = torch.device("cpu"), + y_down: bool = True, +) -> torch.Tensor: + """Unproject pixel grid into a camera-frame direction vector using the UCM.""" + fx_, fy_, cx_, cy_, xi_ = fx, fy, cx, cy, xi + + def to_tensor_flatten(x): + if torch.is_tensor(x): + return x.to(device=device, dtype=dtype).reshape(-1) + return torch.tensor([x], dtype=dtype, device=device) + + fx, fy, cx, cy, xi = map(to_tensor_flatten, (fx, fy, cx, cy, xi)) + B = max(fx.shape[0], fy.shape[0], cx.shape[0], cy.shape[0], xi.shape[0]) + fx = fx.expand(B) + fy = fy.expand(B) + cx = cx.expand(B) + cy = cy.expand(B) + xi = xi.expand(B) + + grid = create_grid(height=height, width=width, batch=B, dtype=dtype, device=device) + u = grid[..., 0] + v = grid[..., 1] + fx = fx[:, None, None] + fy = fy[:, None, None] + cx = cx[:, None, None] + cy = cy[:, None, None] + xi = xi[:, None, None] + x = (u - cx) / fx + y = (v - cy) / fy + if not y_down: + y = -y + r2 = x * x + y * y + alpha = xi + torch.sqrt(1 + (1 - xi * xi) * r2) + gamma = alpha / (1 + r2) + X = gamma * x + Y = gamma * y + Z = gamma - xi + d_cam = torch.stack([X, Y, Z], dim=-1) + is_scalar_input = all(not torch.is_tensor(p) for p in (fx_, fy_, cx_, cy_, xi_)) + if is_scalar_input: + return d_cam[0] + else: + return d_cam + + +def compute_fx_from_fov_xi( + x_fov: Union[torch.Tensor, float], + xi: Union[torch.Tensor, float], + width: int, + device: Union[torch.device, str] = "cpu", + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Recover focal length ``fx`` from horizontal FoV (degrees) + UCM xi.""" + + def to_tensor_flatten(x): + if torch.is_tensor(x): + return x.to(device=device, dtype=dtype).view(-1) + return torch.tensor([x], dtype=dtype, device=device) + + x_fov = to_tensor_flatten(x_fov) + xi = to_tensor_flatten(xi) + B = max(x_fov.shape[0], xi.shape[0]) + x_fov = x_fov.expand(B) + xi = xi.expand(B) + theta = torch.deg2rad(0.5 * x_fov) + eps = torch.finfo(dtype).eps + denom = torch.sin(theta).clamp_min(eps) + fx = (width * 0.5) * (torch.cos(theta) + xi) / denom + return fx + + +def project_ucm_points(X, Y, Z, fx, fy, cx, cy, xi): + """Project 3D points in camera frame to UCM image plane.""" + r = torch.sqrt(X * X + Y * Y + Z * Z) + + def reshape_param(p, target): + if torch.is_tensor(p): + if p.numel() == 1: + return p + if p.ndim == 1 and target.ndim == 4: + return p.view(target.shape[0], target.shape[1], 1, 1) + while p.ndim < target.ndim: + p = p.unsqueeze(-1) + return p + + xi = reshape_param(xi, X) + fx = reshape_param(fx, X) + fy = reshape_param(fy, X) + cx = reshape_param(cx, X) + cy = reshape_param(cy, X) + + alpha = Z + xi * r + du = fx * (X / alpha) + cx + dv = fy * (Y / alpha) + cy + return du, dv + + +def project_ucm_points_fov(X, Y, Z, x_fov, y_fov, xi, height, width, cx, cy): + """Project 3D points in camera frame to UCM image plane using FoV-based intrinsics.""" + fx = compute_fx_from_fov_xi(x_fov, xi, width, X.device, X.dtype) + fy = compute_fx_from_fov_xi(y_fov, xi, height, X.device, X.dtype) + return project_ucm_points(X, Y, Z, fx, fy, cx, cy, xi) + + +def compute_up_lat_map( + R: torch.Tensor, + x_fov: torch.Tensor, + y_fov: torch.Tensor, + xi: torch.Tensor, + height: int, + width: int, + cx: torch.Tensor, + cy: torch.Tensor, + device: torch.device = torch.device("cpu"), + delta: float = 0.1, +): + """Compute UCPE absolute embedding maps ``(up_map, lat_map)``. + + ``up_map`` is a 2-channel projected up-direction; ``lat_map`` is a 1-channel latitude. Concatenated they form the + 3-channel absmap consumed by the camera branch. + """ + B, T, _, _ = R.shape + dtype = R.dtype + R = R.float() + d_cam = ucm_unproject_grid_fov( + x_fov=x_fov, + y_fov=y_fov, + xi=xi, + height=height, + width=width, + cx=cx, + cy=cy, + device=device, + dtype=torch.float32, + ) + + if d_cam.ndim == 3: + # (H, W, C) -> (B, T, H, W, C) + d_cam_exp = d_cam[None, None].expand(B, T, -1, -1, -1) + elif d_cam.ndim == 4: + if d_cam.shape[0] == B * T: + d_cam_exp = d_cam.view(B, T, height, width, 3) + else: + # (B, H, W, C) -> (B, T, H, W, C) + d_cam_exp = d_cam.unsqueeze(1).expand(-1, T, -1, -1, -1) + else: + d_cam_exp = d_cam + + mask_exp = d_cam_exp.isnan().any(dim=-1, keepdim=True) + d_world = torch.einsum("btij,bthwj->bthwi", R, d_cam_exp) + d_world = d_world / torch.clamp_min(d_world.norm(dim=-1, keepdim=True), 1e-8) + Xw, Yw, Zw = d_world[..., 0], d_world[..., 1], d_world[..., 2] + lat_map = torch.atan2(-Yw, torch.sqrt(Xw**2 + Zw**2)).unsqueeze(-1) + v = d_world + up_world = torch.tensor([0, -1, 0], device=device, dtype=torch.float32) + k = torch.cross(v, up_world.unsqueeze(0).unsqueeze(0).unsqueeze(0).expand_as(v), dim=-1) + k = k / torch.clamp_min(k.norm(dim=-1, keepdim=True), 1e-8) + delta_t = torch.tensor(delta, device=device, dtype=torch.float32) + cos_eps = torch.cos(delta_t) + sin_eps = torch.sin(delta_t) + v_rot = ( + v * cos_eps + torch.cross(k, v, dim=-1) * sin_eps + k * (k * (v * 1).sum(dim=-1, keepdim=True)) * (1 - cos_eps) + ) + dirs_cam = torch.einsum("btij,bthwj->bthwi", R.transpose(-1, -2), v_rot) + Xs, Ys, Zs = dirs_cam[..., 0], dirs_cam[..., 1], dirs_cam[..., 2] + du, dv = project_ucm_points_fov( + Xs, + Ys, + Zs, + x_fov=x_fov.float(), + y_fov=y_fov.float(), + xi=xi.float(), + height=height, + width=width, + cx=cx.float(), + cy=cy.float(), + ) + grid = create_grid( + height=height, + width=width, + batch=B, + dtype=torch.float32, + device=device, + ) + grid_x = grid[..., 0].unsqueeze(1) + grid_y = grid[..., 1].unsqueeze(1) + up_map = torch.stack((du - grid_x, dv - grid_y), dim=-1) + up_map = up_map / torch.clamp_min(up_map.norm(dim=-1, keepdim=True), 1e-8) + up_map = up_map.to(dtype=dtype) + lat_map = lat_map.to(dtype=dtype) + up_map = up_map.masked_fill(mask_exp, 0.0) + lat_map = lat_map.masked_fill(mask_exp, 0.0) + return up_map, lat_map + + +def _process_camera_conditions_ucpe(camera_conditions, B, HW, patch_size): + """Convert ``(B, F, 20)`` camera conditions (C2W flat + fx,fy,cx,cy) into + ``(raymats, absmap)``. + + ``raymats`` is ``(B, F, H, W, 4, 4)`` ``ray<-world`` transforms; ``absmap`` is ``(B, F, H, W, 3)`` (up_map 2-ch + + lat_map 1-ch). + """ + F_dim = camera_conditions.shape[1] + c2w_flat = camera_conditions[..., :16] + C_to_W = c2w_flat.view(B, F_dim, 4, 4) + + fx = camera_conditions[..., 16] + fy = camera_conditions[..., 17] + cx = camera_conditions[..., 18] + cy = camera_conditions[..., 19] + H_dim, W_dim = HW[1], HW[2] + image_width = W_dim * patch_size[2] + image_height = H_dim * patch_size[1] + + # xi is fixed at 0 (pinhole) in this stack. + xi = torch.zeros((B, F_dim), device=camera_conditions.device, dtype=camera_conditions.dtype) + x_fov = compute_fov_from_fx_xi( + fx, xi, image_width, device=camera_conditions.device, dtype=camera_conditions.dtype + ).view(B, F_dim) + y_fov = compute_fov_from_fx_xi( + fy, xi, image_height, device=camera_conditions.device, dtype=camera_conditions.dtype + ).view(B, F_dim) + + d_cam = ucm_unproject_grid_fov( + x_fov, + y_fov, + xi, + H_dim, + W_dim, + cx / patch_size[2], + cy / patch_size[1], + device=camera_conditions.device, + dtype=camera_conditions.dtype, + ) + if d_cam.ndim == 4 and d_cam.shape[0] == B * F_dim: + d_cam = d_cam.view(B, F_dim, H_dim, W_dim, 3) + + raymats = world_to_ray_mats(d_cam, C_to_W) # [B, F, H, W, 4, 4] + + up_map, lat_map = compute_up_lat_map( + R=C_to_W[..., :3, :3], + x_fov=x_fov, + y_fov=y_fov, + xi=xi, + height=image_height, + width=image_width, + cx=cx, + cy=cy, + device=camera_conditions.device, + ) + absmap = torch.cat([up_map, lat_map], dim=-1) # (B, F, H, W, 3) + + return raymats, absmap + + +# --------------------------------------------------------------------------- +# Block-diagonal apply primitives shared by camera and main branches +# --------------------------------------------------------------------------- + + +def _apply_ucpe_transform( + feats: torch.Tensor, + matrix: torch.Tensor, + rotary_emb: Optional[torch.Tensor] = None, + inverse_rope: bool = False, +) -> torch.Tensor: + """Apply the block-diagonal UCPE transform to per-token features. + + The channel axis is split in half: the first half is rotated by the per-token 4x4 ray matrix (applied to channels + grouped by 4), the second half gets complex RoPE. + + Args: + feats (`torch.Tensor`): Features of shape `(batch, heads, seq_len, head_dim)`. + matrix (`torch.Tensor`): Per-token 4x4 transform of shape `(batch, seq_len, 4, 4)`. + rotary_emb (`torch.Tensor`, *optional*): Complex RoPE frequencies; `None` leaves the second half unchanged. + inverse_rope (`bool`, defaults to `False`): Conjugate the frequencies (inverse rotation), used on the output. + + Returns: + `torch.Tensor`: Transformed features with the same shape as `feats`. + """ + batch, num_heads, seq_len, head_dim = feats.shape + half_dim = head_dim // 2 + projected, rotated = feats.split(half_dim, dim=-1) + + matrix_dim = matrix.shape[-1] + projected = torch.einsum( + "bnij,bhnkj->bhnki", + matrix, + projected.reshape(batch, num_heads, seq_len, -1, matrix_dim), + ).reshape(batch, num_heads, seq_len, half_dim) + + if rotary_emb is not None: + rotated_fp32 = rotated.to(torch.float32) + if rotated_fp32.stride(-1) != 1: + rotated_fp32 = rotated_fp32.contiguous() + freqs = rotary_emb.conj() if inverse_rope else rotary_emb + rotated_complex = torch.view_as_complex(rotated_fp32.unflatten(-1, (-1, 2))) + rotated = torch.view_as_real(rotated_complex * freqs).flatten(-2, -1).type_as(rotated) + + return torch.cat([projected, rotated], dim=-1) + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + """Closed-form inverse of a 4x4 SE(3) batch.""" + if not (transforms.shape[-2:] == (4, 4)): + raise ValueError(f"`transforms` must have shape (..., 4, 4), got {tuple(transforms.shape)}.") + Rinv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = Rinv + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3]) + out[..., 3, 3] = 1.0 + return out + + +# --------------------------------------------------------------------------- +# UCPE ray-transform preparation +# --------------------------------------------------------------------------- + + +def _slice_rope_for_cam( + rotary_emb: Optional[torch.Tensor], + head_dim: int, + rope_dim: int, +) -> Optional[torch.Tensor]: + """Re-slice WAN RoPE frequencies for a smaller rope_dim using the same (T, H, W) split.""" + if rotary_emb is None: + return None + orig_t_size = head_dim // 2 - 2 * (head_dim // 6) + orig_h_size = head_dim // 6 + new_t_size = rope_dim // 2 - 2 * (rope_dim // 6) + new_h_size = rope_dim // 6 + new_w_size = rope_dim // 6 + t_part = rotary_emb[..., :new_t_size] + h_part = rotary_emb[..., orig_t_size : orig_t_size + new_h_size] + w_part = rotary_emb[..., orig_t_size + orig_h_size : orig_t_size + orig_h_size + new_w_size] + return torch.cat([t_part, h_part, w_part], dim=-1) + + +def _prepare_ucpe_ray_transforms( + head_dim: int, + camera_conditions: torch.Tensor, + HW: Tuple[int, int, int], + patch_size: Tuple[int, int, int], + rotary_emb: Optional[torch.Tensor] = None, + raymats: Optional[torch.Tensor] = None, + cam_pos_embeds: Optional[dict] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Precompute the UCPE ray matrices once for a batch, shared across all blocks. + + Accepts either precomputed matrices (`cam_pos_embeds` with `P`, `P_inv`, `pos_embeds_cam`) or raw camera conditions + plus optional `raymats`. + + Returns: + `Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]`: `(P, P_T, P_inv, rotary_emb_cam)`, + where `P` is the `ray<-world` transform used on the output and `P_T` / `P_inv` are used on Q and K/V. + """ + batch_size = camera_conditions.shape[0] + + # Priority 1: use precomputed matrices. + if cam_pos_embeds is not None: + P = cam_pos_embeds.get("P") + P_inv = cam_pos_embeds.get("P_inv") + rotary_emb_cam = cam_pos_embeds.get("pos_embeds_cam") + + if P is not None and P_inv is not None: + if P.ndim == 3: + P = P.unsqueeze(0).repeat(batch_size, 1, 1, 1) + if P_inv.ndim == 3: + P_inv = P_inv.unsqueeze(0).repeat(batch_size, 1, 1, 1) + + if rotary_emb_cam is not None and rotary_emb_cam.ndim == 3: + rotary_emb_cam = rotary_emb_cam.unsqueeze(0).repeat(batch_size, 1, 1, 1) + elif rotary_emb_cam is None and rotary_emb is not None: + rotary_emb_cam = _slice_rope_for_cam(rotary_emb, head_dim, head_dim // 2) + elif rotary_emb_cam is None: + rotary_emb_cam = rotary_emb + + return P, P.transpose(-1, -2), P_inv, rotary_emb_cam + + # Priority 2: online path. + if raymats is None: + raymats, _ = _process_camera_conditions_ucpe(camera_conditions, batch_size, HW, patch_size) + P = raymats.reshape(batch_size, -1, 4, 4) + rotary_emb_cam = _slice_rope_for_cam(rotary_emb, head_dim, head_dim // 2) + + return P, P.transpose(-1, -2), _invert_SE3(P), rotary_emb_cam + + +def flip_and_shift(x, dim=2, shift_val=0.0): + """Flip a sequence and shift it right by one step. + + The operation reverses the sequence, drops the last element, and pads the front with ``shift_val``. + + Example: + [x0, x1, x2, x3] -> flip [x3, x2, x1, x0] -> shift [v, x3, x2, x1] + + Args: + x: Input tensor with a time dimension at ``dim``. + dim: Dimension to flip and shift. + shift_val: Value used for the padded step. + + Returns: + Tensor with the same shape as ``x``. + """ + x_flip = torch.flip(x, dims=[dim]) + x_shifted = x_flip.narrow(dim, 0, x.shape[dim] - 1) + pad_shape = list(x.shape) + pad_shape[dim] = 1 + padding = torch.full(pad_shape, shift_val, device=x.device, dtype=x.dtype) + return torch.cat([padding, x_shifted], dim=dim) + + +def torch_chunk_sana_gdn( + q, + k, + v, + q_rot, + k_rot, + beta, + decay, + recall_gate=None, + chunk_size: int | None = 21, + eps: float = 1e-6, + return_components: bool = False, +): + del recall_gate # Accepted so the chunk and fused scan share one signature; unused by this rule. + + B, H, D, N = q.shape + if beta.ndim not in (3, 4): + raise ValueError(f"Expected beta.ndim in (3, 4), got {beta.ndim}.") + T = beta.shape[2] + if T <= 0: + raise ValueError(f"Expected T > 0, got T={T}.") + if N % T != 0: + raise ValueError(f"Expected N divisible by T, got N={N}, T={T}.") + S = N // T + + target_z = 1.0 + scale = 1.0 + + def to_frame_seq(x): + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q, k, v = to_frame_seq(q), to_frame_seq(k), to_frame_seq(v) + q_rot, k_rot = to_frame_seq(q_rot), to_frame_seq(k_rot) + + if beta.ndim == 4: + beta = beta.unsqueeze(3) + else: + beta = beta.view(B, H, T, 1, 1) + + decay = decay.view(B, H, T, 1, 1) + + # ========================================================================= + # 1. PARALLEL PRE-PROCESSING + # ========================================================================= + + I = torch.eye(D, device=q.device, dtype=q.dtype).view(1, 1, 1, D, D) + + # KV State Matrices: W = g * (I - c * K @ K^T) + k_rot_beta = k_rot * beta + W_kv = decay * (I - scale * torch.matmul(k_rot_beta, k_rot.transpose(-1, -2))) + U_kv = torch.matmul(v * beta, k_rot.transpose(-1, -2)) + + # Z State Matrices: W = g * (I - c * K @ K^T) + k_beta = k * beta + W_z = decay * (I - scale * torch.matmul(k_beta, k.transpose(-1, -2))) + U_z = target_z * k_beta.sum(dim=-1, keepdim=True) # Equivalent to Kt @ bt^T over spatial dim + + # ========================================================================= + # 2. CHUNKING LOGIC + # ========================================================================= + + # Uniform chunk boundaries over the temporal axis. A small trailing remainder is absorbed + # into the last chunk, since `causal_conv1d` crashes on length-1 sequences. + boundaries = list(range(0, T, chunk_size)) or [0] + if len(boundaries) > 1 and (T - boundaries[-1]) < chunk_size: + boundaries.pop() + if boundaries[-1] != T: + boundaries.append(T) + split_sizes = [boundaries[i + 1] - boundaries[i] for i in range(len(boundaries) - 1)] + + W_kv_c = W_kv.split(split_sizes, dim=2) + U_kv_c = U_kv.split(split_sizes, dim=2) + W_z_c = W_z.split(split_sizes, dim=2) + U_z_c = U_z.split(split_sizes, dim=2) + + # ========================================================================= + # 3. FAST INTRA-CHUNK SCAN OVER DxD SPACE + # ========================================================================= + + S_kv = torch.zeros(B, H, D, D, device=q.device, dtype=q.dtype) + S_z = torch.zeros(B, H, D, 1, device=q.device, dtype=q.dtype) + + out_S_kv = [] + out_S_z = [] + + def _chunk_scan(w_kv, u_kv, w_z, u_z, s_kv, s_z): + c_len = w_kv.shape[2] + s_kv_list, s_z_list = [], [] + for t in range(c_len): + s_kv = torch.matmul(s_kv, w_kv[:, :, t]) + u_kv[:, :, t] + s_z = torch.matmul(w_z[:, :, t], s_z) + u_z[:, :, t] + s_kv_list.append(s_kv) + s_z_list.append(s_z) + return torch.stack(s_kv_list, dim=2), s_kv, torch.stack(s_z_list, dim=2), s_z + + for i in range(len(split_sizes)): + s_kv_all, S_kv, s_z_all, S_z = _chunk_scan(W_kv_c[i], U_kv_c[i], W_z_c[i], U_z_c[i], S_kv, S_z) + out_S_kv.append(s_kv_all) + out_S_z.append(s_z_all) + + S_kv_all = torch.cat(out_S_kv, dim=2) + S_z_all = torch.cat(out_S_z, dim=2) + + # ========================================================================= + # 4. PARALLEL OUTPUT PROJECTION + # ========================================================================= + + out_num = torch.matmul(S_kv_all, q_rot) + out_den = torch.matmul(S_z_all.transpose(-1, -2), q) + + final_num = out_num.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + final_den = out_den.permute(0, 1, 3, 2, 4).reshape(B, H, 1, N) + + if return_components: + return final_num, final_den + + return final_num / (final_den + eps) + + +# --------------------------------------------------------------------------- +# Helpers for hot-path operations +# --------------------------------------------------------------------------- + + +def _apply_rotary_emb( + hidden_states: torch.Tensor, + freqs: torch.Tensor, +) -> torch.Tensor: + """Apply rotary embeddings to `(batch, heads, dim, seq_len)` features.""" + x_rotated = torch.view_as_complex( + hidden_states.permute(0, 1, 3, 2).to(torch.float32).unflatten(3, (-1, 2)), + ) + x_out = torch.view_as_real(x_rotated * freqs).flatten(3, 4).permute(0, 1, 3, 2) + return x_out.type_as(hidden_states) + + +def torch_chunk_cam_single_path_delta_rule( + q_rot: torch.Tensor, + k_rot: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + chunk_size: int | None = 21, +) -> torch.Tensor: + """Parallel chunk-scan version of the single-path delta-rule recurrence. + + Restructured as a linear recurrence in D x D state space so that Phases 1 (transition-matrix construction) and 3 + (output projection) are fully parallel over T, while Phase 2 (the D x D state scan) is chunked. + + The recurrence: + state[t] = state[t-1] * g[t] + delta_v[t] @ k_rot[t]^T + where delta_v[t] = (v[t] - state[t-1]*g[t] @ k_rot[t]) * beta[t] + + is equivalent to: + state[t] = state[t-1] @ W[t] + U[t] + with: + W[t] = g[t] * (I - beta[t] * k_rot[t] @ k_rot[t]^T) U[t] = beta[t] * v[t] @ k_rot[t]^T + """ + B, H, D, N = q_rot.shape + if beta.ndim not in (3, 4): + raise ValueError(f"Expected beta.ndim in (3, 4), got {beta.ndim}.") + T = beta.shape[2] + if T <= 0: + raise ValueError(f"Expected T > 0, got T={T}.") + if N % T != 0: + raise ValueError(f"Expected N divisible by T, got N={N}, T={T}.") + S = N // T + + def to_frame_seq(x: torch.Tensor) -> torch.Tensor: + return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4) + + q_rot = to_frame_seq(q_rot) + k_rot = to_frame_seq(k_rot) + v = to_frame_seq(v) + + if beta.ndim == 4: + beta = beta.unsqueeze(3) + else: + beta = beta.view(B, H, T, 1, 1) + decay = decay.view(B, H, T, 1, 1) + + # ========================================================================= + # Phase 1: PARALLEL PRE-PROCESSING (fully parallel over T) + # ========================================================================= + I = torch.eye(D, device=q_rot.device, dtype=q_rot.dtype).view(1, 1, 1, D, D) + + k_rot_beta = k_rot * beta + W_kv = decay * (I - torch.matmul(k_rot_beta, k_rot.transpose(-1, -2))) + U_kv = torch.matmul(v * beta, k_rot.transpose(-1, -2)) + + # ========================================================================= + # Phase 2: CHUNKED SCAN over D x D state space + # ========================================================================= + # Uniform chunk boundaries over the temporal axis. A small trailing remainder is absorbed + # into the last chunk, since `causal_conv1d` crashes on length-1 sequences. + boundaries = list(range(0, T, chunk_size)) or [0] + if len(boundaries) > 1 and (T - boundaries[-1]) < chunk_size: + boundaries.pop() + if boundaries[-1] != T: + boundaries.append(T) + split_sizes = [boundaries[i + 1] - boundaries[i] for i in range(len(boundaries) - 1)] + + W_kv_c = W_kv.split(split_sizes, dim=2) + U_kv_c = U_kv.split(split_sizes, dim=2) + + S_kv = torch.zeros(B, H, D, D, device=q_rot.device, dtype=q_rot.dtype) + out_S_kv: list[torch.Tensor] = [] + + def _chunk_scan_kv( + w_kv: torch.Tensor, u_kv: torch.Tensor, s_kv: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + c_len = w_kv.shape[2] + s_kv_list: list[torch.Tensor] = [] + for t in range(c_len): + s_kv = torch.matmul(s_kv, w_kv[:, :, t]) + u_kv[:, :, t] + s_kv_list.append(s_kv) + return torch.stack(s_kv_list, dim=2), s_kv + + for i in range(len(split_sizes)): + s_kv_all, S_kv = _chunk_scan_kv(W_kv_c[i], U_kv_c[i], S_kv) + out_S_kv.append(s_kv_all) + + S_kv_all = torch.cat(out_S_kv, dim=2) + + # ========================================================================= + # Phase 3: PARALLEL OUTPUT PROJECTION (no denominator) + # ========================================================================= + out = torch.matmul(S_kv_all, q_rot) # (B, H, T, D, S) + + return out.permute(0, 1, 3, 2, 4).reshape(B, H, D, N) + + +def _prepare_frame_valid_masks( + frame_valid_mask: torch.Tensor | None, + *, + batch_size: int, + num_frames: int, + spatial_size: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Convert a frame-valid mask to the token / beta / decay masks the attention branches use. + + Args: + frame_valid_mask (`torch.Tensor`, *optional*): + Per-frame validity mask shaped `(B, 1, T, 1, 1)`, `(B, 1, T)` or `(B, T)`. `None` disables all masking. + batch_size (`int`): Batch size `B`. + num_frames (`int`): Number of frames `T`. + spatial_size (`int`): Tokens per frame `S`. + device (`torch.device`): Device of the returned masks. + dtype (`torch.dtype`): Dtype of the returned masks. + + Returns: + `tuple`: `(token_valid_mask, beta_valid_mask, decay_valid_mask)` shaped `(B, T * S)`, `(B, 1, T, 1)` and `(B, + 1, T)`, or three `None` when `frame_valid_mask` is `None`. + """ + if frame_valid_mask is None: + return None, None, None + + mask = frame_valid_mask + if mask.ndim == 5: + # (B, 1, T, 1, 1) + mask = mask[:, 0, :, 0, 0] + elif mask.ndim == 3 and mask.shape[1] == 1: + # (B, 1, T) + mask = mask[:, 0, :] + elif mask.ndim != 2: + raise ValueError( + "frame_valid_mask must be shaped (B, 1, T, 1, 1), (B, 1, T), or (B, T); " + f"got shape={list(frame_valid_mask.shape)}" + ) + + if mask.shape[0] != batch_size or mask.shape[1] != num_frames: + raise ValueError( + f"frame_valid_mask shape mismatch: expected (B={batch_size}, T={num_frames}), got {list(mask.shape)}" + ) + + mask = mask.to(device=device, dtype=dtype) + token_valid_mask = mask[:, :, None].expand(batch_size, num_frames, spatial_size).reshape(batch_size, -1) + beta_valid_mask = mask.view(batch_size, 1, num_frames, 1) + decay_valid_mask = mask.view(batch_size, 1, num_frames) + return token_valid_mask, beta_valid_mask, decay_valid_mask + + +def _downscale_to_reference_rms( + reference: torch.Tensor, + transformed: torch.Tensor, + eps: float = 1e-6, +) -> torch.Tensor: + """Downscale a UCPE-transformed tensor if its channel RMS exceeds the reference. + + Args: + reference (`torch.Tensor`): Pre-UCPE tensor carrying the target magnitude, shaped `(B, H, D, N)`. + transformed (`torch.Tensor`): Tensor to stabilize, same shape as `reference`. + eps (`float`, defaults to 1e-6): Numerical epsilon of the RMS. + + Returns: + `torch.Tensor`: Stabilized tensor whose per-`(B, H, N)` channel RMS is not larger than the reference's. + """ + reference_rms = reference.square().mean(dim=2, keepdim=True).add(eps).sqrt() + transformed_rms = transformed.square().mean(dim=2, keepdim=True).add(eps).sqrt() + scale = (reference_rms / transformed_rms.clamp_min(eps)).clamp(max=1.0) + return transformed * scale + + +class SanaWMBidirectionalGDNAttention(nn.Module): + """Bidirectional gated-delta-net linear attention over the temporal axis. + + The delta rule runs twice -- forwards (frames ``1..t``) and backwards (frames ``t+1..T``) -- and the numerator / + denominator streams of both passes are summed before the final normalization. RoPE is applied to the numerator + stream only, so the denominator (``Z``) stream keeps unrotated queries/keys and mass is conserved. + + This module holds no parameters. The projections, short convolutions, norms and gates feeding it live on + [`BidirectionalGDNUCPESinglePathLiteLA`], whose `forward` issues every layer call and passes the resulting tensors + in. + + Args: + eps (`float`, defaults to 1e-15): Denominator epsilon of the linear-attention normalization. + chunk_size (`int`, defaults to 21): Temporal chunk length of the state scan. + """ + + def __init__(self, eps: float = 1e-15, chunk_size: int = 21) -> None: + super().__init__() + self.eps = eps + self.chunk_size = chunk_size + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + recall_gate: torch.Tensor, + num_frames: int, + rotary_emb: torch.Tensor | None = None, + token_valid_mask: torch.Tensor | None = None, + beta_valid_mask: torch.Tensor | None = None, + decay_valid_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run the bidirectional delta rule. + + Args: + query (`torch.Tensor`): Queries of shape `(B, N, H, D)`, already normalized and passed through the kernel. + key (`torch.Tensor`): Keys of shape `(B, N, H, D)`, same preprocessing as `query`. + value (`torch.Tensor`): Values of shape `(B, N, H, D)`. + beta (`torch.Tensor`): Per-frame delta-rule gate of shape `(B, H, T, S)`. + decay (`torch.Tensor`): Per-frame decay gate of shape `(B, H, T)`. + recall_gate (`torch.Tensor`): Recall gate buffer forwarded to the scan. + num_frames (`int`): Number of frames `T` the sequence axis is split into. + rotary_emb (`torch.Tensor`, *optional*): Rotary embeddings applied to the numerator stream. + token_valid_mask (`torch.Tensor`, *optional*): Token mask of shape `(B, N)`. + beta_valid_mask (`torch.Tensor`, *optional*): Frame mask of shape `(B, 1, T, 1)` applied to `beta`. + decay_valid_mask (`torch.Tensor`, *optional*): Frame mask of shape `(B, 1, T)` applied to `decay`. + + Returns: + `torch.Tensor`: Raw attention output of shape `(B, N, H * D)`; the shared output gate and projection are + applied by the caller. + """ + batch_size, seq_len, num_heads, head_dim = query.shape + spatial_size = seq_len // num_frames + dtype_orig = value.dtype + + key = key * ((head_dim**-0.5) * (spatial_size**-0.5)) + + # Permute to (B, H, D, N) for processing. + query = query.permute(0, 2, 3, 1) + key = key.permute(0, 2, 3, 1) + value = value.permute(0, 2, 3, 1) + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(batch_size, 1, 1, seq_len) + query = query * token_mask_qkv + key = key * token_mask_qkv + value = value * token_mask_qkv + + # RoPE preparation (numerator only). + if rotary_emb is not None: + query_rot = _apply_rotary_emb(query, rotary_emb) + key_rot = _apply_rotary_emb(key, rotary_emb) + else: + query_rot = query + key_rot = key + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(batch_size, 1, 1, seq_len) + query_rot = query_rot * token_mask_qkv + key_rot = key_rot * token_mask_qkv + + if beta_valid_mask is not None: + beta = beta * beta_valid_mask.to(beta.dtype) + if decay_valid_mask is not None: + decay_mask = decay_valid_mask.to(decay.dtype) + decay = decay * decay_mask + (1.0 - decay_mask) + + # Force FP32 to preserve recurrent stability. + query = query.float() + key = key.float() + value = value.float() + query_rot = query_rot.float() + key_rot = key_rot.float() + beta = beta.float() + decay = decay.float() + recall_gate = recall_gate.float() + + # Forward pass (inclusive: 1..t). + num_fwd, den_fwd = torch_chunk_sana_gdn( + query, + key, + value, + query_rot, + key_rot, + beta, + decay, + recall_gate=recall_gate, + chunk_size=self.chunk_size, + eps=self.eps, + return_components=True, + ) + + # Backward pass (exclusive: t+1..T). + def to_time_structure(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(batch_size, num_heads, head_dim, num_frames, spatial_size).permute(0, 1, 3, 2, 4) + + def from_time_structure(tensor: torch.Tensor) -> torch.Tensor: + return tensor.permute(0, 1, 3, 2, 4).reshape(batch_size, num_heads, head_dim, seq_len) + + query_bwd = torch.flip(to_time_structure(query), dims=[2]) + query_rot_bwd = torch.flip(to_time_structure(query_rot), dims=[2]) + key_bwd = flip_and_shift(to_time_structure(key), dim=2, shift_val=0.0) + value_bwd = flip_and_shift(to_time_structure(value), dim=2, shift_val=0.0) + key_rot_bwd = flip_and_shift(to_time_structure(key_rot), dim=2, shift_val=0.0) + beta_bwd = flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = flip_and_shift(decay, dim=2, shift_val=1.0) + + num_bwd_flipped, den_bwd_flipped = torch_chunk_sana_gdn( + from_time_structure(query_bwd), + from_time_structure(key_bwd), + from_time_structure(value_bwd), + from_time_structure(query_rot_bwd), + from_time_structure(key_rot_bwd), + beta_bwd, + decay_bwd, + recall_gate=recall_gate, + chunk_size=self.chunk_size, + eps=self.eps, + return_components=True, + ) + + def flip_back(tensor: torch.Tensor) -> torch.Tensor: + # The denominator stream carries a single channel, hence the runtime `dim` lookup. + dim = tensor.shape[2] + tensor = tensor.view(batch_size, num_heads, dim, num_frames, spatial_size) + return torch.flip(tensor, dims=[3]).reshape(batch_size, num_heads, dim, seq_len) + + total_num = num_fwd + flip_back(num_bwd_flipped) + total_den = den_fwd + flip_back(den_bwd_flipped) + + hidden_states = total_num / (total_den + self.eps) + + if dtype_orig != torch.float32: + hidden_states = hidden_states.to(dtype_orig) + + hidden_states = hidden_states.permute(0, 3, 1, 2).reshape(batch_size, seq_len, num_heads * head_dim) + if token_valid_mask is not None: + hidden_states = hidden_states * token_valid_mask.view(batch_size, seq_len, 1).to(hidden_states.dtype) + return hidden_states + + +class SanaWMBidirectionalGDNCamAttention(nn.Module): + """Camera-control counterpart of [`SanaWMBidirectionalGDNAttention`]. + + The recurrence is the same bidirectional delta rule, but the queries/keys/values are positionally encoded with the + UCPE per-ray transforms instead of RoPE, and the rule is reduced to its numerator ("single path") stream. The + transformed tensors are downscaled back to their pre-UCPE RMS envelope, and the energy the transform still adds is + discounted from ``beta``. + + This module holds no parameters; [`BidirectionalGDNUCPESinglePathLiteLA.forward`] issues every layer call and + passes the resulting tensors in. + + Args: + chunk_size (`int`, defaults to 21): Temporal chunk length of the state scan. + """ + + def __init__(self, chunk_size: int = 21) -> None: + super().__init__() + self.chunk_size = chunk_size + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ray_transforms: tuple, + beta: torch.Tensor, + decay: torch.Tensor, + num_frames: int, + token_valid_mask: torch.Tensor | None = None, + beta_valid_mask: torch.Tensor | None = None, + decay_valid_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run the bidirectional single-path delta rule in UCPE ray space. + + Args: + query (`torch.Tensor`): Camera queries of shape `(B, N, H, D)`, already normalized and passed through the + kernel. + key (`torch.Tensor`): Camera keys of shape `(B, N, H, D)`, same preprocessing as `query`. + value (`torch.Tensor`): Camera values of shape `(B, N, H, D)`. + ray_transforms (`tuple`): `(P, P_T, P_inv, rotary_emb_cam)` UCPE transforms; `P_T` encodes the queries, + `P_inv` the keys/values and `P` decodes the output. + beta (`torch.Tensor`): Per-frame delta-rule gate of shape `(B, H, T, S)`. + decay (`torch.Tensor`): Per-frame decay gate of shape `(B, H, T)`. + num_frames (`int`): Number of frames `T` the sequence axis is split into. + token_valid_mask (`torch.Tensor`, *optional*): Token mask of shape `(B, N)`. + beta_valid_mask (`torch.Tensor`, *optional*): Frame mask of shape `(B, 1, T, 1)` applied to `beta`. + decay_valid_mask (`torch.Tensor`, *optional*): Frame mask of shape `(B, 1, T)` applied to `decay`. + + Returns: + `torch.Tensor`: Raw camera attention output of shape `(B, N, H * D)`, in world space; the caller projects + it back into the residual stream. + """ + batch_size, seq_len, num_heads, head_dim = query.shape + spatial_size = seq_len // num_frames + dtype_orig = value.dtype + + key = key * ((head_dim**-0.5) * (spatial_size**-0.5)) + + # Permute to (B, H, D, N) for processing. + query = query.permute(0, 2, 3, 1).contiguous() + key = key.permute(0, 2, 3, 1).contiguous() + value = value.permute(0, 2, 3, 1).contiguous() + + # Measure the safe geometric norm before UCPE applies translations. + pre_ucpe_key_norm = torch.linalg.vector_norm(key, dim=2, keepdim=True).clamp_min(1e-6) + + # UCPE expects (B, h, N, d); our tensors are (B, h, d, N). Avoid eager contiguous copies before the + # transforms, and fuse the K/V transform (both use P_inv) into one call, then split back. + P, P_T, P_inv, rotary_emb_cam = ray_transforms + query_ucpe = _apply_ucpe_transform(query.transpose(-1, -2), P_T, rotary_emb_cam).transpose(-1, -2).contiguous() + key_value = torch.cat([key, value], dim=1) + key_value_ucpe = ( + _apply_ucpe_transform(key_value.transpose(-1, -2), P_inv, rotary_emb_cam).transpose(-1, -2).contiguous() + ) + key_ucpe, value_ucpe = torch.chunk(key_value_ucpe, chunks=2, dim=1) + + # Downscale the transformed tensors back to their pre-UCPE RMS envelope. + query_ucpe = _downscale_to_reference_rms(query, query_ucpe) + key_ucpe = _downscale_to_reference_rms(key, key_ucpe) + value_ucpe = _downscale_to_reference_rms(value, value_ucpe) + + # Measure the inflated geometric norm after UCPE, for the beta discount below. + post_ucpe_key_norm = torch.linalg.vector_norm(key_ucpe, dim=2, keepdim=True).clamp_min(1e-6) + inflation_sq = (post_ucpe_key_norm / pre_ucpe_key_norm) ** 2 + + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(batch_size, 1, 1, seq_len) + value_ucpe = value_ucpe * token_mask_qkv + query_ucpe = query_ucpe * token_mask_qkv + key_ucpe = key_ucpe * token_mask_qkv + + # Dynamic beta discounting: scale beta by the UCPE inflation factor. + frame_inflation_sq = inflation_sq.view(batch_size, num_heads, num_frames, spatial_size).mean(dim=-1) + if beta.ndim == 3: + beta = beta / frame_inflation_sq.clamp_min(1.0) + elif beta.ndim == 4: + beta = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0) + + if beta_valid_mask is not None: + beta = beta * beta_valid_mask.to(beta.dtype) + if decay_valid_mask is not None: + decay_mask = decay_valid_mask.to(decay.dtype) + decay = decay * decay_mask + (1.0 - decay_mask) + + # Forward pass (inclusive: 1..t). Force FP32 to preserve recurrent stability. + out_fwd = torch_chunk_cam_single_path_delta_rule( + query_ucpe.float(), + key_ucpe.float(), + value_ucpe.float(), + beta.float(), + decay.float(), + chunk_size=self.chunk_size, + ) + + # Backward pass (exclusive: t+1..T). + def to_time_structure(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(batch_size, num_heads, head_dim, num_frames, spatial_size).permute(0, 1, 3, 2, 4) + + def from_time_structure(tensor: torch.Tensor) -> torch.Tensor: + return tensor.permute(0, 1, 3, 2, 4).reshape(batch_size, num_heads, head_dim, seq_len) + + query_bwd = torch.flip(to_time_structure(query_ucpe), dims=[2]) + key_bwd = flip_and_shift(to_time_structure(key_ucpe), dim=2, shift_val=0.0) + value_bwd = flip_and_shift(to_time_structure(value_ucpe), dim=2, shift_val=0.0) + beta_bwd = flip_and_shift(beta, dim=2, shift_val=0.0) + decay_bwd = flip_and_shift(decay, dim=2, shift_val=1.0) + + out_bwd_flipped = torch_chunk_cam_single_path_delta_rule( + from_time_structure(query_bwd).float(), + from_time_structure(key_bwd).float(), + from_time_structure(value_bwd).float(), + beta_bwd.float(), + decay_bwd.float(), + chunk_size=self.chunk_size, + ) + out_bwd = torch.flip( + out_bwd_flipped.view(batch_size, num_heads, head_dim, num_frames, spatial_size), + dims=[3], + ).reshape(batch_size, num_heads, head_dim, seq_len) + + hidden_states = out_fwd + out_bwd + + if dtype_orig != torch.float32: + hidden_states = hidden_states.to(dtype_orig) + if token_valid_mask is not None: + hidden_states = hidden_states * token_valid_mask.view(batch_size, 1, 1, seq_len).to(hidden_states.dtype) + + # Decode back from ray space to world space. + hidden_states = ( + _apply_ucpe_transform(hidden_states.transpose(-1, -2), P, rotary_emb_cam, inverse_rope=True) + .transpose(-1, -2) + .contiguous() + ) + hidden_states = hidden_states.reshape(batch_size, num_heads * head_dim, seq_len).permute(0, 2, 1) + if token_valid_mask is not None: + hidden_states = hidden_states * token_valid_mask.view(batch_size, seq_len, 1).to(hidden_states.dtype) + return hidden_states + + +class BidirectionalGDNUCPESinglePathLiteLA(nn.Module): + """Bidirectional Gated-Delta-Net attention with a UCPE camera-control branch. + + This is the attention block used by every non-softmax layer of the released SANA-WM checkpoint. Two branches run + over the same tokens and are summed before a single shared output gate + projection: + + - **Main branch** ([`SanaWMBidirectionalGDNAttention`]) -- bidirectional linear attention with a gated delta rule + over the temporal axis. A ReLU kernel is applied to Q/K, RoPE is applied to the numerator stream only, and the + denominator (Z) stream keeps unrotated Q/K so mass is conserved. The gates (``beta`` / ``decay``) are computed + per frame and shared spatially, while the states are maintained per pixel. + - **Camera branch** ([`SanaWMBidirectionalGDNCamAttention`]) -- the same bidirectional recurrence, but positionally + encoded with UCPE per-ray transforms instead of RoPE and reduced to a numerator-only ("single path") delta rule. + The transformed camera tensors are downscaled back to their pre-UCPE RMS envelope before entering the recurrence. + + Both branch modules are parameter-free: every layer call happens in this class's `forward`, which owns the + projections, short convolutions, norms and gates the two branches consume. + + Camera-specific parameters: ``q_proj_cam``, ``k_proj_cam``, ``v_proj_cam``, ``out_proj_cam``, ``q_norm_cam``, + ``k_norm_cam`` and ``conv_k_cam``. The GDN gates (``beta_proj`` / ``gate_proj`` / ``dt_bias`` / ``A_log`` / + ``recall_gate``), the output gate and the output projection are shared by both branches. + + Args: + in_dim (`int`): Input channels. + out_dim (`int`): Output channels. + cam_dim (`int`): Camera-branch width; must equal `in_dim` so the shared parameters line up. + cam_heads (`int`): Camera-branch heads; must equal `heads` and divide `cam_dim` into multiples of 4. + patch_size (`tuple[int, int, int]`, defaults to `(1, 2, 2)`): Latent patch size, used to map camera + intrinsics onto the token grid. + heads (`int`, *optional*): Number of attention heads; derived from `out_dim // dim * heads_ratio` when `None`. + heads_ratio (`float`, defaults to 1.0): Head-count multiplier used when `heads` is `None`. + dim (`int`, defaults to 32): Head dimension used when `heads` is `None`. + eps (`float`, defaults to 1e-15): Denominator epsilon of the linear-attention normalization. + use_bias (`bool`, defaults to `False`): Whether the fused QKV projection has a bias. + qk_norm (`bool`, defaults to `False`): Apply RMSNorm to Q/K. + norm_eps (`float`, defaults to 1e-5): Epsilon of the Q/K RMSNorm. + use_output_gate (`bool`, defaults to `True`): Apply the shared silu output gate. + chunk_gdn_chunk_size (`int`, defaults to 21): Temporal chunk length of the state scan. + conv_kernel_size (`int`, defaults to 4): Temporal short-convolution width; `0` disables the convolutions. + k_conv_only (`bool`, defaults to `True`): Apply the short convolution to K only. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + *, + cam_dim: int, + cam_heads: int, + patch_size: tuple[int, int, int] = (1, 2, 2), + heads: int | None = None, + heads_ratio: float = 1.0, + dim: int = 32, + eps: float = 1e-15, + use_bias: bool = False, + qk_norm: bool = False, + norm_eps: float = 1e-5, + use_output_gate: bool = True, + chunk_gdn_chunk_size: int = 21, + conv_kernel_size: int = 4, + k_conv_only: bool = True, + **kwargs: object, + ) -> None: + heads = heads or int(out_dim // dim * heads_ratio) + super().__init__() + + # Fused QKV projection and output projection (the `q_norm` / `k_norm` + # attributes are set further down, depending on `qk_norm`). + self.num_heads = heads + self.qkv = nn.Linear(in_dim, in_dim * 3, bias=use_bias) + self.proj = nn.Linear(in_dim, in_dim) + + self.in_dim = in_dim + self.out_dim = out_dim + self.heads = heads + self.dim = out_dim // heads + self.eps = eps + self.k_conv_only = k_conv_only + + self.kernel_func = nn.ReLU(inplace=False) + + if qk_norm: + self.q_norm = RMSNorm(self.in_dim, eps=norm_eps) + self.k_norm = RMSNorm(self.in_dim, eps=norm_eps) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + # Gate projections operate on pooled frame features (B, T, D) -> (B, T, H). + self.beta_proj = nn.Linear(in_dim, heads, bias=True) + self.gate_proj = nn.Linear(in_dim, heads, bias=True) + + A = torch.zeros(self.heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.heads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + + # `recall_gate` is unused by the forward; kept as a buffer for checkpoint compatibility. + self.register_buffer("recall_gate", torch.zeros(1)) + + self.use_output_gate = use_output_gate + if use_output_gate: + self.output_gate = nn.Linear(in_dim, out_dim, bias=True) + else: + self.output_gate = None + + self.chunk_gdn_chunk_size = chunk_gdn_chunk_size + + # Short convolutions (depthwise Conv1d along T). + self.conv_kernel_size = conv_kernel_size + if conv_kernel_size > 0: + self.conv_k = SanaWMTemporalShortConvolution( + hidden_size=out_dim, + kernel_size=conv_kernel_size, + ) + if k_conv_only: + self.conv_q = None + self.conv_v = None + else: + self.conv_q = SanaWMTemporalShortConvolution( + hidden_size=out_dim, + kernel_size=conv_kernel_size, + ) + self.conv_v = SanaWMTemporalShortConvolution( + hidden_size=out_dim, + kernel_size=conv_kernel_size, + ) + else: + self.conv_q = None + self.conv_k = None + self.conv_v = None + + self.patch_size = patch_size + self.cam_dim = cam_dim + self.cam_heads = cam_heads + self.cam_head_dim = cam_dim // cam_heads + + if cam_dim != in_dim: + raise ValueError(f"Parameter sharing requires cam_dim == in_dim, got cam_dim={cam_dim}, in_dim={in_dim}.") + if cam_heads != self.heads: + raise ValueError( + f"Parameter sharing requires cam_heads == heads, got cam_heads={cam_heads}, heads={self.heads}." + ) + if self.cam_head_dim % 4 != 0: + raise ValueError( + "UCPE camera branch requires cam_head_dim divisible by 4, " + f"got {self.cam_head_dim} (cam_dim={cam_dim}, cam_heads={cam_heads})." + ) + + # ---- Camera-specific: QKV + output projections only ---- + self.q_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.k_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.v_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.out_proj_cam = nn.Linear(cam_dim, out_dim, bias=True) + + # Keep branch-specific Q/K norms so camera statistics do not disturb the + # main branch (and vice versa). Start from identical weights. + self.q_norm_cam = deepcopy(self.q_norm) + self.k_norm_cam = deepcopy(self.k_norm) + + # Short convolutions for the camera branch (matching the main branch). + if self.conv_kernel_size > 0: + self.conv_k_cam = SanaWMTemporalShortConvolution( + hidden_size=cam_dim, + kernel_size=self.conv_kernel_size, + ) + if self.k_conv_only: + self.conv_q_cam = None + self.conv_v_cam = None + else: + self.conv_q_cam = SanaWMTemporalShortConvolution( + hidden_size=cam_dim, + kernel_size=self.conv_kernel_size, + ) + self.conv_v_cam = SanaWMTemporalShortConvolution( + hidden_size=cam_dim, + kernel_size=self.conv_kernel_size, + ) + else: + self.conv_q_cam = None + self.conv_k_cam = None + self.conv_v_cam = None + + # Branch compute modules. They own no parameters -- so the checkpoint layout is untouched -- and only carry + # the recurrence configuration; `forward` runs every layer call and hands them the resulting tensors. + self.attn = SanaWMBidirectionalGDNAttention(eps=eps, chunk_size=chunk_gdn_chunk_size) + self.cam_attn = SanaWMBidirectionalGDNCamAttention(chunk_size=chunk_gdn_chunk_size) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + HW: tuple[int, int, int] | None = None, + rotary_emb: torch.Tensor | None = None, + block_mask: torch.Tensor | None = None, + camera_conditions: torch.Tensor | None = None, + chunk_size: int | None = None, + *, + frame_valid_mask: torch.Tensor | None = None, + ucpe_ray_transforms: tuple | None = None, + ) -> torch.Tensor: + """Dual-branch forward: bidirectional GDN main branch + UCPE camera branch. + + Flow: + 1. attn_output = GDN attention (no gate/proj) + 2. cam_output = GDN+UCPE attention (no gate/proj) + 3. combined = attn_output + out_proj_cam(cam_output) [zero at init] + 4. output = proj(output_gate(combined)) [shared, once] + + Args: + x: Input tensor of shape ``(B, N, C)``. + mask: Unused attention mask (kept for API compatibility). + HW: Tuple of ``(T, H, W)`` describing the token layout. + rotary_emb: Optional rotary embeddings for q/k. + block_mask: Unused block mask (kept for API compatibility). + camera_conditions: Raw ``(B, T, 20)`` camera conditions enabling the camera branch. + chunk_size: Unused chunk length (kept for API compatibility). + frame_valid_mask: Optional per-frame validity mask used to zero out padded frames, shaped + ``(B, 1, T, 1, 1)``, ``(B, 1, T)`` or ``(B, T)``. + ucpe_ray_transforms: Optional pre-computed UCPE transforms shared across blocks. + + Returns: + Tensor of shape ``(B, N, C)`` after attention and projection. + """ + del mask, block_mask, chunk_size + + if HW is None: + raise ValueError("HW (T, H, W) must be provided for GDN attention.") + + batch_size, seq_len, channels = x.shape + num_frames, height, width = HW + spatial_size = height * width + + token_valid_mask, beta_valid_mask, decay_valid_mask = _prepare_frame_valid_masks( + frame_valid_mask, + batch_size=batch_size, + num_frames=num_frames, + spatial_size=spatial_size, + device=x.device, + dtype=x.dtype, + ) + + # Per-frame beta / decay gates, computed once from the unmasked input and shared by both branches. `beta` is + # broadcast over the spatial axis; `decay` is a per-frame scalar per head. + beta = ( + self.beta_proj(x).sigmoid().reshape(batch_size, num_frames, spatial_size, self.heads).permute(0, 3, 1, 2) + ) + frame_hidden_states = x.reshape(batch_size, num_frames, spatial_size, channels).mean(dim=2) + gate = self.gate_proj(frame_hidden_states).float() + dt = self.dt_bias.float().view(1, 1, -1) + decay_rate = self.A_log.float().exp().view(1, 1, -1) + decay = (-decay_rate * F.softplus(gate + dt)).exp().transpose(1, 2) + + # ---- Main branch: fused QKV -> short conv -> Q/K norm -> ReLU kernel -> bidirectional GDN ---- + hidden_states = x + if token_valid_mask is not None: + hidden_states = hidden_states * token_valid_mask.view(batch_size, seq_len, 1) + + query, key, value = self.qkv(hidden_states).reshape(batch_size, seq_len, 3, self.heads, self.dim).unbind(2) + if token_valid_mask is not None: + token_mask = token_valid_mask.view(batch_size, seq_len, 1, 1) + query = query * token_mask + key = key * token_mask + value = value * token_mask + + # Short convolution along T (before norm / kernel activation). + if self.conv_q is not None: + query = self.conv_q(query.reshape(batch_size, seq_len, channels), num_frames).reshape( + batch_size, seq_len, self.heads, self.dim + ) + if self.conv_k is not None: + key = self.conv_k(key.reshape(batch_size, seq_len, channels), num_frames).reshape( + batch_size, seq_len, self.heads, self.dim + ) + if self.conv_v is not None: + value = self.conv_v(value.reshape(batch_size, seq_len, channels), num_frames).reshape( + batch_size, seq_len, self.heads, self.dim + ) + + # Q/K norm runs on the flattened channels (B, N, C), then the tensors go back to (B, N, H, D). + query = self.q_norm(query.reshape(batch_size, seq_len, channels)).reshape( + batch_size, seq_len, self.heads, self.dim + ) + key = self.k_norm(key.reshape(batch_size, seq_len, channels)).reshape( + batch_size, seq_len, self.heads, self.dim + ) + query = self.kernel_func(query) + key = self.kernel_func(key) + + attn_output = self.attn( + query, + key, + value, + beta, + decay, + self.recall_gate, + num_frames, + rotary_emb=rotary_emb, + token_valid_mask=token_valid_mask, + beta_valid_mask=beta_valid_mask, + decay_valid_mask=decay_valid_mask, + ) + + # ---- Camera branch: same pipeline on the camera projections, positionally encoded with UCPE ---- + if camera_conditions is not None: + cam_hidden_states = x + if token_valid_mask is not None: + cam_hidden_states = cam_hidden_states * token_valid_mask.view(batch_size, seq_len, 1) + + query_cam = self.q_proj_cam(cam_hidden_states) + key_cam = self.k_proj_cam(cam_hidden_states) + value_cam = self.v_proj_cam(cam_hidden_states) + if token_valid_mask is not None: + token_mask = token_valid_mask.view(batch_size, seq_len, 1) + query_cam = query_cam * token_mask + key_cam = key_cam * token_mask + value_cam = value_cam * token_mask + + if self.conv_q_cam is not None: + query_cam = self.conv_q_cam(query_cam, num_frames) + if self.conv_k_cam is not None: + key_cam = self.conv_k_cam(key_cam, num_frames) + if self.conv_v_cam is not None: + value_cam = self.conv_v_cam(value_cam, num_frames) + + query_cam = self.q_norm_cam(query_cam).reshape(batch_size, seq_len, self.cam_heads, self.cam_head_dim) + key_cam = self.k_norm_cam(key_cam).reshape(batch_size, seq_len, self.cam_heads, self.cam_head_dim) + value_cam = value_cam.reshape(batch_size, seq_len, self.cam_heads, self.cam_head_dim) + query_cam = self.kernel_func(query_cam) + key_cam = self.kernel_func(key_cam) + + # Reuse the model-level cache when available, to avoid recomputing the ray transforms per block. + ray_transforms = ucpe_ray_transforms + if ray_transforms is None: + ray_transforms = _prepare_ucpe_ray_transforms( + head_dim=self.cam_head_dim, + camera_conditions=camera_conditions, + HW=HW, + patch_size=self.patch_size, + rotary_emb=rotary_emb, + ) + + cam_output = self.cam_attn( + query_cam, + key_cam, + value_cam, + ray_transforms, + beta, + decay, + num_frames, + token_valid_mask=token_valid_mask, + beta_valid_mask=beta_valid_mask, + decay_valid_mask=decay_valid_mask, + ) + attn_output = attn_output + self.out_proj_cam(cam_output) + + # Shared output gate + projection, applied once over both branches. + if self.use_output_gate and self.output_gate is not None: + attn_output = attn_output * F.silu(self.output_gate(x).to(torch.float32)) + return self.proj(attn_output.to(x.dtype)) + + +class SanaWMSoftmaxAttention(nn.Module): + """Softmax counterpart of [`SanaWMBidirectionalGDNAttention`]. + + Replaces the bidirectional recurrence with a full (non-causal) `F.scaled_dot_product_attention` over the whole + token sequence. RoPE is applied to the queries and keys, and no linear-attention kernel or key scaling is needed + (softmax attention brings its own ``1 / sqrt(d_k)``). + + This module holds no parameters; [`_SoftmaxUCPESinglePathLiteLA.forward`] issues every layer call and passes the + resulting tensors in. + """ + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + rotary_emb: torch.Tensor | None = None, + token_valid_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run softmax attention over the token sequence. + + Args: + query (`torch.Tensor`): Queries of shape `(B, N, H, D)`, already normalized. + key (`torch.Tensor`): Keys of shape `(B, N, H, D)`, already normalized. + value (`torch.Tensor`): Values of shape `(B, N, H, D)`. + rotary_emb (`torch.Tensor`, *optional*): Rotary embeddings applied to `query` and `key`. + token_valid_mask (`torch.Tensor`, *optional*): Token mask of shape `(B, N)`. + + Returns: + `torch.Tensor`: Raw attention output of shape `(B, N, H * D)`; the shared output gate and projection are + applied by the caller. + """ + batch_size, seq_len, num_heads, head_dim = query.shape + dtype_orig = value.dtype + + # `_apply_rotary_emb` works on (B, H, D, N) features. + if rotary_emb is not None: + query = _apply_rotary_emb(query.permute(0, 2, 3, 1), rotary_emb).permute(0, 3, 1, 2) + key = _apply_rotary_emb(key.permute(0, 2, 3, 1), rotary_emb).permute(0, 3, 1, 2) + + if token_valid_mask is not None: + token_mask = token_valid_mask.view(batch_size, seq_len, 1, 1) + query = query * token_mask + key = key * token_mask + value = value * token_mask + + query = query.transpose(1, 2) # (B, H, N, D) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + # SDPA / FlashAttention only supports bf16/fp16; fp32 falls back to the math backend. + if query.dtype == torch.float32: + query, key, value = query.bfloat16(), key.bfloat16(), value.bfloat16() + + hidden_states = F.scaled_dot_product_attention(query, key, value) + return hidden_states.transpose(1, 2).reshape(batch_size, seq_len, num_heads * head_dim).to(dtype_orig) + + +class SanaWMSoftmaxCamAttention(nn.Module): + """Softmax counterpart of [`SanaWMBidirectionalGDNCamAttention`]. + + Keeps the UCPE per-ray encode/decode of the camera branch but replaces the single-path delta rule with a full + (non-causal) `F.scaled_dot_product_attention`. Padded frames are masked with an additive logit bias on the keys + instead of being dropped from the recurrence, and no linear-attention kernel, key scaling or gate is involved. + + This module holds no parameters; [`_SoftmaxUCPESinglePathLiteLA.forward`] issues every layer call and passes the + resulting tensors in. + """ + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ray_transforms: tuple, + token_valid_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run softmax attention in UCPE ray space. + + Args: + query (`torch.Tensor`): Camera queries of shape `(B, N, H, D)`, already normalized. + key (`torch.Tensor`): Camera keys of shape `(B, N, H, D)`, already normalized. + value (`torch.Tensor`): Camera values of shape `(B, N, H, D)`. + ray_transforms (`tuple`): `(P, P_T, P_inv, rotary_emb_cam)` UCPE transforms; `P_T` encodes the queries, + `P_inv` the keys/values and `P` decodes the output. + token_valid_mask (`torch.Tensor`, *optional*): Token mask of shape `(B, N)`. + + Returns: + `torch.Tensor`: Raw camera attention output of shape `(B, N, H * D)`, in world space; the caller projects + it back into the residual stream. + """ + batch_size, seq_len, num_heads, head_dim = query.shape + dtype_orig = value.dtype + + # Permute to (B, H, D, N), matching the GDN camera branch. + query = query.permute(0, 2, 3, 1).contiguous() + key = key.permute(0, 2, 3, 1).contiguous() + value = value.permute(0, 2, 3, 1).contiguous() + + # UCPE expects (B, h, N, d); our tensors are (B, h, d, N). Avoid eager contiguous copies before the + # transforms, and fuse the K/V transform (both use P_inv) into one call, then split back. + P, P_T, P_inv, rotary_emb_cam = ray_transforms + query_ucpe = _apply_ucpe_transform(query.transpose(-1, -2), P_T, rotary_emb_cam).transpose(-1, -2).contiguous() + key_value = torch.cat([key, value], dim=1) + key_value_ucpe = ( + _apply_ucpe_transform(key_value.transpose(-1, -2), P_inv, rotary_emb_cam).transpose(-1, -2).contiguous() + ) + key_ucpe, value_ucpe = torch.chunk(key_value_ucpe, chunks=2, dim=1) + + # Downscale the transformed tensors back to their pre-UCPE RMS envelope. + query_ucpe = _downscale_to_reference_rms(query, query_ucpe) + key_ucpe = _downscale_to_reference_rms(key, key_ucpe) + value_ucpe = _downscale_to_reference_rms(value, value_ucpe) + + if token_valid_mask is not None: + token_mask_qkv = token_valid_mask.view(batch_size, 1, 1, seq_len) + query_ucpe = query_ucpe * token_mask_qkv + value_ucpe = value_ucpe * token_mask_qkv + + query_sdpa = query_ucpe.transpose(-1, -2) + key_sdpa = key_ucpe.transpose(-1, -2) + value_sdpa = value_ucpe.transpose(-1, -2) + + query_sdpa, key_sdpa, value_sdpa = query_sdpa.float(), key_sdpa.float(), value_sdpa.float() + # SDPA / FlashAttention only supports bf16/fp16; fp32 falls back to math backend. + if query_sdpa.dtype == torch.float32: + query_sdpa, key_sdpa, value_sdpa = query_sdpa.bfloat16(), key_sdpa.bfloat16(), value_sdpa.bfloat16() + + # Invalid frames are masked out of the keys with an additive bias rather than being zeroed. + invalid_kv_logit_bias = None + if token_valid_mask is not None and not bool(token_valid_mask.all()): + invalid_kv_logit_bias = torch.where( + token_valid_mask.bool().view(batch_size, 1, 1, -1), + torch.zeros((), dtype=query_sdpa.dtype, device=query_sdpa.device), + torch.full((), -1e9, dtype=query_sdpa.dtype, device=query_sdpa.device), + ) + + # FlashAttention-2 only supports head_dim in {32, 64, 128, 256}. + need_pad = head_dim not in (32, 64, 128, 256) and head_dim < 256 + if need_pad: + pad_size = (128 if head_dim <= 128 else 256) - head_dim + query_sdpa = F.pad(query_sdpa, (0, pad_size)) + key_sdpa = F.pad(key_sdpa, (0, pad_size)) + value_sdpa = F.pad(value_sdpa, (0, pad_size)) + hidden_states = F.scaled_dot_product_attention( + query_sdpa, key_sdpa, value_sdpa, attn_mask=invalid_kv_logit_bias + ) + if need_pad: + hidden_states = hidden_states[..., :head_dim] + + hidden_states = hidden_states.transpose(-1, -2) + if hidden_states.dtype != dtype_orig: + hidden_states = hidden_states.to(dtype_orig) + if token_valid_mask is not None: + hidden_states = hidden_states * token_valid_mask.view(batch_size, 1, 1, seq_len).to(hidden_states.dtype) + + # Decode back from ray space to world space. + hidden_states = ( + _apply_ucpe_transform(hidden_states.transpose(-1, -2), P, rotary_emb_cam, inverse_rope=True) + .transpose(-1, -2) + .contiguous() + ) + hidden_states = hidden_states.reshape(batch_size, num_heads * head_dim, seq_len).permute(0, 2, 1) + if token_valid_mask is not None: + hidden_states = hidden_states * token_valid_mask.view(batch_size, seq_len, 1).to(hidden_states.dtype) + return hidden_states + + +class _SoftmaxUCPESinglePathLiteLA(nn.Module): + """Softmax counterpart of [`BidirectionalGDNUCPESinglePathLiteLA`]. + + The released checkpoint uses this block for every ``softmax_every_n``-th layer. It keeps the exact parameter layout + of [`BidirectionalGDNUCPESinglePathLiteLA`] -- so both variants load from the same checkpoint -- but replaces the + main-branch recurrence with [`SanaWMSoftmaxAttention`] and the camera-branch recurrence with + [`SanaWMSoftmaxCamAttention`], both a full (non-causal) ``F.scaled_dot_product_attention``. Short convolutions are + never built for this variant, and the GDN-only parameters (``beta_proj`` / ``gate_proj`` / ``dt_bias`` / ``A_log`` + / ``recall_gate``) exist only so the shared checkpoint loads: they are created in the same order, under the same + names, and are unused by the forward. + + Both branch modules are parameter-free: every layer call happens in this class's `forward`, which owns the + projections and norms the two branches consume. + + Args: + in_dim (`int`): Input channels. + out_dim (`int`): Output channels. + cam_dim (`int`): Camera-branch width; must equal `in_dim` so the shared parameters line up. + cam_heads (`int`): Camera-branch heads; must equal `heads` and divide `cam_dim` into multiples of 4. + patch_size (`tuple[int, int, int]`, defaults to `(1, 2, 2)`): Latent patch size, used to map camera + intrinsics onto the token grid. + heads (`int`, *optional*): Number of attention heads; derived from `out_dim // dim * heads_ratio` when `None`. + heads_ratio (`float`, defaults to 1.0): Head-count multiplier used when `heads` is `None`. + dim (`int`, defaults to 32): Head dimension used when `heads` is `None`. + use_bias (`bool`, defaults to `False`): Whether the fused QKV projection has a bias. + qk_norm (`bool`, defaults to `False`): Apply RMSNorm to Q/K. + norm_eps (`float`, defaults to 1e-5): Epsilon of the Q/K RMSNorm. + use_output_gate (`bool`, defaults to `True`): Apply the shared silu output gate. + kwargs: The GDN-only options (`eps`, `chunk_gdn_chunk_size`, `conv_kernel_size`, `k_conv_only`) are accepted + and ignored, so both attention variants can be built from the same block keywords. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + *, + cam_dim: int, + cam_heads: int, + patch_size: tuple[int, int, int] = (1, 2, 2), + heads: int | None = None, + heads_ratio: float = 1.0, + dim: int = 32, + use_bias: bool = False, + qk_norm: bool = False, + norm_eps: float = 1e-5, + use_output_gate: bool = True, + **kwargs: object, + ) -> None: + heads = heads or int(out_dim // dim * heads_ratio) + super().__init__() + + # Fused QKV projection and output projection (the `q_norm` / `k_norm` + # attributes are set further down, depending on `qk_norm`). + self.qkv = nn.Linear(in_dim, in_dim * 3, bias=use_bias) + self.proj = nn.Linear(in_dim, in_dim) + + self.heads = heads + self.dim = out_dim // heads + + if qk_norm: + self.q_norm = RMSNorm(in_dim, eps=norm_eps) + self.k_norm = RMSNorm(in_dim, eps=norm_eps) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + # GDN-only gates. Softmax attention never reads them, but they are part of the + # shared checkpoint, so they are built here in the checkpoint's order. + self.beta_proj = nn.Linear(in_dim, heads, bias=True) + self.gate_proj = nn.Linear(in_dim, heads, bias=True) + + A = torch.zeros(heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(heads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + + # `recall_gate` is unused by the forward; kept as a buffer for checkpoint compatibility. + self.register_buffer("recall_gate", torch.zeros(1)) + + self.use_output_gate = use_output_gate + if use_output_gate: + self.output_gate = nn.Linear(in_dim, out_dim, bias=True) + else: + self.output_gate = None + + self.patch_size = patch_size + self.cam_dim = cam_dim + self.cam_heads = cam_heads + self.cam_head_dim = cam_dim // cam_heads + + if cam_dim != in_dim: + raise ValueError(f"Parameter sharing requires cam_dim == in_dim, got cam_dim={cam_dim}, in_dim={in_dim}.") + if cam_heads != self.heads: + raise ValueError( + f"Parameter sharing requires cam_heads == heads, got cam_heads={cam_heads}, heads={self.heads}." + ) + if self.cam_head_dim % 4 != 0: + raise ValueError( + "UCPE camera branch requires cam_head_dim divisible by 4, " + f"got {self.cam_head_dim} (cam_dim={cam_dim}, cam_heads={cam_heads})." + ) + + # ---- Camera-specific: QKV + output projections only ---- + self.q_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.k_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.v_proj_cam = nn.Linear(in_dim, cam_dim, bias=True) + self.out_proj_cam = nn.Linear(cam_dim, out_dim, bias=True) + + # Keep branch-specific Q/K norms so camera statistics do not disturb the + # main branch (and vice versa). Start from identical weights. + self.q_norm_cam = deepcopy(self.q_norm) + self.k_norm_cam = deepcopy(self.k_norm) + + # Branch compute modules. They own no parameters -- so the checkpoint layout is untouched -- and `forward` + # runs every layer call and hands them the resulting tensors. + self.attn = SanaWMSoftmaxAttention() + self.cam_attn = SanaWMSoftmaxCamAttention() + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + HW: tuple[int, int, int] | None = None, + rotary_emb: torch.Tensor | None = None, + block_mask: torch.Tensor | None = None, + camera_conditions: torch.Tensor | None = None, + chunk_size: int | None = None, + *, + frame_valid_mask: torch.Tensor | None = None, + ucpe_ray_transforms: tuple | None = None, + ) -> torch.Tensor: + """Dual-branch forward: softmax main branch + softmax UCPE camera branch. + + Args: + x: Input tensor of shape ``(B, N, C)``. + mask: Unused attention mask (kept for API compatibility). + HW: Tuple of ``(T, H, W)`` describing the token layout. + rotary_emb: Optional rotary embeddings for q/k. + block_mask: Unused block mask (kept for API compatibility). + camera_conditions: Raw ``(B, T, 20)`` camera conditions enabling the camera branch. + chunk_size: Unused chunk length (kept for API compatibility). + frame_valid_mask: Optional per-frame validity mask used to zero out padded frames, shaped + ``(B, 1, T, 1, 1)``, ``(B, 1, T)`` or ``(B, T)``. + ucpe_ray_transforms: Optional pre-computed UCPE transforms shared across blocks. + + Returns: + Tensor of shape ``(B, N, C)`` after attention and projection. + """ + del mask, block_mask, chunk_size + + if HW is None: + raise ValueError("HW (T, H, W) must be provided for softmax attention.") + + batch_size, seq_len, channels = x.shape + num_frames, height, width = HW + spatial_size = height * width + + token_valid_mask, _, _ = _prepare_frame_valid_masks( + frame_valid_mask, + batch_size=batch_size, + num_frames=num_frames, + spatial_size=spatial_size, + device=x.device, + dtype=x.dtype, + ) + + # ---- Main branch: fused QKV -> Q/K norm -> softmax attention ---- + hidden_states = x + if token_valid_mask is not None: + hidden_states = hidden_states * token_valid_mask.view(batch_size, seq_len, 1) + + query, key, value = self.qkv(hidden_states).reshape(batch_size, seq_len, 3, self.heads, self.dim).unbind(2) + if token_valid_mask is not None: + token_mask = token_valid_mask.view(batch_size, seq_len, 1, 1) + query = query * token_mask + key = key * token_mask + value = value * token_mask + + # Q/K norm runs on the flattened channels (B, N, C), then the tensors go back to (B, N, H, D). + query = self.q_norm(query.reshape(batch_size, seq_len, channels)).reshape( + batch_size, seq_len, self.heads, self.dim + ) + key = self.k_norm(key.reshape(batch_size, seq_len, channels)).reshape( + batch_size, seq_len, self.heads, self.dim + ) + + attn_output = self.attn(query, key, value, rotary_emb=rotary_emb, token_valid_mask=token_valid_mask) + + # ---- Camera branch: same pipeline on the camera projections, positionally encoded with UCPE ---- + if camera_conditions is not None: + cam_hidden_states = x + if token_valid_mask is not None: + cam_hidden_states = cam_hidden_states * token_valid_mask.view(batch_size, seq_len, 1) + + query_cam = self.q_proj_cam(cam_hidden_states) + key_cam = self.k_proj_cam(cam_hidden_states) + value_cam = self.v_proj_cam(cam_hidden_states) + if token_valid_mask is not None: + token_mask = token_valid_mask.view(batch_size, seq_len, 1) + query_cam = query_cam * token_mask + key_cam = key_cam * token_mask + value_cam = value_cam * token_mask + + query_cam = self.q_norm_cam(query_cam).reshape(batch_size, seq_len, self.cam_heads, self.cam_head_dim) + key_cam = self.k_norm_cam(key_cam).reshape(batch_size, seq_len, self.cam_heads, self.cam_head_dim) + value_cam = value_cam.reshape(batch_size, seq_len, self.cam_heads, self.cam_head_dim) + + # Reuse the model-level cache when available, to avoid recomputing the ray transforms per block. + ray_transforms = ucpe_ray_transforms + if ray_transforms is None: + ray_transforms = _prepare_ucpe_ray_transforms( + head_dim=self.cam_head_dim, + camera_conditions=camera_conditions, + HW=HW, + patch_size=self.patch_size, + rotary_emb=rotary_emb, + ) + + cam_output = self.cam_attn( + query_cam, key_cam, value_cam, ray_transforms, token_valid_mask=token_valid_mask + ) + attn_output = attn_output + self.out_proj_cam(cam_output) + + # Shared output gate + projection, applied once over both branches. + if self.use_output_gate and self.output_gate is not None: + attn_output = attn_output * F.silu(self.output_gate(x).to(torch.float32)) + return self.proj(attn_output.to(x.dtype)) + + +# ============================================================================ +# DiT base + SANA-WM camera-controlled transformer + public wrapper +# ============================================================================ + + +class SanaVideoMSCamCtrlBlock(nn.Module): + """ + A Sana block with global shared adaptive layer norm zero (adaLN-Zero) conditioning. + """ + + def __init__( + self, + hidden_size, + num_heads, + attn_cls, + mlp_ratio=4.0, + qk_norm=False, + ffn_type="mlp", + mlp_acts=("silu", "silu", None), + linear_head_dim=32, + cross_norm=False, + t_kernel_size=3, + patch_size=(1, 2, 2), + cam_attn_compress=2, + chunk_size=10, + use_chunk_plucker_post_attn=False, + **block_kwargs, + ): + super().__init__() + self.hidden_size = hidden_size + self.chunk_size = chunk_size + + if use_chunk_plucker_post_attn: + self.plucker_proj = nn.Linear(hidden_size, hidden_size, bias=True) + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + # Camera-conditioned (UCPE) attention: either the GDN or the softmax variant. + self_num_heads = hidden_size // linear_head_dim + self.attn = attn_cls( + hidden_size, + hidden_size, + heads=self_num_heads, + cam_dim=hidden_size // cam_attn_compress, + cam_heads=max(1, self_num_heads // cam_attn_compress), + eps=1e-8, + qk_norm=qk_norm, + patch_size=patch_size, + **block_kwargs, + ) + + self.cross_attn = MultiHeadCrossAttention(hidden_size, num_heads, qk_norm=cross_norm, **block_kwargs) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + + # MLP + if ffn_type == "GLUMBConvTemp": + self.mlp = GLUMBConvTemp( + in_features=hidden_size, + hidden_features=int(hidden_size * mlp_ratio), + use_bias=(True, True, False), + act=mlp_acts, + t_kernel_size=t_kernel_size, + ) + elif ffn_type == "mlp": + + def approx_gelu(): + return nn.GELU(approximate="tanh") + + self.mlp = Mlp( + in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0 + ) + else: + self.mlp = None + + self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size**0.5) + + @staticmethod + def _build_frame_token_mask( + frame_valid_mask: Optional[torch.Tensor], + *, + B: int, + T: int, + N: int, + device: torch.device, + dtype: torch.dtype, + ) -> Optional[torch.Tensor]: + """Convert frame-valid mask to token mask shaped ``(B, N, 1)``.""" + if frame_valid_mask is None: + return None + + m = frame_valid_mask + if m.ndim == 5: + m = m[:, 0, :, 0, 0] + elif m.ndim == 3 and m.shape[1] == 1: + m = m[:, 0, :] + elif m.ndim != 2: + raise ValueError( + "frame_valid_mask must be shaped (B, 1, T, 1, 1), (B, 1, T), or (B, T); " + f"got shape={list(frame_valid_mask.shape)}" + ) + + if m.shape[0] != B or m.shape[1] != T: + raise ValueError(f"frame_valid_mask shape mismatch: expected (B={B}, T={T}), got {list(m.shape)}") + if T <= 0 or N % T != 0: + raise ValueError(f"Invalid token/frame layout: N={N}, T={T}") + + S = N // T + return m.to(device=device, dtype=dtype).view(B, T, 1).expand(B, T, S).reshape(B, N, 1) + + def forward( + self, + x, + y, + t, + mask=None, + THW=None, + rotary_emb=None, + block_mask=None, + *, + camera_conditions=None, + ucpe_ray_transforms=None, + plucker_emb=None, + frame_valid_mask=None, + chunk_size=None, + ): + """Run one adaLN-Zero block: self-attention -> cross-attention -> FFN. + + Args: + x: ``(B, N, C)`` token sequence. + y: ``(B, 1, L, C)`` text embeddings for cross-attention. + t: ``(B, 1, T, 6 * C)`` adaLN modulation input. + mask: Text padding mask for cross-attention. + THW: ``(T, H, W)`` token layout. + rotary_emb: Rotary embeddings for the self-attention branch. + block_mask: Optional block mask forwarded to the attention. + camera_conditions: Raw camera conditions enabling the camera branch. + ucpe_ray_transforms: Pre-computed UCPE transforms shared across blocks. + plucker_emb: Optional post-attention Plucker embedding. + frame_valid_mask: Optional per-frame validity mask. + chunk_size: Chunk length override; falls back to ``self.chunk_size``. + """ + B, N, C = x.shape + num_frames = t.shape[2] + frame_token_mask = self._build_frame_token_mask( + frame_valid_mask, + B=B, + T=num_frames, + N=N, + device=x.device, + dtype=x.dtype, + ) + if frame_token_mask is not None: + x = x * frame_token_mask + + t = t.reshape(B, num_frames, 6, -1) # B,F,6,D + # scale_shift_table: 6, hidden_size -> 1,1,6,hidden_size + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None, :, :] + t + ).chunk(6, dim=-2) # each chunk: B,F,1,D + if chunk_size is None: + chunk_size = self.chunk_size + + x_norm1 = self.norm1(x).reshape(B, num_frames, -1, C) + x_msa_in = (x_norm1 * (1 + scale_msa) + shift_msa).reshape(B, N, C) + if frame_token_mask is not None: + x_msa_in = x_msa_in * frame_token_mask + # Camera-conditioned (UCPE) attention: dual-branch (main + camera) forward. + attn_out = self.attn( + x_msa_in, + HW=THW, + rotary_emb=rotary_emb, + block_mask=block_mask, + camera_conditions=camera_conditions, + chunk_size=chunk_size, + frame_valid_mask=frame_valid_mask, + ucpe_ray_transforms=ucpe_ray_transforms, + ) + attn_out = attn_out.reshape(B, num_frames, -1, C) + attn_out = (gate_msa * attn_out).reshape(B, N, C) + if frame_token_mask is not None: + attn_out = attn_out * frame_token_mask + x = x + attn_out + if frame_token_mask is not None: + x = x * frame_token_mask + + if plucker_emb is not None and hasattr(self, "plucker_proj"): + x = x + self.plucker_proj(plucker_emb) + + x = x + self.cross_attn(x, y, mask=mask) + if frame_token_mask is not None: + x = x * frame_token_mask + + x_norm2 = self.norm2(x).reshape(B, num_frames, -1, C) + x_mlp_in = (x_norm2 * (1 + scale_mlp) + shift_mlp).reshape(B, N, C) + if frame_token_mask is not None: + x_mlp_in = x_mlp_in * frame_token_mask + mlp_out = self.mlp(x_mlp_in, HW=THW).reshape(B, num_frames, -1, C) + mlp_out = (gate_mlp * mlp_out).reshape(B, N, C) + if frame_token_mask is not None: + mlp_out = mlp_out * frame_token_mask + x = x + mlp_out + if frame_token_mask is not None: + x = x * frame_token_mask + + return x + + +class SanaWMTransformer3DModel(ModelMixin, ConfigMixin): + r""" + SANA-WM 1600M bidirectional camera-controlled DiT. + + A single-class DiT (depth=20, hidden_size=2240, patch_size=(1,1,1), num_heads=20 — i.e. the public + ``Efficient-Large-Model/SANA-WM_bidirectional`` release). ``save_pretrained`` / ``from_pretrained`` work out of the + box via :class:`~diffusers.configuration_utils.ConfigMixin`. + + Every block runs a camera-conditioned (UCPE) attention: [`BidirectionalGDNUCPESinglePathLiteLA`], except every + ``softmax_every_n``-th block which runs its softmax counterpart [`_SoftmaxUCPESinglePathLiteLA`]. + + Args: + in_channels (`int`, defaults to 128): VAE latent channels (LTX-2). + softmax_every_n (`int`, defaults to 4): Use a softmax attention block every N blocks. + linear_head_dim (`int`, defaults to 112): GDN head dimension. + ffn_type (`str`, defaults to ``"GLUMBConvTemp"``): FFN. + t_kernel_size (`int`, defaults to 3): Temporal conv kernel. + conv_kernel_size (`int`, defaults to 4): Spatial conv kernel inside attention. + k_conv_only (`bool`, defaults to True): Apply conv only on K. + pos_embed_type (`str`, defaults to ``"wan_rope"``): Position embedding. + qk_norm (`bool`, defaults to True): RMSNorm on Q/K. + cross_norm (`bool`, defaults to True): RMSNorm on cross-attention K. + y_norm (`bool`, defaults to True): Apply ``attention_y_norm`` to text embeddings. + init_cam_from_base (`bool`, defaults to True): Unused; the camera branch is loaded from the checkpoint. + Kept so released `config.json` files load. + use_chunk_plucker_post_attn (`bool`, defaults to True). + chunk_plucker_channels (`int`, defaults to 48): ``6 dims * temporal_stride 8``. + chunk_plucker_post_attn_blocks (`int`, defaults to 20): All blocks. + fp32_attention (`bool`, defaults to True): Unused; attention always runs in fp32. Kept so released + `config.json` files load. + image_size (`int`, defaults to 720): Nominal image size. + caption_channels (`int`, defaults to 2304): Gemma-2 hidden size. + model_max_length (`int`, defaults to 300): Max prompt tokens. + + The state-dict is identical to the public sana checkpoint apart from the intentionally-removed ``pos_embed`` + buffer. + """ + + _supports_gradient_checkpointing = False + _no_split_modules = ["SanaVideoMSCamCtrlBlock"] + _repeated_blocks = ["SanaVideoMSCamCtrlBlock"] + _skip_layerwise_casting_patterns = ["x_embedder", "plucker_embedder", "norm"] + # NOTE: `_keep_in_fp32_modules` is intentionally unset. SANA-WM's blocks apply the + # timestep modulation inline, so holding `t_embedder` / `t_block` / + # `scale_shift_table` in fp32 would upcast the hidden states and feed fp32 activations + # to bf16 weights. Supporting it needs explicit casts in the block forward first. + + @register_to_config + def __init__( + self, + in_channels: int = 128, + num_layers: int = 20, + hidden_size: int = 2240, + num_attention_heads: int = 20, + patch_size: tuple[int, int, int] = (1, 1, 1), + softmax_every_n: int = 4, + linear_head_dim: int = 112, + ffn_type: str = "GLUMBConvTemp", + t_kernel_size: int = 3, + conv_kernel_size: int = 4, + k_conv_only: bool = True, + pos_embed_type: str = "wan_rope", + qk_norm: bool = True, + cross_norm: bool = True, + y_norm: bool = True, + cam_attn_compress: int = 1, + init_cam_from_base: bool = True, + use_chunk_plucker_post_attn: bool = True, + chunk_plucker_channels: int = 48, + chunk_plucker_post_attn_blocks: int = 20, + fp32_attention: bool = True, + image_size: int = 720, + caption_channels: int = 2304, + model_max_length: int = 300, + mlp_ratio: float = 3.0, + mlp_acts: tuple = ("silu", "silu", None), + use_pe: bool = True, + learn_sigma: bool = False, + pred_sigma: bool = False, + mixed_precision: str = "bf16", + ) -> None: + super().__init__() + + # The defaults describe the public SANA-WM_bidirectional release; they are + # configurable so a small variant can be built (e.g. for tests). + depth = num_layers + num_heads = num_attention_heads + patch_size = tuple(patch_size) + + # Remaining SanaMSVideoCamCtrl.__init__ defaults not exposed by the config signature. + mlp_acts = list(mlp_acts) + pe_interpolation = 1.0 + norm_eps = 1e-5 + patch_embed_kernel = None + cfg_embed = False + timestep_norm_scale_factor = 1.0 + rope_fhw_dim = None + pack_latents = False + chunk_size = 10 + use_chunk_plucker_input = False + + # --- Base DiT config attributes (from Sana.__init__) --- + self.pred_sigma = pred_sigma + self.in_channels = in_channels + self.out_channels = in_channels * 2 if pred_sigma else in_channels + self.hidden_size = hidden_size + self.num_heads = num_heads + self.linear_head_dim = linear_head_dim + self.pe_interpolation = pe_interpolation + self.depth = depth + self.use_pe = use_pe + self.pos_embed_type = pos_embed_type + self.y_norm = y_norm + # NOTE: ``self.config`` is provided (read-only) by ConfigMixin via @register_to_config. + self.timestep_norm_scale_factor = timestep_norm_scale_factor + + self.t_embedder = TimestepEmbedder(hidden_size) + self.cfg_embedder = None + if cfg_embed: + self.cfg_embedder = TimestepEmbedder(hidden_size) + + if self.y_norm: + self.attention_y_norm = RMSNorm(hidden_size, eps=norm_eps) + + # --- Video camera-controlled DiT modules (from SanaMSVideoCamCtrl.__init__) --- + self.chunk_size = chunk_size + self.patch_size = patch_size + + def approx_gelu(): + return nn.GELU(approximate="tanh") + + self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True)) + self.pos_embed_ms = None + self.pack_latents = pack_latents + self.cam_attn_compress = cam_attn_compress + + kernel_size = patch_embed_kernel or patch_size + x_embedder_in_channels = in_channels + if self.pack_latents: + x_embedder_in_channels = x_embedder_in_channels * 2 * 2 + self.out_channels = in_channels * 2 * 2 + + self.x_embedder = PatchEmbedMS3D( + patch_size, x_embedder_in_channels, hidden_size, kernel_size=kernel_size, bias=True + ) + + self.y_embedder = CaptionEmbedder( + in_channels=caption_channels, + hidden_size=hidden_size, + act_layer=approx_gelu, + token_num=model_max_length, + ) + + self.use_chunk_plucker_input = use_chunk_plucker_input + self.use_chunk_plucker_post_attn = use_chunk_plucker_post_attn + if self.use_chunk_plucker_input or self.use_chunk_plucker_post_attn: + self.plucker_embedder = PatchEmbedMS3D( + patch_size, chunk_plucker_channels, hidden_size, kernel_size=kernel_size, bias=True + ) + nn.init.zeros_(self.plucker_embedder.proj.weight) + nn.init.zeros_(self.plucker_embedder.proj.bias) + + # UCPE-style camera branch uses a 3-channel absmap (up_map + lat_map). + self.raymap_embedder = PatchEmbedMS3D(patch_size, 3, hidden_size, kernel_size=kernel_size, bias=True) + + if use_pe: + if pos_embed_type != "wan_rope": + raise ValueError(f'`pos_embed_type` must be "wan_rope", got {pos_embed_type!r}.') + self.rope = SanaWMRotaryPosEmbed( + attention_head_dim=linear_head_dim, patch_size=patch_size, max_seq_len=1024, fhw_dim=rope_fhw_dim + ) + + # Every ``softmax_every_n``-th block swaps the GDN recurrence for softmax attention; both variants share the + # same parameter layout. + self.softmax_every_n = softmax_every_n + attn_cls_list = [] + for i in range(depth): + if softmax_every_n > 0 and (i + 1) % softmax_every_n == 0: + attn_cls = _SoftmaxUCPESinglePathLiteLA + else: + attn_cls = BidirectionalGDNUCPESinglePathLiteLA + attn_cls_list.append(attn_cls) + + self.blocks = nn.ModuleList( + [ + SanaVideoMSCamCtrlBlock( + hidden_size, + num_heads, + attn_cls=attn_cls_list[i], + mlp_ratio=mlp_ratio, + qk_norm=qk_norm, + ffn_type=ffn_type, + mlp_acts=mlp_acts, + linear_head_dim=linear_head_dim, + cross_norm=cross_norm, + t_kernel_size=t_kernel_size, + patch_size=patch_size, + cam_attn_compress=self.cam_attn_compress, + chunk_size=chunk_size, + conv_kernel_size=conv_kernel_size, + k_conv_only=k_conv_only, + use_chunk_plucker_post_attn=( + use_chunk_plucker_post_attn + and (chunk_plucker_post_attn_blocks < 0 or i < chunk_plucker_post_attn_blocks) + ), + ) + for i in range(depth) + ] + ) + self.final_layer = T2IFinalLayer(hidden_size, patch_size, self.out_channels) + + if ffn_type == "GLUMBConvTemp": + logger.info(f"{ffn_type} Temporal kernal: {t_kernel_size}") + + self.in_channels = self.out_channels = in_channels + + @staticmethod + def _pack_latents(latents, batch_size, num_channels_latents, height, width, frame): + latents = latents.view(batch_size, num_channels_latents, frame, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 4, 6, 2, 3, 5) + latents = latents.reshape(batch_size, num_channels_latents * 4, frame, height // 2, width // 2) + + return latents + + @staticmethod + def _unpack_latents(latents, height, width, frame): + batch_size, channels, frame, H, W = latents.shape + + if not (height % 2 == 0 and width % 2 == 0): + raise ValueError(f"Latent height and width must be divisible by 2, got {height}x{width}.") + # latent height and width to be divisible by 2. + latents = latents.view(batch_size, channels // 4, 2, 2, frame, height // 2, width // 2) + latents = latents.permute(0, 1, 4, 5, 2, 6, 3) + latents = latents.reshape(batch_size, channels // (2 * 2), frame, height, width) + + return latents + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + return_dict: bool = True, + data_info: Optional[dict] = None, + camera_conditions: torch.Tensor | None = None, + chunk_plucker: torch.Tensor | None = None, + cam_pos_embeds: Optional[dict] = None, + pos_embeds: torch.Tensor | None = None, + raymats: torch.Tensor | None = None, + frame_valid_mask: torch.Tensor | None = None, + chunk_size: int | None = None, + ): + """Run the SANA-WM DiT. + + Args: + hidden_states: ``(B, C, T, H, W)`` latents. + timestep: ``(B, 1, T)`` per-frame diffusion timesteps (LTX style). + encoder_hidden_states: ``(B, 1, L, D_caption)`` text embeddings. + encoder_attention_mask: ``(B, L)`` text attention mask (diffusers convention). + mask: Alias for ``encoder_attention_mask`` matching the sana DiT's + kwarg name. If both are passed, ``mask`` takes precedence. + return_dict: If ``True`` (default), returns a :class:`Transformer2DModelOutput`; + otherwise returns a one-tuple ``(sample,)``. + data_info: Extra conditioning; ``data_info["image_vae_embeds"]`` is + concatenated to the latents along the channel axis when present. + camera_conditions: ``(B, T, 20)`` raw camera conditions driving the + camera-control (UCPE) branch. + chunk_plucker: Plucker ray embeddings ``(B, C, T, H, W)``, consumed when + the model is configured with ``use_chunk_plucker_input`` / ``use_chunk_plucker_post_attn``. + cam_pos_embeds: Optional pre-computed camera positional embeddings + (``"absmap"`` / ``"P"`` entries) reused instead of recomputing them. + pos_embeds: Optional pre-computed rotary position embeddings; when ``None`` + they are built from the latent shape. + raymats: Optional pre-computed UCPE ray matrices (used only when + ``cam_pos_embeds`` does not carry ``"P"``). + frame_valid_mask: Optional per-frame validity mask, shaped + ``(B, 1, T, 1, 1)``, ``(B, 1, T)`` or ``(B, T)``. + chunk_size: Chunk length override forwarded to the blocks; falls back + to each block's configured ``chunk_size``. + + Returns: + :class:`Transformer2DModelOutput` with ``sample`` of shape ``(B, C, T, H, W)``. + """ + # The sana DiT names its text mask kwarg ``mask``. + # Accept both ``mask=`` (sana convention) and ``encoder_attention_mask=`` + # (diffusers convention); the former wins if both are provided. + if mask is None: + mask = encoder_attention_mask + x = hidden_states + y = encoder_hidden_states + + bs = x.shape[0] + x = x.to(self.dtype) + if self.timestep_norm_scale_factor != 1.0: + timestep = (timestep.float() / self.timestep_norm_scale_factor).to(torch.float32) + else: + timestep = timestep.long().to(torch.float32) + y = y.to(self.dtype) + post_patch_num_frames, post_patch_height, post_patch_width = ( + x.shape[-3] // self.patch_size[0], + x.shape[-2] // self.patch_size[1], + x.shape[-1] // self.patch_size[2], + ) + + if data_info is None: + data_info = {} + if data_info.get("image_vae_embeds", None) is not None: + x = torch.cat([x, data_info["image_vae_embeds"].to(self.dtype)], dim=1) + cam_embeds = camera_conditions + if self.pack_latents: + x = self._pack_latents(x, bs, self.in_channels, post_patch_height, post_patch_width, post_patch_num_frames) + if cam_embeds is not None: + cam_embeds = cam_embeds.to(self.dtype) + + post_patch_height = post_patch_height // 2 + post_patch_width = post_patch_width // 2 + + if self.x_embedder.patch_size != self.x_embedder.kernel_size and self.x_embedder.kernel_size == (1, 2, 2): + x = F.pad(x, (0, 1, 0, 1, 0, 0)) + if cam_embeds is not None: + cam_embeds = F.pad(cam_embeds, (0, 1, 0, 1, 0, 0)) + + x = self.x_embedder(x) + if cam_embeds is not None: + # Both attention variants are UCPE-style: build raymats + 3-channel absmap + # (up_map + lat_map) from the raw (B,F,20) camera conditions. + raw_cam_conditions = cam_embeds + if cam_pos_embeds is not None and "absmap" in cam_pos_embeds: + cam_embeds = cam_pos_embeds["absmap"] + if "P" in cam_pos_embeds: + raymats = cam_pos_embeds["P"] + else: + raymats, cam_embeds = _process_camera_conditions_ucpe( + raw_cam_conditions, + bs, + (post_patch_num_frames, post_patch_height, post_patch_width), + self.patch_size, + ) + cam_embeds = cam_embeds.permute(0, 4, 1, 2, 3).to(self.dtype) + if not (self.use_chunk_plucker_input or self.use_chunk_plucker_post_attn): + cam_embeds = self.raymap_embedder(cam_embeds) + x = x + cam_embeds + camera_conditions = raw_cam_conditions + + post_attn_plucker_emb = None + if self.use_chunk_plucker_input and chunk_plucker is not None: + plucker_input = chunk_plucker.to(self.dtype) + plucker_emb = self.plucker_embedder(plucker_input) + x = x + plucker_emb + + if self.use_chunk_plucker_post_attn and chunk_plucker is not None: + plucker_input = chunk_plucker.to(self.dtype) + post_attn_plucker_emb = self.plucker_embedder(plucker_input) + + image_pos_embed = pos_embeds + if self.use_pe and image_pos_embed is None: + image_pos_embed = self.rope((post_patch_num_frames, post_patch_height, post_patch_width)) + elif image_pos_embed is not None: + image_pos_embed = image_pos_embed.to(x.device) + while image_pos_embed.ndim > 4: + image_pos_embed = image_pos_embed.squeeze(1) + + t = self.t_embedder(timestep.flatten()) # (N, D) + t0 = self.t_block(t) + t = t.unflatten(dim=0, sizes=timestep.shape) + t0 = t0.unflatten(dim=0, sizes=timestep.shape) + + y = self.y_embedder(y) # (N, D) + if self.y_norm: + y = self.attention_y_norm(y) + + if mask is None: + raise ValueError( + "`mask` is required: SANA-WM's cross-attention needs the text padding mask to build its attention " + "bias. Pass the prompt attention mask returned by the pipeline's `encode_prompt`." + ) + mask = mask.to(torch.int16) + mask = mask.repeat(y.shape[0] // mask.shape[0], 1) if mask.shape[0] != y.shape[0] else mask + mask = mask.squeeze(1).squeeze(1) + y_lens = mask + + block_mask = None + + ucpe_ray_transforms = None + if camera_conditions is not None: + # Pre-compute the UCPE ray matrices once and share them across blocks + # (both attention variants are UCPE-style). + head_dim = self.linear_head_dim + + if cam_pos_embeds is not None: + for k, v in cam_pos_embeds.items(): + if isinstance(v, torch.Tensor): + v = v.to(x.device) + if k == "absmap": + while v.ndim > 5: + v = v.squeeze(1) + else: + while v.ndim > 4: + v = v.squeeze(1) + cam_pos_embeds[k] = v + + ucpe_ray_transforms = _prepare_ucpe_ray_transforms( + head_dim=head_dim, + camera_conditions=camera_conditions, + HW=(post_patch_num_frames, post_patch_height, post_patch_width), + patch_size=self.patch_size, + rotary_emb=image_pos_embed, + raymats=raymats, + cam_pos_embeds=cam_pos_embeds, + ) + + for i, block in enumerate(self.blocks): + x = block( + x, + y, + t0, + y_lens, + (post_patch_num_frames, post_patch_height, post_patch_width), + image_pos_embed, + block_mask=block_mask if i > 1 else None, + camera_conditions=camera_conditions, + ucpe_ray_transforms=ucpe_ray_transforms, + plucker_emb=post_attn_plucker_emb, + frame_valid_mask=frame_valid_mask, + chunk_size=chunk_size, + ) # (N, T, D) + + x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels) + x = self.unpatchify(x, post_patch_num_frames, post_patch_height, post_patch_width) # (N, out_channels, H, W) + if self.pack_latents: + x = self._unpack_latents(x, post_patch_height * 2, post_patch_width * 2, post_patch_num_frames) + + return Transformer2DModelOutput(sample=x) if return_dict else (x,) + + def unpatchify(self, x, post_patch_num_frames, post_patch_height, post_patch_width): + """ + x: (N, T, patch_size**2 * C) imgs: (N, H, W, C) + """ + c = self.out_channels + p_f, p_h, p_w = self.x_embedder.patch_size + if post_patch_num_frames * post_patch_height * post_patch_width != x.shape[1]: + raise ValueError( + f"Expected {post_patch_num_frames * post_patch_height * post_patch_width} tokens for a " + f"({post_patch_num_frames}, {post_patch_height}, {post_patch_width}) latent, but got {x.shape[1]}." + ) + + x = x.reshape(shape=(x.shape[0], post_patch_num_frames, post_patch_height, post_patch_width, p_f, p_h, p_w, c)) + x = torch.einsum("nfhwopqc->ncfohpwq", x) + imgs = x.reshape( + shape=(x.shape[0], c, post_patch_num_frames * p_f, post_patch_height * p_h, post_patch_width * p_w) + ) + + return imgs diff --git a/src/diffusers/models/transformers/transformer_sana_wm_refiner.py b/src/diffusers/models/transformers/transformer_sana_wm_refiner.py new file mode 100644 index 000000000000..faa0b039dbae --- /dev/null +++ b/src/diffusers/models/transformers/transformer_sana_wm_refiner.py @@ -0,0 +1,1556 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...utils import apply_lora_scale, is_torch_version, logging +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..cache_utils import CacheMixin +from ..embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings, PixArtAlphaTextProjection +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import RMSNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Copied from diffusers.models.transformers.transformer_ltx2.apply_interleaved_rotary_emb +def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + cos, sin = freqs + x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) # [B, S, C // 2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + return out + + +# Copied from diffusers.models.transformers.transformer_ltx2.apply_split_rotary_emb +def apply_split_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + cos, sin = freqs + + x_dtype = x.dtype + needs_reshape = False + if x.ndim != 4 and cos.ndim == 4: + # cos is (b, h, t, r) -> reshape x to (b, h, t, dim_per_head) + b, h, t, _ = cos.shape + x = x.reshape(b, t, h, -1).swapaxes(1, 2) + needs_reshape = True + + # Split last dim (2*r) into (d=2, r) + last = x.shape[-1] + if last % 2 != 0: + raise ValueError(f"Expected x.shape[-1] to be even for split rotary, got {last}.") + r = last // 2 + + # (..., 2, r) + split_x = x.reshape(*x.shape[:-1], 2, r).float() # Explicitly upcast to float + first_x = split_x[..., :1, :] # (..., 1, r) + second_x = split_x[..., 1:, :] # (..., 1, r) + + cos_u = cos.unsqueeze(-2) # broadcast to (..., 1, r) against (..., 2, r) + sin_u = sin.unsqueeze(-2) + + out = split_x * cos_u + first_out = out[..., :1, :] + second_out = out[..., 1:, :] + + first_out.addcmul_(-sin_u, second_x) + second_out.addcmul_(sin_u, first_x) + + out = out.reshape(*out.shape[:-2], last) + + if needs_reshape: + out = out.swapaxes(1, 2).reshape(b, t, -1) + + out = out.to(dtype=x_dtype) + return out + + +# Copied from diffusers.models.transformers.transformer_ltx2.LTX2AdaLayerNormSingle +class SanaWMLTX2AdaLayerNormSingle(nn.Module): + r""" + Norm layer adaptive layer norm single (adaLN-single). + + As proposed in PixArt-Alpha (see: https://huggingface.co/papers/2310.00426; Section 2.3) and adapted by the LTX-2.0 + model. In particular, the number of modulation parameters to be calculated is now configurable. + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + num_mod_params (`int`, *optional*, defaults to `6`): + The number of modulation parameters which will be calculated in the first return argument. The default of 6 + is standard, but sometimes we may want to have a different (usually smaller) number of modulation + parameters. + use_additional_conditions (`bool`, *optional*, defaults to `False`): + Whether to use additional conditions for normalization or not. + """ + + def __init__(self, embedding_dim: int, num_mod_params: int = 6, use_additional_conditions: bool = False): + super().__init__() + self.num_mod_params = num_mod_params + + self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim, size_emb_dim=embedding_dim // 3, use_additional_conditions=use_additional_conditions + ) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, self.num_mod_params * embedding_dim, bias=True) + + def forward( + self, + timestep: torch.Tensor, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + batch_size: int | None = None, + hidden_dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # No modulation happening here. + added_cond_kwargs = added_cond_kwargs or {"resolution": None, "aspect_ratio": None} + embedded_timestep = self.emb(timestep, **added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_dtype) + return self.linear(self.silu(embedded_timestep)), embedded_timestep + + +# Copied from diffusers.models.transformers.transformer_ltx2.LTX2AudioVideoAttnProcessor with LTX2->SanaWMLTX2 +class SanaWMLTX2AudioVideoAttnProcessor: + r""" + Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model. + Compared to the LTX-1.0 model, we allow the RoPE embeddings for the queries and keys to be separate so that we can + support audio-to-video (a2v) and video-to-audio (v2a) cross attention. + """ + + _attention_backend = None + _parallel_config = None + + def __init__(self): + if is_torch_version("<", "2.0"): + raise ValueError( + "LTX attention processors require a minimum PyTorch version of 2.0. Please upgrade your PyTorch installation." + ) + + def __call__( + self, + attn: "SanaWMLTX2Attention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + if attn.to_gate_logits is not None: + # Calculate gate logits on original hidden_states + gate_logits = attn.to_gate_logits(hidden_states) + + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if query_rotary_emb is not None: + if attn.rope_type == "interleaved": + query = apply_interleaved_rotary_emb(query, query_rotary_emb) + key = apply_interleaved_rotary_emb( + key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb + ) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb) + key = apply_split_rotary_emb(key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + if attn.to_gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) # [B, T, H, D] + # The factor of 2.0 is so that if the gates logits are zero-initialized the initial gates are all 1 + gates = 2.0 * torch.sigmoid(gate_logits) # [B, T, H] + hidden_states = hidden_states * gates.unsqueeze(-1) + hidden_states = hidden_states.flatten(2, 3) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +# Copied from diffusers.models.transformers.transformer_ltx2.LTX2PerturbedAttnProcessor with LTX2->SanaWMLTX2 +class SanaWMLTX2PerturbedAttnProcessor: + r""" + Processor which implements attention with perturbation masking and per-head gating for LTX-2.X models. + """ + + _attention_backend = None + _parallel_config = None + + def __init__(self): + if is_torch_version("<", "2.0"): + raise ValueError( + "LTX attention processors require a minimum PyTorch version of 2.0. Please upgrade your PyTorch installation." + ) + + def __call__( + self, + attn: "SanaWMLTX2Attention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + perturbation_mask: torch.Tensor | None = None, + all_perturbed: bool | None = None, + ) -> torch.Tensor: + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + + if attn.to_gate_logits is not None: + # Calculate gate logits on original hidden_states + gate_logits = attn.to_gate_logits(hidden_states) + + value = attn.to_v(encoder_hidden_states) + if all_perturbed is None: + all_perturbed = torch.all(perturbation_mask == 0) if perturbation_mask is not None else False + + if all_perturbed: + # Skip attention, use the value projection value + hidden_states = value + else: + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if query_rotary_emb is not None: + if attn.rope_type == "interleaved": + query = apply_interleaved_rotary_emb(query, query_rotary_emb) + key = apply_interleaved_rotary_emb( + key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb + ) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb) + key = apply_split_rotary_emb( + key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb + ) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + if perturbation_mask is not None: + value = value.flatten(2, 3) + hidden_states = torch.lerp(value, hidden_states, perturbation_mask) + + if attn.to_gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) # [B, T, H, D] + # The factor of 2.0 is so that if the gates logits are zero-initialized the initial gates are all 1 + gates = 2.0 * torch.sigmoid(gate_logits) # [B, T, H] + hidden_states = hidden_states * gates.unsqueeze(-1) + hidden_states = hidden_states.flatten(2, 3) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +# Copied from diffusers.models.transformers.transformer_ltx2.LTX2Attention with LTX2->SanaWMLTX2 +class SanaWMLTX2Attention(torch.nn.Module, AttentionModuleMixin): + r""" + Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key + RoPE embeddings separately for audio-to-video (a2v) and video-to-audio (v2a) cross-attention. + """ + + _default_processor_cls = SanaWMLTX2AudioVideoAttnProcessor + _available_processors = [SanaWMLTX2AudioVideoAttnProcessor, SanaWMLTX2PerturbedAttnProcessor] + + def __init__( + self, + query_dim: int, + heads: int = 8, + kv_heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = True, + cross_attention_dim: int | None = None, + out_bias: bool = True, + qk_norm: str = "rms_norm_across_heads", + norm_eps: float = 1e-6, + norm_elementwise_affine: bool = True, + rope_type: str = "interleaved", + apply_gated_attention: bool = False, + processor=None, + ): + super().__init__() + if qk_norm != "rms_norm_across_heads": + raise NotImplementedError("Only 'rms_norm_across_heads' is supported as a valid value for `qk_norm`.") + + self.head_dim = dim_head + self.inner_dim = dim_head * heads + self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads + self.query_dim = query_dim + self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim + self.use_bias = bias + self.dropout = dropout + self.out_dim = query_dim + self.heads = heads + self.rope_type = rope_type + + self.norm_q = torch.nn.RMSNorm(dim_head * heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine) + self.norm_k = torch.nn.RMSNorm(dim_head * kv_heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine) + self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = torch.nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) + self.to_v = torch.nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias) + self.to_out = torch.nn.ModuleList([]) + self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(torch.nn.Dropout(dropout)) + + if apply_gated_attention: + # Per head gate values + self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True) + else: + self.to_gate_logits = None + + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters] + if len(unused_kwargs) > 0: + logger.warning( + f"attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored." + ) + kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters} + hidden_states = self.processor( + self, hidden_states, encoder_hidden_states, attention_mask, query_rotary_emb, key_rotary_emb, **kwargs + ) + return hidden_states + + +# Copied from diffusers.models.transformers.transformer_ltx2.LTX2AudioVideoRotaryPosEmbed +class SanaWMLTX2AudioVideoRotaryPosEmbed(nn.Module): + """ + Video and audio rotary positional embeddings (RoPE) for the LTX-2.0 model. + + Args: + causal_offset (`int`, *optional*, defaults to `1`): + Offset in the temporal axis for causal VAE modeling. This is typically 1 (for causal modeling where the VAE + treats the very first frame differently), but could also be 0 (for non-causal modeling). + """ + + def __init__( + self, + dim: int, + patch_size: int = 1, + patch_size_t: int = 1, + base_num_frames: int = 20, + base_height: int = 2048, + base_width: int = 2048, + sampling_rate: int = 16000, + hop_length: int = 160, + scale_factors: tuple[int, ...] = (8, 32, 32), + theta: float = 10000.0, + causal_offset: int = 1, + modality: str = "video", + double_precision: bool = True, + rope_type: str = "interleaved", + num_attention_heads: int = 32, + ) -> None: + super().__init__() + + self.dim = dim + self.patch_size = patch_size + self.patch_size_t = patch_size_t + + if rope_type not in ["interleaved", "split"]: + raise ValueError(f"{rope_type=} not supported. Choose between 'interleaved' and 'split'.") + self.rope_type = rope_type + + self.base_num_frames = base_num_frames + self.num_attention_heads = num_attention_heads + + # Video-specific + self.base_height = base_height + self.base_width = base_width + + # Audio-specific + self.sampling_rate = sampling_rate + self.hop_length = hop_length + self.audio_latents_per_second = float(sampling_rate) / float(hop_length) / float(scale_factors[0]) + + self.scale_factors = scale_factors + self.theta = theta + self.causal_offset = causal_offset + + self.modality = modality + if self.modality not in ["video", "audio"]: + raise ValueError(f"Modality {modality} is not supported. Supported modalities are `video` and `audio`.") + self.double_precision = double_precision + + def prepare_video_coords( + self, + batch_size: int, + num_frames: int, + height: int, + width: int, + device: torch.device, + fps: float = 24.0, + ) -> torch.Tensor: + """ + Create per-dimension bounds [inclusive start, exclusive end) for each patch with respect to the original pixel + space video grid (num_frames, height, width). This will ultimately have shape (batch_size, 3, num_patches, 2) + where + - axis 1 (size 3) enumerates (frame, height, width) dimensions (e.g. idx 0 corresponds to frames) + - axis 3 (size 2) stores `[start, end)` indices within each dimension + + Args: + batch_size (`int`): + Batch size of the video latents. + num_frames (`int`): + Number of latent frames in the video latents. + height (`int`): + Latent height of the video latents. + width (`int`): + Latent width of the video latents. + device (`torch.device`): + Device on which to create the video grid. + + Returns: + `torch.Tensor`: + Per-dimension patch boundaries tensor of shape [batch_size, 3, num_patches, 2]. + """ + + # 1. Generate grid coordinates for each spatiotemporal dimension (frames, height, width) + # Always compute rope in fp32 + grid_f = torch.arange(start=0, end=num_frames, step=self.patch_size_t, dtype=torch.float32, device=device) + grid_h = torch.arange(start=0, end=height, step=self.patch_size, dtype=torch.float32, device=device) + grid_w = torch.arange(start=0, end=width, step=self.patch_size, dtype=torch.float32, device=device) + # indexing='ij' ensures that the dimensions are kept in order as (frames, height, width) + grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij") + grid = torch.stack(grid, dim=0) # [3, N_F, N_H, N_W], where e.g. N_F is the number of temporal patches + + # 2. Get the patch boundaries with respect to the latent video grid + patch_size = (self.patch_size_t, self.patch_size, self.patch_size) + patch_size_delta = torch.tensor(patch_size, dtype=grid.dtype, device=grid.device) + patch_ends = grid + patch_size_delta.view(3, 1, 1, 1) + + # Combine the start (grid) and end (patch_ends) coordinates along new trailing dimension + latent_coords = torch.stack([grid, patch_ends], dim=-1) # [3, N_F, N_H, N_W, 2] + # Reshape to (batch_size, 3, num_patches, 2) + latent_coords = latent_coords.flatten(1, 3) + latent_coords = latent_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1) + + # 3. Calculate the pixel space patch boundaries from the latent boundaries. + scale_tensor = torch.tensor(self.scale_factors, device=latent_coords.device) + # Broadcast the VAE scale factors such that they are compatible with latent_coords's shape + broadcast_shape = [1] * latent_coords.ndim + broadcast_shape[1] = -1 # This is the (frame, height, width) dim + # Apply per-axis scaling to convert latent coordinates to pixel space coordinates + pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) + + # As the VAE temporal stride for the first frame is 1 instead of self.vae_scale_factors[0], we need to shift + # and clamp to keep the first-frame timestamps causal and non-negative. + pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + self.causal_offset - self.scale_factors[0]).clamp(min=0) + + # Scale the temporal coordinates by the video FPS + pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps + + return pixel_coords + + def prepare_audio_coords( + self, + batch_size: int, + num_frames: int, + device: torch.device, + shift: int = 0, + ) -> torch.Tensor: + """ + Create per-dimension bounds [inclusive start, exclusive end) of start and end timestamps for each latent frame. + This will ultimately have shape (batch_size, 3, num_patches, 2) where + - axis 1 (size 1) represents the temporal dimension + - axis 3 (size 2) stores `[start, end)` indices within each dimension + + Args: + batch_size (`int`): + Batch size of the audio latents. + num_frames (`int`): + Number of latent frames in the audio latents. + device (`torch.device`): + Device on which to create the audio grid. + shift (`int`, *optional*, defaults to `0`): + Offset on the latent indices. Different shift values correspond to different overlapping windows with + respect to the same underlying latent grid. + + Returns: + `torch.Tensor`: + Per-dimension patch boundaries tensor of shape [batch_size, 1, num_patches, 2]. + """ + + # 1. Generate coordinates in the frame (time) dimension. + # Always compute rope in fp32 + grid_f = torch.arange( + start=shift, end=num_frames + shift, step=self.patch_size_t, dtype=torch.float32, device=device + ) + + # 2. Calculate start timestamps in seconds with respect to the original spectrogram grid + audio_scale_factor = self.scale_factors[0] + # Scale back to mel spectrogram space + grid_start_mel = grid_f * audio_scale_factor + # Handle first frame causal offset, ensuring non-negative timestamps + grid_start_mel = (grid_start_mel + self.causal_offset - audio_scale_factor).clip(min=0) + # Convert mel bins back into seconds + grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate + + # 3. Calculate start timestamps in seconds with respect to the original spectrogram grid + grid_end_mel = (grid_f + self.patch_size_t) * audio_scale_factor + grid_end_mel = (grid_end_mel + self.causal_offset - audio_scale_factor).clip(min=0) + grid_end_s = grid_end_mel * self.hop_length / self.sampling_rate + + audio_coords = torch.stack([grid_start_s, grid_end_s], dim=-1) # [num_patches, 2] + audio_coords = audio_coords.unsqueeze(0).expand(batch_size, -1, -1) # [batch_size, num_patches, 2] + audio_coords = audio_coords.unsqueeze(1) # [batch_size, 1, num_patches, 2] + return audio_coords + + def prepare_coords(self, *args, **kwargs): + if self.modality == "video": + return self.prepare_video_coords(*args, **kwargs) + elif self.modality == "audio": + return self.prepare_audio_coords(*args, **kwargs) + + def forward( + self, coords: torch.Tensor, device: str | torch.device | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + device = device or coords.device + + # Number of spatiotemporal dimensions (3 for video, 1 (temporal) for audio and cross attn) + num_pos_dims = coords.shape[1] + + # 1. If the coords are patch boundaries [start, end), use the midpoint of these boundaries as the patch + # position index + if coords.ndim == 4: + coords_start, coords_end = coords.chunk(2, dim=-1) + coords = (coords_start + coords_end) / 2.0 + coords = coords.squeeze(-1) # [B, num_pos_dims, num_patches] + + # 2. Get coordinates as a fraction of the base data shape + if self.modality == "video": + max_positions = (self.base_num_frames, self.base_height, self.base_width) + elif self.modality == "audio": + max_positions = (self.base_num_frames,) + # [B, num_pos_dims, num_patches] --> [B, num_patches, num_pos_dims] + grid = torch.stack([coords[:, i] / max_positions[i] for i in range(num_pos_dims)], dim=-1).to(device) + # Number of spatiotemporal dimensions (3 for video, 1 for audio and cross attn) times 2 for cos, sin + num_rope_elems = num_pos_dims * 2 + + # 3. Create a 1D grid of frequencies for RoPE + freqs_dtype = torch.float64 if self.double_precision else torch.float32 + pow_indices = torch.pow( + self.theta, + torch.linspace(start=0.0, end=1.0, steps=self.dim // num_rope_elems, dtype=freqs_dtype, device=device), + ) + freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) + + # 4. Tensor-vector outer product between pos ids tensor of shape (B, 3, num_patches) and freqs vector of shape + # (self.dim // num_elems,) + freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs # [B, num_patches, num_pos_dims, self.dim // num_elems] + freqs = freqs.transpose(-1, -2).flatten(2) # [B, num_patches, self.dim // 2] + + # 5. Get real, interleaved (cos, sin) frequencies, padded to self.dim + # TODO: consider implementing this as a utility and reuse in `connectors.py`. + # src/diffusers/pipelines/ltx2/connectors.py + if self.rope_type == "interleaved": + cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) + sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) + + if self.dim % num_rope_elems != 0: + cos_padding = torch.ones_like(cos_freqs[:, :, : self.dim % num_rope_elems]) + sin_padding = torch.zeros_like(cos_freqs[:, :, : self.dim % num_rope_elems]) + cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) + sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) + + elif self.rope_type == "split": + expected_freqs = self.dim // 2 + current_freqs = freqs.shape[-1] + pad_size = expected_freqs - current_freqs + cos_freq = freqs.cos() + sin_freq = freqs.sin() + + if pad_size != 0: + cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) + sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) + + cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) + sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) + + # Reshape freqs to be compatible with multi-head attention + b = cos_freq.shape[0] + t = cos_freq.shape[1] + + cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1) + sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1) + + cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) + sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) + + return cos_freqs, sin_freqs + + +# ``kv_cache_mode`` values accepted by [`SanaWMLTX2RefinerTransformer3DModel`] and +# [`SanaWMLTX2RefinerTransformerBlock`]. See [`SanaWMRefinerKVCache`] for the AR contract they implement. +class SanaWMRefinerKVLayerCache: + r""" + Per-layer KV cache for the SANA-WM stage-2 chunk-causal AR refiner. + + Holds the two halves of the sliding-window prefix that the refiner's self-attention attends to, plus a slot for + reading back the K/V that the last forward captured. All tensors are `(batch_size, num_tokens, inner_dim)` (i.e. + before the per-head unflatten), matching the layout the refiner's self-attention concatenates in. + + * ``sink_k_pre`` / ``sink_v``: **pre**-RoPE K/V of the attention-sink frames, captured once from the raw stage-1 + latents. They are stored pre-RoPE so each AR window can re-apply RoPE at its own shifted sink offset + (``SanaWMRefinerKVCache.sink_pe``). + * ``history_k`` / ``history_v``: **post**-RoPE K/V of the already refined recent frames, ready to be concatenated + as-is. + * ``captured_k_pre`` / ``captured_v_pre`` and ``captured_k_post`` / ``captured_v_post``: readback slots written by + the capture ``kv_cache_mode``s. + """ + + def __init__(self): + self.sink_k_pre: torch.Tensor | None = None + self.sink_v: torch.Tensor | None = None + self.history_k: torch.Tensor | None = None + self.history_v: torch.Tensor | None = None + self.captured_k_pre: torch.Tensor | None = None + self.captured_v_pre: torch.Tensor | None = None + self.captured_k_post: torch.Tensor | None = None + self.captured_v_post: torch.Tensor | None = None + + def store_sink(self, sink_k_pre: torch.Tensor, sink_v: torch.Tensor) -> None: + """Store the pre-RoPE sink K/V.""" + self.sink_k_pre = sink_k_pre + self.sink_v = sink_v + + def get_sink(self) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return the pre-RoPE sink K/V, or `None` if it has not been captured (or is empty).""" + if self.sink_k_pre is None or self.sink_v is None or self.sink_k_pre.shape[1] == 0: + return None + return self.sink_k_pre, self.sink_v + + def store_history(self, history_k: torch.Tensor, history_v: torch.Tensor) -> None: + """Store the post-RoPE recent-history K/V.""" + self.history_k = history_k + self.history_v = history_v + + def get_history(self) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return the post-RoPE recent-history K/V, or `None` if empty.""" + if self.history_k is None or self.history_v is None or self.history_k.shape[1] == 0: + return None + return self.history_k, self.history_v + + def store_captured_pre_rope(self, key: torch.Tensor, value: torch.Tensor) -> None: + """Store the pre-RoPE K/V produced by the current forward.""" + self.captured_k_pre = key + self.captured_v_pre = value + + def get_captured_pre_rope(self) -> tuple[torch.Tensor, torch.Tensor]: + """Pop the pre-RoPE K/V captured by the last forward.""" + if self.captured_k_pre is None: + raise RuntimeError("No pre-RoPE K/V was captured. Run a forward with `kv_cache_mode='capture_pre_rope'`.") + key, value = self.captured_k_pre, self.captured_v_pre + # Release the references so the caller owns the only handle. + self.captured_k_pre = self.captured_v_pre = None + return key, value + + def store_captured_post_rope(self, key: torch.Tensor, value: torch.Tensor) -> None: + """Store the post-RoPE K/V produced by the current forward.""" + self.captured_k_post = key + self.captured_v_post = value + + def get_captured_post_rope(self) -> tuple[torch.Tensor, torch.Tensor]: + """Pop the post-RoPE K/V captured by the last forward.""" + if self.captured_k_post is None: + raise RuntimeError( + "No post-RoPE K/V was captured. Run a forward with `kv_cache_mode='inject_and_capture_post_rope'`." + ) + key, value = self.captured_k_post, self.captured_v_post + self.captured_k_post = self.captured_v_post = None + return key, value + + def clear(self) -> None: + self.sink_k_pre = None + self.sink_v = None + self.history_k = None + self.history_v = None + self.captured_k_pre = None + self.captured_v_pre = None + self.captured_k_post = None + self.captured_v_post = None + + +class SanaWMRefinerKVCache: + r""" + Container holding one [`SanaWMRefinerKVLayerCache`] per transformer block, plus the shared sink RoPE. + + This implements the ``rf_shifted_sink`` KV-cache contract the SANA-WM stage-2 refiner was trained with. Refinement + is chunk-causal: `block_size` latent frames are denoised at a time while attending to a bounded window of + `[attention sink + recent history + active block]` K/V. + + * ``sink_pe``: the `(cos, sin)` RoPE tuple for the sink frames, rebuilt per AR window at the sliding + ``sink_rope_offset`` so the sink sits immediately before the bounded working cache. Shared across layers because + RoPE does not depend on the layer. + + Args: + num_layers (`int`): + Number of transformer blocks to allocate a per-layer cache for. + """ + + def __init__(self, num_layers: int): + self.layer_caches = [SanaWMRefinerKVLayerCache() for _ in range(num_layers)] + self.sink_pe: tuple[torch.Tensor, torch.Tensor] | None = None + + def __len__(self) -> int: + return len(self.layer_caches) + + def get(self, layer_idx: int) -> SanaWMRefinerKVLayerCache: + return self.layer_caches[layer_idx] + + def clear(self) -> None: + for layer_cache in self.layer_caches: + layer_cache.clear() + self.sink_pe = None + + +class SanaWMLTX2RefinerAttnProcessor: + """Self-attention over `[sink + history + current]` K/V for the sliding-window AR refiner. + + Unlike the plain LTX-2 processors this one is cache-aware: `kv_cache_mode` decides whether the layer cache's sink + and recent-history K/V are prepended before the single SDPA call, and whether the current block's K/V is written + back for the next window. + """ + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: SanaWMLTX2Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + sink_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + kv_cache: SanaWMRefinerKVLayerCache | None = None, + kv_cache_mode: str | None = None, + ) -> torch.Tensor: + """LTX-2 self-attention over `[sink + history + current]` K/V. + + The queries always come from the active block only. Depending on `kv_cache_mode`, the layer cache's pre-RoPE + sink K/V (re-RoPE'd here with `sink_rotary_emb`) and post-RoPE recent-history K/V are prepended to the current + K/V before a single SDPA call, and/or the current K/V is written back to the cache. + """ + del encoder_hidden_states, attention_mask, key_rotary_emb + + gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + # Capture PRE-RoPE (post-norm) K/V so a future window can re-apply RoPE at its shifted sink offset. + if kv_cache_mode == "capture_pre_rope": + kv_cache.store_captured_pre_rope(key.detach().clone(), value.detach().clone()) + + if attn.rope_type == "interleaved": + query = apply_interleaved_rotary_emb(query, query_rotary_emb) + key = apply_interleaved_rotary_emb(key, query_rotary_emb) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb) + key = apply_split_rotary_emb(key, query_rotary_emb) + else: + raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}") + + # Capture POST-RoPE K/V so the next window can concatenate the recent history directly. Deliberately taken + # before the prefix is prepended, so only the current block's tokens are recorded. + if kv_cache_mode == "inject_and_capture_post_rope": + kv_cache.store_captured_post_rope(key.detach().clone(), value.detach().clone()) + + if kv_cache_mode in ("inject", "inject_and_capture_post_rope"): + prefix_k_parts: list[torch.Tensor] = [] + prefix_v_parts: list[torch.Tensor] = [] + sink_kv = kv_cache.get_sink() + if sink_kv is not None: + if sink_rotary_emb is None: + raise ValueError("Injecting the attention sink requires the `sink_pe` RoPE tuple on the KV cache.") + sink_k_pre, sink_v = sink_kv + sink_k_pre = sink_k_pre.to(key.dtype) + if attn.rope_type == "interleaved": + sink_k = apply_interleaved_rotary_emb(sink_k_pre, sink_rotary_emb) + else: + sink_k = apply_split_rotary_emb(sink_k_pre, sink_rotary_emb) + prefix_k_parts.append(sink_k) + prefix_v_parts.append(sink_v.to(value.dtype)) + history_kv = kv_cache.get_history() + if history_kv is not None: + prefix_k_parts.append(history_kv[0].to(key.dtype)) + prefix_v_parts.append(history_kv[1].to(value.dtype)) + if prefix_k_parts: + key = torch.cat([*prefix_k_parts, key], dim=1) + value = torch.cat([*prefix_v_parts, value], dim=1) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + if gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) + gates = 2.0 * torch.sigmoid(gate_logits) + hidden_states = hidden_states * gates.unsqueeze(-1) + hidden_states = hidden_states.flatten(2, 3) + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class SanaWMLTX2RefinerTransformerBlock(nn.Module): + r""" + Video-only, streaming-attention variant of [`LTX2VideoTransformerBlock`] used by the SANA-WM stage-2 refiner. + + The submodule structure is copied verbatim from [`LTX2VideoTransformerBlock`] (so LTX-2 checkpoints load as-is); + only [`~SanaWMLTX2RefinerTransformerBlock.forward`] differs. It runs the video stream only (self-attn -> prompt + cross-attn -> feed-forward), skipping the audio and audio/video cross-attention branches, and routes the + self-attention through a KV-cached sliding window instead of plain full self-attention. + """ + + # Copied from diffusers.models.transformers.transformer_ltx2.LTX2VideoTransformerBlock.__init__ with LTX2->SanaWMLTX2 + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + cross_attention_dim: int, + audio_dim: int, + audio_num_attention_heads: int, + audio_attention_head_dim, + audio_cross_attention_dim: int, + video_gated_attn: bool = False, + video_cross_attn_adaln: bool = False, + audio_gated_attn: bool = False, + audio_cross_attn_adaln: bool = False, + qk_norm: str = "rms_norm_across_heads", + activation_fn: str = "gelu-approximate", + attention_bias: bool = True, + attention_out_bias: bool = True, + eps: float = 1e-6, + elementwise_affine: bool = False, + rope_type: str = "interleaved", + perturbed_attn: bool = False, + ff_bias: bool = True, + audio_ff_bias: bool = True, + ): + super().__init__() + + self.perturbed_attn = perturbed_attn + if perturbed_attn: + attn_processor_cls = SanaWMLTX2PerturbedAttnProcessor + else: + attn_processor_cls = SanaWMLTX2AudioVideoAttnProcessor + + # 1. Self-Attention (video and audio) + self.norm1 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) + self.attn1 = SanaWMLTX2Attention( + query_dim=dim, + heads=num_attention_heads, + kv_heads=num_attention_heads, + dim_head=attention_head_dim, + bias=attention_bias, + cross_attention_dim=None, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + apply_gated_attention=video_gated_attn, + processor=attn_processor_cls(), + ) + + self.audio_norm1 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) + self.audio_attn1 = SanaWMLTX2Attention( + query_dim=audio_dim, + heads=audio_num_attention_heads, + kv_heads=audio_num_attention_heads, + dim_head=audio_attention_head_dim, + bias=attention_bias, + cross_attention_dim=None, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + apply_gated_attention=audio_gated_attn, + processor=attn_processor_cls(), + ) + + # 2. Prompt Cross-Attention + self.norm2 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) + self.attn2 = SanaWMLTX2Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + kv_heads=num_attention_heads, + dim_head=attention_head_dim, + bias=attention_bias, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + apply_gated_attention=video_gated_attn, + processor=attn_processor_cls(), + ) + + self.audio_norm2 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) + self.audio_attn2 = SanaWMLTX2Attention( + query_dim=audio_dim, + cross_attention_dim=audio_cross_attention_dim, + heads=audio_num_attention_heads, + kv_heads=audio_num_attention_heads, + dim_head=audio_attention_head_dim, + bias=attention_bias, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + apply_gated_attention=audio_gated_attn, + processor=attn_processor_cls(), + ) + + # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention + # Audio-to-Video (a2v) Attention --> Q: Video; K,V: Audio + self.audio_to_video_norm = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) + self.audio_to_video_attn = SanaWMLTX2Attention( + query_dim=dim, + cross_attention_dim=audio_dim, + heads=audio_num_attention_heads, + kv_heads=audio_num_attention_heads, + dim_head=audio_attention_head_dim, + bias=attention_bias, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + apply_gated_attention=video_gated_attn, + processor=attn_processor_cls(), + ) + + # Video-to-Audio (v2a) Attention --> Q: Audio; K,V: Video + self.video_to_audio_norm = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) + self.video_to_audio_attn = SanaWMLTX2Attention( + query_dim=audio_dim, + cross_attention_dim=dim, + heads=audio_num_attention_heads, + kv_heads=audio_num_attention_heads, + dim_head=audio_attention_head_dim, + bias=attention_bias, + out_bias=attention_out_bias, + qk_norm=qk_norm, + rope_type=rope_type, + apply_gated_attention=audio_gated_attn, + processor=attn_processor_cls(), + ) + + # 4. Feedforward layers + self.norm3 = RMSNorm(dim, eps=eps, elementwise_affine=elementwise_affine) + self.ff = FeedForward(dim, activation_fn=activation_fn, bias=ff_bias) + + self.audio_norm3 = RMSNorm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) + self.audio_ff = FeedForward(audio_dim, activation_fn=activation_fn, bias=audio_ff_bias) + + # 5. Per-Layer Modulation Parameters + # Self-Attention (attn1) / Feedforward AdaLayerNorm-Zero mod params + # 6 base mod params for text cross-attn K,V; if cross_attn_adaln, also has mod params for Q + self.video_cross_attn_adaln = video_cross_attn_adaln + self.audio_cross_attn_adaln = audio_cross_attn_adaln + video_mod_param_num = 9 if self.video_cross_attn_adaln else 6 + audio_mod_param_num = 9 if self.audio_cross_attn_adaln else 6 + self.scale_shift_table = nn.Parameter(torch.randn(video_mod_param_num, dim) / dim**0.5) + self.audio_scale_shift_table = nn.Parameter(torch.randn(audio_mod_param_num, audio_dim) / audio_dim**0.5) + + # Prompt cross-attn (attn2) additional modulation params + self.cross_attn_adaln = video_cross_attn_adaln or audio_cross_attn_adaln + if self.cross_attn_adaln: + self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim)) + self.audio_prompt_scale_shift_table = nn.Parameter(torch.randn(2, audio_dim)) + + # Per-layer a2v, v2a Cross-Attention mod params + self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim)) + self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, audio_dim)) + + @staticmethod + # Copied from diffusers.models.transformers.transformer_ltx2.LTX2VideoTransformerBlock.get_mod_params + def get_mod_params( + scale_shift_table: torch.Tensor, temb: torch.Tensor, batch_size: int + ) -> tuple[torch.Tensor, ...]: + num_ada_params = scale_shift_table.shape[0] + ada_values = scale_shift_table[None, None].to(temb.device) + temb.reshape( + batch_size, temb.shape[1], num_ada_params, -1 + ) + ada_params = ada_values.unbind(dim=2) + return ada_params + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + video_rotary_emb: tuple[torch.Tensor, torch.Tensor], + encoder_attention_mask: torch.Tensor | None = None, + sink_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + kv_cache: SanaWMRefinerKVLayerCache | None = None, + kv_cache_mode: str | None = None, + ) -> torch.Tensor: + batch_size = hidden_states.size(0) + + # 1. Video self-attention over the KV-cached sliding window + norm_hidden_states = self.norm1(hidden_states) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.get_mod_params( + self.scale_shift_table, temb, batch_size + ) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + + attn_hidden_states = self.attn1( + norm_hidden_states, + query_rotary_emb=video_rotary_emb, + sink_rotary_emb=sink_rotary_emb, + kv_cache=kv_cache, + kv_cache_mode=kv_cache_mode, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa + + # 2. Prompt cross-attention + norm_hidden_states = self.norm2(hidden_states) + attn_hidden_states = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + query_rotary_emb=None, + attention_mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states + + # 3. Feed-forward + norm_hidden_states = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp + hidden_states = hidden_states + self.ff(norm_hidden_states) * gate_mlp + return hidden_states + + +class SanaWMLTX2RefinerTransformer3DModel( + ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin, PeftAdapterMixin, CacheMixin +): + r""" + The chunk-causal autoregressive refiner transformer used as SANA-WM stage 2. + + Architecturally identical to [`LTX2VideoTransformer3DModel`] — same config arguments, same submodules, same + parameter names — so a released LTX-2 checkpoint loads into it unchanged. What differs is the forward pass: + + * only the video stream is run (the audio and audio/video cross-attention branches are skipped), + * self-attention runs against an explicit sliding-window KV cache ([`SanaWMRefinerKVCache`]) holding the attention + sink plus recent refined history, so refinement cost is bounded per AR block and scales linearly with video + length, + * the caller supplies the video RoPE, which lets each AR window keep every frame's absolute index in the source + video (see [`~SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions`]). + + Args: + in_channels (`int`, defaults to `128`): + The number of channels in the input. + out_channels (`int`, defaults to `128`): + The number of channels in the output. + patch_size (`int`, defaults to `1`): + The size of the spatial patches to use in the patch embedding layer. + patch_size_t (`int`, defaults to `1`): + The size of the temporal patches to use in the patch embedding layer. + num_attention_heads (`int`, defaults to `32`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, defaults to `128`): + The number of channels in each head. + cross_attention_dim (`int`, defaults to `4096`): + The number of channels for cross attention heads. + num_layers (`int`, defaults to `48`): + The number of layers of Transformer blocks to use. + activation_fn (`str`, defaults to `"gelu-approximate"`): + Activation function to use in feed-forward. + qk_norm (`str`, defaults to `"rms_norm_across_heads"`): + The normalization layer to use. + rope_type (`str`, defaults to `"interleaved"`): + Which RoPE application to use (`"interleaved"` or `"split"`). + + The remaining arguments mirror [`LTX2VideoTransformer3DModel`] one-for-one. The audio-side arguments and submodules + are kept purely so the checkpoint's audio weights round-trip; they are not used by the refiner forward. + """ + + _skip_layerwise_casting_patterns = ["norm"] + _repeated_blocks = ["SanaWMLTX2RefinerTransformerBlock"] + _skip_keys = ["kv_cache"] + + @register_to_config + def __init__( + self, + in_channels: int = 128, # Video Arguments + out_channels: int | None = 128, + patch_size: int = 1, + patch_size_t: int = 1, + num_attention_heads: int = 32, + attention_head_dim: int = 128, + cross_attention_dim: int = 4096, + vae_scale_factors: tuple[int, int, int] = (8, 32, 32), + pos_embed_max_pos: int = 20, + base_height: int = 2048, + base_width: int = 2048, + gated_attn: bool = False, + cross_attn_mod: bool = False, + audio_in_channels: int = 128, # Audio Arguments + audio_out_channels: int | None = 128, + audio_patch_size: int = 1, + audio_patch_size_t: int = 1, + audio_num_attention_heads: int = 32, + audio_attention_head_dim: int = 64, + audio_cross_attention_dim: int = 2048, + audio_scale_factor: int = 4, + audio_pos_embed_max_pos: int = 20, + audio_sampling_rate: int = 16000, + audio_hop_length: int = 160, + audio_gated_attn: bool = False, + audio_cross_attn_mod: bool = False, + num_layers: int = 48, # Shared arguments + activation_fn: str = "gelu-approximate", + qk_norm: str = "rms_norm_across_heads", + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + caption_channels: int = 3840, + attention_bias: bool = True, + attention_out_bias: bool = True, + rope_theta: float = 10000.0, + rope_double_precision: bool = True, + causal_offset: int = 1, + timestep_scale_multiplier: int = 1000, + cross_attn_timestep_scale_multiplier: int = 1000, + rope_type: str = "interleaved", + use_prompt_embeddings=True, + perturbed_attn: bool = False, + ff_bias: bool = True, + audio_ff_bias: bool = True, + use_prompt_adaln_single: bool = True, + use_keyframes_abs_pos_embedding: bool = False, + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + audio_out_channels = audio_out_channels or audio_in_channels + inner_dim = num_attention_heads * attention_head_dim + audio_inner_dim = audio_num_attention_heads * audio_attention_head_dim + + # 1. Patchification input projections + self.proj_in = nn.Linear(in_channels, inner_dim) + self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim) + + if use_keyframes_abs_pos_embedding: + self.keyframes_abs_pos_embedding = nn.Parameter(torch.zeros(1, inner_dim)) + + # 2. Prompt embeddings + if use_prompt_embeddings: + self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) + self.audio_caption_projection = PixArtAlphaTextProjection( + in_features=caption_channels, hidden_size=audio_inner_dim + ) + + # 3. Timestep Modulation Params and Embedding + self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod + + # 3.1. Global Timestep Modulation Parameters (except for cross-attention) and timestep + size embedding + video_time_emb_mod_params = 9 if cross_attn_mod else 6 + audio_time_emb_mod_params = 9 if audio_cross_attn_mod else 6 + self.time_embed = SanaWMLTX2AdaLayerNormSingle( + inner_dim, num_mod_params=video_time_emb_mod_params, use_additional_conditions=False + ) + self.audio_time_embed = SanaWMLTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=audio_time_emb_mod_params, use_additional_conditions=False + ) + + # 3.2. Global Cross Attention Modulation Parameters + self.av_cross_attn_video_scale_shift = SanaWMLTX2AdaLayerNormSingle( + inner_dim, num_mod_params=4, use_additional_conditions=False + ) + self.av_cross_attn_audio_scale_shift = SanaWMLTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=4, use_additional_conditions=False + ) + self.av_cross_attn_video_a2v_gate = SanaWMLTX2AdaLayerNormSingle( + inner_dim, num_mod_params=1, use_additional_conditions=False + ) + self.av_cross_attn_audio_v2a_gate = SanaWMLTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=1, use_additional_conditions=False + ) + + # 3.3. Output Layer Scale/Shift Modulation parameters + self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) + self.audio_scale_shift_table = nn.Parameter(torch.randn(2, audio_inner_dim) / audio_inner_dim**0.5) + + # 3.4. Prompt Scale/Shift Modulation parameters (LTX-2.3) + if self.prompt_modulation and use_prompt_adaln_single: + self.prompt_adaln = SanaWMLTX2AdaLayerNormSingle( + inner_dim, num_mod_params=2, use_additional_conditions=False + ) + self.audio_prompt_adaln = SanaWMLTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=2, use_additional_conditions=False + ) + + # 4. Rotary Positional Embeddings (RoPE) + self.rope = SanaWMLTX2AudioVideoRotaryPosEmbed( + dim=inner_dim, + patch_size=patch_size, + patch_size_t=patch_size_t, + base_num_frames=pos_embed_max_pos, + base_height=base_height, + base_width=base_width, + scale_factors=vae_scale_factors, + theta=rope_theta, + causal_offset=causal_offset, + modality="video", + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=num_attention_heads, + ) + self.audio_rope = SanaWMLTX2AudioVideoRotaryPosEmbed( + dim=audio_inner_dim, + patch_size=audio_patch_size, + patch_size_t=audio_patch_size_t, + base_num_frames=audio_pos_embed_max_pos, + sampling_rate=audio_sampling_rate, + hop_length=audio_hop_length, + scale_factors=[audio_scale_factor], + theta=rope_theta, + causal_offset=causal_offset, + modality="audio", + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=audio_num_attention_heads, + ) + + # Audio-to-Video, Video-to-Audio Cross-Attention + cross_attn_pos_embed_max_pos = max(pos_embed_max_pos, audio_pos_embed_max_pos) + self.cross_attn_rope = SanaWMLTX2AudioVideoRotaryPosEmbed( + dim=audio_cross_attention_dim, + patch_size=patch_size, + patch_size_t=patch_size_t, + base_num_frames=cross_attn_pos_embed_max_pos, + base_height=base_height, + base_width=base_width, + theta=rope_theta, + causal_offset=causal_offset, + modality="video", + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=num_attention_heads, + ) + self.cross_attn_audio_rope = SanaWMLTX2AudioVideoRotaryPosEmbed( + dim=audio_cross_attention_dim, + patch_size=audio_patch_size, + patch_size_t=audio_patch_size_t, + base_num_frames=cross_attn_pos_embed_max_pos, + sampling_rate=audio_sampling_rate, + hop_length=audio_hop_length, + theta=rope_theta, + causal_offset=causal_offset, + modality="audio", + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=audio_num_attention_heads, + ) + + # 5. Transformer Blocks + self.transformer_blocks = nn.ModuleList( + [ + SanaWMLTX2RefinerTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + cross_attention_dim=cross_attention_dim, + audio_dim=audio_inner_dim, + audio_num_attention_heads=audio_num_attention_heads, + audio_attention_head_dim=audio_attention_head_dim, + audio_cross_attention_dim=audio_cross_attention_dim, + video_gated_attn=gated_attn, + video_cross_attn_adaln=cross_attn_mod, + audio_gated_attn=audio_gated_attn, + audio_cross_attn_adaln=audio_cross_attn_mod, + qk_norm=qk_norm, + activation_fn=activation_fn, + attention_bias=attention_bias, + attention_out_bias=attention_out_bias, + eps=norm_eps, + elementwise_affine=norm_elementwise_affine, + rope_type=rope_type, + perturbed_attn=perturbed_attn, + ff_bias=ff_bias, + audio_ff_bias=audio_ff_bias, + ) + for _ in range(num_layers) + ] + ) + # The blocks are built by LTX-2's `__init__`, which installs LTX-2's own self-attention processor. Swap in + # the KV-cached sliding-window one the AR refiner needs; the module and its weights are otherwise identical. + for block in self.transformer_blocks: + block.attn1.set_processor(SanaWMLTX2RefinerAttnProcessor()) + + # 6. Output layers + self.norm_out = nn.LayerNorm(inner_dim, eps=1e-6, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels) + + self.audio_norm_out = nn.LayerNorm(audio_inner_dim, eps=1e-6, elementwise_affine=False) + self.audio_proj_out = nn.Linear(audio_inner_dim, audio_out_channels) + + self.gradient_checkpointing = False + + def build_rotary_emb_for_absolute_positions( + self, + batch_size: int, + frame_positions: list[int], + height: int, + width: int, + device: torch.device, + fps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Build the video RoPE for an explicit list of absolute latent-frame indices. + + [`SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords`] assumes a contiguous `torch.arange(num_frames)`, + which is fine for bidirectional inference. The sliding-window AR refiner instead needs to keep each frame's + absolute index in the source video, so RoPE captures the correct temporal phase across the `[sink + recent + + active]` window. + + Args: + batch_size (`int`): + Batch size to broadcast the coordinates to. + frame_positions (`list[int]`): + Absolute latent-frame indices covered by this window. + height (`int`), width (`int`): + Latent spatial resolution. + device (`torch.device`): + Device to build the coordinates on. + fps (`float`): + Video frame rate, which drives LTX-2's temporal RoPE scaling. + + Returns: + `tuple[torch.Tensor, torch.Tensor]`: the `(cos, sin)` RoPE tuple. + """ + rope = self.rope + patch_size_t = int(rope.patch_size_t) + patch_size = int(rope.patch_size) + f_positions = torch.tensor(frame_positions, dtype=torch.float32, device=device) + if patch_size_t > 1: + # Each patch covers ``patch_size_t`` latent frames; pick the start of each patch. + f_positions = f_positions[::patch_size_t] + grid_h = torch.arange(start=0, end=height, step=patch_size, dtype=torch.float32, device=device) + grid_w = torch.arange(start=0, end=width, step=patch_size, dtype=torch.float32, device=device) + grid = torch.meshgrid(f_positions, grid_h, grid_w, indexing="ij") + grid = torch.stack(grid, dim=0) + + patch_size_delta = torch.tensor((patch_size_t, patch_size, patch_size), dtype=grid.dtype, device=device) + patch_ends = grid + patch_size_delta.view(3, 1, 1, 1) + latent_coords = torch.stack([grid, patch_ends], dim=-1) + latent_coords = latent_coords.flatten(1, 3).unsqueeze(0).repeat(batch_size, 1, 1, 1) + + scale_tensor = torch.tensor(rope.scale_factors, device=device) + broadcast_shape = [1] * latent_coords.ndim + broadcast_shape[1] = -1 + pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) + pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + rope.causal_offset - rope.scale_factors[0]).clamp(min=0) + pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / float(fps) + return rope(pixel_coords, device=device) + + @apply_lora_scale("attention_kwargs") + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + video_rotary_emb: tuple[torch.Tensor, torch.Tensor], + encoder_attention_mask: torch.Tensor | None = None, + kv_cache: SanaWMRefinerKVCache | None = None, + kv_cache_mode: str | None = None, + attention_kwargs: dict | None = None, + return_dict: bool = True, + ): + r""" + Video-only forward pass over a single AR block. + + Args: + hidden_states (`torch.Tensor`): + Patchified video latents of the active block, of shape `(batch_size, num_video_tokens, in_channels)`. + encoder_hidden_states (`torch.Tensor`): + Text embeddings of shape `(batch_size, text_seq_len, caption_channels)`. + timestep (`torch.Tensor`): + Timestep of shape `(batch_size, num_video_tokens)`, already scaled by + `self.config.timestep_scale_multiplier`. + video_rotary_emb (`tuple[torch.Tensor, torch.Tensor]`): + The `(cos, sin)` RoPE for the active block's absolute frame positions, as returned by + [`~SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions`]. + encoder_attention_mask (`torch.Tensor`, *optional*): + Multiplicative text attention mask of shape `(batch_size, text_seq_len)`. + kv_cache (`SanaWMRefinerKVCache`, *optional*): + Sliding-window KV cache holding the per-layer attention sink and recent refined history, plus the + shared `sink_pe` RoPE. Required whenever `kv_cache_mode` is set. + kv_cache_mode (`str`, *optional*): + One of: + + - `"inject"`: attend to `[sink + history + current]` K/V (the denoising steps). + - `"capture_pre_rope"`: no prefix; record the pre-RoPE K/V of this forward into the cache (used once to + seed the attention sink from the raw stage-1 latents). + - `"inject_and_capture_post_rope"`: attend to `[sink + history + current]` K/V and record this block's + post-RoPE K/V into the cache so it can be appended to the history. + + When `None`, the block runs plain full self-attention over the current tokens only. + attention_kwargs (`dict`, *optional*): + Optional kwargs forwarded to the LoRA scale handling. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a plain tuple. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: the predicted velocity for the active + block, of shape `(batch_size, num_video_tokens, out_channels)`. + """ + if kv_cache_mode is not None: + if kv_cache_mode not in ("inject", "capture_pre_rope", "inject_and_capture_post_rope"): + raise ValueError( + "`kv_cache_mode` must be one of 'inject', 'capture_pre_rope', " + f"'inject_and_capture_post_rope' or `None`, got {kv_cache_mode!r}." + ) + if kv_cache is None: + raise ValueError(f"`kv_cache_mode={kv_cache_mode!r}` requires a `SanaWMRefinerKVCache`.") + if len(kv_cache) != len(self.transformer_blocks): + raise ValueError( + f"`kv_cache` holds {len(kv_cache)} layer caches but the model has " + f"{len(self.transformer_blocks)} transformer blocks." + ) + + batch_size = hidden_states.size(0) + + # Convert encoder_attention_mask to an additive bias. + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: + encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 1. Patchification input projection + hidden_states = self.proj_in(hidden_states) + + # 2. Timestep embedding and modulation parameters + temb, embedded_timestep = self.time_embed( + timestep.flatten(), + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(batch_size, -1, temb.size(-1)) + embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) + + # 3. Prompt embeddings + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) + + # 4. Transformer blocks + sink_rotary_emb = kv_cache.sink_pe if kv_cache is not None else None + for i, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + sink_rotary_emb=sink_rotary_emb, + kv_cache=kv_cache.get(i) if kv_cache is not None else None, + kv_cache_mode=kv_cache_mode, + ) + + # 5. Output norm and projection + scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None] + shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] + hidden_states = self.norm_out(hidden_states) + hidden_states = hidden_states * (1 + scale) + shift + output = self.proj_out(hidden_states) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 32f193a03080..5934cc2952f3 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -388,6 +388,11 @@ "SanaVideoPipeline", "SanaImageToVideoPipeline", ] + _import_structure["sana_wm"] = [ + "SanaWMPipeline", + "SanaWMLTX2Refiner", + "SanaWMPipelineOutput", + ] _import_structure["shap_e"] = ["ShapEImg2ImgPipeline", "ShapEPipeline"] _import_structure["stable_audio"] = [ "StableAudioProjectionModel", @@ -873,6 +878,11 @@ SanaSprintPipeline, ) from .sana_video import SanaImageToVideoPipeline, SanaVideoPipeline + from .sana_wm import ( + SanaWMLTX2Refiner, + SanaWMPipeline, + SanaWMPipelineOutput, + ) from .shap_e import ShapEImg2ImgPipeline, ShapEPipeline from .stable_audio import StableAudioPipeline, StableAudioProjectionModel from .stable_audio_3 import ( diff --git a/src/diffusers/pipelines/sana_wm/__init__.py b/src/diffusers/pipelines/sana_wm/__init__.py new file mode 100644 index 000000000000..c33e4f615751 --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/__init__.py @@ -0,0 +1,49 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_output"] = ["SanaWMPipelineOutput"] + _import_structure["pipeline_sana_wm"] = ["SanaWMPipeline"] + _import_structure["refiner"] = ["SanaWMLTX2Refiner"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_output import SanaWMPipelineOutput + from .pipeline_sana_wm import SanaWMPipeline + from .refiner import SanaWMLTX2Refiner +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/sana_wm/cam_utils.py b/src/diffusers/pipelines/sana_wm/cam_utils.py new file mode 100644 index 000000000000..8865e46acd56 --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/cam_utils.py @@ -0,0 +1,321 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Camera + image utilities for the SANA-WM pipeline. + +* Action-string DSL → camera-to-world trajectory. +* Resize-and-center-crop to (704, 1280) with intrinsics adjustment. +* Plücker / raymap packing for the DiT camera-control branch. +""" + +from __future__ import annotations + +import math + +import numpy as np +import torch +from PIL import Image + + +TARGET_HEIGHT = 704 +TARGET_WIDTH = 1280 + +DEFAULT_TRANSLATION_SPEED = 0.05 +DEFAULT_ROTATION_SPEED_DEG = 1.2 +DEFAULT_PITCH_LIMIT_DEG = 85.0 +ALLOWED_ACTION_KEYS: frozenset[str] = frozenset("wasdijkl") + + +# --------------------------------------------------------------------------- +# Action DSL → camera-to-world trajectory +# --------------------------------------------------------------------------- + + +def _rot_x(angle_rad: float) -> np.ndarray: + c, s = np.cos(angle_rad), np.sin(angle_rad) + return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float64) + + +def _rot_y(angle_rad: float) -> np.ndarray: + c, s = np.cos(angle_rad), np.sin(angle_rad) + return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64) + + +def _parse_action_string(action: str) -> list[list[str]]: + cleaned = "".join(action.replace("，", ",").split()) + if not cleaned: + raise ValueError("action string is empty") + per_frame: list[list[str]] = [] + for segment in cleaned.split(","): + if not segment or "-" not in segment: + raise ValueError(f"Invalid action segment {segment!r}: expected '-'.") + keys_part, dur_str = segment.rsplit("-", 1) + if not dur_str.isdigit() or int(dur_str) <= 0: + raise ValueError(f"Action segment {segment!r} has a non-positive duration {dur_str!r}.") + n = int(dur_str) + keys_lower = keys_part.lower() + if keys_lower == "none": + keys: list[str] = [] + else: + bad = sorted({c for c in keys_lower if c not in ALLOWED_ACTION_KEYS}) + if bad: + raise ValueError( + f"Action segment {segment!r} contains unknown keys {bad}; " + f"allowed: {''.join(sorted(ALLOWED_ACTION_KEYS))}." + ) + keys = sorted(set(keys_lower)) + per_frame.extend([list(keys) for _ in range(n)]) + return per_frame + + +def action_string_to_c2w( + action: str, + *, + translation_speed: float = DEFAULT_TRANSLATION_SPEED, + rotation_speed_deg: float = DEFAULT_ROTATION_SPEED_DEG, + pitch_limit_deg: float = DEFAULT_PITCH_LIMIT_DEG, +) -> np.ndarray: + """Roll out a ``(N+1, 4, 4)`` c2w trajectory from a WASD+IJKL action DSL. + + Coordinate convention: OpenCV (``+X right, +Y down, +Z forward``). WASD translates on the world XZ plane; IJKL + applies pitch / yaw. + """ + per_frame = _parse_action_string(action) + rotate_rad = math.radians(rotation_speed_deg) + pitch_limit_rad = math.radians(pitch_limit_deg) + current = np.eye(4, dtype=np.float64) + poses = [current.copy()] + current_pitch = 0.0 + + for keys in per_frame: + held = set(keys) + R = current[:3, :3] + T_ = current[:3, 3] + + pitch_delta = (rotate_rad if "i" in held else 0.0) - (rotate_rad if "k" in held else 0.0) + new_pitch = current_pitch + pitch_delta + if not (-pitch_limit_rad <= new_pitch <= pitch_limit_rad): + pitch_delta = 0.0 + else: + current_pitch = new_pitch + + yaw_delta = (rotate_rad if "l" in held else 0.0) - (rotate_rad if "j" in held else 0.0) + R_new = _rot_y(yaw_delta) @ R @ _rot_x(pitch_delta) + + forward = R_new[:, 2].copy() + forward[1] = 0.0 + right = R_new[:, 0].copy() + right[1] = 0.0 + if (fn := float(np.linalg.norm(forward))) > 0: + forward /= fn + if (rn := float(np.linalg.norm(right))) > 0: + right /= rn + move = np.zeros(3, dtype=np.float64) + if "w" in held: + move += forward * translation_speed + if "s" in held: + move -= forward * translation_speed + if "d" in held: + move += right * translation_speed + if "a" in held: + move -= right * translation_speed + + current = np.eye(4, dtype=np.float64) + current[:3, :3] = R_new + current[:3, 3] = T_ + move + poses.append(current.copy()) + + return np.stack(poses, axis=0).astype(np.float32) + + +# --------------------------------------------------------------------------- +# Intrinsics handling +# --------------------------------------------------------------------------- + + +def transform_intrinsics_for_crop( + intrinsics_vec4: np.ndarray, + src_size: tuple[int, int], + resized_size: tuple[int, int], + crop_offset: tuple[int, int], +) -> np.ndarray: + """Adjust ``[fx, fy, cx, cy]`` to match a resize-then-center-crop image.""" + src_w, src_h = src_size + rw, rh = resized_size + cl, ct = crop_offset + sx, sy = rw / src_w, rh / src_h + out = intrinsics_vec4.copy() + out[..., 0] *= sx + out[..., 2] = out[..., 2] * sx - cl + out[..., 1] *= sy + out[..., 3] = out[..., 3] * sy - ct + return out + + +# --------------------------------------------------------------------------- +# Image preprocessing +# --------------------------------------------------------------------------- + + +def resize_and_center_crop( + image: Image.Image, + target_h: int = TARGET_HEIGHT, + target_w: int = TARGET_WIDTH, +) -> tuple[Image.Image, tuple[int, int], tuple[int, int], tuple[int, int]]: + """Aspect-preserving resize then center-crop to ``(target_h, target_w)``.""" + src_w, src_h = image.size + scale = max(target_h / src_h, target_w / src_w) + rw = max(target_w, int(round(src_w * scale))) + rh = max(target_h, int(round(src_h * scale))) + resized = image.resize((rw, rh), Image.LANCZOS) + left = (rw - target_w) // 2 + top = (rh - target_h) // 2 + cropped = resized.crop((left, top, left + target_w, top + target_h)) + return cropped, (src_w, src_h), (rw, rh), (left, top) + + +# --------------------------------------------------------------------------- +# Camera condition packing — Plücker + raymap +# --------------------------------------------------------------------------- + + +def compute_raymap( + intrinsics: torch.Tensor, + poses: torch.Tensor, + H: int, + W: int, + *, + use_plucker: bool = True, +) -> torch.Tensor: + """Compute a per-pixel ray geometry map. + + Args: + intrinsics: ``(T, 4)`` ``[fx, fy, cx, cy]`` per frame. + poses: ``(T, 4, 4)`` camera-to-world poses (OpenCV convention). + H: spatial height. + W: spatial width. + use_plucker: if True returns Plücker coordinates ``(d, m)``; otherwise + returns ``(origin, direction)``. + + Returns: + ``(T, H, W, 6)`` tensor. + """ + T = intrinsics.shape[0] + device = intrinsics.device + dtype = intrinsics.dtype + y_grid, x_grid = torch.meshgrid( + torch.arange(H, device=device, dtype=dtype), + torch.arange(W, device=device, dtype=dtype), + indexing="ij", + ) + x_grid = x_grid[None].expand(T, -1, -1) + y_grid = y_grid[None].expand(T, -1, -1) + fx = intrinsics[:, 0].view(T, 1, 1) + fy = intrinsics[:, 1].view(T, 1, 1) + cx = intrinsics[:, 2].view(T, 1, 1) + cy = intrinsics[:, 3].view(T, 1, 1) + dirs_cam = torch.stack( + [(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], + dim=-1, + ) + R = poses[:, :3, :3] + t = poses[:, :3, 3] + dirs_world = torch.einsum("tij,thwj->thwi", R, dirs_cam) + dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True) + origins = t.view(T, 1, 1, 3).expand_as(dirs_world) + if use_plucker: + moments = torch.cross(origins, dirs_world, dim=-1) + return torch.cat([dirs_world, moments], dim=-1) + return torch.cat([origins, dirs_world], dim=-1) + + +def _pose_inverse(T44: torch.Tensor) -> torch.Tensor: + R = T44[..., :3, :3] + t = T44[..., :3, 3:] + Rt = R.transpose(-1, -2) + out = torch.zeros_like(T44) + out[..., :3, :3] = Rt + out[..., :3, 3:] = -Rt @ t + out[..., 3, 3] = 1.0 + return out + + +def prepare_camera( + poses_c2w: np.ndarray, + intrinsics_vec4: np.ndarray, + *, + target_size: tuple[int, int], + vae_stride: tuple[int, int, int], +) -> dict[str, torch.Tensor]: + """Build the DiT-input camera tensors. + + Returns a dict with: + + * ``raymap`` ``(T_lat, 20)`` — flattened (rel-pose, intrinsics) per latent frame + * ``chunk_plucker`` ``(6 * vae_time_stride, T_lat, H_lat, W_lat)`` — Plücker coordinates packed by chunk. + """ + num_frames = poses_c2w.shape[0] + vae_time_stride, vae_spatial_stride = vae_stride[0], vae_stride[-1] + H_pixel, W_pixel = target_size + latent_h = H_pixel // vae_spatial_stride + latent_w = W_pixel // vae_spatial_stride + latent_frames = (num_frames - 1) // vae_time_stride + 1 + + poses = torch.from_numpy(poses_c2w).float() + first_inv = _pose_inverse(poses[0:1]).squeeze(0) + poses_rel = torch.matmul(first_inv, poses[1:]) + poses = torch.cat([torch.eye(4).unsqueeze(0), poses_rel], dim=0) + + intrinsics = torch.from_numpy(intrinsics_vec4).float() + intrinsics_latent = intrinsics.clone() + intrinsics_latent[:, [0, 2]] *= latent_w / float(W_pixel) + intrinsics_latent[:, [1, 3]] *= latent_h / float(H_pixel) + + time_indices = torch.arange(0, num_frames, vae_time_stride) + if len(time_indices) > latent_frames: + time_indices = time_indices[:latent_frames] + + raymap = torch.cat( + [poses[time_indices].reshape(len(time_indices), -1), intrinsics_latent[time_indices]], + dim=-1, + ) + + chunk_starts = time_indices - (vae_time_stride - 1) + chunks = [] + for start in chunk_starts: + s = max(0, int(start)) + e = s + vae_time_stride + chunk_poses, chunk_intrs = poses[s:e], intrinsics_latent[s:e] + if chunk_poses.shape[0] < vae_time_stride: + pad = vae_time_stride - chunk_poses.shape[0] + chunk_poses = torch.cat([chunk_poses, chunk_poses[-1:].repeat(pad, 1, 1)], dim=0) + chunk_intrs = torch.cat([chunk_intrs, chunk_intrs[-1:].repeat(pad, 1)], dim=0) + plucker = compute_raymap(chunk_intrs, chunk_poses, latent_h, latent_w, use_plucker=True) + chunks.append(plucker.permute(0, 3, 1, 2).reshape(-1, latent_h, latent_w)) + chunk_plucker = torch.stack(chunks).permute(1, 0, 2, 3) + return {"raymap": raymap, "chunk_plucker": chunk_plucker} + + +def snap_num_frames(n: int, stride: int = 8, *, upper_bound: int | None = None) -> int: + """Snap ``n`` to the nearest ``stride*k + 1`` (LTX-2 VAE constraint).""" + if n < 1: + return 1 + if (n - 1) % stride == 0: + return n + floor_cand = n - ((n - 1) % stride) + ceil_cand = floor_cand + stride + snapped = floor_cand if (n - floor_cand) < (ceil_cand - n) else ceil_cand + if upper_bound is not None and snapped > upper_bound: + snapped = floor_cand + return max(snapped, 1) diff --git a/src/diffusers/pipelines/sana_wm/image_processor.py b/src/diffusers/pipelines/sana_wm/image_processor.py new file mode 100644 index 000000000000..c18326cd33fd --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/image_processor.py @@ -0,0 +1,67 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import numpy as np +import PIL.Image +import torch + +from ...configuration_utils import register_to_config +from ...image_processor import VaeImageProcessor +from .cam_utils import TARGET_HEIGHT, TARGET_WIDTH, resize_and_center_crop, transform_intrinsics_for_crop + + +class SanaWMImageProcessor(VaeImageProcessor): + r""" + Image processor for SANA-WM's first-frame input. + + SANA-WM was trained at a fixed 704×1280 resolution with an aspect-preserving *resize + center-crop* transform. The + pipeline also needs to rescale the per-frame camera intrinsics ``[fx, fy, cx, cy]`` to match the crop — + ``preprocess_with_intrinsics`` does both in one call so the two stay in lockstep. + + Args: + vae_scale_factor (`int`, defaults to `32`): + LTX-2 VAE spatial stride. + do_normalize (`bool`, defaults to `True`): + Standard `VaeImageProcessor` [-1, 1] normalization. + """ + + @register_to_config + def __init__(self, vae_scale_factor: int = 32, do_normalize: bool = True) -> None: + super().__init__(vae_scale_factor=vae_scale_factor, do_normalize=do_normalize) + + def preprocess_with_intrinsics( + self, + image: PIL.Image.Image, + intrinsics: np.ndarray, + height: int = TARGET_HEIGHT, + width: int = TARGET_WIDTH, + ) -> tuple[torch.Tensor, np.ndarray]: + """Resize + center-crop the image and rescale ``intrinsics`` to match. + + Args: + image: RGB PIL image (any size). + intrinsics: ``(F, 4)`` ``[fx, fy, cx, cy]`` per frame in original-image pixel coordinates. + height / width: Target crop size (defaults to SANA-WM's training resolution). + + Returns: + ``(pixel_values, intrinsics_cropped)``: + * ``pixel_values`` — ``(1, 3, H, W)`` tensor in `[-1, 1]` (VaeImageProcessor convention). + * ``intrinsics_cropped`` — ``(F, 4)`` array rescaled for the resize + crop. + """ + cropped, src_size, resized_size, crop_offset = resize_and_center_crop(image, height, width) + pixel_values = self.preprocess(cropped, height=height, width=width) + intr = transform_intrinsics_for_crop(intrinsics, src_size, resized_size, crop_offset) + return pixel_values, intr diff --git a/src/diffusers/pipelines/sana_wm/pipeline_output.py b/src/diffusers/pipelines/sana_wm/pipeline_output.py new file mode 100644 index 000000000000..9a007071c32f --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/pipeline_output.py @@ -0,0 +1,43 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +import numpy as np +import PIL.Image +import torch + +from ...utils import BaseOutput + + +@dataclass +class SanaWMPipelineOutput(BaseOutput): + """ + Output class for SANA-WM image-to-video pipeline. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or `list[PIL.Image.Image]`): + Generated video. Shape ``(T, H, W, 3)`` as a float ``np.ndarray`` / ``torch.Tensor`` in ``[0, 1]`` when + ``output_type="np"`` / ``"latent"``, or a list of ``PIL.Image`` of length ``T`` when ``output_type="pil"``. + c2w (`np.ndarray`): + Camera-to-world poses ``(T, 4, 4)`` aligned with ``frames`` (the refiner drops the sink anchor frame; this + array is realigned accordingly when the refiner ran). + latent (`torch.Tensor`, optional): + Latent tensor in LTX-2 VAE space, shape ``(B, C, T_lat, H_lat, W_lat)``. Returned when + ``output_type="latent"``. + """ + + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] + c2w: np.ndarray | None = None + latent: torch.Tensor | None = None diff --git a/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py b/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py new file mode 100644 index 000000000000..d14ec03901c2 --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py @@ -0,0 +1,631 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import inspect +from pathlib import Path +from typing import Literal + +import numpy as np +import PIL.Image +import torch +from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast + +from ...models import AutoencoderKLLTX2Video, SanaWMTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .cam_utils import ( + TARGET_HEIGHT, + TARGET_WIDTH, + action_string_to_c2w, + prepare_camera, + snap_num_frames, +) +from .image_processor import SanaWMImageProcessor +from .pipeline_output import SanaWMPipelineOutput + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from PIL import Image + >>> from diffusers import SanaWMPipeline + + >>> pipe = SanaWMPipeline.from_pretrained( + ... "Efficient-Large-Model/SANA-WM_bidirectional-diffusers", torch_dtype=torch.bfloat16 + ... ).to("cuda") + + >>> output = 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", + ... intrinsics=[800.0, 800.0, 845.0, 464.0], # fx, fy, cx, cy in original-image pixels + ... num_inference_steps=60, + ... ) + >>> # output.frames is (T, H, W, 3) float np.ndarray in [0, 1] (diffusers convention). + ``` +""" + + +# Default instruction prefix prepended to the user prompt before Gemma-2 encoding. +# SANA-WM was trained with this prefix, so changing it degrades prompt adherence. +DEFAULT_CHI_PROMPT: list[str] = [ + 'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:', + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.", + "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.", + "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", + "User Prompt: ", +] + + +class SanaWMPipeline(DiffusionPipeline): + r""" + SANA-WM camera-controlled image-to-video pipeline. + + Generates a video from a first-frame image, a text prompt, and a camera trajectory (explicit ``c2w`` poses or a + WASD/IJKL action string). Uses the 1600M bidirectional SANA DiT for stage-1 sampling and the LTX-2 + sink-bidirectional Euler refiner for stage-2 polish; both decode through the LTX-2 VAE. + + Args: + tokenizer ([`GemmaTokenizer`] or [`GemmaTokenizerFast`]): + The Gemma-2 tokenizer. + text_encoder ([`Gemma2PreTrainedModel`]): + The Gemma-2 text encoder. + vae ([`AutoencoderKLLTX2Video`]): + The LTX-2 VAE. + transformer ([`SanaWMTransformer3DModel`]): + The 1600M bidirectional SANA-WM DiT. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + Flow-matching Euler scheduler (LTX-style per-token timesteps). latents directly. + """ + + # the offload sequence; it manages its own sub-module device placement. + model_cpu_offload_seq = "text_encoder->transformer->vae" + + def __init__( + self, + tokenizer: GemmaTokenizer | GemmaTokenizerFast, + text_encoder: Gemma2PreTrainedModel, + vae: AutoencoderKLLTX2Video, + transformer: SanaWMTransformer3DModel, + scheduler: FlowMatchEulerDiscreteScheduler, + ) -> None: + super().__init__() + self.register_modules( + tokenizer=tokenizer, + text_encoder=text_encoder, + vae=vae, + transformer=transformer, + scheduler=scheduler, + ) + # Read VAE strides from the registered component (LTX2Pipeline pattern). + # Fall back to the LTX-2 defaults (32 spatial / 8 temporal) if the VAE + # hasn't been registered yet — matches SANA-WM's training config. + self.vae_spatial_compression_ratio = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.vae_temporal_compression_ratio = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + # ``image_processor`` handles first-frame input (resize + center-crop + # + [-1, 1] normalization + intrinsics rescale for the crop); + # ``video_processor`` handles the decoded [-1, 1] video -> user-chosen + # ``output_type`` conversion. + self.image_processor = SanaWMImageProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) + + # ------------------------------------------------------------------ + # Prompt encoding + # ------------------------------------------------------------------ + + def encode_prompt( + self, + prompt: str, + negative_prompt: str = "", + *, + device: torch.device, + max_sequence_length: int = 300, + chi_prompt: list[str] | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode prompt + negative prompt through Gemma-2. + + Mirrors the SANA chi-prompt-prefix trick: the chi prompt is prepended to the user prompt, then a ``select_index + = [0, -L+1, ..., -1]`` slice takes the BOS token plus the last ``max_sequence_length - 1`` tokens. + + Returns: + ``(cond, cond_mask, neg, neg_mask)`` where ``cond`` and ``neg`` are ``(1, 1, L, D)``-shaped Gemma hidden + states and the masks are ``(1, L)``. + """ + if (prompt_embeds is None) != (prompt_attention_mask is None): + raise ValueError("`prompt_embeds` and `prompt_attention_mask` must be passed together.") + if (negative_prompt_embeds is None) != (negative_prompt_attention_mask is None): + raise ValueError("`negative_prompt_embeds` and `negative_prompt_attention_mask` must be passed together.") + if prompt_embeds is not None and negative_prompt_embeds is not None: + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + chi = "\n".join(chi_prompt) if chi_prompt else "" + if chi: + full_prompt = chi + prompt + max_length_all = len(self.tokenizer.encode(chi)) + max_sequence_length - 2 + else: + full_prompt = prompt + max_length_all = max_sequence_length + + def _encode(text: str, length: int) -> tuple[torch.Tensor, torch.Tensor]: + # SANA was trained with right-padded prompts; Gemma's tokenizer defaults to "left". + tok = self.tokenizer( + [text], + max_length=length, + padding="max_length", + padding_side="right", + truncation=True, + return_tensors="pt", + ).to(device) + # Go through the outer ``Gemma2ForCausalLM`` so the CPU-offload + # hook moves the encoder to GPU; grab the final-layer hidden + # states (== ``Gemma2Model.last_hidden_state``). + out = self.text_encoder( + input_ids=tok.input_ids, + attention_mask=tok.attention_mask, + output_hidden_states=True, + return_dict=True, + ) + return out.hidden_states[-1], tok.attention_mask + + if prompt_embeds is None: + cond, cond_mask = _encode(full_prompt, max_length_all) + select = [0] + list(range(-max_sequence_length + 1, 0)) + prompt_embeds = cond[:, None][:, :, select] + prompt_attention_mask = cond_mask[:, select] + + if negative_prompt_embeds is None: + neg, neg_mask = _encode(negative_prompt, max_sequence_length) + negative_prompt_embeds, negative_prompt_attention_mask = neg[:, None], neg_mask + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + # ------------------------------------------------------------------ + # First-frame VAE encode (deterministic — uses posterior mode) + # ------------------------------------------------------------------ + + def _encode_first_frame( + self, pixel_values: torch.Tensor, device: torch.device, dtype: torch.dtype + ) -> torch.Tensor: + # ``pixel_values`` is ``(1, 3, H, W)`` in [-1, 1] (from ``SanaWMImageProcessor``). + # Add the temporal axis to match the LTX-2 VAE input shape ``(B, C, 1, H, W)``. + img = pixel_values.unsqueeze(2).to(device, dtype=self.vae.dtype) + z = self.vae.encode(img).latent_dist.mode() + latents_mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(z) + latents_std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(z) + z = (z - latents_mean) * self.vae.config.scaling_factor / latents_std + return z.to(dtype) + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + """Decode latents to a `(B, C, F, H, W)` tensor in `[-1, 1]` (the VAE's native output range). + + Post-processing (e.g. `[-1, 1]` → PIL frames / `np.ndarray` in `[0, 1]`) is handled by + `self.video_processor.postprocess_video` at the call site so callers get the diffusers convention that the + `VideoProcessor` / `export_to_video` helpers assume. + """ + latents = latents.to(self.vae.device, dtype=self.vae.dtype) + latents_mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents) + latents_std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents) + latents = latents / self.vae.config.scaling_factor * latents_std + latents_mean + return self.vae.decode(latents, return_dict=False)[0] + + # ------------------------------------------------------------------ + # Camera conditioning packing + # ------------------------------------------------------------------ + + def _build_camera_kwargs( + self, + c2w: np.ndarray, + intrinsics_vec4: np.ndarray, + target_size: tuple[int, int], + *, + device: torch.device, + dtype: torch.dtype, + do_cfg: bool, + ) -> dict[str, torch.Tensor]: + cam = prepare_camera( + c2w, + intrinsics_vec4, + target_size=target_size, + vae_stride=( + self.vae_temporal_compression_ratio, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ), + ) + raymap = cam["raymap"].unsqueeze(0).to(device, dtype=dtype) + chunk_plucker = cam["chunk_plucker"].unsqueeze(0).to(device, dtype=dtype) + if do_cfg: + raymap = torch.cat([raymap, raymap], dim=0) + chunk_plucker = torch.cat([chunk_plucker, chunk_plucker], dim=0) + return {"camera_conditions": raymap, "chunk_plucker": chunk_plucker} + + def check_inputs( + self, + c2w: np.ndarray | None, + action: str | None, + intrinsics: np.ndarray | list[float] | None, + ) -> None: + """Validate `__call__` inputs. Raises on bad input and returns nothing.""" + if (c2w is None) == (action is None): + raise ValueError("Provide exactly one of `c2w` or `action`.") + if c2w is not None: + poses = np.asarray(c2w, dtype=np.float32) + if poses.ndim != 3 or poses.shape[1:] != (4, 4): + raise ValueError(f"`c2w` must be `(F, 4, 4)`; got {poses.shape}.") + if intrinsics is None: + raise ValueError( + "Pass `intrinsics` as either `[fx, fy, cx, cy]`, a 3x3 K matrix, " + "an `(F, 4)` per-frame [fx,fy,cx,cy], or `(F, 3, 3)` per-frame K — " + "all in original-image pixel coordinates. Use " + "the `Efficient-Large-Model/pi3x-intrinsics-estimator` modular block for an automatic " + "estimate." + ) + + def prepare_camera_trajectory( + self, + c2w: np.ndarray | None, + action: str | None, + intrinsics: np.ndarray | list[float], + num_frames: int, + ) -> tuple[np.ndarray, np.ndarray]: + """Normalize the camera inputs to ``(c2w_(F, 4, 4), intrinsics_(F, 4))``. + + Also snaps ``num_frames`` to the VAE-friendly ``8k+1`` and trims both arrays to match. The cropped image and + rescaled intrinsics come later, once the target resolution is known. + """ + if action is not None: + c2w = action_string_to_c2w(action) + c2w = np.asarray(c2w, dtype=np.float32) + + num_frames = min(num_frames, c2w.shape[0]) + num_frames = snap_num_frames(num_frames, stride=self.vae_temporal_compression_ratio, upper_bound=c2w.shape[0]) + c2w = c2w[:num_frames] + + intr = np.asarray(intrinsics, dtype=np.float32) + # Accept (3, 3), (F, 3, 3), (4,) and (F, 4) — normalize to (F, 4). + if intr.shape == (3, 3): + intr = np.array([intr[0, 0], intr[1, 1], intr[0, 2], intr[1, 2]], dtype=np.float32) + elif intr.ndim == 3 and intr.shape[-2:] == (3, 3): + intr = np.stack([intr[:, 0, 0], intr[:, 1, 1], intr[:, 0, 2], intr[:, 1, 2]], axis=-1) + if intr.shape == (4,): + intr = np.broadcast_to(intr, (num_frames, 4)).copy() + if intr.ndim == 2 and intr.shape[1] == 4 and intr.shape[0] >= num_frames: + # Caller may pass a full-trajectory intrinsics array; trim to match. + intr = intr[:num_frames] + if intr.shape != (num_frames, 4): + raise ValueError( + f"`intrinsics` must be `(4,)`, `(F>={num_frames}, 4)`, `(3, 3)`, or " + f"`(F>={num_frames}, 3, 3)`; got shape {np.asarray(intrinsics).shape}." + ) + return c2w, intr + + def prepare_latents( + self, + first_latent: torch.Tensor, + num_frames: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Sample initial latents and pin the first frame as the conditioning anchor. + + Returns ``(latents, condition_mask)`` where ``condition_mask`` has ones on the first frame's tokens (they are + held clean throughout sampling) and zeros elsewhere. + """ + latent_T = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_h = height // self.vae_spatial_compression_ratio + latent_w = width // self.vae_spatial_compression_ratio + latent_channels = first_latent.shape[1] + latents = randn_tensor( + (1, latent_channels, latent_T, latent_h, latent_w), + generator=generator, + device=device, + dtype=dtype, + ) + latents[:, :, :1] = first_latent + condition_mask = torch.zeros_like(latents) + condition_mask[:, :, :1] = 1.0 + return latents, condition_mask + + # ------------------------------------------------------------------ + # __call__ + # ------------------------------------------------------------------ + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + image: PIL.Image.Image | str | Path, + prompt: str, + *, + c2w: np.ndarray | None = None, + action: str | None = None, + intrinsics: np.ndarray | list[float] | None = None, + height: int = TARGET_HEIGHT, + width: int = TARGET_WIDTH, + num_frames: int = 161, + fps: int = 16, + num_inference_steps: int = 60, + guidance_scale: float = 5.0, + negative_prompt: str = "", + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int = 300, + chi_prompt: list[str] | None = None, + output_type: Literal["np", "pil", "latent"] = "np", + return_dict: bool = True, + ) -> SanaWMPipelineOutput | tuple: + r""" + Generate a SANA-WM camera-controlled video. + + Args: + image (`PIL.Image.Image` or `str`): + First-frame image (PIL or path). + prompt (`str`): + Text prompt. + c2w (`np.ndarray`, *optional*): + ``(F, 4, 4)`` camera-to-world poses. Mutually exclusive with `action`. + action (`str`, *optional*): + Action-DSL string e.g. ``"w-80,jw-40,w-40"``. Mutually exclusive with `c2w`. + intrinsics (`np.ndarray` or `list[float]`): + ``[fx, fy, cx, cy]`` in **original-image** pixel coordinates. The pipeline applies the resize+crop + transform internally. + height (`int`, defaults to 704): + Output frame height (fixed for the public model). + width (`int`, defaults to 1280): + Output frame width (fixed for the public model). + num_frames (`int`, defaults to 161): + Target frame count; snapped to ``8k+1`` (LTX-2 VAE constraint). + fps (`int`, defaults to 16): + Output frame rate (also fed to the refiner). + num_inference_steps (`int`, defaults to 60): + Number of stage-1 DiT sampling steps. + guidance_scale (`float`, defaults to 5.0): + Classifier-free guidance scale. + negative_prompt (`str`, defaults to ""): + Optional negative prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or more torch generators to make the noise sampling deterministic. If `None`, sampling is + non-deterministic. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-computed text embeddings, to skip the text encoder. Must be passed together with + `prompt_attention_mask`. + prompt_attention_mask (`torch.Tensor`, *optional*): + Attention mask for `prompt_embeds`. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-computed negative text embeddings. Must be passed together with `negative_prompt_attention_mask`. + negative_prompt_attention_mask (`torch.Tensor`, *optional*): + Attention mask for `negative_prompt_embeds`. + max_sequence_length (`int`, defaults to 300): + Max prompt tokens. + chi_prompt (`list[str]`, *optional*): + Override the chi-prompt prefix (default mirrors the public release). + output_type (`"np"`, `"pil"`, or `"latent"`, defaults to `"np"`): + Output format. + return_dict (`bool`, defaults to True): + Return [`SanaWMPipelineOutput`] vs tuple. + + Returns: + [`SanaWMPipelineOutput`] with `.frames` of shape ``(T, H, W, 3)``, float ``np.ndarray`` in ``[0, 1]`` for + `output_type="np"`, a list of ``PIL.Image.Image`` of length ``T`` for `"pil"`, or the raw latent tensor for + `"latent"`. + + Examples: + """ + self.check_inputs(c2w, action, intrinsics) + c2w, intr = self.prepare_camera_trajectory(c2w, action, intrinsics, num_frames) + num_frames = c2w.shape[0] + pixel_values, intr = self.image_processor.preprocess_with_intrinsics(image, intr, height, width) + + device = self._execution_device + dtype = self.transformer.dtype + + cond, cond_mask, neg, neg_mask = self.encode_prompt( + prompt, + negative_prompt, + device=device, + max_sequence_length=max_sequence_length, + chi_prompt=chi_prompt or DEFAULT_CHI_PROMPT, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + + first_latent = self._encode_first_frame(pixel_values, device, dtype) + cam_kwargs = self._build_camera_kwargs( + c2w, intr, (height, width), device=device, dtype=dtype, do_cfg=guidance_scale > 1.0 + ) + + self._guidance_scale = guidance_scale + self._current_timestep = None + self._interrupt = False + do_cfg = self.do_classifier_free_guidance + + # Stage-1 denoising — LTX-style flow-matching Euler with per-token + # timesteps. The first latent frame is the conditioning anchor: its + # per-token timestep is pinned to 0 so it is never denoised away. + latents, condition_mask = self.prepare_latents( + first_latent, num_frames, height, width, dtype, device, generator + ) + timesteps, _ = retrieve_timesteps(self.scheduler, num_inference_steps, device, None) + self._num_timesteps = len(timesteps) + + prompt_embeds = torch.cat([neg, cond], dim=0) if do_cfg else cond + mask_cfg = torch.cat([neg_mask, cond_mask], dim=0) if do_cfg else cond_mask + model_kwargs = { + "data_info": { + "img_hw": torch.tensor([[height, width]], dtype=torch.float, device=device), + }, + "mask": mask_cfg, + **cam_kwargs, + } + + for t in self.progress_bar(timesteps): + if self.interrupt: + continue + self._current_timestep = t + cond_mask_input = torch.cat([condition_mask] * 2) if do_cfg else condition_mask + latent_model_input = torch.cat([latents] * 2) if do_cfg else latents + timestep = t.expand(cond_mask_input.shape).float() + timestep = torch.min(timestep, (1.0 - cond_mask_input) * 1000.0) + + noise_pred = self.transformer( + latent_model_input, + timestep[:, :1, :, 0, 0], # (B, 1, T) + prompt_embeds, + return_dict=False, + **model_kwargs, + )[0] + + if do_cfg: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + timestep = timestep.chunk(2)[0] + + B, C, F, H, W = latents.shape + denoised = self.scheduler.step( + -noise_pred.reshape(B, C, -1).transpose(1, 2), + t, + latents.reshape(B, C, -1).transpose(1, 2), + per_token_timesteps=timestep.reshape(B, C, -1)[:, 0], + return_dict=False, + )[0] + denoised = denoised.transpose(1, 2).reshape(B, C, F, H, W) + keep_clean = t / 1000.0 - 1e-6 < (1.0 - condition_mask) + latents = torch.where(keep_clean, denoised, latents).to(dtype) + + if output_type == "latent": + if not return_dict: + return (latents, c2w, latents) + return SanaWMPipelineOutput(frames=latents, c2w=c2w, latent=latents) + + self._current_timestep = None + + decoded = self._decode_latents(latents) # (B=1, C=3, F, H, W) in [-1, 1] + video_c2w = c2w[:num_frames] + + # ``VideoProcessor.postprocess_video`` handles the standard [-1, 1] -> + # requested output_type conversion (uint8 PIL frames, float np.ndarray + # in [0, 1], or the raw pt tensor). + frames = self.video_processor.postprocess_video(decoded, output_type=output_type)[0] + + if not return_dict: + return (frames, video_c2w, latents) + return SanaWMPipelineOutput(frames=frames, c2w=video_c2w, latent=latents) diff --git a/src/diffusers/pipelines/sana_wm/refiner.py b/src/diffusers/pipelines/sana_wm/refiner.py new file mode 100644 index 000000000000..7b3bd228bed6 --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/refiner.py @@ -0,0 +1,515 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LTX-2 chunk-causal AR refiner used as SANA-WM stage 2. + +Wraps [`SanaWMLTX2RefinerTransformer3DModel`] (an LTX-2 DiT with a sliding-window KV cache and a video-only forward) +plus ``LTX2TextConnectors`` and a Gemma-3 text encoder. + +Refinement is chunk-causal / autoregressive (``block_size=3``, ``kv_max_frames=11``): ``block_size`` latent frames are +processed at a time over a sliding window of ``[source_sink + recent_history + active_block]`` K/V. The model was +trained with this contract; per-block compute is bounded by the window size, so total cost scales linearly with video +length. +""" + +from __future__ import annotations + +import torch +from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast + +from ...models.autoencoders import AutoencoderKLLTX2Video +from ...models.transformers.transformer_sana_wm_refiner import ( + SanaWMLTX2RefinerTransformer3DModel, + SanaWMRefinerKVCache, +) +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor + +# TODO: `LTX2TextConnectors` lives in the LTX-2 pipeline folder, so stage 2 has to reach across +# pipelines for it. Once https://github.com/huggingface/diffusers/issues/14749 moves the connector +# to a shared home (e.g. `models/`), import it from there and drop this cross-pipeline import. +from ..ltx2.connectors import LTX2TextConnectors +from ..pipeline_utils import DiffusionPipeline + + +# Sigma schedule for the 3-step distilled refiner (matches the public release). +STAGE_2_DISTILLED_SIGMA_VALUES: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0) + + +class SanaWMLTX2Refiner(DiffusionPipeline): + r""" + LTX-2 sink-bidirectional Euler refiner — SANA-WM stage 2, as a standalone pipeline. + + Wraps the LTX-2 components (refiner transformer + text connectors + Gemma-3 text encoder + tokenizer) plus a + [`FlowMatchEulerDiscreteScheduler`] that carries the distilled sigma schedule and performs the Euler steps. It is + registered as an optional component of [`SanaWMPipeline`] and can also be used on its own to refine stage-1 + latents. + + Args: + transformer ([`SanaWMLTX2RefinerTransformer3DModel`]): + The LTX-2 video DiT with the chunk-causal sliding-window KV cache. + connectors ([`LTX2TextConnectors`]): + LTX-2 text connectors. + tokenizer: + Gemma-3 tokenizer. + text_encoder: + Gemma-3 text encoder. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + Flow-matching Euler scheduler. Constructed with ``shift=1.0`` so the distilled sigmas pass through + unmodified. + vae ([`AutoencoderKLLTX2Video`], *optional*): + The same VAE used by [`SanaWMPipeline`]; pass `vae=pipe.vae` to share the weights. When given, the refiner + decodes to video, otherwise it returns refined latents. + text_max_sequence_length (`int`, defaults to 1024): + Maximum tokens passed to the Gemma-3 tokenizer. + """ + + model_cpu_offload_seq = "text_encoder->connectors->transformer->vae" + _optional_components = ["vae"] + + def __init__( + self, + transformer: SanaWMLTX2RefinerTransformer3DModel, + connectors: LTX2TextConnectors, + tokenizer: GemmaTokenizer | GemmaTokenizerFast, + text_encoder: Gemma3ForConditionalGeneration, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLLTX2Video | None = None, + text_max_sequence_length: int = 1024, + ) -> None: + super().__init__() + self.register_modules( + transformer=transformer, + connectors=connectors, + tokenizer=tokenizer, + text_encoder=text_encoder, + scheduler=scheduler, + vae=vae, + ) + self.video_processor = VideoProcessor( + vae_scale_factor=self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.register_to_config(text_max_sequence_length=int(text_max_sequence_length)) + self.text_max_sequence_length = int(text_max_sequence_length) + + # ------------------------------------------------------------------ + # forward + # ------------------------------------------------------------------ + + @torch.no_grad() + def __call__( + self, + sana_latent: torch.Tensor, + prompt: str, + *, + fps: float, + sink_size: int = 1, + generator: torch.Generator | None = None, + block_size: int = 3, + kv_max_frames: int = 11, + sigmas: tuple[float, ...] = STAGE_2_DISTILLED_SIGMA_VALUES, + output_type: str = "np", + ) -> torch.Tensor: + """Run the LTX-2 refiner and return refined VAE latents. + + Uses the chunk-causal AR recipe the model was trained on (``block_size=3``, ``kv_max_frames=11``): a sliding + window of ``[source_sink + recent_history + active_block]`` K/V is fed to the transformer one block at a time, + so per-block compute is bounded and total refinement cost scales linearly with video length. + + Args: + sana_latent: ``(B, C, F, H, W)`` stage-1 latent. + prompt: text prompt. + fps: video frame rate (drives LTX-2 RoPE temporal scaling). + sink_size: how many leading raw ``z_sana`` frames to anchor as the + attention sink (canonical: 1). + generator: torch.Generator for the FM endpoint noise. Defaults to a generator seeded with 42 + so results are reproducible out of the box. + block_size: latent frames per AR block (canonical: 3). + kv_max_frames: maximum context+active frames retained in the + sliding window (canonical: 11 = 1 sink + 10 recent). + sigmas: descending Euler schedule terminating at 0.0 (canonical + 3-step distilled: ``(0.909375, 0.725, 0.421875, 0.0)``). Fed to ``self.scheduler`` (minus the trailing + 0.0, which the scheduler appends itself). + output_type: `"latent"` returns the refined latents. Anything else decodes through `self.vae` and + post-processes to that type (`"np"`, `"pt"`, `"pil"`); without a `vae` the latents are returned + regardless. + + Returns: + `torch.Tensor`: Refined VAE latents of shape ``(B, C, F, H, W)`` — the first ``sink_size`` frames carry the + raw stage-1 sink latents unchanged, the rest carry the refined output. + """ + if sana_latent.shape[2] <= sink_size: + raise ValueError(f"Stage-1 latent has {sana_latent.shape[2]} frames but sink_size={sink_size}.") + + # Stage 2 is memory hungry (a Gemma-3 text encoder plus a 48-layer DiT), so it is meant to be + # run under `enable_model_cpu_offload()`: `model_cpu_offload_seq` walks + # `text_encoder -> connectors -> transformer -> vae`, which is exactly the order below, so each + # sub-model is on the accelerator only while it runs. + device = self._execution_device + dtype = next(self.transformer.parameters()).dtype + transformer_config = self.transformer.config + sink_size = int(sink_size) + block_size = int(block_size) + if generator is None: + generator = torch.Generator(device=device).manual_seed(42) + + # 1. Load the distilled sigma schedule into the scheduler. Drop the trailing + # 0.0 — ``FlowMatchEulerDiscreteScheduler.set_timesteps`` appends the + # terminal 0.0 itself, so ``self.scheduler.sigmas`` reproduces ``sigmas``. + self.scheduler.set_timesteps(sigmas=list(sigmas[:-1]), device=device) + sigmas_t = self.scheduler.sigmas.to(device=device, dtype=torch.float32) + sigma_max = float(sigmas_t[0]) + + # 2. Encode the prompt. + prompt_embeds, prompt_attention_mask = self._encode_prompt(prompt, device=device, dtype=dtype) + + # 3. Prepare the latents. The output keeps the raw sink prefix verbatim; the AR blocks fill + # frames [sink_size, num_frames). + z = sana_latent.to(device=device, dtype=dtype) + latents = z.clone() + batch_size, _, num_frames, height, width = z.shape + num_blocks = (num_frames - sink_size + block_size - 1) // block_size + + # 4. Chunk-causal AR refinement implementing the canonical `rf_shifted_sink` KV-cache contract: + # + # a. Pre-capture **pre-RoPE** sink K/V from raw `z_sana[:sink_size]` at sigma=0. The sink frames themselves + # are never refined — they sit unchanged in the output volume. + # b. AR blocks cover frames `[sink_size, num_frames)` in `block_size`-frame chunks. For each block: + # - Initialize `x_t = (1-sigma_0) * z_sana_block + sigma_0 * eps` (single eps per block). + # - 3-step deterministic Euler. Each step injects the per-layer prefix + # `{sink_k_pre, sink_v, sink_pe, history_k, history_v}`, where `sink_pe` is rebuilt at + # `sink_rope_offset = block_start - history_frames - sink_size` so the sink slides to sit immediately + # before the bounded working cache. + # - Capture **post-RoPE** K/V from the refined block under the same prefix, append it to the history, + # and trim the history to `kv_max_frames - sink_size` frames. + num_layers = len(self.transformer.transformer_blocks) + max_history_frames = int(kv_max_frames) - sink_size + # ``_pack_latents`` emits ``(T // patch_size_t) * (H // p) * (W // p)`` tokens, so a single latent + # frame contributes ``(H // p) * (W // p) / patch_size_t`` tokens. (No-op for LTX-2, which uses + # ``patch_size_t=1``.) + tokens_per_frame = ( + (height // transformer_config.patch_size) + * (width // transformer_config.patch_size) + // transformer_config.patch_size_t + ) + history_frames = 0 + + kv_cache = SanaWMRefinerKVCache(num_layers) + self._capture_block_kv( + clean_block=z[:, :, :sink_size].contiguous(), + frame_positions=list(range(sink_size)), + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + fps=fps, + kv_cache=kv_cache, + kv_cache_mode="capture_pre_rope", + device=device, + ) + for layer_idx in range(num_layers): + layer_cache = kv_cache.get(layer_idx) + layer_cache.store_sink(*layer_cache.get_captured_pre_rope()) + + with self.progress_bar(total=num_blocks) as progress_bar: + for block_idx in range(num_blocks): + block_start = sink_size + block_idx * block_size + block_end = min(block_start + block_size, num_frames) + clean_block = z[:, :, block_start:block_end] + frame_positions = list(range(block_start, block_end)) + + # Slide the sink's RoPE so it sits immediately before the bounded working cache. + sink_rope_offset = block_start - history_frames - sink_size + kv_cache.sink_pe = self.transformer.build_rotary_emb_for_absolute_positions( + batch_size=batch_size, + frame_positions=list(range(sink_rope_offset, sink_rope_offset + sink_size)), + height=height, + width=width, + device=device, + fps=float(fps), + ) + + # FM endpoint at sigma_max: a single epsilon per block. + noise = randn_tensor(clean_block.shape, generator=generator, device=device, dtype=dtype) + latent_block = ((1.0 - sigma_max) * clean_block.float() + sigma_max * noise.float()).to(dtype) + + # Reset the shared scheduler to step 0 for this block's Euler run (blocks are processed + # sequentially, so re-seeding the schedule per block is safe). + self.scheduler.set_timesteps(sigmas=[float(s) for s in sigmas_t[:-1]], device=device) + timesteps = self.scheduler.timesteps + + for i, t in enumerate(timesteps): + sigma = float(sigmas_t[i].item()) + + # Only the active block is forwarded; its queries attend to the `[sink, history, current]` + # K/V supplied by `kv_cache`. All active tokens carry the same sigma. + latent_tokens = _pack_latents( + latent_block, + patch_size=transformer_config.patch_size, + patch_size_t=transformer_config.patch_size_t, + ) + seq_len = latent_tokens.shape[1] + timestep = torch.full( + (batch_size, seq_len), + sigma * float(transformer_config.timestep_scale_multiplier), + dtype=torch.float32, + device=device, + ) + video_rotary_emb = self.transformer.build_rotary_emb_for_absolute_positions( + batch_size=batch_size, + frame_positions=frame_positions, + height=height, + width=width, + device=device, + fps=float(fps), + ) + velocity_pred = self.transformer( + hidden_states=latent_tokens, + encoder_hidden_states=prompt_embeds, + timestep=timestep, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=prompt_attention_mask, + kv_cache=kv_cache, + kv_cache_mode="inject", + return_dict=False, + )[0] + + # FM x0 prediction: x_t - σ_cur · v. + raw_sigma = torch.full((batch_size, seq_len, 1), sigma, dtype=torch.float32, device=device) + denoised_tokens = latent_tokens.float() - velocity_pred.float() * raw_sigma + pred_x0 = _unpack_latents( + denoised_tokens.to(dtype), + num_frames=block_end - block_start, + height=height, + width=width, + patch_size=transformer_config.patch_size, + patch_size_t=transformer_config.patch_size_t, + ) + + if sigma <= 1.0e-6: + latent_block = pred_x0.to(dtype) + else: + # FM velocity from x0; the scheduler applies the Euler update. + velocity = (latent_block.float() - pred_x0.float()) / sigma + latent_block = self.scheduler.step(velocity, t, latent_block.float(), return_dict=False)[0].to( + dtype + ) + + # Capture POST-RoPE K/V for this refined block under the same prefix and append it to the history. + self._capture_block_kv( + clean_block=latent_block, + frame_positions=frame_positions, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + fps=fps, + kv_cache=kv_cache, + kv_cache_mode="inject_and_capture_post_rope", + device=device, + ) + for layer_idx in range(num_layers): + layer_cache = kv_cache.get(layer_idx) + new_key, new_value = layer_cache.get_captured_post_rope() + history = layer_cache.get_history() + if history is None: + layer_cache.store_history(new_key, new_value) + else: + layer_cache.store_history( + torch.cat([history[0], new_key], dim=1), + torch.cat([history[1], new_value], dim=1), + ) + history_frames += block_end - block_start + + # Trim the history so the sliding window stays bounded. + if max_history_frames > 0 and history_frames > max_history_frames: + keep_tokens = max_history_frames * tokens_per_frame + for layer_idx in range(num_layers): + layer_cache = kv_cache.get(layer_idx) + history = layer_cache.get_history() + if history is not None: + layer_cache.store_history(history[0][:, -keep_tokens:], history[1][:, -keep_tokens:]) + history_frames = max_history_frames + + latents[:, :, block_start:block_end] = latent_block + progress_bar.update() + + if self.vae is None or output_type == "latent": + self.maybe_free_model_hooks() + return latents + + # The sink frames are carried through unrefined, so drop the anchor before decoding. + decoded = self._decode_latents(latents)[:, :, sink_size:] + video = self.video_processor.postprocess_video(decoded, output_type=output_type)[0] + + self.maybe_free_model_hooks() + return video + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + """Decode latents to a `(B, C, F, H, W)` tensor in `[-1, 1]` (the VAE's native output range).""" + latents = latents.to(self._execution_device, dtype=self.vae.dtype) + latents_mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents) + latents_std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents) + latents = latents / self.vae.config.scaling_factor * latents_std + latents_mean + return self.vae.decode(latents, return_dict=False)[0] + + def _capture_block_kv( + self, + *, + clean_block: torch.Tensor, + frame_positions: list[int], + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + fps: float, + kv_cache: SanaWMRefinerKVCache, + kv_cache_mode: str, + device: torch.device, + ) -> None: + """Run one forward at σ=0 in a capturing ``kv_cache_mode``; the K/V lands in ``kv_cache``. + + ``'capture_pre_rope'`` saves PRE-RoPE K/V (so a future window can re-RoPE the sink to its shifted offset) and + injects no prefix. ``'inject_and_capture_post_rope'`` attends to the current window's prefix and saves the + block's POST-RoPE K/V, ready to be appended to the recent history. + """ + latent_tokens = _pack_latents( + clean_block, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + batch_size, seq_len, _ = latent_tokens.shape + model_timestep = torch.zeros(batch_size, seq_len, dtype=torch.float32, device=device) + + video_rotary_emb = self.transformer.build_rotary_emb_for_absolute_positions( + batch_size=batch_size, + frame_positions=frame_positions, + height=int(clean_block.shape[3]), + width=int(clean_block.shape[4]), + device=device, + fps=float(fps), + ) + + self.transformer( + hidden_states=latent_tokens, + encoder_hidden_states=prompt_embeds, + timestep=model_timestep, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=prompt_attention_mask, + kv_cache=kv_cache, + kv_cache_mode=kv_cache_mode, + return_dict=False, + ) + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + + def _encode_prompt( + self, prompt: str, *, device: torch.device, dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + tokenizer = self.tokenizer + text_inputs = tokenizer( + [prompt.strip()], + padding="max_length", + padding_side="left", + max_length=self.text_max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + input_ids = text_inputs.input_ids.to(device) + attention_mask = text_inputs.attention_mask.to(device) + + # Call the top-level `text_encoder` (not its inner backbone) so that the model CPU offload + # hook installed on it by `enable_model_cpu_offload()` fires and onloads it first. + outputs = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) + hidden_states = torch.stack(outputs.hidden_states, dim=-1) + sequence_lengths = attention_mask.sum(dim=-1) + prompt_embeds = _pack_text_embeds( + hidden_states, + sequence_lengths, + device=device, + padding_side="left", + ).to(dtype=dtype) + + connector_prompt_embeds, _, connector_attention_mask = self.connectors(prompt_embeds, attention_mask) + return ( + connector_prompt_embeds.to(device=device, dtype=dtype), + connector_attention_mask.to(device=device), + ) + + +# ------------------------------------------------------------------------- +# private helpers (text embedding + latent packing) +# ------------------------------------------------------------------------- + + +def _pack_text_embeds( + text_hidden_states: torch.Tensor, + sequence_lengths: torch.Tensor, + device: str | torch.device, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, +) -> torch.Tensor: + batch_size, seq_len, hidden_dim, _ = text_hidden_states.shape + original_dtype = text_hidden_states.dtype + + token_indices = torch.arange(seq_len, device=device).unsqueeze(0) + if padding_side == "right": + mask = token_indices < sequence_lengths[:, None] + elif padding_side == "left": + start_indices = seq_len - sequence_lengths[:, None] + mask = token_indices >= start_indices + else: + raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") + mask = mask[:, :, None, None] + + masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) + num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) + masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) + + x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) + x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) + + normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) + normalized_hidden_states = normalized_hidden_states * scale_factor + normalized_hidden_states = normalized_hidden_states.flatten(2) + mask_flat = mask.squeeze(-1).expand(-1, -1, normalized_hidden_states.shape[-1]) + normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) + return normalized_hidden_states.to(dtype=original_dtype) + + +def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: + batch_size, _, num_frames, height, width = latents.shape + latents = latents.reshape( + batch_size, + -1, + num_frames // patch_size_t, + patch_size_t, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + return latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + + +def _unpack_latents( + latents: torch.Tensor, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, +) -> torch.Tensor: + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + return latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 3434c6416cce..fd615dc24b16 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2164,6 +2164,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class SanaWMLTX2RefinerTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class SanaWMTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class SD3ControlNetModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index ed724e7de751..cc0626f9919d 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -3977,6 +3977,51 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class SanaWMLTX2Refiner(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class SanaWMPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class SanaWMPipelineOutput(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class SemanticStableDiffusionPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_sana_wm.py b/tests/models/transformers/test_models_transformer_sana_wm.py new file mode 100644 index 000000000000..29c6c53ebc6e --- /dev/null +++ b/tests/models/transformers/test_models_transformer_sana_wm.py @@ -0,0 +1,184 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import SanaWMTransformer3DModel +from diffusers.models.transformers.transformer_sana_wm import SanaWMTemporalShortConvolution +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class SanaWMTransformer3DTesterConfig(BaseModelTesterConfig): + # Tiny stand-in for the public `Efficient-Large-Model/SANA-WM_bidirectional` release + # (depth 20 / hidden 2240 / 20 heads). `num_layers=2` together with `softmax_every_n=2` + # keeps both camera-branch variants covered: block 0 is the GDN one and block 1 is the + # softmax one the model's per-layer `attn_cls` loop selects. + num_layers = 2 + in_channels = 4 + caption_channels = 8 + chunk_plucker_channels = 8 + sequence_length = 16 + + num_frames = 4 + height = 8 + width = 8 + + @property + def model_class(self): + return SanaWMTransformer3DModel + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def input_shape(self) -> tuple: + return (self.in_channels, self.num_frames, self.height, self.width) + + @property + def output_shape(self) -> tuple: + return (self.in_channels, self.num_frames, self.height, self.width) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "in_channels": self.in_channels, + "num_layers": self.num_layers, + "hidden_size": 32, + "num_attention_heads": 2, + "patch_size": (1, 1, 1), + "softmax_every_n": 2, + "linear_head_dim": 16, + "t_kernel_size": 3, + "conv_kernel_size": 4, + "caption_channels": self.caption_channels, + "model_max_length": self.sequence_length, + "mlp_ratio": 2.0, + "chunk_plucker_channels": self.chunk_plucker_channels, + "chunk_plucker_post_attn_blocks": self.num_layers, + } + + def get_dummy_camera_conditions(self, batch_size: int = 1) -> torch.Tensor: + """Build the `(B, F, 20)` camera conditioning: a flat 4x4 c2w followed by `[fx, fy, cx, cy]`. + + The trajectory is a pure forward translation with an identity rotation, and the intrinsics are a + pinhole camera centred on the latent grid, which keeps the UCPE ray maps well conditioned. + """ + c2w = torch.eye(4).repeat(batch_size, self.num_frames, 1, 1) + c2w[..., 2, 3] = torch.arange(self.num_frames, dtype=torch.float32) * 0.1 + intrinsics = torch.tensor([float(self.width), float(self.height), self.width / 2, self.height / 2]) + intrinsics = intrinsics.expand(batch_size, self.num_frames, 4) + return torch.cat([c2w.flatten(start_dim=-2), intrinsics], dim=-1).to(torch_device) + + def get_dummy_inputs(self, batch_size: int = 1) -> dict[str, torch.Tensor]: + shape = (batch_size, self.in_channels, self.num_frames, self.height, self.width) + plucker_shape = (batch_size, self.chunk_plucker_channels, self.num_frames, self.height, self.width) + + return { + "hidden_states": randn_tensor(shape, generator=self.generator, device=torch_device), + "timestep": torch.randint(0, 1000, size=(batch_size, 1, self.num_frames), generator=self.generator).to( + device=torch_device, dtype=torch.float32 + ), + "encoder_hidden_states": randn_tensor( + (batch_size, 1, self.sequence_length, self.caption_channels), + generator=self.generator, + device=torch_device, + ), + # SANA-WM's cross-attention needs the text padding mask to build its attention bias. + "encoder_attention_mask": torch.ones( + batch_size, self.sequence_length, dtype=torch.long, device=torch_device + ), + "camera_conditions": self.get_dummy_camera_conditions(batch_size), + "chunk_plucker": randn_tensor(plucker_shape, generator=self.generator, device=torch_device), + } + + +class TestSanaWMTransformer3D(SanaWMTransformer3DTesterConfig, ModelTesterMixin): + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): + # Skip: fp16/bf16 require very high atol to pass, providing little signal. + # Dtype preservation is already tested by test_from_save_pretrained_dtype. + pytest.skip("Tolerance requirements too high for meaningful test") + + +class TestSanaWMTransformer3DGDNBranch(SanaWMTransformer3DTesterConfig): + """Guards against the GDN branch silently contributing nothing on a from-scratch model. + + `SanaWMTemporalShortConvolution` allocates its weight with `torch.zeros`, which `from_pretrained` overwrites but + a freshly constructed model does not. With a zero `conv_k` the GDN key is identically zero and the whole branch + returns zeros, so every other test in this file would keep passing while exercising none of the GDN maths. + """ + + def _build(self): + return self.model_class(**self.get_init_dict()).to(torch_device).eval() + + def test_short_conv_weights_start_at_zero(self): + model = self._build() + conv_weights = [m.weight for m in model.modules() if isinstance(m, SanaWMTemporalShortConvolution)] + assert conv_weights, "expected the GDN blocks to build temporal short convolutions" + assert all(torch.count_nonzero(w) == 0 for w in conv_weights) + + def test_gdn_branch_is_live_once_the_short_convs_are_populated(self): + inputs = self.get_dummy_inputs() + + model = self._build() + with torch.no_grad(): + zero_conv_out = model(**inputs, return_dict=False)[0] + + generator = torch.Generator("cpu").manual_seed(0) + for module in model.modules(): + if isinstance(module, SanaWMTemporalShortConvolution): + module.weight.data = randn_tensor( + tuple(module.weight.shape), generator=generator, device=module.weight.device + ).to(module.weight.dtype) + with torch.no_grad(): + live_conv_out = model(**inputs, return_dict=False)[0] + + assert not torch.allclose(zero_conv_out, live_conv_out), ( + "populating the short convolutions did not change the output, so the GDN branch is still dead" + ) + + +class TestSanaWMTransformer3DMemory(SanaWMTransformer3DTesterConfig, MemoryTesterMixin): + pass + + +class TestSanaWMTransformer3DCompile(SanaWMTransformer3DTesterConfig, TorchCompileTesterMixin): + def test_torch_compile_repeated_blocks(self): + # The repeated `SanaVideoMSCamCtrlBlock` runs two attention variants (`softmax_every_n` swaps a + # softmax block in for the GDN one), so the shared block forward compiles once per variant. + super().test_torch_compile_repeated_blocks(recompile_limit=2) + + +class TestSanaWMTransformer3DTraining(SanaWMTransformer3DTesterConfig, TrainingTesterMixin): + # `SanaWMTransformer3DModel._supports_gradient_checkpointing` is `False`, so the + # gradient-checkpointing tests of this mixin skip themselves. + pass diff --git a/tests/pipelines/sana_wm/__init__.py b/tests/pipelines/sana_wm/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/sana_wm/test_sana_wm.py b/tests/pipelines/sana_wm/test_sana_wm.py new file mode 100644 index 000000000000..e7d6ed7c18a2 --- /dev/null +++ b/tests/pipelines/sana_wm/test_sana_wm.py @@ -0,0 +1,204 @@ +# Copyright 2025 The HuggingFace Team and SANA-WM Authors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SANA-WM CPU unit tests. + +Covers the standalone helpers (action DSL, intrinsics math, resize-and-crop) and the public-surface registration. +""" + +import inspect + +import numpy as np +import pytest +import torch +from PIL import Image + +import diffusers +from diffusers import DiffusionPipeline, SanaWMPipeline, SanaWMPipelineOutput +from diffusers.pipelines.sana_wm import SanaWMLTX2Refiner +from diffusers.pipelines.sana_wm.cam_utils import ( + TARGET_HEIGHT, + TARGET_WIDTH, + action_string_to_c2w, + resize_and_center_crop, + snap_num_frames, + transform_intrinsics_for_crop, +) + +from ...testing_utils import enable_full_determinism + + +enable_full_determinism() + + +class TestSanaWMCamUtils: + """Pure-numpy/PIL helpers — no torch.cuda required.""" + + def test_action_dsl_forward_only(self): + c2w = action_string_to_c2w("w-5", translation_speed=0.1) + # 5 action frames + leading identity = 6 total + assert c2w.shape == (6, 4, 4) + assert c2w.dtype == np.float32 + # First frame is identity (the anchor). + np.testing.assert_allclose(c2w[0], np.eye(4, dtype=np.float32), atol=1e-6) + # 'w' moves forward (+Z in OpenCV convention). + assert float(c2w[-1, 2, 3]) == pytest.approx(0.5, abs=1e-5) + # No yaw / pitch -> rotation is identity throughout. + for i in range(c2w.shape[0]): + np.testing.assert_allclose(c2w[i, :3, :3], np.eye(3), atol=1e-6) + + def test_action_dsl_concat_segments(self): + c2w = action_string_to_c2w("w-3,a-2", translation_speed=0.1) + assert c2w.shape == (6, 4, 4) # 3 + 2 + identity anchor + + @pytest.mark.parametrize( + "action", + ["", "x-5", "w-0"], + ids=["empty", "unknown-key", "zero-length-segment"], + ) + def test_action_dsl_rejects_bad_input(self, action): + with pytest.raises(ValueError): + action_string_to_c2w(action) + + def test_action_dsl_none_segment_is_idle(self): + c2w = action_string_to_c2w("none-3", translation_speed=0.1) + assert c2w.shape == (4, 4, 4) + # No motion -> all frames are identity. + for i in range(c2w.shape[0]): + np.testing.assert_allclose(c2w[i], np.eye(4), atol=1e-6) + + def test_transform_intrinsics_for_crop_scalar(self): + # (fx, fy, cx, cy) for a 1000x500 source, resized to 1280x704, then + # center-cropped to 1280x704 (no extra crop offset). + intr = np.array([800.0, 800.0, 500.0, 250.0], dtype=np.float32) + out = transform_intrinsics_for_crop(intr, src_size=(1000, 500), resized_size=(1280, 704), crop_offset=(0, 0)) + assert float(out[0]) == pytest.approx(800.0 * 1280 / 1000, abs=1e-4) # fx scales with x + assert float(out[1]) == pytest.approx(800.0 * 704 / 500, abs=1e-4) + assert float(out[2]) == pytest.approx(500.0 * 1280 / 1000, abs=1e-4) + assert float(out[3]) == pytest.approx(250.0 * 704 / 500, abs=1e-4) + + def test_transform_intrinsics_for_crop_with_offset(self): + intr = np.array([800.0, 800.0, 500.0, 250.0], dtype=np.float32) + # After resize, an extra crop offset shifts the principal point. + out = transform_intrinsics_for_crop( + intr, src_size=(1000, 500), resized_size=(2000, 1000), crop_offset=(360, 148) + ) + assert float(out[2]) == pytest.approx(500.0 * 2.0 - 360.0, abs=1e-4) + assert float(out[3]) == pytest.approx(250.0 * 2.0 - 148.0, abs=1e-4) + + def test_resize_and_center_crop_default_target(self): + src = Image.new("RGB", (1691, 930)) + cropped, src_size, resized_size, crop_offset = resize_and_center_crop(src) + assert cropped.size == (TARGET_WIDTH, TARGET_HEIGHT) + assert src_size == (1691, 930) + # Resize preserves aspect; one of the resized dimensions equals the target. + resized_width, resized_height = resized_size + assert resized_width >= TARGET_WIDTH + assert resized_height >= TARGET_HEIGHT + crop_left, crop_top = crop_offset + assert crop_left >= 0 + assert crop_top >= 0 + # Center crop produces 0 offset on the dimension that hit the target exactly. + assert crop_left == 0 or crop_top == 0 + + # The LTX-2 VAE requires a (8k + 1)-shaped temporal dim, so ``snap_num_frames`` rounds to + # the nearest such value (ties break to the ceil). + @pytest.mark.parametrize("num_frames", [1, 9, 17, 81, 161, 321, 801]) + def test_snap_num_frames_is_a_noop_on_8k_plus_1(self, num_frames): + assert snap_num_frames(num_frames) == num_frames + + @pytest.mark.parametrize( + ("num_frames", "expected"), + [ + (2, 1), + (10, 9), # 10 is closer to 9 than 17 + (80, 81), # 80 is closer to 81 than 73 + (100, 97), # 100 is closer to 97 than 105 + ], + ) + def test_snap_num_frames_to_8k_plus_1(self, num_frames, expected): + assert snap_num_frames(num_frames) == expected + + def test_snap_num_frames_respects_upper_bound(self): + # ``upper_bound`` caps the result (the snap falls back to the floor). + assert snap_num_frames(100, upper_bound=100) <= 100 + assert snap_num_frames(100, upper_bound=100) == 97 + + +class TestSanaWMRegistration: + """Verify the SANA-WM symbols are reachable through the public diffusers surface.""" + + @pytest.mark.parametrize( + "name", ["SanaWMPipeline", "SanaWMTransformer3DModel", "SanaWMLTX2Refiner", "SanaWMPipelineOutput"] + ) + def test_top_level_symbols(self, name): + assert hasattr(diffusers, name), f"{name!r} not exported from diffusers top-level" + + def test_pipeline_output_dataclass(self): + frames = np.zeros((3, 8, 8, 3), dtype=np.float32) + c2w = np.broadcast_to(np.eye(4, dtype=np.float32), (3, 4, 4)).copy() + latent = torch.zeros(1, 16, 1, 4, 4) + output = SanaWMPipelineOutput(frames=frames, c2w=c2w, latent=latent) + assert tuple(output.frames.shape) == (3, 8, 8, 3) + assert tuple(output.c2w.shape) == (3, 4, 4) + assert tuple(output.latent.shape) == (1, 16, 1, 4, 4) + + def test_refiner_is_pipeline_with_ar_call_defaults(self): + # The refiner is a standalone DiffusionPipeline. + assert issubclass(SanaWMLTX2Refiner, DiffusionPipeline) + + # Its denoising entry point is ``__call__`` with the canonical AR defaults. + params = inspect.signature(SanaWMLTX2Refiner.__call__).parameters + assert "block_size" in params + assert "kv_max_frames" in params + # AR mode is on by default. + assert params["block_size"].default == 3 + assert params["kv_max_frames"].default == 11 + + @pytest.mark.parametrize( + "name", + [ + "intrinsics", + "c2w", + "action", + # Standard diffusers pipeline arguments. + "generator", + "prompt_embeds", + "prompt_attention_mask", + "negative_prompt_embeds", + "negative_prompt_attention_mask", + ], + ) + def test_pipeline_call_intrinsics_signature(self, name): + params = inspect.signature(SanaWMPipeline.__call__).parameters + assert name in params + + def test_pipeline_call_takes_generator_not_seed(self): + # Pipelines take a `generator`; `seed` shortcuts are not part of the diffusers interface. + params = inspect.signature(SanaWMPipeline.__call__).parameters + assert "seed" not in params + assert "refiner_seed" not in params + + def test_refiner_is_not_a_component_of_the_base_pipeline(self): + # The two stages run as separate pipelines, so the base one neither holds a + # refiner nor exposes a switch for it. + assert "refiner" not in inspect.signature(SanaWMPipeline.__init__).parameters + assert "use_refiner" not in inspect.signature(SanaWMPipeline.__call__).parameters + + def test_refiner_takes_an_optional_vae_for_decoding(self): + # With a `vae` the refiner returns video; without one, refined latents. + params = inspect.signature(SanaWMLTX2Refiner.__init__).parameters + assert "vae" in params + assert params["vae"].default is None + assert "output_type" in inspect.signature(SanaWMLTX2Refiner.__call__).parameters