diff --git a/.ai/skills/custom-blocks/SKILL.md b/.ai/skills/custom-blocks/SKILL.md index b98aa18b61..261b4cdcd8 100644 --- a/.ai/skills/custom-blocks/SKILL.md +++ b/.ai/skills/custom-blocks/SKILL.md @@ -57,7 +57,8 @@ diffusers-cli custom_blocks [--block_module_name ] [--block_class_name ### What it does 1. **AST scan**: parses `` without executing it, walks top-level `ClassDef` nodes, and collects every - class whose `bases` include `ModularPipelineBlocks`. + class whose `bases` include `ModularPipelineBlocks` or one of its composite subclasses + (`SequentialPipelineBlocks`, `AutoPipelineBlocks`, `ConditionalPipelineBlocks`, `LoopSequentialPipelineBlocks`). 2. **Pick a class**: uses `--block_class_name` if given, else the first found. Errors with the list of available classes if your name doesn't match. 3. **Load and save**: imports the file via `importlib.util.spec_from_file_location` (this does execute the @@ -137,9 +138,9 @@ diffusers-cli run --model my-user/my-denoise-block --trust-remote-code \ - **`block_class_name could not be retrieved. Available classes from : [ClassA, ClassB]`** — your `--block_class_name` doesn't match any `ModularPipelineBlocks` subclass found. Pick from the list shown. - **No classes found**: silent — the command will try to use the first entry in an empty list and raise - `IndexError`. If you hit that, double-check your class actually inherits from `ModularPipelineBlocks` - (the AST scan looks for that literal base-class name; aliased imports like `from diffusers import ... - as MPB` won't be picked up). + `IndexError`. If you hit that, double-check your class actually inherits from `ModularPipelineBlocks` or one + of the composite classes listed above (the AST scan looks for those literal class names; aliased imports like + `from diffusers import ... as MPB` won't be picked up). - **Block requires constructor args**: the command calls `()` with no args. If your block needs `__init__` parameters, refactor to take them from `state`/`components` at `__call__` time instead, or hardcode defaults in `__init__`. diff --git a/.github/workflows/pr_labeler.yml b/.github/workflows/pr_labeler.yml index d657bed951..9df5dc517a 100644 --- a/.github/workflows/pr_labeler.yml +++ b/.github/workflows/pr_labeler.yml @@ -11,38 +11,33 @@ jobs: label: permissions: contents: read + issues: read pull-requests: write runs-on: ubuntu-latest + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + PR_NUMBER: ${{ github.event.pull_request.number }} + REPO: ${{ github.repository }} steps: + # label - uses: actions/labeler@8558fd74291d67161a8a78ce36a881fa63b766a9 # v5 with: sync-labels: true - missing-tests: - permissions: - contents: read - issues: read - pull-requests: write - runs-on: ubuntu-latest - steps: + # missing-tests - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + id: checkout + if: ${{ !cancelled() }} with: ref: ${{ github.event.pull_request.base.sha }} - name: Check for missing tests id: check - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - PR_NUMBER: ${{ github.event.pull_request.number }} - REPO: ${{ github.repository }} + if: ${{ !cancelled() && steps.checkout.outcome == 'success' }} run: | gh api --paginate "repos/${REPO}/pulls/${PR_NUMBER}/files" \ | python utils/check_test_missing.py - name: Add or remove missing-tests label if: always() - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - PR_NUMBER: ${{ github.event.pull_request.number }} - REPO: ${{ github.repository }} run: | HAS_LABEL=$(gh api "repos/${REPO}/issues/${PR_NUMBER}/labels" --jq 'any(.[]; .name == "missing-tests")') if [ "${{ steps.check.outcome }}" = "failure" ]; then @@ -55,17 +50,9 @@ jobs: fi fi - fixes-issue: - permissions: - issues: read - pull-requests: write - runs-on: ubuntu-latest - steps: + # fixes-issue - name: Check for linked closing issues - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - PR_NUMBER: ${{ github.event.pull_request.number }} - REPO: ${{ github.repository }} + if: ${{ !cancelled() }} run: | OWNER="${REPO%/*}" NAME="${REPO#*/}" @@ -93,17 +80,9 @@ jobs: fi fi - size-label: - permissions: - issues: read - pull-requests: write - runs-on: ubuntu-latest - steps: + # size-label - name: Label PR by diff size - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - PR_NUMBER: ${{ github.event.pull_request.number }} - REPO: ${{ github.repository }} + if: ${{ !cancelled() }} run: | DIFF_SIZE=$(gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.additions + .deletions') if [ "$DIFF_SIZE" -lt 50 ]; then diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 4b4e424719..ab98a52430 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -114,6 +114,8 @@ title: Intel Gaudi - local: optimization/neuron title: AWS Neuron + - local: optimization/tpu + title: TPU title: Hardware-specific acceleration - isExpanded: false sections: @@ -409,6 +411,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 +657,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 0000000000..b5251db213 --- /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 0000000000..82c8fa3397 --- /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 0000000000..efbb64e3e8 --- /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/docs/source/en/optimization/tpu.md b/docs/source/en/optimization/tpu.md new file mode 100644 index 0000000000..e5f75480f8 --- /dev/null +++ b/docs/source/en/optimization/tpu.md @@ -0,0 +1,157 @@ + + +# TorchTPU + +[TorchTPU](https://github.com/google-pytorch/torch_tpu/) is a PyTorch backend for Google's Tensor Processing Units +(TPUs), which lets you run Diffusers pipelines on Cloud TPUs (v6e, v7x) with minimal code changes. + +Two execution modes are available: + +| Mode | Constant | How to activate | Notes | +|---|---|---|---| +| Strict eager (default) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Operations dispatched one at a time, asynchronous | +| Compile | — | `torch.compile(module, backend="tpu")` | AOT compilation with `TpuBackend` | + +Follow the [TorchTPU installation guide](https://github.com/google-pytorch/torch_tpu/). Once installed, `import torch` +loads it automatically and registers the `"tpu"` device, so `pipe.to("tpu")` is the only change needed. Add +`import torch_tpu` only if you disabled backend autoloading with `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. + +## Eager mode + +FLUX.1-schnell doesn't fit on a single v6e chip all at once, so use [`~DiffusionPipeline.enable_model_cpu_offload`] to +move each model to the TPU only while it runs. It detects the `"tpu"` device automatically. + +```python +import torch + +from diffusers import FluxPipeline + +pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16) +pipe.enable_model_cpu_offload() + +image = pipe( + prompt="a golden retriever surfing a wave, photorealistic", + height=1024, + width=1024, + num_inference_steps=4, + guidance_scale=0.0, +).images[0] + +image.save("output.png") +``` + +If a model is too large for a single chip, or you have several chips and want lower latency, shard the models +across chips instead. See the [Tensor parallelism](#tensor-parallelism) section. + +## Compiled mode + +TorchTPU registers `"tpu"` as a `torch.compile` backend name (`TpuBackend` under the hood), so +components compile like any other `torch.compile` target. The first +call (warmup) is slow because it compiles; later calls with the same shapes reuse the compiled graph. + +> [!IMPORTANT] +> TorchTPU requires **static shapes**, so pass `dynamic=False`. A new `height` or `width` compiles again for that +> shape, once; shapes already seen are reused. Changing `num_inference_steps` doesn't recompile. + +When the whole pipeline fits on one chip, move it to the TPU and compile the full transformer. Stable Diffusion 3.5 +Medium (~15GB in bf16) fits on a single v6e chip. + +```python +import torch + +from diffusers import StableDiffusion3Pipeline + +pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16) +pipe.to("tpu") +pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False) + +# Warmup — triggers static graph compilation. +pipe(prompt="warmup", height=1024, width=1024, num_inference_steps=40, guidance_scale=4.5) + +# Later calls with the same shapes reuse the compiled graph. +image = pipe( + prompt="a golden retriever surfing a wave, photorealistic", + height=1024, + width=1024, + num_inference_steps=40, + guidance_scale=4.5, +).images[0] + +image.save("output.png") +``` + +If the pipeline doesn't fit on one chip: + +- With several chips, shard it with [tensor parallelism](#tensor-parallelism) instead. Everything stays on the TPU, + and `pipe.transformer.compile(...)` works the same way on the sharded transformer. +- With [`~DiffusionPipeline.enable_model_cpu_offload`], the offload hooks can't be traced by `torch.compile`. Compile + only the transformer's repeated blocks instead, with + `pipe.transformer.compile_repeated_blocks(backend="tpu", fullgraph=True, dynamic=False)`. + +## Tensor parallelism + +Shard models too large for one chip across several. FLUX.2-dev's text encoder (~48GB) and transformer (~64GB) each +exceed a single chip, so the example below shards both: + +- the transformer with [`TensorParallelConfig`], passed to the `parallel_config` argument of [`~ModelMixin.from_pretrained`]. Each rank reads only its own slice of every sharded weight, so the full model is never materialized. For general TP details (`_tp_plan`, colwise/rowwise), see the [Tensor parallelism](../training/distributed_inference#tensor-parallelism) guide. +- the text encoder with Transformers' own [tensor parallelism](https://huggingface.co/docs/transformers/perf_infer_gpu_multi), passing a [`~transformers.DistributedConfig`] and the same mesh. + +On TPU, initialize the process group with `backend="tpu_dist"` and build the mesh with `DeviceMesh("tpu", ...)`. + +```python +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from transformers import DistributedConfig, Mistral3ForConditionalGeneration + +from diffusers import Flux2Pipeline, Flux2Transformer2DModel, TensorParallelConfig + +dist.init_process_group(backend="tpu_dist") +mesh = DeviceMesh("tpu", list(range(dist.get_world_size()))) + +repo_id = "black-forest-labs/FLUX.2-dev" +text_encoder = Mistral3ForConditionalGeneration.from_pretrained( + repo_id, + subfolder="text_encoder", + dtype=torch.bfloat16, + distributed_config=DistributedConfig(tp_plan="auto"), + device_mesh=mesh, +) +transformer = Flux2Transformer2DModel.from_pretrained( + repo_id, subfolder="transformer", dtype=torch.bfloat16, parallel_config=TensorParallelConfig(mesh=mesh) +) +pipe = Flux2Pipeline.from_pretrained( + repo_id, text_encoder=text_encoder, transformer=transformer, dtype=torch.bfloat16 +) +pipe.vae.to("tpu") + +image = pipe( + prompt="a golden retriever surfing a wave, photorealistic", + num_inference_steps=28, + generator=torch.Generator("cpu").manual_seed(0), +).images[0] +if dist.get_rank() == 0: + image.save("output.png") +``` + +Launch one process per chip. Set `--nproc_per_node` to use all the number of TPU chips on your host. + +```bash +eval $(python -m torch_tpu._internal.distributed.launchers.singlehost_wrapper | sed 's/^/export /') +torchrun --nproc_per_node=8 flux2_tp.py +``` + +> [!WARNING] +> The launch command above uses a TorchTPU internal module (`torch_tpu._internal`), not a stable API, and is expected +> to change in a future TorchTPU release. Check the [TorchTPU repository](https://github.com/google-pytorch/torch_tpu/) +> for the current way to launch multi-chip jobs. diff --git a/docs/source/en/using-diffusers/cli.md b/docs/source/en/using-diffusers/cli.md index 01339f4628..8625001cef 100644 --- a/docs/source/en/using-diffusers/cli.md +++ b/docs/source/en/using-diffusers/cli.md @@ -290,8 +290,9 @@ hf sandbox kill ## `custom_blocks` Package a local `ModularPipelineBlocks` subclass for upload to the Hub. Reads a Python file, AST-scans it for -subclasses of `ModularPipelineBlocks`, instantiates the chosen one, and calls `save_pretrained` in the current -working directory. +classes deriving from `ModularPipelineBlocks` or one of its composite subclasses (`SequentialPipelineBlocks`, +`AutoPipelineBlocks`, `ConditionalPipelineBlocks`, `LoopSequentialPipelineBlocks`), instantiates the chosen one, and +calls `save_pretrained` in the current working directory. ```bash # Package the first block found in ./block.py diff --git a/scripts/convert_sana_wm_to_diffusers.py b/scripts/convert_sana_wm_to_diffusers.py new file mode 100644 index 0000000000..c8ca2eeb96 --- /dev/null +++ b/scripts/convert_sana_wm_to_diffusers.py @@ -0,0 +1,247 @@ +#!/usr/bin/env python +# Copyright 2026 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) + # `y_embedding` is the training-time null-caption embedding; inference never reads it. + sd.pop("y_embedder.y_embedding", None) + # The timestep MLP is diffusers' `TimestepEmbedding`, which names its layers `linear_1` / `linear_2`. + sd = { + k.replace("t_embedder.mlp.0.", "t_embedder.linear_1.").replace("t_embedder.mlp.2.", "t_embedder.linear_2."): v + for k, v in sd.items() + } + # `.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/setup.py b/setup.py index bdc87e5c19..14eff97197 100644 --- a/setup.py +++ b/setup.py @@ -100,7 +100,7 @@ _deps = [ "filelock", "ftfy", "hf-doc-builder>=0.3.0", - "huggingface-hub>=1.32.0,<2.0", + "huggingface-hub>=1.32.0,<3.0", "requests-mock==1.10.0", "importlib_metadata", "invisible-watermark>=0.2.0", diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 599237a6f2..36cd6d2838 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -348,6 +348,8 @@ else: "SanaControlNetModel", "SanaTransformer2DModel", "SanaVideoTransformer3DModel", + "SanaWMLTX2RefinerTransformer3DModel", + "SanaWMTransformer3DModel", "SD3ControlNetModel", "SD3MultiControlNetModel", "SD3Transformer2DModel", @@ -798,6 +800,9 @@ else: "SanaSprintPipeline", "SanaVideoPipeline", "SanaVideoPipeline", + "SanaWMLTX2Refiner", + "SanaWMPipeline", + "SanaWMPipelineOutput", "SemanticStableDiffusionPipeline", "ShapEImg2ImgPipeline", "ShapEPipeline", @@ -1241,6 +1246,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: SanaControlNetModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, + SanaWMLTX2RefinerTransformer3DModel, + SanaWMTransformer3DModel, SD3ControlNetModel, SD3MultiControlNetModel, SD3Transformer2DModel, @@ -1665,6 +1672,9 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: SanaSprintImg2ImgPipeline, SanaSprintPipeline, SanaVideoPipeline, + SanaWMLTX2Refiner, + SanaWMPipeline, + SanaWMPipelineOutput, SemanticStableDiffusionPipeline, ShapEImg2ImgPipeline, ShapEPipeline, diff --git a/src/diffusers/commands/custom_blocks.py b/src/diffusers/commands/custom_blocks.py index a3649117e0..bc6aa257cc 100644 --- a/src/diffusers/commands/custom_blocks.py +++ b/src/diffusers/commands/custom_blocks.py @@ -28,7 +28,14 @@ from ..utils import logging from . import BaseDiffusersCLICommand -EXPECTED_PARENT_CLASSES = ["ModularPipelineBlocks"] +# the block base class and its composite subclasses; a class deriving from any of them is a packageable block +EXPECTED_PARENT_CLASSES = [ + "ModularPipelineBlocks", + "SequentialPipelineBlocks", + "AutoPipelineBlocks", + "ConditionalPipelineBlocks", + "LoopSequentialPipelineBlocks", +] def conversion_command_factory(args: Namespace): @@ -128,7 +135,10 @@ class CustomBlocksCommand(BaseDiffusersCLICommand): if not isinstance(node, ast.ClassDef): continue - base_names = [bname for b in node.bases if (bname := self._get_base_name(b)) is not None] + # compare on the class name only, so `diffusers.modular_pipelines.SequentialPipelineBlocks` matches too + base_names = [ + bname.rsplit(".", 1)[-1] for b in node.bases if (bname := self._get_base_name(b)) is not None + ] for allowed in EXPECTED_PARENT_CLASSES: if allowed in base_names: diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 43390d8d54..156d2a1066 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -8,7 +8,7 @@ deps = { "filelock": "filelock", "ftfy": "ftfy", "hf-doc-builder": "hf-doc-builder>=0.3.0", - "huggingface-hub": "huggingface-hub>=1.32.0,<2.0", + "huggingface-hub": "huggingface-hub>=1.32.0,<3.0", "requests-mock": "requests-mock==1.10.0", "importlib_metadata": "importlib_metadata", "invisible-watermark": "invisible-watermark>=0.2.0", diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index 0f56c3644e..19a772704e 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/src/diffusers/hooks/tensor_parallel.py @@ -22,7 +22,7 @@ from ..utils import get_logger, is_peft_available logger = get_logger(__name__) # pylint: disable=invalid-name -_SUPPORTED_TP_DEVICES = ("cuda", "neuron") +_SUPPORTED_TP_DEVICES = ("cuda", "neuron", "tpu", "xpu") class PackedColwiseParallel: diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index 4f6f4bd52b..6dccc02dd5 100644 --- a/src/diffusers/image_processor.py +++ b/src/diffusers/image_processor.py @@ -14,6 +14,7 @@ import math import warnings +from typing import Any import numpy as np import PIL.Image @@ -885,7 +886,7 @@ class InpaintProcessor(ConfigMixin): height: int | None = None, width: int | None = None, padding_mask_crop: int | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor | None, dict[str, Any]]: """ Preprocess the image and mask. """ @@ -894,7 +895,13 @@ class InpaintProcessor(ConfigMixin): # if mask is None, same behavior as regular image processor if mask is None: - return self._image_processor.preprocess(image, height=height, width=width) + processed_image = self._image_processor.preprocess(image, height=height, width=width) + postprocessing_kwargs = { + "crops_coords": None, + "original_image": None, + "original_mask": None, + } + return processed_image, None, postprocessing_kwargs if padding_mask_crop is not None: crops_coords = self._image_processor.get_crop_region(mask, width, height, pad=padding_mask_crop) diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index ec23d6e09d..6c383b1c33 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -921,7 +921,7 @@ class LoraBaseMixin: for component in self._lora_loadable_modules: model = getattr(self, component, None) - if model is not None and issubclass(model.__class__, ModelMixin): + if model is not None and issubclass(model.__class__, (ModelMixin, PreTrainedModel)): for module in model.modules(): if isinstance(module, BaseTunerLayer): active_adapters = module.active_adapters diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index cef2e88454..bfaec3b0b5 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2736,6 +2736,160 @@ def _convert_kohya_flux2_lora_to_diffusers(state_dict): return ait_sd +def _bake_lokr_alpha_(state_dict): + """ + Consume `.alpha` keys by baking the LyCORIS `alpha / rank` scaling into the left Kronecker factor. The scaling only + applies when a factor is rank-decomposed (`lokr_w1_a/b` or `lokr_w2_a/b`); when both factors are stored as full + matrices, LoKr applies no alpha scaling and the alpha key is simply dropped. + """ + for alpha_key in [k for k in state_dict if k.endswith(".alpha")]: + alpha = state_dict.pop(alpha_key).item() + module = alpha_key.removesuffix(".alpha") + w1_b = state_dict.get(f"{module}.lokr_w1_b") + w2_b = state_dict.get(f"{module}.lokr_w2_b") + rank = w2_b.shape[0] if w2_b is not None else w1_b.shape[0] if w1_b is not None else None + if rank is None: + continue + w1_key = f"{module}.lokr_w1" if f"{module}.lokr_w1" in state_dict else f"{module}.lokr_w1_a" + state_dict[w1_key] = state_dict[w1_key] * (alpha / rank) + + +def _convert_non_diffusers_lokr_to_diffusers(state_dict): + """ + Convert a non-diffusers LoKr state dict whose module paths already match the diffusers model (e.g. ai-toolkit + Z-Image checkpoints with keys like `diffusion_model.layers.0.attention.to_q.lokr_w1`) to the peft-loadable format: + the `diffusion_model.` prefix is replaced with `transformer.` and the `.alpha` keys are consumed. + """ + state_dict = {k.removeprefix("diffusion_model."): v for k, v in state_dict.items()} + _bake_lokr_alpha_(state_dict) + + non_lokr_keys = [k for k in state_dict if ".lokr_" not in k] + if non_lokr_keys: + raise ValueError(f"`state_dict` contains unexpected non-LoKr keys: {non_lokr_keys}.") + + return {f"transformer.{k}": v for k, v in state_dict.items()} + + +_LOKR_SUFFIXES = ("lokr_w1", "lokr_w1_a", "lokr_w1_b", "lokr_w2", "lokr_w2_a", "lokr_w2_b") + + +def _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict): + """ + Convert a BFL-format Flux2 LoKr state dict (e.g. trained with ai-toolkit, keys like + `diffusion_model.double_blocks.0.img_attn.qkv.lokr_w1`) to the peft-loadable diffusers format. + + BFL checkpoints apply LoKr to the fused QKV projections of the double blocks. Unlike a LoRA delta, a Kronecker + product delta over the fused projection cannot be split exactly into separate Q/K/V factors, so these are mapped to + the model's fused `to_qkv`/`to_added_qkv` projections instead; `Flux2LoraLoaderMixin.load_lora_weights` fuses the + model's projections before injecting such an adapter. + """ + original_state_dict = {k.removeprefix("diffusion_model."): v for k, v in state_dict.items()} + _bake_lokr_alpha_(original_state_dict) + + converted_state_dict = {} + + num_double_layers = 0 + num_single_layers = 0 + for key in original_state_dict.keys(): + if key.startswith("single_blocks."): + num_single_layers = max(num_single_layers, int(key.split(".")[1]) + 1) + elif key.startswith("double_blocks."): + num_double_layers = max(num_double_layers, int(key.split(".")[1]) + 1) + + def _remap(bfl_path, diffusers_path): + for suffix in _LOKR_SUFFIXES: + weight = original_state_dict.pop(f"{bfl_path}.{suffix}", None) + if weight is not None: + converted_state_dict[f"{diffusers_path}.{suffix}"] = weight + + for sl in range(num_single_layers): + _remap(f"single_blocks.{sl}.linear1", f"single_transformer_blocks.{sl}.attn.to_qkv_mlp_proj") + _remap(f"single_blocks.{sl}.linear2", f"single_transformer_blocks.{sl}.attn.to_out") + + for dl in range(num_double_layers): + tb = f"transformer_blocks.{dl}" + db = f"double_blocks.{dl}" + + _remap(f"{db}.img_attn.qkv", f"{tb}.attn.to_qkv") + _remap(f"{db}.txt_attn.qkv", f"{tb}.attn.to_added_qkv") + + _remap(f"{db}.img_attn.proj", f"{tb}.attn.to_out.0") + _remap(f"{db}.txt_attn.proj", f"{tb}.attn.to_add_out") + + _remap(f"{db}.img_mlp.0", f"{tb}.ff.linear_in") + _remap(f"{db}.img_mlp.2", f"{tb}.ff.linear_out") + _remap(f"{db}.txt_mlp.0", f"{tb}.ff_context.linear_in") + _remap(f"{db}.txt_mlp.2", f"{tb}.ff_context.linear_out") + + extra_mappings = { + "img_in": "x_embedder", + "txt_in": "context_embedder", + "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1", + "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2", + "guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1", + "guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2", + "final_layer.linear": "proj_out", + "final_layer.adaLN_modulation.1": "norm_out.linear", + "single_stream_modulation.lin": "single_stream_modulation.linear", + "double_stream_modulation_img.lin": "double_stream_modulation_img.linear", + "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear", + } + for bfl_key, diffusers_key in extra_mappings.items(): + _remap(bfl_key, diffusers_key) + + if len(original_state_dict) > 0: + raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.") + + return {f"transformer.{k}": v for k, v in converted_state_dict.items()} + + +# Mapping from LyCORIS underscore-encoded sub-paths to dotted Flux2 module paths. +_LYCORIS_FLUX2_SUBPATH_MAP = { + "attn_to_q": "attn.to_q", + "attn_to_k": "attn.to_k", + "attn_to_v": "attn.to_v", + "attn_to_out_0": "attn.to_out.0", + "attn_to_add_out": "attn.to_add_out", + "attn_add_q_proj": "attn.add_q_proj", + "attn_add_k_proj": "attn.add_k_proj", + "attn_add_v_proj": "attn.add_v_proj", + "attn_to_qkv_mlp_proj": "attn.to_qkv_mlp_proj", + "attn_to_out": "attn.to_out", + "ff_linear_in": "ff.linear_in", + "ff_linear_out": "ff.linear_out", + "ff_context_linear_in": "ff_context.linear_in", + "ff_context_linear_out": "ff_context.linear_out", +} + + +def _convert_lycoris_flux2_lokr_to_diffusers(state_dict): + """ + Convert a LyCORIS-format Flux2 LoKr state dict (keys like `lycoris_transformer_blocks_0_attn_to_q.lokr_w1`) to the + peft-loadable diffusers format. LyCORIS wraps the diffusers model directly and encodes each module path with + underscores, which are decoded through a lookup of the known block sub-paths. + """ + state_dict = dict(state_dict) + _bake_lokr_alpha_(state_dict) + + lycoris_key_pattern = re.compile(r"^lycoris_((?:single_)?transformer_blocks)_(\d+)_(.+)\.(.+)$") + + converted_state_dict = {} + unrecognized_keys = [] + for key, value in state_dict.items(): + match = lycoris_key_pattern.match(key) + diffusers_sub_path = _LYCORIS_FLUX2_SUBPATH_MAP.get(match.group(3)) if match is not None else None + if diffusers_sub_path is None: + unrecognized_keys.append(key) + continue + container, block_idx, _, suffix = match.groups() + converted_state_dict[f"transformer.{container}.{block_idx}.{diffusers_sub_path}.{suffix}"] = value + + if unrecognized_keys: + raise ValueError(f"These keys are not LyCORIS Flux2 LoKr keys: {unrecognized_keys}.") + + return converted_state_dict + + def _convert_non_diffusers_z_image_lora_to_diffusers(state_dict): """ Convert non-diffusers ZImage LoRA state dict to diffusers format. diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 1003aa57c4..46867b2257 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -45,13 +45,16 @@ from .lora_conversion_utils import ( _convert_hunyuan_video_lora_to_diffusers, _convert_kohya_flux2_lora_to_diffusers, _convert_kohya_flux_lora_to_diffusers, + _convert_lycoris_flux2_lokr_to_diffusers, _convert_musubi_wan_lora_to_diffusers, _convert_non_diffusers_ace_step_lora_to_diffusers, _convert_non_diffusers_anima_lora_to_diffusers, + _convert_non_diffusers_flux2_lokr_to_diffusers, _convert_non_diffusers_flux2_lora_to_diffusers, _convert_non_diffusers_hidream_lora_to_diffusers, _convert_non_diffusers_ideogram4_lora_to_diffusers, _convert_non_diffusers_krea2_lora_to_diffusers, + _convert_non_diffusers_lokr_to_diffusers, _convert_non_diffusers_lora_to_diffusers, _convert_non_diffusers_ltx2_lora_to_diffusers, _convert_non_diffusers_ltxv_lora_to_diffusers, @@ -5443,14 +5446,19 @@ class ZImageLoraLoaderMixin(LoraBaseMixin): has_lora_unet = any(k.startswith("lora_unet_") for k in state_dict) has_diffusion_model = any(k.startswith("diffusion_model.") for k in state_dict) has_default = any("default." in k for k in state_dict) - if has_alphas_in_sd or has_lora_unet or has_diffusion_model or has_default: + is_lokr = any(".lokr_" in k for k in state_dict) + if is_lokr: + # ai-toolkit Z-Image LoKr checkpoints store module paths that already match the diffusers model. + if has_diffusion_model or has_alphas_in_sd: + state_dict = _convert_non_diffusers_lokr_to_diffusers(state_dict) + elif has_alphas_in_sd or has_lora_unet or has_diffusion_model or has_default: state_dict = _convert_non_diffusers_z_image_lora_to_diffusers(state_dict) out = (state_dict, metadata) if return_lora_metadata else state_dict return out @require_peft_backend - # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights + # Copied from diffusers.loaders.lora_pipeline.Flux2LoraLoaderMixin.load_lora_weights def load_lora_weights( self, pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], @@ -5471,9 +5479,9 @@ class ZImageLoraLoaderMixin(LoraBaseMixin): kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError("Invalid adapter checkpoint. We currently support LoRA and LoKr.") self.load_lora_into_transformer( state_dict, @@ -5831,15 +5839,23 @@ class Flux2LoraLoaderMixin(LoraBaseMixin): if is_peft_format: state_dict = {k.replace("base_model.model.", "diffusion_model."): v for k, v in state_dict.items()} + is_lokr = any(".lokr_" in k for k in state_dict) is_ai_toolkit = any(k.startswith("diffusion_model.") for k in state_dict) - if is_ai_toolkit: + if is_lokr: + if any(k.startswith("lycoris_") for k in state_dict): + state_dict = _convert_lycoris_flux2_lokr_to_diffusers(state_dict) + elif is_ai_toolkit: + state_dict = _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict) + elif not any(k.startswith("transformer.") for k in state_dict): + # Bare dotted diffusers module paths (e.g. SimpleTuner exports), possibly with alpha keys. + state_dict = _convert_non_diffusers_lokr_to_diffusers(state_dict) + elif is_ai_toolkit: state_dict = _convert_non_diffusers_flux2_lora_to_diffusers(state_dict) out = (state_dict, metadata) if return_lora_metadata else state_dict return out @require_peft_backend - # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], @@ -5860,9 +5876,9 @@ class Flux2LoraLoaderMixin(LoraBaseMixin): kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError("Invalid adapter checkpoint. We currently support LoRA and LoKr.") self.load_lora_into_transformer( state_dict, diff --git a/src/diffusers/loaders/peft.py b/src/diffusers/loaders/peft.py index 5ab34f4f0a..3f0652c979 100644 --- a/src/diffusers/loaders/peft.py +++ b/src/diffusers/loaders/peft.py @@ -37,7 +37,12 @@ from ..utils import ( set_adapter_layers, set_weights_and_activate_adapters, ) -from ..utils.peft_utils import _create_lora_config, _maybe_warn_for_unhandled_keys +from ..utils.peft_utils import ( + _create_lokr_config, + _create_lora_config, + _maybe_fuse_qkv_projections_for_lokr, + _maybe_warn_for_unhandled_keys, +) from .lora_base import _fetch_state_dict, _func_optionally_disable_offloading from .unet_loader_utils import _maybe_expand_lora_scales @@ -217,56 +222,65 @@ class PeftAdapterMixin: "Please choose an existing adapter name or set `hotswap=False` to prevent hotswapping." ) - # check with first key if is not in peft format - first_key = next(iter(state_dict.keys())) - if "lora_A" not in first_key: - state_dict = convert_unet_state_dict_to_peft(state_dict) - - # Control LoRA from SAI is different from BFL Control LoRA - # https://huggingface.co/stabilityai/control-lora - # https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors - is_sai_sd_control_lora = "lora_controlnet" in state_dict - if is_sai_sd_control_lora: - state_dict = convert_sai_sd_control_lora_state_dict_to_peft(state_dict) - - rank = {} - for key, val in state_dict.items(): - # Cannot figure out rank from lora layers that don't have at least 2 dimensions. - # Bias layers in LoRA only have a single dimension - if "lora_B" in key and val.ndim > 1: - # Check out https://github.com/huggingface/peft/pull/2419 for the `^` symbol. - # We may run into some ambiguous configuration values when a model has module - # names, sharing a common prefix (`proj_out.weight` and `blocks.transformer.proj_out.weight`, - # for example) and they have different LoRA ranks. - rank[f"^{key}"] = val.shape[1] - - if network_alphas is not None and len(network_alphas) >= 1: - alpha_keys = [k for k in network_alphas.keys() if k.startswith(f"{prefix}.")] - network_alphas = { - k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys - } - - # adapter_name if adapter_name is None: adapter_name = get_adapter_name(self) - # create LoraConfig - lora_config = _create_lora_config( - state_dict, - network_alphas, - metadata, - rank, - model_state_dict=self.state_dict(), - adapter_name=adapter_name, - ) + # LoKr adapters (Kronecker product factors) use a different peft config and state dict layout + # (`{module}.lokr_w1` etc.) than LoRA; detect them before any LoRA-specific key handling. + is_lokr = any(".lokr_" in k for k in state_dict) - # Adjust LoRA config for Control LoRA - if is_sai_sd_control_lora: - lora_config.lora_alpha = lora_config.r - lora_config.alpha_pattern = lora_config.rank_pattern - lora_config.bias = "all" - lora_config.modules_to_save = lora_config.exclude_modules - lora_config.exclude_modules = None + if is_lokr: + if hotswap: + raise ValueError("Hotswapping LoKr adapters is not supported.") + _maybe_fuse_qkv_projections_for_lokr(self, state_dict) + adapter_config = _create_lokr_config(state_dict, metadata) + else: + # check with first key if is not in peft format + first_key = next(iter(state_dict.keys())) + if "lora_A" not in first_key: + state_dict = convert_unet_state_dict_to_peft(state_dict) + + # Control LoRA from SAI is different from BFL Control LoRA + # https://huggingface.co/stabilityai/control-lora + # https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors + is_sai_sd_control_lora = "lora_controlnet" in state_dict + if is_sai_sd_control_lora: + state_dict = convert_sai_sd_control_lora_state_dict_to_peft(state_dict) + + rank = {} + for key, val in state_dict.items(): + # Cannot figure out rank from lora layers that don't have at least 2 dimensions. + # Bias layers in LoRA only have a single dimension + if "lora_B" in key and val.ndim > 1: + # Check out https://github.com/huggingface/peft/pull/2419 for the `^` symbol. + # We may run into some ambiguous configuration values when a model has module + # names, sharing a common prefix (`proj_out.weight` and `blocks.transformer.proj_out.weight`, + # for example) and they have different LoRA ranks. + rank[f"^{key}"] = val.shape[1] + + if network_alphas is not None and len(network_alphas) >= 1: + alpha_keys = [k for k in network_alphas.keys() if k.startswith(f"{prefix}.")] + network_alphas = { + k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys + } + + # create LoraConfig + adapter_config = _create_lora_config( + state_dict, + network_alphas, + metadata, + rank, + model_state_dict=self.state_dict(), + adapter_name=adapter_name, + ) + + # Adjust LoRA config for Control LoRA + if is_sai_sd_control_lora: + adapter_config.lora_alpha = adapter_config.r + adapter_config.alpha_pattern = adapter_config.rank_pattern + adapter_config.bias = "all" + adapter_config.modules_to_save = adapter_config.exclude_modules + adapter_config.exclude_modules = None # 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) -> torch.Tensor: + 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(self, x, t) -> torch.Tensor: + # 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 + + +class SanaWMCaptionEmbedder(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"), + ): + super().__init__() + self.y_proj = SanaWMMlp( + in_features=in_channels, hidden_features=hidden_size, out_features=hidden_size, act_layer=act_layer, drop=0 + ) + + def forward(self, caption) -> torch.Tensor: + return self.y_proj(caption) + + +class SanaWMPatchEmbed(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) -> torch.Tensor: + 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 rather than being split into real cos/sin buffers. + + The complex table is held as a plain attribute rather than a registered buffer. `nn.Module.to(dtype)` casts complex + buffers to the requested *real* dtype, silently discarding the imaginary (sine) half of every rotation, so + `transformer.to(torch.bfloat16)` would corrupt it. Keeping it out of the buffer registry makes dtype casts leave it + alone; it is moved to the activations' device on each call. + """ + + 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.freqs = torch.cat(freqs, dim=1) + + def forward(self, fhw: Tuple[int, int, int], device: torch.device | None = None) -> torch.Tensor: + ppf, pph, ppw = fhw + + freqs = self.freqs.to(device).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). + """ + # Camera geometry runs in the input dtype. `torch.cross` has no autocast rule and raises under autocast, so + # keep autocast off here: without autocast these ops already run in the input dtype, so results are unchanged. + with torch.autocast(device_type=camera_conditions.device.type, enabled=False): + 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 = SanaWMMlp( + 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, + ) -> torch.Tensor: + """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. + + The state-dict is identical to the public sana checkpoint apart from the intentionally-removed ``pos_embed`` + buffer. + """ + + _supports_gradient_checkpointing = True + _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, + 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 + 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.gradient_checkpointing = False + + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.t_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=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 = SanaWMPatchEmbed( + patch_size, x_embedder_in_channels, hidden_size, kernel_size=kernel_size, bias=True + ) + + self.y_embedder = SanaWMCaptionEmbedder( + in_channels=caption_channels, + hidden_size=hidden_size, + act_layer=approx_gelu, + ) + + 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 = SanaWMPatchEmbed( + 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 = SanaWMPatchEmbed(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, + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: + """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), device=x.device) + 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(self.time_proj(timestep.flatten()).to(self.dtype)) # (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): + block_args = ( + x, + y, + t0, + y_lens, + (post_patch_num_frames, post_patch_height, post_patch_width), + image_pos_embed, + block_mask if i > 1 else None, + camera_conditions, + ucpe_ray_transforms, + post_attn_plucker_emb, + frame_valid_mask, + chunk_size, + ) + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = self._gradient_checkpointing_func(block, *block_args) + else: + x = block(*block_args) # (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 0000000000..892eca5eb8 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_sana_wm_refiner.py @@ -0,0 +1,1556 @@ +# Copyright 2026 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, + ) -> tuple[torch.Tensor] | Transformer2DModelOutput: + 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/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 913075ff9e..37a9e7c2fe 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -13,6 +13,7 @@ # limitations under the License. import math +from typing import Any import torch import torch.nn as nn @@ -24,6 +25,7 @@ from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...models.attention_processor import Attention from ...models.modeling_utils import ModelMixin from ...models.normalization import RMSNorm +from ...utils import apply_lora_scale from ...utils.torch_utils import maybe_allow_in_graph from ..attention_dispatch import dispatch_attention_fn from ..modeling_outputs import Transformer2DModelOutput @@ -891,6 +893,7 @@ class ZImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOr return unified, unified_freqs, attn_mask, noise_mask_tensor + @apply_lora_scale("joint_attention_kwargs") def forward( self, x: list[torch.Tensor, list[list[torch.Tensor]]], @@ -902,6 +905,7 @@ class ZImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOr image_noise_mask: list[list[int]] | None = None, patch_size: int = 2, f_patch_size: int = 1, + joint_attention_kwargs: dict[str, Any] | None = None, ) -> Transformer2DModelOutput | tuple[torch.Tensor]: """ The [`ZImageTransformer2DModel`] forward method. @@ -930,6 +934,8 @@ class ZImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOr Spatial patch size used to patchify the input latents. f_patch_size (`int`, *optional*, defaults to 1): Temporal patch size used to patchify the input latents. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary. Only the `scale` entry is used, to scale the LoRA layers. Returns: If `return_dict` is True, a [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a diff --git a/src/diffusers/modular_pipelines/anima/encoders.py b/src/diffusers/modular_pipelines/anima/encoders.py index 726a68280d..2f86334fd5 100644 --- a/src/diffusers/modular_pipelines/anima/encoders.py +++ b/src/diffusers/modular_pipelines/anima/encoders.py @@ -58,41 +58,49 @@ class AnimaTextEncoderStep(ModularPipelineBlocks): return [ OutputParam( "qwen_prompt_embeds", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Qwen prompt embeddings to be consumed by the Anima text conditioner.", ), OutputParam( "qwen_attention_mask", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Qwen prompt attention mask to be consumed by the Anima text conditioner.", ), OutputParam( "t5_input_ids", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="T5 prompt token ids to be consumed by the Anima text conditioner.", ), OutputParam( "t5_attention_mask", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="T5 prompt attention mask to be consumed by the Anima text conditioner.", ), OutputParam( "negative_qwen_prompt_embeds", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Negative Qwen prompt embeddings to be consumed by the Anima text conditioner.", ), OutputParam( "negative_qwen_attention_mask", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Negative Qwen prompt attention mask to be consumed by the Anima text conditioner.", ), OutputParam( "negative_t5_input_ids", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Negative T5 prompt token ids to be consumed by the Anima text conditioner.", ), OutputParam( "negative_t5_attention_mask", + kwargs_type="denoiser_input_fields", type_hint=torch.Tensor, description="Negative T5 prompt attention mask to be consumed by the Anima text conditioner.", ), diff --git a/src/diffusers/modular_pipelines/anima/modular_pipeline.py b/src/diffusers/modular_pipelines/anima/modular_pipeline.py index 1430ac5558..a7136bb97c 100644 --- a/src/diffusers/modular_pipelines/anima/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/anima/modular_pipeline.py @@ -39,7 +39,7 @@ class AnimaModularPipeline(ModularPipeline, AnimaLoraLoaderMixin): @property def vae_scale_factor(self): vae_scale_factor = 8 - if self.vae is not None: + if getattr(self, "vae", None) is not None: vae_scale_factor = 2 ** len(self.vae.temperal_downsample) return vae_scale_factor diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index b8cb626d82..4f1f600e4f 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -388,6 +388,11 @@ else: "SanaVideoPipeline", "SanaImageToVideoPipeline", ] + _import_structure["sana_wm"] = [ + "SanaWMPipeline", + "SanaWMLTX2Refiner", + "SanaWMPipelineOutput", + ] _import_structure["shap_e"] = ["ShapEImg2ImgPipeline", "ShapEPipeline"] _import_structure["stable_audio"] = [ "StableAudioProjectionModel", @@ -885,6 +890,11 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: 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 0000000000..c33e4f6157 --- /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 0000000000..b7b41815be --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/cam_utils.py @@ -0,0 +1,321 @@ +# Copyright 2026 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 0000000000..0d4bc196af --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/image_processor.py @@ -0,0 +1,67 @@ +# Copyright 2026 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 0000000000..3664ab652f --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/pipeline_output.py @@ -0,0 +1,44 @@ +# Copyright 2026 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: a float ``np.ndarray`` of shape ``(T, H, W, 3)`` in ``[0, 1]`` for ``output_type="np"``, a + ``(T, 3, H, W)`` tensor for ``"pt"``, a list of ``T`` ``PIL.Image`` for ``"pil"``, or the latent tensor + ``(B, C, T_lat, H_lat, W_lat)`` for ``"latent"``. + 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 0000000000..9a06f8e7eb --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py @@ -0,0 +1,659 @@ +# Copyright 2026 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 Callable, 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" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + 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 | None = None, + 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 + + device = device or self._execution_device + 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._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] + + # ------------------------------------------------------------------ + # 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, + callback_on_step_end_tensor_inputs: list[str] | None = None, + ) -> None: + """Validate `__call__` inputs. Raises on bad input and returns nothing.""" + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found " + f"{[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + 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", "pt", "latent"] = "np", + callback_on_step_end: Callable[[DiffusionPipeline, int, int, dict], dict] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + 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"`, `"pt"`, or `"latent"`, defaults to `"np"`): + Output format. + callback_on_step_end (`Callable`, *optional*): + Called at the end of each denoising step as `callback_on_step_end(self, step, timestep, + callback_kwargs)`, where `callback_kwargs` holds the tensors named in + `callback_on_step_end_tensor_inputs`. Tensors in the returned dict replace the pipeline's. + callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`): + Which tensors to pass to `callback_on_step_end`; must be a subset of `._callback_tensor_inputs`. + 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"`, a ``(T, 3, H, W)`` tensor + for `"pt"`, or the raw latent tensor for `"latent"`. + + Examples: + """ + self.check_inputs(c2w, action, intrinsics, callback_on_step_end_tensor_inputs) + 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 i, t in enumerate(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 + self.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 callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + self._current_timestep = None + + if output_type == "latent": + self.maybe_free_model_hooks() + if not return_dict: + return (latents, c2w, latents) + return SanaWMPipelineOutput(frames=latents, c2w=c2w, latent=latents) + + 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] + self.maybe_free_model_hooks() + + 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 0000000000..b0dcc96c8e --- /dev/null +++ b/src/diffusers/pipelines/sana_wm/refiner.py @@ -0,0 +1,515 @@ +# Copyright 2026 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/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py index 9509adde74..8a0407f678 100644 --- a/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py +++ b/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py @@ -1067,6 +1067,9 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latent_model_input.shape[0]) + # on TPU, a view into `timesteps` makes `torch.compile` recompile every step + if timestep.device.type == "tpu": + timestep = timestep.clone() noise_pred = self.transformer( hidden_states=latent_model_input, @@ -1089,6 +1092,8 @@ class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingle ) if skip_guidance_layers is not None and should_skip_layers: timestep = t.expand(latents.shape[0]) + if timestep.device.type == "tpu": + timestep = timestep.clone() latent_model_input = latents noise_pred_skip_layers = self.transformer( hidden_states=latent_model_input, diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index b97bdda170..f7aa3a6d28 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image.py @@ -531,7 +531,11 @@ class ZImagePipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMix latent_model_input_list = list(latent_model_input.unbind(dim=0)) model_out_list = self.transformer( - latent_model_input_list, timestep_model_input, prompt_embeds_model_input, return_dict=False + latent_model_input_list, + timestep_model_input, + prompt_embeds_model_input, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, )[0] if apply_cfg: diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py index 6ba698a30d..0653186e26 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py @@ -672,6 +672,7 @@ class ZImageControlNetPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSin timestep_model_input, prompt_embeds_model_input, controlnet_block_samples=controlnet_block_samples, + joint_attention_kwargs=self.joint_attention_kwargs, )[0] if apply_cfg: diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py index cb72c80770..33ff876ea2 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py @@ -699,6 +699,7 @@ class ZImageControlNetInpaintPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, timestep_model_input, prompt_embeds_model_input, controlnet_block_samples=controlnet_block_samples, + joint_attention_kwargs=self.joint_attention_kwargs, )[0] if apply_cfg: diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py b/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py index efa6f50962..48e3fae94a 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py @@ -647,6 +647,7 @@ class ZImageImg2ImgPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingle latent_model_input_list, timestep_model_input, prompt_embeds_model_input, + joint_attention_kwargs=self.joint_attention_kwargs, )[0] if apply_cfg: diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py index 62a53ea261..64e0cca28c 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py @@ -852,6 +852,7 @@ class ZImageInpaintPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingle latent_model_input_list, timestep_model_input, prompt_embeds_model_input, + joint_attention_kwargs=self.joint_attention_kwargs, )[0] if apply_cfg: diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py b/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py index c73e0d7d18..46da8b135d 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py +++ b/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py @@ -680,6 +680,7 @@ class ZImageOmniPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFil cap_feats=prompt_embeds_model_input, siglip_feats=condition_siglip_embeds_model_input, image_noise_mask=image_noise_mask, + joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index b3051dfcb9..5f89eb318d 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -116,6 +116,7 @@ from .import_utils import ( is_torch_mlu_available, is_torch_neuronx_available, is_torch_npu_available, + is_torch_tpu_available, is_torch_version, is_torch_xla_available, is_torch_xla_version, diff --git a/src/diffusers/utils/accelerate_utils.py b/src/diffusers/utils/accelerate_utils.py index ae6a8ca747..f7deeb6e39 100644 --- a/src/diffusers/utils/accelerate_utils.py +++ b/src/diffusers/utils/accelerate_utils.py @@ -15,6 +15,8 @@ Accelerate utilities: Utilities related to accelerate """ +import functools + from packaging import version from .import_utils import is_accelerate_available @@ -40,6 +42,7 @@ def apply_forward_hook(method): if version.parse(accelerate_version) < version.parse("0.17.0"): return method + @functools.wraps(method) def wrapper(self, *args, **kwargs): if hasattr(self, "_hf_hook") and hasattr(self._hf_hook, "pre_forward"): self._hf_hook.pre_forward(self) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 2d89f6f955..9e1f5f027c 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2254,6 +2254,36 @@ class SanaVideoTransformer3DModel(metaclass=DummyObject): 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 9ede7a7543..a068ef8672 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -4037,6 +4037,51 @@ class SanaVideoPipeline(metaclass=DummyObject): 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/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index d2cf394cd9..af765021f4 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -178,6 +178,7 @@ except importlib_metadata.PackageNotFoundError: _torch_xla_available, _torch_xla_version = _is_package_available("torch_xla") _torch_npu_available, _torch_npu_version = _is_package_available("torch_npu") _torch_mlu_available, _torch_mlu_version = _is_package_available("torch_mlu") +_torch_tpu_available, _torch_tpu_version = _is_package_available("torch_tpu") _torch_neuronx_available, _torch_neuronx_version = _is_package_available("torch_neuronx") _transformers_available, _transformers_version = _is_package_available("transformers") _hf_hub_available, _hf_hub_version = _is_package_available("huggingface_hub") @@ -238,6 +239,10 @@ def is_torch_mlu_available(): return _torch_mlu_available +def is_torch_tpu_available(): + return _torch_tpu_available + + def is_torch_neuronx_available(): return _torch_neuronx_available diff --git a/src/diffusers/utils/peft_utils.py b/src/diffusers/utils/peft_utils.py index ea6f867981..f20e3f32b0 100644 --- a/src/diffusers/utils/peft_utils.py +++ b/src/diffusers/utils/peft_utils.py @@ -368,6 +368,126 @@ def check_peft_version(min_version: str) -> None: ) +def _maybe_fuse_qkv_projections_for_lokr(model, state_dict) -> None: + """ + Fuse the model's QKV projections when a peft-format LoKr state dict targets fused `to_qkv` / `to_added_qkv` + projections that the model does not have yet. + + BFL-format Flux2 LoKr checkpoints apply LoKr to the fused QKV projections. Unlike a LoRA delta, a Kronecker product + delta over the fused projection cannot be split exactly into separate Q/K/V factors, so the model's projections are + fused instead and the adapter maps 1:1. + """ + fused_targets = { + module + for module in (k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k) + if module.rsplit(".", 1)[-1] in ("to_qkv", "to_added_qkv") + } + named_modules = dict(model.named_modules()) + if all(module in named_modules for module in fused_targets) or not hasattr(model, "fuse_qkv_projections"): + return + + if getattr(model, "is_quantized", False): + raise ValueError( + "This LoKr checkpoint targets fused QKV projections. Fusing concatenates the Q/K/V weights into a new " + "`nn.Linear`, which is not possible with quantized weights. Please load the transformer without " + "quantization." + ) + + # Fusing replaces to_q/to_k/to_v (and the add_*_proj) with a single projection, which would orphan any adapter + # already injected on the unfused ones. + from peft.tuners.tuners_utils import BaseTunerLayer + + unfused_projections = {"to_q", "to_k", "to_v", "add_q_proj", "add_k_proj", "add_v_proj"} + adapted = [ + name + for name, module in named_modules.items() + if isinstance(module, BaseTunerLayer) and name.rsplit(".", 1)[-1] in unfused_projections + ] + if adapted: + raise ValueError( + "This LoKr checkpoint targets fused QKV projections, but an adapter is already loaded on the unfused " + f"projections (e.g. `{adapted[0]}`). Unload it with `unload_lora_weights()` before loading this checkpoint." + ) + + logger.info("The LoKr checkpoint targets fused QKV projections; calling `fuse_qkv_projections()` on the model.") + model.fuse_qkv_projections() + + +def _create_lokr_config(state_dict, metadata): + """ + Create a `LoKrConfig` from a peft-format LoKr state dict (keys like `{module}.lokr_w1`). + + Without metadata, the config is inferred from the tensor shapes. The checkpoint alpha is expected to be already + baked into the weights by the state dict conversion, so `alpha` is set equal to the rank (runtime scaling 1.0). + peft re-derives each module's Kronecker factorization from `decompose_factor` and only creates rank-decomposed + factors when the rank is small compared to the factorized dimensions, so for modules whose checkpoint factors are + full matrices the rank is set to `max(lokr_w2.shape)` to make peft create full matrices as well. + """ + from peft import LoKrConfig + from peft.tuners.lokr.layer import factorization + + if metadata is not None: + try: + return LoKrConfig(**metadata) + except TypeError as e: + raise TypeError("`LoKrConfig` class could not be instantiated.") from e + + modules = sorted({k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k}) + + # Reconstruct each module's factorized dimensions, (out_l, out_k) x (in_m, in_n), from the checkpoint. A factor + # is either stored as a full matrix (`lokr_w1`) or rank-decomposed into `lokr_w1_a @ lokr_w1_b` (same for w2). + factorizations = {} + rank_dict = {} + for module in modules: + w1, w1_a, w1_b = (state_dict.get(f"{module}.lokr_w1{s}") for s in ("", "_a", "_b")) + w2, w2_a, w2_b = (state_dict.get(f"{module}.lokr_w2{s}") for s in ("", "_a", "_b")) + out_l, in_m = w1.shape if w1 is not None else (w1_a.shape[0], w1_b.shape[1]) + out_k, in_n = w2.shape if w2 is not None else (w2_a.shape[0], w2_b.shape[1]) + factorizations[module] = ((out_l, out_k), (in_m, in_n)) + if w2_a is not None: + rank_dict[module] = w2_a.shape[1] + elif w1_a is not None: + rank_dict[module] = w1_a.shape[1] + else: + rank_dict[module] = max(w2.shape) + + # Find the `decompose_factor` under which peft reproduces the checkpoint factorizations. A fixed factor shows up + # as the left dimension of the modules it divides (modules it does not divide fall back to a near-square + # factorization, like with factor -1), so every observed left dimension is a candidate. + left_dims = {dims[0][0] for dims in factorizations.values()} + decompose_factor = None + for candidate in sorted(left_dims) + [-1]: + if all( + factorization(out_l * out_k, candidate) == (out_l, out_k) + and factorization(in_m * in_n, candidate) == (in_m, in_n) + for (out_l, out_k), (in_m, in_n) in factorizations.values() + ): + decompose_factor = candidate + break + if decompose_factor is None: + raise ValueError( + "Could not infer a `decompose_factor` that reproduces the Kronecker factorizations of this LoKr " + "state dict. Please open an issue: https://github.com/huggingface/diffusers/issues/new" + ) + + r = collections.Counter(rank_dict.values()).most_common(1)[0][0] + rank_pattern = {k: v for k, v in rank_dict.items() if v != r} + + lokr_config_kwargs = { + "r": r, + "alpha": r, + "rank_pattern": rank_pattern, + "alpha_pattern": dict(rank_pattern), + "target_modules": modules, + "decompose_both": any(".lokr_w1_a" in k for k in state_dict), + "decompose_factor": decompose_factor, + } + try: + return LoKrConfig(**lokr_config_kwargs) + except TypeError as e: + raise TypeError("`LoKrConfig` class could not be instantiated.") from e + + def _create_lora_config( state_dict, network_alphas, metadata, rank_pattern_dict, is_unet=True, model_state_dict=None, adapter_name=None ): @@ -397,7 +517,7 @@ def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name: str) -> None # Check only for unexpected keys. unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) if unexpected_keys: - lora_unexpected_keys = [k for k in unexpected_keys if ".lora_" in k] + lora_unexpected_keys = [k for k in unexpected_keys if ".lora_" in k or "lokr_" in k] if lora_unexpected_keys: warn_msg = ( f"Loading adapter weights from state_dict led to unexpected keys found in the model:" @@ -407,7 +527,7 @@ def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name: str) -> None # Filter missing keys specific to the current adapter. missing_keys = getattr(incompatible_keys, "missing_keys", None) if missing_keys: - lora_missing_keys = [k for k in missing_keys if ".lora_" in k and adapter_name in k] + lora_missing_keys = [k for k in missing_keys if (".lora_" in k or "lokr_" in k) and adapter_name in k] if lora_missing_keys: warn_msg += ( f"Loading adapter weights from state_dict led to missing keys in the model:" diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 2d7d5ae232..ee0ab21e87 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -17,12 +17,14 @@ from .cache import ( from .common import BaseModelTesterConfig, ModelTesterMixin from .compile import TorchCompileTesterMixin from .ip_adapter import IPAdapterTesterMixin +from .lokr import LoKrTesterMixin from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin from .parallelism import ( ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, TensorParallelTesterMixin, + TensorParallelTPUTesterMixin, ) from .quantization import ( AutoRoundCompileTesterMixin, @@ -67,6 +69,7 @@ __all__ = [ "ContextParallelTesterMixin", "ContextParallelAttentionBackendsTesterMixin", "TensorParallelTesterMixin", + "TensorParallelTPUTesterMixin", "CPUOffloadTesterMixin", "FasterCacheConfigMixin", "FasterCacheTesterMixin", @@ -80,6 +83,7 @@ __all__ = [ "GroupOffloadTesterMixin", "IPAdapterTesterMixin", "LayerwiseCastingTesterMixin", + "LoKrTesterMixin", "LoraHotSwappingForModelTesterMixin", "LoraTesterMixin", "MemoryTesterMixin", diff --git a/tests/models/testing_utils/lokr.py b/tests/models/testing_utils/lokr.py new file mode 100644 index 0000000000..d1589d28fa --- /dev/null +++ b/tests/models/testing_utils/lokr.py @@ -0,0 +1,139 @@ +# 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 +import torch.nn as nn + +from diffusers.utils.import_utils import is_peft_available + +from ...testing_utils import assert_tensors_close, is_lora, require_peft_backend, torch_device +from .common import BaseModelOutputMixin + + +if is_peft_available(): + from peft.tuners.lokr.layer import LoKrLayer, factorization + + from diffusers.loaders.peft import PeftAdapterMixin + + +def make_lokr_factors(out_features, in_features, factor=4, rank=None): + """ + Random LoKr factors for an `out_features x in_features` layer, factorized like peft does with `decompose_factor`. + With `rank`, the right factor is stored rank-decomposed (`lokr_w2_a @ lokr_w2_b`). + + Returns the factors keyed by their state dict suffix, and the delta weight they encode. + """ + out_l, out_k = factorization(out_features, factor) + in_m, in_n = factorization(in_features, factor) + w1 = torch.randn(out_l, in_m) + if rank is None: + w2 = torch.randn(out_k, in_n) + return {"lokr_w1": w1, "lokr_w2": w2}, torch.kron(w1, w2) + w2_a, w2_b = torch.randn(out_k, rank), torch.randn(rank, in_n) + return {"lokr_w1": w1, "lokr_w2_a": w2_a, "lokr_w2_b": w2_b}, torch.kron(w1, w2_a @ w2_b) + + +def check_lokr_deltas(model, expected_deltas, adapter_name="default", atol=1e-5): + """Check that exactly the expected modules carry a LoKr adapter, each with the expected delta weight.""" + named_modules = dict(model.named_modules()) + adapted = {name for name, module in named_modules.items() if isinstance(module, LoKrLayer)} + assert adapted == set(expected_deltas) + for name, expected in expected_deltas.items(): + delta = named_modules[name].get_delta_weight(adapter_name).cpu() + assert_tensors_close(delta, expected, atol=atol, rtol=0, msg=f"Wrong LoKr delta on {name}") + + +@is_lora +@require_peft_backend +class LoKrTesterMixin(BaseModelOutputMixin): + """ + Mixin class for testing loading LoKr (LyCORIS Kronecker product) adapters with `load_lora_adapter`. + + Expected from config mixin: + - model_class: The model class to test + + Expected methods from config mixin: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: lora + Use `pytest -m "not lora"` to skip these tests + """ + + def setup_method(self): + if not issubclass(self.model_class, PeftAdapterMixin): + pytest.skip(f"PEFT is not supported for this model ({self.model_class.__name__}).") + + def _flatten_output(self, output): + # Some models (e.g. Z-Image) return a list of per-sample tensors. + if isinstance(output, (list, tuple)): + return torch.cat([t.flatten() for t in output]) + return output + + def _model_output(self, model, inputs_dict): + return self._flatten_output(model(**inputs_dict, return_dict=False)[0]) + + def get_lokr_state_dict(self, model, rank=None): + """ + A peft-format LoKr state dict on every attention `to_q` and `to_v`, with the expected delta of each. With + `rank`, the `to_v` factors are rank-decomposed, so the config inference has to mix both kinds. + """ + state_dict, expected_deltas = {}, {} + for name, module in model.named_modules(): + projection = name.rsplit(".", 1)[-1] + if isinstance(module, nn.Linear) and projection in ("to_q", "to_v"): + factors, delta = make_lokr_factors( + module.out_features, module.in_features, rank=rank if projection == "to_v" else None + ) + state_dict.update({f"{name}.{suffix}": weight for suffix, weight in factors.items()}) + expected_deltas[name] = delta + return state_dict, expected_deltas + + @pytest.mark.parametrize("rank", [None, 1], ids=["full_factors", "rank_decomposed_factors"]) + @torch.no_grad() + def test_lokr_adapter_loads_exact_kronecker_deltas(self, base_model_output, rank): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + state_dict, expected_deltas = self.get_lokr_state_dict(model, rank=rank) + + model.load_lora_adapter(state_dict, prefix=None, adapter_name="default") + + check_lokr_deltas(model, expected_deltas) + output = self._model_output(model, self.get_dummy_inputs()) + base_output = self._flatten_output(base_model_output) + assert not torch.allclose(output, base_output, atol=1e-4, rtol=1e-4), "Output should differ with LoKr" + + @torch.no_grad() + def test_lokr_unload_restores_base_output(self, base_model_output): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + state_dict, _ = self.get_lokr_state_dict(model) + + model.load_lora_adapter(state_dict, prefix=None, adapter_name="default") + model.unload_lora() + + assert not any(isinstance(module, LoKrLayer) for module in model.modules()) + output = self._model_output(model, self.get_dummy_inputs()) + assert_tensors_close(output, self._flatten_output(base_model_output), atol=1e-4, rtol=1e-4) + + def test_lokr_hotswap_raises(self): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + state_dict, _ = self.get_lokr_state_dict(model) + model.load_lora_adapter(state_dict, prefix=None, adapter_name="default") + + with pytest.raises(ValueError, match="Hotswapping LoKr adapters is not supported"): + model.load_lora_adapter(state_dict, prefix=None, adapter_name="default", hotswap=True) diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index b525637c95..eca6bff5d2 100644 --- a/tests/models/testing_utils/parallelism.py +++ b/tests/models/testing_utils/parallelism.py @@ -31,6 +31,7 @@ from ...testing_utils import ( is_kernels_available, is_tensor_parallel, require_torch_multi_accelerator, + require_torch_tpu, torch_device, ) from .common import calculate_expected_num_shards, compute_module_persistent_sizes @@ -41,6 +42,7 @@ from .utils import _maybe_cast_to_bf16 DEVICE_CONFIG = { "cuda": {"backend": "nccl", "module": torch.cuda}, "xpu": {"backend": "xccl", "module": torch.xpu}, + "tpu": {"backend": "tpu_dist", "module": None}, } @@ -245,7 +247,15 @@ def _custom_mesh_worker( def _tensor_parallel_worker( - rank, world_size, master_port, model_class, init_dict, inputs_dict, return_dict, state_dict + rank, + world_size, + master_port, + model_class, + init_dict, + inputs_dict, + return_dict, + state_dict, + device_type=torch_device, ): """Worker function for tensor parallel inference testing. @@ -258,16 +268,19 @@ def _tensor_parallel_worker( os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(master_port) os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) - device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"]) - backend = device_config["backend"] - device_module = device_config["module"] + device_config = DEVICE_CONFIG.get(device_type, DEVICE_CONFIG["cuda"]) + dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size) - dist.init_process_group(backend=backend, rank=rank, world_size=world_size) - - device_module.set_device(rank) - device = torch.device(f"{torch_device}:{rank}") + if device_type == "tpu": + # Each spawned process is bound to one chip. Avoid bf16 matmuls to keep the tolerance tight. + device = torch.device("tpu") + torch.set_float32_matmul_precision("highest") + else: + device_config["module"].set_device(rank) + device = torch.device(f"{device_type}:{rank}") model = model_class(**init_dict) model.load_state_dict(state_dict) @@ -303,8 +316,8 @@ def _tensor_parallel_from_pretrained_worker( """Worker for `from_pretrained(..., parallel_config=...)`, i.e. sharding while reading the checkpoint. Each rank loads only its own slice of every `_tp_plan` parameter straight into a `DTensor` and runs a forward - pass. Rank 0 reports its output and the local/global shapes of one sharded weight so the caller can check both the - numerics and that sharding actually happened. + pass. Rank 0 checks that exactly the parameters `_tp_plan` covers were loaded as `DTensor`s, each with the + placement and local shape its shard spec implies, and reports its output so the caller can check the numerics. """ try: os.environ["MASTER_ADDR"] = "localhost" @@ -316,7 +329,9 @@ def _tensor_parallel_from_pretrained_worker( dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size) device_config["module"].set_device(rank) - from torch.distributed.tensor import DTensor + from torch.distributed.tensor import DTensor, Replicate, Shard + + from diffusers.hooks.tensor_parallel import resolve_tp_shard_specs model = model_class.from_pretrained( checkpoint_dir, parallel_config=TensorParallelConfig(tp_degree=world_size) @@ -330,12 +345,30 @@ def _tensor_parallel_from_pretrained_worker( output = output.full_tensor() if rank == 0: - sharded = {k: v for k, v in model.state_dict().items() if isinstance(v, DTensor)} - assert sharded, "No parameter was sharded into a DTensor by the streaming load." - name, param = next(iter(sharded.items())) + specs = resolve_tp_shard_specs(model, model_class._tp_plan, world_size) + state_dict = model.state_dict() + + # The numerics check alone wouldn't catch a planned parameter left unsharded. + for name, spec in specs.items(): + param = state_dict[name] + assert isinstance(param, DTensor), ( + f"'{name}' is covered by `_tp_plan` but was not loaded as a DTensor." + ) + placement = Replicate() if spec.dim is None else Shard(spec.dim) + assert param.placements == (placement,), ( + f"'{name}' has placements {param.placements}, not {placement}." + ) + expected_local_shape = list(param.shape) + if spec.dim is not None: + expected_local_shape[spec.dim] //= world_size + assert list(param.to_local().shape) == expected_local_shape, ( + f"'{name}' has local shape {list(param.to_local().shape)}, not {expected_local_shape}." + ) + + unplanned = sorted(k for k, v in state_dict.items() if isinstance(v, DTensor) and k not in specs) + assert not unplanned, f"Parameters not covered by `_tp_plan` were loaded as DTensors: {unplanned}" + return_dict["status"] = "success" - return_dict["num_sharded"] = len(sharded) - return_dict["shard_example"] = (name, list(param.to_local().shape), list(param.shape)) return_dict["output"] = output.float().cpu().tolist() except Exception as e: @@ -456,15 +489,75 @@ class TensorParallelTesterMixin: f"Tensor parallel `from_pretrained` failed: {return_dict.get('error', 'Unknown error')}" ) - name, local_shape, global_shape = return_dict["shard_example"] - assert local_shape != global_shape, ( - f"'{name}' has local shape {local_shape} equal to its global shape, so it was not sharded." - ) - # Sharded matmuls + all-reduce reorder the summation, so allow a small tolerance over the reference. torch.testing.assert_close(reference, torch.tensor(return_dict["output"]), atol=1e-3, rtol=1e-3) +@is_tensor_parallel +@require_torch_tpu +class TensorParallelTPUTesterMixin: + """Same check as `TensorParallelTesterMixin`, spawning one process per TPU chip. + + Run these tests in their own pytest process (e.g. `-k TensorParallelTPU`): `torch.manual_seed` or a backward pass + in an earlier test binds every TPU chip to the pytest process, and the spawned workers then fail. + """ + + # Smallest slice every TPU generation supports, so the test is the same on any CI host. + tp_world_size = 4 + + def test_tensor_parallel_tpu_inference(self, atol=1e-3, rtol=1e-3): + from torch_tpu._internal.distributed.launchers.singlehost_wrapper import prepare_tpu_environment + from torch_tpu._internal.utils import hardware + + if getattr(self.model_class, "_tp_plan", None) is None: + pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.") + + world_size = self.tp_world_size + if hardware.get_tpu_device_count() < world_size: + pytest.skip(f"Needs at least {world_size} TPU chips.") + init_dict = self.get_init_dict() + num_heads = init_dict.get("num_attention_heads") + if num_heads is not None and num_heads % world_size != 0: + pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).") + + # Reference on CPU: touching the TPU here would bind every chip to this process. + inputs_dict = self.get_dummy_inputs(device="cpu") + model = self.model_class(**init_dict).eval() + with torch.no_grad(): + ref_output = model(**inputs_dict, return_dict=False)[0].float() + + tpu_env = ("TORCH_TPU_TOPOLOGY", "TORCH_TPU_SLICEBUILDER_ADDRESSES") + for key in tpu_env: + os.environ.pop(key, None) + prepare_tpu_environment(world_size) + return_dict = mp.Manager().dict() + try: + mp.spawn( + _tensor_parallel_worker, + args=( + world_size, + _find_free_port(), + self.model_class, + init_dict, + inputs_dict, + return_dict, + model.state_dict(), + "tpu", + ), + nprocs=world_size, + join=True, + ) + finally: + for key in tpu_env: + os.environ.pop(key, None) + + assert return_dict.get("status") == "success", ( + f"Tensor parallel inference failed: {return_dict.get('error', 'Unknown error')}" + ) + tp_output = torch.tensor(return_dict["output"]) + torch.testing.assert_close(ref_output, tp_output, atol=atol, rtol=rtol) + + @is_context_parallel @require_torch_multi_accelerator class ContextParallelTesterMixin: diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index b08fde2963..04542253c7 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -53,6 +53,7 @@ from ..testing_utils import ( SingleFileTesterMixin, TaylorSeerCacheTesterMixin, TensorParallelTesterMixin, + TensorParallelTPUTesterMixin, TorchAoCompileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, @@ -276,6 +277,14 @@ class TestFluxTransformerTensorParallel(FluxTransformerTesterConfig, TensorParal """Tensor Parallel inference tests for Flux Transformer (CUDA/XPU multi-accelerator).""" +class TestFluxTransformerTensorParallelTPU(FluxTransformerTesterConfig, TensorParallelTPUTesterMixin): + """Tensor Parallel inference test for Flux Transformer on TPU.""" + + def get_init_dict(self): + # One head per chip. + return {**super().get_init_dict(), "num_attention_heads": self.tp_world_size} + + def make_neuron_tp_spec(): """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`). diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index 3263ce6820..31d47f1153 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -17,9 +17,11 @@ import os import subprocess import sys +import pytest import torch from diffusers import Flux2Transformer2DModel +from diffusers.loaders.lora_pipeline import Flux2LoraLoaderMixin from diffusers.models.transformers.transformer_flux2 import ( Flux2KVAttnProcessor, Flux2KVCache, @@ -36,17 +38,20 @@ from ..testing_utils import ( ContextParallelTesterMixin, GGUFCompileTesterMixin, GGUFTesterMixin, + LoKrTesterMixin, LoraHotSwappingForModelTesterMixin, LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, SingleFileTesterMixin, TensorParallelTesterMixin, + TensorParallelTPUTesterMixin, TorchAoCompileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) +from ..testing_utils.lokr import check_lokr_deltas, make_lokr_factors enable_full_determinism() @@ -176,6 +181,14 @@ def make_neuron_tp_spec(): return Flux2Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu") +class TestFlux2TransformerTensorParallelTPU(Flux2TransformerTesterConfig, TensorParallelTPUTesterMixin): + """Tensor Parallel inference test for Flux2 Transformer on TPU.""" + + def get_init_dict(self): + # One head per chip. + return {**super().get_init_dict(), "num_attention_heads": self.tp_world_size} + + @is_tensor_parallel @require_torch_neuron class TestFlux2TransformerTensorParallelNeuron: @@ -202,6 +215,116 @@ class TestFlux2TransformerLoRA(Flux2TransformerTesterConfig, LoraTesterMixin): """LoRA adapter tests for Flux2 Transformer.""" +class TestFlux2TransformerLoKr(Flux2TransformerTesterConfig, LoKrTesterMixin): + """LoKr adapter tests for Flux2 Transformer, including the Flux2 LoKr checkpoint formats.""" + + # ai-toolkit stores a placeholder alpha for full-matrix factors, where LoKr applies no scaling. + placeholder_alpha = torch.tensor(9999220736.0) + + def get_bfl_qkv_state_dict(self, model): + """A BFL-format LoKr state dict on the fused QKV projections of the first double block.""" + to_q = model.transformer_blocks[0].attn.to_q + state_dict, expected_deltas = {}, {} + for bfl_path, diffusers_path in [ + ("double_blocks.0.img_attn.qkv", "transformer_blocks.0.attn.to_qkv"), + ("double_blocks.0.txt_attn.qkv", "transformer_blocks.0.attn.to_added_qkv"), + ]: + factors, expected_deltas[diffusers_path] = make_lokr_factors(3 * to_q.out_features, to_q.in_features) + state_dict.update({f"diffusion_model.{bfl_path}.{k}": v for k, v in factors.items()}) + state_dict[f"diffusion_model.{bfl_path}.alpha"] = self.placeholder_alpha + return state_dict, expected_deltas + + @torch.no_grad() + def test_lokr_bfl_checkpoint(self): + # BFL checkpoints (e.g. ai-toolkit) apply LoKr to the fused QKV projections. A Kronecker product delta cannot + # be split exactly into Q/K/V, so loading fuses the model's projections and maps the adapter 1:1. + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + state_dict, expected_deltas = self.get_bfl_qkv_state_dict(model) + for bfl_path, diffusers_path in [ + ("single_blocks.0.linear1", "single_transformer_blocks.0.attn.to_qkv_mlp_proj"), + ("double_blocks.0.img_attn.proj", "transformer_blocks.0.attn.to_out.0"), + ("double_blocks.0.img_mlp.0", "transformer_blocks.0.ff.linear_in"), + ]: + linear = model.get_submodule(diffusers_path) + factors, expected_deltas[diffusers_path] = make_lokr_factors(linear.out_features, linear.in_features) + state_dict.update({f"diffusion_model.{bfl_path}.{k}": v for k, v in factors.items()}) + state_dict[f"diffusion_model.{bfl_path}.alpha"] = self.placeholder_alpha + + converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict) + model.load_lora_adapter(converted, prefix="transformer", adapter_name="default") + + assert model.transformer_blocks[0].attn.fused_projections + check_lokr_deltas(model, expected_deltas) + + def test_lokr_fused_qkv_checkpoint_refuses_when_unfused_projections_are_adapted(self): + # Fusing would replace to_q and orphan the adapter already injected there. + from peft import LoraConfig + + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + model.add_adapter(LoraConfig(r=2, target_modules=["to_q"]), adapter_name="lora") + state_dict, _ = self.get_bfl_qkv_state_dict(model) + converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict) + + with pytest.raises(ValueError, match="already loaded on the unfused projections"): + model.load_lora_adapter(converted, prefix="transformer", adapter_name="lokr") + assert not model.transformer_blocks[0].attn.fused_projections + + @torch.no_grad() + def test_lokr_lycoris_checkpoint(self): + # LyCORIS wraps the diffusers model and encodes module paths with underscores under a `lycoris_` prefix. + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + state_dict, expected_deltas = {}, {} + for diffusers_path in [ + "single_transformer_blocks.0.attn.to_qkv_mlp_proj", + "transformer_blocks.0.attn.to_q", + "transformer_blocks.0.attn.to_out.0", + "transformer_blocks.0.ff.linear_in", + ]: + linear = model.get_submodule(diffusers_path) + factors, expected_deltas[diffusers_path] = make_lokr_factors(linear.out_features, linear.in_features) + lycoris_path = "lycoris_" + diffusers_path.replace(".", "_") + state_dict.update({f"{lycoris_path}.{k}": v for k, v in factors.items()}) + state_dict[f"{lycoris_path}.alpha"] = torch.tensor(16.0) + + converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict) + model.load_lora_adapter(converted, prefix="transformer", adapter_name="default") + + check_lokr_deltas(model, expected_deltas) + + def test_lokr_lycoris_checkpoint_with_unknown_keys_raises(self): + state_dict = { + "lycoris_transformer_blocks_0_attn_to_q.lokr_w1": torch.randn(4, 4), + "lycoris_transformer_blocks_0_attn_norm_q.lokr_w1": torch.randn(4, 4), + } + with pytest.raises(ValueError, match="lycoris_transformer_blocks_0_attn_norm_q.lokr_w1"): + Flux2LoraLoaderMixin.lora_state_dict(state_dict) + + @torch.no_grad() + def test_lokr_diffusers_names_checkpoint(self): + # Checkpoints that store the diffusers module paths directly, with alpha keys and no prefix (e.g. SimpleTuner, + # `bghira/flux2-klein-9b-distillation-lokr`). Alpha scales the rank-decomposed factors only. + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + rank, alpha = 2, 1.0 + state_dict, expected_deltas = {}, {} + for diffusers_path, factor_rank in [ + ("single_transformer_blocks.0.attn.to_out", None), + ("transformer_blocks.0.attn.to_k", rank), + ]: + linear = model.get_submodule(diffusers_path) + factors, delta = make_lokr_factors(linear.out_features, linear.in_features, rank=factor_rank) + state_dict.update({f"{diffusers_path}.{k}": v for k, v in factors.items()}) + state_dict[f"{diffusers_path}.alpha"] = torch.tensor(alpha) + expected_deltas[diffusers_path] = delta if factor_rank is None else (alpha / rank) * delta + + converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict) + model.load_lora_adapter(converted, prefix="transformer", adapter_name="default") + + check_lokr_deltas(model, expected_deltas) + + class TestFlux2TransformerLoRAHotSwap(Flux2TransformerTesterConfig, LoraHotSwappingForModelTesterMixin): """LoRA hot-swapping tests for Flux2 Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 5fcf37f6ff..6210aecff8 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -37,6 +37,7 @@ from ..testing_utils import ( MemoryTesterMixin, ModelTesterMixin, TensorParallelTesterMixin, + TensorParallelTPUTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, @@ -307,6 +308,18 @@ class TestQwenImageTransformerTensorParallel(QwenImageTransformerTesterConfig, T """Tensor Parallel inference tests for QwenImage Transformer (CUDA/XPU multi-accelerator).""" +class TestQwenImageTransformerTensorParallelTPU(QwenImageTransformerTesterConfig, TensorParallelTPUTesterMixin): + """Tensor Parallel inference test for QwenImage Transformer on TPU.""" + + def get_init_dict(self): + # One head per chip. + return {**super().get_init_dict(), "num_attention_heads": self.tp_world_size} + + def test_tensor_parallel_tpu_inference(self): + # TPU numerics differ by ~1e-2 between sharded and unsharded QwenImage (~1e-7 on CPU). + super().test_tensor_parallel_tpu_inference(atol=2e-2, rtol=2e-2) + + def make_neuron_tp_spec(): """Model spec consumed by the generic Neuron TP worker (``_neuron_tp_worker.py``). 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 0000000000..149a174787 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_sana_wm.py @@ -0,0 +1,149 @@ +# 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 torch + +from diffusers import SanaWMTransformer3DModel +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, + "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): + def test_dtype_cast_keeps_complex_rotary_frequencies(self): + # The RoPE table is complex. A registered buffer would be cast to the *real* dtype by `model.to(dtype)`, + # silently dropping the sine half of every rotation, so a same-dtype round trip must be a no-op. + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + with torch.no_grad(): + reference = model(**inputs, return_dict=False)[0] + model.to(torch.float32) + assert model.rope.freqs.is_complex() + output = model(**inputs, return_dict=False)[0] + assert torch.equal(reference, output) + + +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): + def test_gradient_checkpointing_is_applied(self): + expected_set = {"SanaWMTransformer3DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) diff --git a/tests/models/transformers/test_models_transformer_sana_wm_refiner.py b/tests/models/transformers/test_models_transformer_sana_wm_refiner.py new file mode 100644 index 0000000000..c28fef4a1a --- /dev/null +++ b/tests/models/transformers/test_models_transformer_sana_wm_refiner.py @@ -0,0 +1,236 @@ +# 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 torch + +from diffusers import SanaWMLTX2RefinerTransformer3DModel +from diffusers.models.transformers.transformer_sana_wm_refiner import ( + SanaWMLTX2AudioVideoRotaryPosEmbed, + SanaWMRefinerKVCache, +) +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 SanaWMLTX2RefinerTransformer3DTesterConfig(BaseModelTesterConfig): + # Tiny stand-in for the LTX-2 based SANA-WM stage-2 refiner. The refiner forward is video-only, but the audio + # submodules are still built (so LTX-2 checkpoints round-trip), hence the small audio arguments below. + batch_size = 2 + num_frames = 2 + height = 4 + width = 4 + in_channels = 4 + num_attention_heads = 2 + attention_head_dim = 8 + caption_channels = 16 + sequence_length = 8 + fps = 24.0 + + @property + def model_class(self): + return SanaWMLTX2RefinerTransformer3DModel + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def input_shape(self) -> tuple: + return (self.num_frames * self.height * self.width, self.in_channels) + + @property + def output_shape(self) -> tuple: + return (self.num_frames * self.height * self.width, self.in_channels) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "in_channels": self.in_channels, + "out_channels": self.in_channels, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": self.num_attention_heads, + "attention_head_dim": self.attention_head_dim, + "cross_attention_dim": 16, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_cross_attention_dim": 8, + "num_layers": 2, + "qk_norm": "rms_norm_across_heads", + "caption_channels": self.caption_channels, + "rope_double_precision": False, + } + + def get_video_rotary_emb(self, batch_size: int, num_frames: int) -> tuple[torch.Tensor, torch.Tensor]: + """Build the caller-supplied video RoPE for frames `[0, num_frames)` with the model's default RoPE config.""" + rope = SanaWMLTX2AudioVideoRotaryPosEmbed( + dim=self.num_attention_heads * self.attention_head_dim, + patch_size=1, + patch_size_t=1, + base_num_frames=20, + base_height=2048, + base_width=2048, + scale_factors=(8, 32, 32), + modality="video", + double_precision=False, + num_attention_heads=self.num_attention_heads, + ) + coords = rope.prepare_video_coords( + batch_size, num_frames, self.height, self.width, device=torch_device, fps=self.fps + ) + return rope(coords, device=torch_device) + + def get_dummy_inputs(self, batch_size: int | None = None) -> dict: + batch_size = batch_size or self.batch_size + num_tokens = self.num_frames * self.height * self.width + + return { + "hidden_states": randn_tensor( + (batch_size, num_tokens, self.in_channels), generator=self.generator, device=torch_device + ), + "encoder_hidden_states": randn_tensor( + (batch_size, self.sequence_length, self.caption_channels), + generator=self.generator, + device=torch_device, + ), + # The refiner takes a per-token timestep that is already scaled by `timestep_scale_multiplier`. + "timestep": torch.rand((batch_size, num_tokens), generator=self.generator).to(torch_device) * 1000, + "video_rotary_emb": self.get_video_rotary_emb(batch_size, self.num_frames), + "encoder_attention_mask": torch.ones((batch_size, self.sequence_length), device=torch_device), + } + + +class TestSanaWMLTX2RefinerTransformer3D(SanaWMLTX2RefinerTransformer3DTesterConfig, ModelTesterMixin): + def test_kv_cache_sink_injection_changes_output(self): + """Drive the KV cache the way `SanaWMLTX2Refiner` does: seed the attention sink from a clean block with + `capture_pre_rope`, then denoise the next block with `inject` and check the sink is actually attended to.""" + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + num_layers = len(model.transformer_blocks) + + batch_size = self.batch_size + sink_size = 1 + tokens_per_frame = self.height * self.width + encoder_hidden_states = randn_tensor( + (batch_size, self.sequence_length, self.caption_channels), generator=self.generator, device=torch_device + ) + encoder_attention_mask = torch.ones((batch_size, self.sequence_length), device=torch_device) + + # Sink block: frames `[0, sink_size)` at sigma=0. + sink_tokens = randn_tensor( + (batch_size, sink_size * tokens_per_frame, self.in_channels), generator=self.generator, device=torch_device + ) + sink_inputs = { + "hidden_states": sink_tokens, + "encoder_hidden_states": encoder_hidden_states, + "timestep": torch.zeros((batch_size, sink_tokens.shape[1]), device=torch_device), + "video_rotary_emb": model.build_rotary_emb_for_absolute_positions( + batch_size=batch_size, + frame_positions=list(range(sink_size)), + height=self.height, + width=self.width, + device=torch_device, + fps=self.fps, + ), + "encoder_attention_mask": encoder_attention_mask, + } + + # Active block: frames `[sink_size, sink_size + num_frames)` at some intermediate sigma. + block_start = sink_size + frame_positions = list(range(block_start, block_start + self.num_frames)) + active_tokens = randn_tensor( + (batch_size, self.num_frames * tokens_per_frame, self.in_channels), + generator=self.generator, + device=torch_device, + ) + active_inputs = { + "hidden_states": active_tokens, + "encoder_hidden_states": encoder_hidden_states, + "timestep": torch.full((batch_size, active_tokens.shape[1]), 500.0, device=torch_device), + "video_rotary_emb": model.build_rotary_emb_for_absolute_positions( + batch_size=batch_size, + frame_positions=frame_positions, + height=self.height, + width=self.width, + device=torch_device, + fps=self.fps, + ), + "encoder_attention_mask": encoder_attention_mask, + } + + kv_cache = SanaWMRefinerKVCache(num_layers) + with torch.no_grad(): + # 1. Capture the pre-RoPE sink K/V; capturing injects no prefix, so it must match a plain forward. + sink_plain = model(**sink_inputs, return_dict=False)[0] + sink_captured = model( + **sink_inputs, kv_cache=kv_cache, kv_cache_mode="capture_pre_rope", return_dict=False + )[0] + assert torch.allclose(sink_plain, sink_captured, atol=1e-6) + + inner_dim = self.num_attention_heads * self.attention_head_dim + for layer_idx in range(num_layers): + layer_cache = kv_cache.get(layer_idx) + sink_k_pre, sink_v = layer_cache.get_captured_pre_rope() + assert sink_k_pre.shape == (batch_size, sink_size * tokens_per_frame, inner_dim) + assert sink_v.shape == (batch_size, sink_size * tokens_per_frame, inner_dim) + layer_cache.store_sink(sink_k_pre, sink_v) + + # 2. Slide the sink RoPE so it sits immediately before the active block (no history yet). + history_frames = 0 + sink_rope_offset = block_start - history_frames - sink_size + kv_cache.sink_pe = model.build_rotary_emb_for_absolute_positions( + batch_size=batch_size, + frame_positions=list(range(sink_rope_offset, sink_rope_offset + sink_size)), + height=self.height, + width=self.width, + device=torch_device, + fps=self.fps, + ) + + # 3. Denoise the active block with and without the injected sink. + no_cache = model(**active_inputs, return_dict=False)[0] + injected = model(**active_inputs, kv_cache=kv_cache, kv_cache_mode="inject", return_dict=False)[0] + + assert injected.shape == no_cache.shape == (batch_size, active_tokens.shape[1], self.in_channels) + assert torch.isfinite(injected).all() + assert not torch.allclose(injected, no_cache, atol=1e-4), "Injecting the sink K/V did not change the output." + + +class TestSanaWMLTX2RefinerTransformer3DMemory(SanaWMLTX2RefinerTransformer3DTesterConfig, MemoryTesterMixin): + pass + + +class TestSanaWMLTX2RefinerTransformer3DCompile(SanaWMLTX2RefinerTransformer3DTesterConfig, TorchCompileTesterMixin): + pass + + +class TestSanaWMLTX2RefinerTransformer3DTraining(SanaWMLTX2RefinerTransformer3DTesterConfig, TrainingTesterMixin): + pass diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 35bafa5702..1a6881c131 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -18,6 +18,7 @@ import pytest import torch from diffusers import ZImageTransformer2DModel +from diffusers.loaders.lora_pipeline import ZImageLoraLoaderMixin from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import assert_tensors_close, torch_device @@ -25,6 +26,7 @@ from ..testing_utils import ( AutoRoundCompileTesterMixin, AutoRoundTesterMixin, BaseModelTesterConfig, + LoKrTesterMixin, LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, @@ -32,6 +34,7 @@ from ..testing_utils import ( TorchCompileTesterMixin, TrainingTesterMixin, ) +from ..testing_utils.lokr import check_lokr_deltas, make_lokr_factors # Z-Image requires torch.use_deterministic_algorithms(False) due to complex64 RoPE operations @@ -183,6 +186,36 @@ class TestZImageTransformerLoRA(ZImageTransformerTesterConfig, LoraTesterMixin): """LoRA adapter tests for Z-Image Transformer.""" +class TestZImageTransformerLoKr(ZImageTransformerTesterConfig, LoKrTesterMixin): + """LoKr adapter tests for Z-Image Transformer, including the ai-toolkit Z-Image LoKr checkpoint format.""" + + @torch.no_grad() + def test_lokr_ai_toolkit_checkpoint(self): + # ai-toolkit stores the diffusers module paths under a `diffusion_model.` prefix. Full-matrix factors come with + # a placeholder alpha, where LoKr applies no scaling; rank-decomposed factors are scaled by `alpha / rank`. + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).eval().to(torch_device) + rank = 1 + state_dict, expected_deltas = {}, {} + for module, factor_rank, alpha in [ + ("layers.0.attention.to_q", None, 9999220736.0), + ("layers.0.feed_forward.w1", None, 9999220736.0), + ("layers.0.adaLN_modulation.0", None, 9999220736.0), + ("layers.0.attention.to_v", rank, 0.5), + ]: + linear = model.get_submodule(module) + factors, delta = make_lokr_factors(linear.out_features, linear.in_features, rank=factor_rank) + state_dict.update({f"diffusion_model.{module}.{k}": v for k, v in factors.items()}) + state_dict[f"diffusion_model.{module}.alpha"] = torch.tensor(alpha) + expected_deltas[module] = delta if factor_rank is None else (alpha / rank) * delta + + converted = ZImageLoraLoaderMixin.lora_state_dict(state_dict) + assert all(k.startswith("transformer.") and ".lokr_" in k for k in converted) + model.load_lora_adapter(converted, prefix="transformer", adapter_name="default") + + check_lokr_deltas(model, expected_deltas) + + # TODO: Add pretrained_model_name_or_path once a tiny Z-Image model is available on the Hub # class TestZImageTransformerBitsAndBytes(ZImageTransformerTesterConfig, BitsAndBytesTesterMixin): # """BitsAndBytes quantization tests for Z-Image Transformer.""" diff --git a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index c3ae0d9f2e..ff8a52839d 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py @@ -17,6 +17,7 @@ import unittest import numpy as np import PIL.Image +import pytest import torch from diffusers import ( @@ -25,6 +26,7 @@ from diffusers import ( AnimaTextConditioner, CosmosTransformer3DModel, ) +from diffusers.modular_pipelines.anima.modular_blocks_anima import AnimaCoreDenoiseStep from ...testing_utils import enable_full_determinism, is_lora, require_peft_backend from ..testing_utils import ( @@ -127,6 +129,55 @@ class AnimaModularPipelineTesterConfig(BaseModularPipelineTesterConfig): class TestAnimaModularPipelineFast(AnimaModularPipelineTesterConfig, ModularPipelineTesterMixin): + @pytest.mark.parametrize("negative_prompt", [None, "", "bad quality"]) + def test_split_text_encoder(self, negative_prompt): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["negative_prompt"] = negative_prompt + expected = pipe(**inputs).images + + blocks = self.pipeline_blocks_class() + text_block = blocks.sub_blocks.pop("text_encoder") + text_pipe = text_block.init_pipeline(self.pretrained_model_name_or_path) + text_pipe.load_components(dtype=torch.float32) + denoise_pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) + denoise_pipe.load_components(dtype=torch.float32) + + inputs = self.get_dummy_inputs() + inputs["negative_prompt"] = negative_prompt + text_inputs = {name: inputs.pop(name) for name in ("prompt", "negative_prompt", "max_sequence_length")} + embeddings = text_pipe(**text_inputs).get_by_kwargs("denoiser_input_fields") + output = denoise_pipe(**embeddings, **inputs).images + + torch.testing.assert_close(output, expected, rtol=1e-5, atol=1e-5) + + def test_standalone_core_denoise(self): + text_block = self.pipeline_blocks_class().sub_blocks["text_encoder"] + text_pipe = text_block.init_pipeline(self.pretrained_model_name_or_path) + text_pipe.load_components(dtype=torch.float32) + inputs = self.get_dummy_inputs() + text_inputs = {name: inputs.pop(name) for name in ("prompt", "negative_prompt", "max_sequence_length")} + embeddings = text_pipe(**text_inputs).get( + [ + "qwen_prompt_embeds", + "qwen_attention_mask", + "t5_input_ids", + "t5_attention_mask", + "negative_qwen_prompt_embeds", + "negative_qwen_attention_mask", + "negative_t5_input_ids", + "negative_t5_attention_mask", + ] + ) + + denoise_pipe = AnimaCoreDenoiseStep().init_pipeline(self.pretrained_model_name_or_path) + denoise_pipe.load_components(dtype=torch.float32) + inputs.pop("output_type") + output = denoise_pipe(**embeddings, **inputs).get("latents") + + assert output.shape == (1, 4, 1, 4, 4) + assert torch.isfinite(output).all() + def test_inference_empty_negative_prompt(self): pipe = self.get_pipeline() diff --git a/tests/others/test_cli_commands.py b/tests/others/test_cli_commands.py index 9b731416df..ba30cfb4a9 100644 --- a/tests/others/test_cli_commands.py +++ b/tests/others/test_cli_commands.py @@ -432,9 +432,16 @@ class TestCustomBlocksCommand: "class OtherBase:\n pass\n" "class NotABlock(OtherBase):\n pass\n" "class MyBlock(ModularPipelineBlocks):\n pass\n" + # composite blocks are packageable too, and the base may be referenced through its module + "class MyBlocks(SequentialPipelineBlocks):\n pass\n" + "class MyAutoBlocks(diffusers.modular_pipelines.AutoPipelineBlocks):\n pass\n" ) cmd = CustomBlocksCommand() - assert cmd._get_class_names(block_py) == [("MyBlock", "ModularPipelineBlocks")] + assert cmd._get_class_names(block_py) == [ + ("MyBlock", "ModularPipelineBlocks"), + ("MyBlocks", "SequentialPipelineBlocks"), + ("MyAutoBlocks", "AutoPipelineBlocks"), + ] broken = tmp_path / "broken.py" broken.write_text("class Broken(:\n pass\n") @@ -457,6 +464,25 @@ class TestCustomBlocksCommand: assert (tmp_path / "modular_config.json").exists() assert (tmp_path / "modular_model_index.json").exists() + def test_packaging_composite_block(self, tmp_path, monkeypatch): + # a `SequentialPipelineBlocks` subclass is the usual thing to publish; it must be selectable by name + block_py = tmp_path / "block.py" + block_py.write_text( + "from diffusers.modular_pipelines import ModularPipelineBlocks, SequentialPipelineBlocks\n" + "\n" + "class MyStep(ModularPipelineBlocks):\n" + " model_name = 'test'\n" + "\n" + "class MyBlocks(SequentialPipelineBlocks):\n" + " model_name = 'test'\n" + " block_classes = [MyStep]\n" + " block_names = ['step']\n" + ) + monkeypatch.chdir(tmp_path) + CustomBlocksCommand(str(block_py), "MyBlocks").run() + config = json.loads((tmp_path / "modular_config.json").read_text()) + assert config["auto_map"] == {"ModularPipelineBlocks": "block.MyBlocks"} + class TestCli: def test_toplevel_help_lists_all_commands(self): diff --git a/tests/others/test_image_processor.py b/tests/others/test_image_processor.py index 0d358699f1..05cbb8a27a 100644 --- a/tests/others/test_image_processor.py +++ b/tests/others/test_image_processor.py @@ -15,9 +15,10 @@ import numpy as np import PIL.Image +import pytest import torch -from diffusers.image_processor import VaeImageProcessor +from diffusers.image_processor import InpaintProcessor, VaeImageProcessor class TestImageProcessor: @@ -306,3 +307,38 @@ class TestImageProcessor: assert out_np.shape == exp_np_shape, ( f"resized image output shape '{out_np.shape}' didn't match expected shape '{exp_np_shape}'." ) + + @pytest.mark.parametrize( + "has_mask, padding_mask_crop", + [ + (True, None), + (False, None), + (True, 8), + ], + ) + def test_inpaint_processor_preprocess(self, has_mask, padding_mask_crop): + processor = InpaintProcessor() + image = PIL.Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8)) + mask = PIL.Image.fromarray(np.ones((64, 64), dtype=np.uint8) * 255) if has_mask else None + + kwargs = {"height": 64, "width": 64} + if padding_mask_crop is not None: + kwargs["padding_mask_crop"] = padding_mask_crop + + out_img, out_mask, postprocessing_kwargs = processor.preprocess(image, mask=mask, **kwargs) + + assert isinstance(out_img, torch.Tensor) + if has_mask: + assert isinstance(out_mask, torch.Tensor) + else: + assert out_mask is None + assert isinstance(postprocessing_kwargs, dict) + + if padding_mask_crop is not None: + assert postprocessing_kwargs["crops_coords"] is not None + assert postprocessing_kwargs["original_image"] == image + assert postprocessing_kwargs["original_mask"] == mask + else: + assert postprocessing_kwargs["crops_coords"] is None + assert postprocessing_kwargs["original_image"] is None + assert postprocessing_kwargs["original_mask"] is None diff --git a/tests/others/test_utils.py b/tests/others/test_utils.py index bb7e5c2989..0341623d87 100755 --- a/tests/others/test_utils.py +++ b/tests/others/test_utils.py @@ -271,6 +271,18 @@ class TestRandnTensor: assert "moved to" in cuda_out, f"Non-MPS target should still emit the CPU-fallback info log, got: {cuda_out}" +class TestApplyForwardHook: + def test_preserves_wrapped_function_metadata(self): + from diffusers.utils.accelerate_utils import apply_forward_hook + + @apply_forward_hook + def example(self): + """Example method docstring.""" + + assert example.__name__ == "example" + assert example.__doc__ == "Example method docstring." + + class TestTorchDeviceBackend: """Tests for :class:`diffusers.utils.torch_utils.TorchDeviceBackend` on the CPU backend.""" diff --git a/tests/pipelines/sana_wm/__init__.py b/tests/pipelines/sana_wm/__init__.py new file mode 100644 index 0000000000..e69de29bb2 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 0000000000..478e3b78a1 --- /dev/null +++ b/tests/pipelines/sana_wm/test_sana_wm.py @@ -0,0 +1,250 @@ +# Copyright 2026 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 numpy as np +import PIL.Image +import pytest +import torch +from transformers import Gemma2Config, Gemma2Model, GemmaTokenizer + +from diffusers import ( + AutoencoderKLLTX2Video, + FlowMatchEulerDiscreteScheduler, + SanaWMPipeline, + SanaWMPipelineOutput, + SanaWMTransformer3DModel, +) +from diffusers.pipelines.sana_wm.cam_utils import action_string_to_c2w + +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +enable_full_determinism() + + +SINGLE_VIDEO_SKIP_REASON = ( + "SanaWMPipeline generates exactly one video per call (one first-frame image + one camera trajectory); " + "`prompt`/`image`/`c2w` take no batch dimension and there is no `num_videos_per_prompt`." +) + + +class SanaWMPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = SanaWMPipeline + required_input_params_in_call_signature = frozenset( + [ + "image", + "prompt", + "negative_prompt", + "c2w", + "action", + "intrinsics", + "height", + "width", + "num_frames", + "guidance_scale", + "prompt_embeds", + "prompt_attention_mask", + "negative_prompt_embeds", + "negative_prompt_attention_mask", + ] + ) + # One first-frame image + one camera trajectory per call: nothing is batched. + batch_input_params = frozenset() + # The pipeline has no `num_images_per_prompt` (single video per call) and samples its own noise around the + # VAE-encoded first frame, so it takes no user `latents` either. + optional_input_params = frozenset(["num_inference_steps", "generator", "output_type", "return_dict"]) + # `frames` is the un-batched `(num_frames, channels, height, width)` video for `output_type="pt"`, so a single + # element of `pipe(...)[0]` is one `(channels, height, width)` frame. + output_shape = (3, 32, 32) + + # Dummy video geometry. The tiny LTX-2 VAE below compresses 2x spatially and 2x temporally. + num_frames = 5 + height = 32 + width = 32 + + def get_dummy_components(self): + tokenizer = GemmaTokenizer.from_pretrained("hf-internal-testing/dummy-gemma") + + torch.manual_seed(0) + text_encoder_config = Gemma2Config( + head_dim=16, + hidden_size=8, + initializer_range=0.02, + intermediate_size=16, + max_position_embeddings=512, + num_attention_heads=2, + num_hidden_layers=1, + num_key_value_heads=1, + # Match the tokenizer so real prompts (and the chi-prompt prefix) can be embedded. + vocab_size=len(tokenizer), + ) + text_encoder = Gemma2Model(text_encoder_config) + + torch.manual_seed(0) + vae = AutoencoderKLLTX2Video( + in_channels=3, + out_channels=3, + latent_channels=4, + block_out_channels=(8,), + decoder_block_out_channels=(8,), + layers_per_block=(1,), + decoder_layers_per_block=(1, 1), + spatio_temporal_scaling=(True,), + decoder_spatio_temporal_scaling=(True,), + decoder_inject_noise=(False, False), + downsample_type=("spatial",), + upsample_residual=(False,), + upsample_factor=(1,), + timestep_conditioning=False, + patch_size=1, + patch_size_t=1, + encoder_causal=True, + decoder_causal=False, + ) + vae.use_framewise_encoding = False + vae.use_framewise_decoding = False + + # `num_layers=2` with `softmax_every_n=2` covers both block variants (GDN and softmax camera attention). + torch.manual_seed(0) + transformer = SanaWMTransformer3DModel( + in_channels=4, + num_layers=2, + hidden_size=32, + num_attention_heads=2, + softmax_every_n=2, + linear_head_dim=16, + t_kernel_size=3, + conv_kernel_size=4, + caption_channels=text_encoder_config.hidden_size, + mlp_ratio=2.0, + # Plücker rays are packed per VAE temporal chunk: 6 dims * temporal compression ratio. + chunk_plucker_channels=6 * vae.temporal_compression_ratio, + chunk_plucker_post_attn_blocks=2, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + + def get_dummy_image(self): + # A non-square source image so the resize + center-crop (and the matching intrinsics rescale) is exercised. + pixels = np.random.RandomState(0).randint(0, 256, size=(40, 48, 3), dtype=np.uint8) + return PIL.Image.fromarray(pixels) + + def get_dummy_inputs(self): + return { + "image": self.get_dummy_image(), + "prompt": "A car driving across a desert", + # `w-4` rolls out 4 forward steps plus the identity anchor = `num_frames` poses. + "action": "w-4", + # [fx, fy, cx, cy] in the pixel coordinates of the 48x40 source image. + "intrinsics": [40.0, 40.0, 24.0, 20.0], + "height": self.height, + "width": self.width, + "num_frames": self.num_frames, + "generator": self.get_generator(0), + "num_inference_steps": 2, + "guidance_scale": 5.0, + "max_sequence_length": 16, + # A short chi-prompt keeps the tokenized prefix small (the default is the long release prefix). + "chi_prompt": ["Describe the scene: "], + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` videos are `(num_frames, channels, height, width)`. + "output_type": "pt", + } + + +class TestSanaWMPipeline(SanaWMPipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + pipe = self.get_pipeline().to(torch_device) + + output = pipe(**self.get_dummy_inputs()) + + assert isinstance(output, SanaWMPipelineOutput) + assert output.frames.shape == (self.num_frames, *self.output_shape) + assert not torch.isnan(output.frames).any() + assert output.c2w.shape == (self.num_frames, 4, 4) + # (batch, latent channels, (num_frames - 1) // 2 + 1, height // 2, width // 2) + assert output.latent.shape == (1, 4, 3, self.height // 2, self.width // 2) + + def test_latent_output_matches_returned_latent(self): + pipe = self.get_pipeline().to(torch_device) + + latent_output = pipe(**{**self.get_dummy_inputs(), "output_type": "latent"}) + video_output = pipe(**self.get_dummy_inputs()) + + assert_tensors_close(latent_output.frames, video_output.latent, atol=1e-5, msg="Latent outputs differ.") + + def test_c2w_matches_equivalent_action(self): + pipe = self.get_pipeline().to(torch_device) + + inputs = self.get_dummy_inputs() + output_action = pipe(**inputs)[0] + + inputs = self.get_dummy_inputs() + inputs["c2w"] = action_string_to_c2w(inputs.pop("action")) + output_c2w = pipe(**inputs)[0] + + assert_tensors_close(output_c2w, output_action, atol=1e-5, msg="`c2w` and the equivalent `action` differ.") + + def test_intrinsics_matrix_matches_vector(self): + pipe = self.get_pipeline().to(torch_device) + + inputs = self.get_dummy_inputs() + output_vec = pipe(**inputs)[0] + + inputs = self.get_dummy_inputs() + fx, fy, cx, cy = inputs["intrinsics"] + inputs["intrinsics"] = np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], dtype=np.float32) + output_mat = pipe(**inputs)[0] + + assert_tensors_close(output_mat, output_vec, atol=1e-5, msg="3x3 `K` and `[fx, fy, cx, cy]` differ.") + + @pytest.mark.parametrize( + "overrides", + [ + {"action": None}, + {"c2w": np.broadcast_to(np.eye(4, dtype=np.float32), (5, 4, 4)).copy()}, + {"intrinsics": None}, + {"action": None, "c2w": np.eye(4, dtype=np.float32)}, + ], + ids=["no-trajectory", "both-c2w-and-action", "no-intrinsics", "bad-c2w-shape"], + ) + def test_check_inputs_rejects_bad_camera_inputs(self, overrides): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs.update(overrides) + + with pytest.raises(ValueError): + pipe(**inputs) + + @pytest.mark.skip(reason=SINGLE_VIDEO_SKIP_REASON) + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip(reason=SINGLE_VIDEO_SKIP_REASON) + def test_inference_batch_single_identical(self): + pass + + +class TestSanaWMPipelineMemory(SanaWMPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SANA-WM pipeline.""" diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index e0142d6129..c1a3788bb4 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -1189,6 +1189,16 @@ class TestStableDiffusionPipelineDeviceMap: gc.collect() backend_empty_cache(torch_device) + @pytest.fixture + def deterministic_algorithms(self): + enabled = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + yield + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) + def get_inputs(self, generator_device="cpu", seed=0): generator = torch.Generator(device=generator_device).manual_seed(seed) inputs = { @@ -1212,7 +1222,7 @@ class TestStableDiffusionPipelineDeviceMap: return no_device_map_image - def test_forward_pass_balanced_device_map(self): + def test_forward_pass_balanced_device_map(self, deterministic_algorithms): no_device_map_image = self.get_pipeline_output_without_device_map() sd_pipe_with_device_map = StableDiffusionPipeline.from_pretrained( @@ -1232,7 +1242,7 @@ class TestStableDiffusionPipelineDeviceMap: assert len(set(sd_pipe_with_device_map.hf_device_map.values())) >= 2 - def test_max_memory(self): + def test_max_memory(self, deterministic_algorithms): no_device_map_image = self.get_pipeline_output_without_device_map() sd_pipe_with_device_map = StableDiffusionPipeline.from_pretrained( diff --git a/tests/pipelines/testing_utils/lora.py b/tests/pipelines/testing_utils/lora.py index 87d1da8fac..661d2908b9 100644 --- a/tests/pipelines/testing_utils/lora.py +++ b/tests/pipelines/testing_utils/lora.py @@ -700,6 +700,18 @@ class LoraTesterMixin(BaseLoraTesterMixin, BasePipelineOutputMixin): msg="output with no lora and output with lora disabled should give same results", ) + def test_get_active_adapters_text_encoder_only(self): + """ + Tests that get_active_adapters correctly reports active adapters when only the text encoder has a LoRA attached. + """ + if not self.text_encoder_components: + pytest.skip("Skipping test as text encoder LoRAs are not currently supported.") + + pipe = self.get_pipeline().to(torch_device) + self.add_adapters_to_pipeline(pipe, components=self.text_encoder_components, adapter_name="default") + + assert pipe.get_active_adapters() == ["default"] + def test_get_adapters(self): pipe = self.get_pipeline().to(torch_device) diff --git a/tests/pipelines/z_image/test_z_image.py b/tests/pipelines/z_image/test_z_image.py index fac10607cb..866c4eec11 100644 --- a/tests/pipelines/z_image/test_z_image.py +++ b/tests/pipelines/z_image/test_z_image.py @@ -196,14 +196,6 @@ class TestZImagePipelineLoRA(ZImagePipelineTesterConfig, LoraTesterMixin): base_pipe_output, lora_scale, expected_atol=5e-2, expected_rtol=5e-2 ) - @pytest.mark.skip("Needs to be debugged.") - def test_set_adapters_match_attention_kwargs(self, tmp_path, base_pipe_output): - pass - - @pytest.mark.skip("Needs to be debugged.") - def test_simple_inference_with_text_denoiser_lora_and_scale(self, base_pipe_output): - pass - class TestZImagePipelineLoRAMemory(ZImagePipelineTesterConfig, LoraMemoryTesterMixin): """LoRA x memory-optimization tests (group offload, CPU offload) for the Z-Image pipeline.""" diff --git a/tests/testing_utils.py b/tests/testing_utils.py index cce6f15325..a22da6de08 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -46,6 +46,7 @@ from diffusers.utils.import_utils import ( is_timm_available, is_torch_available, is_torch_neuronx_available, + is_torch_tpu_available, is_torch_version, is_torchao_available, is_torchsde_available, @@ -566,6 +567,14 @@ def require_torch_neuron(test_case): )(test_case) +def require_torch_tpu(test_case): + """Decorator marking a test that requires a TPU device (torch_tpu).""" + return pytest.mark.skipif( + not is_torch_tpu_available(), + reason="test requires TPU device (torch_tpu)", + )(test_case) + + def require_torch_multi_gpu(test_case): """ Decorator marking a test that requires a multi-GPU setup (in PyTorch). These tests are skipped on a machine without