From e0415c41f8434f73177127f7bd7a3c96eb9d609f Mon Sep 17 00:00:00 2001 From: Subodh Date: Thu, 8 Oct 2026 11:57:28 +0530 Subject: [PATCH 01/13] Fix active adapter detection for text encoders (#14961) * Fix active adapter detection for text encoders Signed-off-by: Subodh Bansode * test: move text encoder adapter regression test Signed-off-by: Subodh Bansode --------- Signed-off-by: Subodh Bansode Co-authored-by: Sayak Paul --- src/diffusers/loaders/lora_base.py | 2 +- tests/pipelines/testing_utils/lora.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) 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/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) From d961a388fd02e4db38d17350c8dd9b8abe642e05 Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Thu, 8 Oct 2026 10:24:47 +0200 Subject: [PATCH 02/13] fix(image_processor): maintain 3-tuple return contract in InpaintProcessor.preprocess when mask is None (#14481) fix(image_processor): preserve 3-tuple return contract when mask is None (#14470) Ensure InpaintProcessor.preprocess returns (image, None, postprocessing_kwargs) when mask is None, preserving the expected 3-tuple return contract and signature consistency across all preprocess branches. Consolidate unit tests into a parameterized test covering with-mask, without-mask, and crop configurations. --- src/diffusers/image_processor.py | 11 ++++++-- tests/others/test_image_processor.py | 38 +++++++++++++++++++++++++++- 2 files changed, 46 insertions(+), 3 deletions(-) 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/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 From d7d0b9c2c7b60e679926549d6c82fc79d3cab567 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 9 Oct 2026 10:49:19 +0530 Subject: [PATCH 03/13] [ci] consolidate labeling on PRs. (#14911) consolidate labeling on PRs. Co-authored-by: Dhruv Nair --- .github/workflows/pr_labeler.yml | 49 +++++++++----------------------- 1 file changed, 14 insertions(+), 35 deletions(-) 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 From 78bf2c6266668a1c7185b28053b807f786b6f34e Mon Sep 17 00:00:00 2001 From: Christopher Date: Fri, 9 Oct 2026 09:26:56 +0200 Subject: [PATCH 04/13] [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) (#14163) * [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) Adds loading of LoKr (LyCORIS Kronecker product) adapters: - `load_lora_adapter` detects `lokr_` keys and injects a peft `LoKrConfig`, inferred from the tensor shapes via `_create_lokr_config` (decompose factor, per-module rank/alpha patterns). - State dict conversions for the formats in the wild: ai-toolkit Z-Image (dotted diffusers paths under `diffusion_model.`), ai-toolkit BFL Flux2 (fused qkv), LyCORIS underscore format, and bare dotted diffusers paths. - BFL fused-QKV LoKr cannot be split exactly into separate Q/K/V Kronecker factors, so `Flux2LoraLoaderMixin.load_lora_weights` fuses the model's QKV projections and maps the adapter 1:1 (exact). - Alpha follows the LyCORIS convention: scaling applies only to rank-decomposed factors and is baked into the weights at conversion. Fixes #13221 * [LoRA] LoKr review follow-ups: adapter-agnostic error message, refuse fused-QKV load over an unfused adapter - "Invalid adapter checkpoint. We currently support LoRA and LoKr." replaces the message that still said "LoRA checkpoint" and described the substring check. - Loading a fused-QKV LoKr checkpoint now refuses when an adapter is already injected on to_q/to_k/to_v or add_{q,k,v}_proj: fuse_qkv_projections() would replace those modules and orphan it. Covered by a test. Co-Authored-By: Claude Fable 5.1 * [LoRA] rename the LoKr tests to test_lokr_conversion.py The file exercises checkpoint conversion and loading for LoKr adapters, not LoRA. Co-Authored-By: Claude Opus 5.5 * [LoRA] LoKr review follow-ups: scope to Z-Image/Flux2, model-level fused QKV, LoKrTesterMixin - Only the Z-Image and Flux2 loader mixins accept LoKr checkpoints, the two that convert them; the other mixins are back to their LoRA-only check. - The fused-QKV handling moves to `_maybe_fuse_qkv_projections_for_lokr` in peft_utils, called from `load_lora_adapter`, so the model fuses its projections whichever entry point loads the adapter. The quantized-model error explains why. - `_bake_lokr_alpha` becomes `_bake_lokr_alpha_` as it edits the dict in place. - The BFL Flux2 converter no longer accepts expanded diffusers block names under `diffusion_model.`: no LoKr checkpoint uses that layout; diffusers-named checkpoints have no prefix and take the generic converter. - The LyCORIS converter names the keys it does not recognize. - Tests move from tests/lora to a `LoKrTesterMixin` in tests/models/testing_utils, with the checkpoint formats tested on the Flux2 and Z-Image model test classes. Co-Authored-By: Claude Opus 5.5 * [LoRA] mark ZImageLoraLoaderMixin.load_lora_weights as copied from Flux2 --------- Co-authored-by: Claude Fable 5.1 Co-authored-by: Sayak Paul --- .../loaders/lora_conversion_utils.py | 154 ++++++++++++++++++ src/diffusers/loaders/lora_pipeline.py | 32 +++- src/diffusers/loaders/peft.py | 116 +++++++------ src/diffusers/utils/peft_utils.py | 124 +++++++++++++- tests/models/testing_utils/__init__.py | 2 + tests/models/testing_utils/lokr.py | 139 ++++++++++++++++ .../test_models_transformer_flux2.py | 114 +++++++++++++ .../test_models_transformer_z_image.py | 33 ++++ 8 files changed, 653 insertions(+), 61 deletions(-) create mode 100644 tests/models/testing_utils/lokr.py 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 # 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..8bab73f67c 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -17,6 +17,7 @@ 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 ( @@ -80,6 +81,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/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index 3263ce6820..cc4f5cad96 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,6 +38,7 @@ from ..testing_utils import ( ContextParallelTesterMixin, GGUFCompileTesterMixin, GGUFTesterMixin, + LoKrTesterMixin, LoraHotSwappingForModelTesterMixin, LoraTesterMixin, MemoryTesterMixin, @@ -47,6 +50,7 @@ from ..testing_utils import ( TorchCompileTesterMixin, TrainingTesterMixin, ) +from ..testing_utils.lokr import check_lokr_deltas, make_lokr_factors enable_full_determinism() @@ -202,6 +206,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_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.""" From 1d5d056ecff3a8c1d887a7dd1aa66ca4a0323458 Mon Sep 17 00:00:00 2001 From: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com> Date: Fri, 9 Oct 2026 12:46:23 +0200 Subject: [PATCH 05/13] [TPU] TorchTPU backend integration - eager / torch.compile / tp (#14039) * feat: add torchtpu * feat:draft TorchTPU support * fix: wan overflow issue + compile mode error on sdxl * doc: enhance with TorchTPU doc * doc: enhance with TorchTPU doc * style: remove unused imports flagged by ruff Co-Authored-By: Claude Sonnet 4.6 * docs: remove Debug Eager and Fused Eager sections from tpu.md Co-Authored-By: Claude Sonnet 4.6 * feat: add TP for TPU * style: fix import sorting in TPU test scripts Co-Authored-By: Claude Sonnet 4.6 * style: ruff format TPU test scripts Co-Authored-By: Claude Sonnet 4.6 * fix: style * fix: test for native 4 devices * fix: propagate TPU device fixes to Flux/Flux2/Wan-family copies; register tpu.md; drop redundant execution_device check - Propagate the text_encoder.device-based fix (introduced for TPU CPU-offload support) from FluxPipeline/Flux2KleinPipeline/WanPipeline into their `# Copied from` copies (flux/*, flux2_klein_inpaint, visualcloze, anyflow, chronoedit, lucy_edit, skyreels_v2/*). SDXL-family copies of StableDiffusionXLPipeline.encode_prompt are intentionally left untouched; they'll be handled in a follow-up PR that fixes device placement for every pipeline component (not just text encoders). - Register docs/source/en/optimization/tpu.md in _toctree.yml (was breaking the docs build: "not present in the table of contents"). - Remove the redundant "prefer non-CPU, non-meta component" loop from DiffusionPipeline._execution_device: PR #14383 already fixed this in DiffusionPipeline.device, which _execution_device falls back to. Verified on TPU hardware that _execution_device still resolves correctly for a split-placement pipeline after the removal. Co-Authored-By: Claude Sonnet 5 * tests: cleanup * fix: fix compile mode * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * test: remove flux2 e2e test * doc: apply suggestions * doc: apply suggestions * review: remove monkey patch * review: revert neuron-specific changes in the tests * review: apply suggestions * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: add tpu to tp doc * review: remove unnecessary for tpu * test: assert every _tp_plan parameter is sharded after a TP load Co-Authored-By: Claude Opus 5.5 * removal: delete redundant tp shard helpers * removal: drop TPU workarounds no longer needed after recent main changes * doc: keep the FLUX.2-dev text encoder off a single TPU chip in the TP example * doc: encode the prompt on CPU in the TPU TP example so it runs end to end * test: tighten TPU TP tolerance and simplify the TPU TP worker * doc: use enable_model_cpu_offload() in the TPU eager example * doc: shard both FLUX.2-dev text encoder and transformer in the TPU TP example * doc: run the text encoders on TPU in the compiled example * test: drop the redundant TPU sync in the TP worker * test: trim comments in the TPU TP tests * test: run the TPU TP test with mp.spawn like the CUDA one * test: pin the TPU TP test to 4 chips * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: use DistributedConfig and dtype, and drop the torch_tpu import in the TPU examples * test: skip the TPU TP test for models without a _tp_plan * doc: correct the warning of recompilation * fix: clone the SD3 timestep so torch.compile doesn't recompile every step, thus we can use it for the tpu.md compile mode snippet * doc: update advice for compile mode when one chip doesn't fit * test: pass the TPU TP tolerances as test arguments * review: apply suggestions to the doc * review: view compilation patch for tpu-only --------- Co-authored-by: Claude Sonnet 4.6 Co-authored-by: JingyaHuang Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: Sayak Paul --- docs/source/en/_toctree.yml | 2 + docs/source/en/optimization/tpu.md | 157 ++++++++++++++++++ src/diffusers/hooks/tensor_parallel.py | 2 +- src/diffusers/models/_modeling_parallel.py | 2 +- .../pipeline_stable_diffusion_3.py | 5 + src/diffusers/utils/__init__.py | 1 + src/diffusers/utils/import_utils.py | 5 + tests/models/testing_utils/__init__.py | 2 + tests/models/testing_utils/parallelism.py | 135 ++++++++++++--- .../test_models_transformer_flux.py | 9 + .../test_models_transformer_flux2.py | 9 + .../test_models_transformer_qwenimage.py | 13 ++ tests/testing_utils.py | 9 + 13 files changed, 328 insertions(+), 23 deletions(-) create mode 100644 docs/source/en/optimization/tpu.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 4b4e424719..3abfb9052c 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: 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/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index 0f56c3644e..ce0c76f970 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") class PackedColwiseParallel: diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py index b54e86d6b4..58cbc822f8 100644 --- a/src/diffusers/models/_modeling_parallel.py +++ b/src/diffusers/models/_modeling_parallel.py @@ -161,7 +161,7 @@ class TensorParallelConfig: Tensor parallelism shards weight matrices (column-wise and row-wise) across devices. Each device computes a partial result; an AllReduce/AllGather at layer boundaries reconstructs the full output. Uses `torch.distributed.tensor.parallelize_module` with `ColwiseParallel` / `RowwiseParallel` sharding styles. Supported - device types are `"cuda"` and `"neuron"`. + device types are `"cuda"`, `"neuron"` and `"tpu"`. Args: tp_degree (`int`, defaults to `1`): 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/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/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/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 8bab73f67c..ee0ab21e87 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -24,6 +24,7 @@ from .parallelism import ( ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin, TensorParallelTesterMixin, + TensorParallelTPUTesterMixin, ) from .quantization import ( AutoRoundCompileTesterMixin, @@ -68,6 +69,7 @@ __all__ = [ "ContextParallelTesterMixin", "ContextParallelAttentionBackendsTesterMixin", "TensorParallelTesterMixin", + "TensorParallelTPUTesterMixin", "CPUOffloadTesterMixin", "FasterCacheConfigMixin", "FasterCacheTesterMixin", 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 cc4f5cad96..31d47f1153 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -45,6 +45,7 @@ from ..testing_utils import ( ModelTesterMixin, SingleFileTesterMixin, TensorParallelTesterMixin, + TensorParallelTPUTesterMixin, TorchAoCompileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, @@ -180,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: 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/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 From 85c9fa4caee33d1c5b473395784471d2d7957c1b Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Sat, 10 Oct 2026 08:40:11 -1000 Subject: [PATCH 06/13] [cli] detect composite block classes in `custom_blocks` (#15001) `diffusers-cli custom_blocks` only recognized classes whose base was spelled exactly `ModularPipelineBlocks`, so a `SequentialPipelineBlocks` or `AutoPipelineBlocks` subclass (the usual thing to publish) was reported as "could not be retrieved". Accept the composite subclasses as well, and match on the class name so dotted bases like `diffusers.modular_pipelines.SequentialPipelineBlocks` work too. Co-authored-by: Claude Fable 5.1 --- .ai/skills/custom-blocks/SKILL.md | 9 ++++---- docs/source/en/using-diffusers/cli.md | 5 +++-- src/diffusers/commands/custom_blocks.py | 14 +++++++++++-- tests/others/test_cli_commands.py | 28 ++++++++++++++++++++++++- 4 files changed, 47 insertions(+), 9 deletions(-) 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/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/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/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): From e1a6c67dc724e7c1a5ae8eb0c1964c9f47325217 Mon Sep 17 00:00:00 2001 From: Lucas Ruan <138600690+lucasruan1618@users.noreply.github.com> Date: Mon, 12 Oct 2026 04:41:19 +0700 Subject: [PATCH 07/13] Fix Anima text encoder exports and standalone core denoising (#14943) --- .../modular_pipelines/anima/encoders.py | 8 +++ .../anima/modular_pipeline.py | 2 +- .../anima/test_modular_pipeline_anima.py | 51 +++++++++++++++++++ 3 files changed, 60 insertions(+), 1 deletion(-) 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/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() From b1f6cde2ac8396ae5bed19c84f381ccbc041a961 Mon Sep 17 00:00:00 2001 From: Lucain Date: Mon, 12 Oct 2026 06:03:32 +0200 Subject: [PATCH 08/13] Bump huggingface_hub upper bound to <3.0 (diffusers) (#14863) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bump huggingface_hub upper bound to <3.0 Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: célina Co-authored-by: Sayak Paul --- setup.py | 2 +- src/diffusers/dependency_versions_table.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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/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", From 883d51ef02cd7972eeaf158c652d229f3459f419 Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Mon, 12 Oct 2026 12:29:44 +0800 Subject: [PATCH 09/13] Make Stable Diffusion device-map comparisons deterministic (#15010) Description: The balanced device-map tests compare images from single-device and multi-device inference. On XPU, repeated fp16 inference can produce different results even with the same seed, causing the pixel comparison to fail. Enable PyTorch deterministic algorithms for the two image-comparison tests and restore the previous setting afterward. This keeps other tests unaffected. Validation: Both affected tests passed on a multi-XPU runner (RUN_NIGHTLY=1 RUN_SLOW=1). Signed-off-by: Wang, Yi A Co-authored-by: Sayak Paul --- .../stable_diffusion/test_stable_diffusion.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) 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( From 1e004eb6ec9435349c4a6dd6d1c44c7e0fec6b40 Mon Sep 17 00:00:00 2001 From: "Wang, Yi" Date: Mon, 12 Oct 2026 17:09:49 +0800 Subject: [PATCH 10/13] xpu supports TP as well (#15013) xpu support TP as well Signed-off-by: Wang, Yi A --- src/diffusers/hooks/tensor_parallel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index ce0c76f970..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", "tpu") +_SUPPORTED_TP_DEVICES = ("cuda", "neuron", "tpu", "xpu") class PackedColwiseParallel: From 1628cb1c331741b64f67c97209ab49018e71ac70 Mon Sep 17 00:00:00 2001 From: Junsong Chen Date: Tue, 13 Oct 2026 10:30:58 +0800 Subject: [PATCH 11/13] Add SANA-WM camera-controlled image-to-video pipeline (#13881) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(sana-wm): add diffusers-style SANA-WM camera-controlled I2V pipeline Adds the public SANA-WM bidirectional camera-controlled image-to-video model as a first-class diffusers pipeline + transformer. Layout mirrors ``sana_video``: the model lives under ``src/diffusers/models/transformers/`` as a near-single-file (kernels split off so the ``@triton.jit`` decorators don't drown the model body); the pipeline lives under ``src/diffusers/pipelines/sana_wm/``. Files added: src/diffusers/models/transformers/ ├── transformer_sana_wm.py # SanaWMTransformer3DModel + blocks + helpers └── transformer_sana_wm_kernels.py # fused Triton kernels + camera math src/diffusers/pipelines/sana_wm/ ├── __init__.py ├── pipeline_sana_wm.py ├── pipeline_output.py ├── refiner.py └── cam_utils.py Pipeline architecture: * Stage 1: 1600M ``SanaWMTransformer3DModel`` DiT with bidirectional GDN-Triton linear attention + UCPE camera-control branch, LTX-style flow-matching Euler scheduler with per-token timesteps. * Stage 2: LTX-2 sink-bidirectional Euler refiner (3 distilled sigma steps, reuses diffusers' ``LTX2VideoTransformer3DModel`` + ``LTX2TextConnectors`` + Gemma-3 text encoder). * Decode through the LTX-2 VAE (``AutoencoderKLLTX2Video``). One-line usage: pipe = SanaWMPipeline.from_pretrained( "Efficient-Large-Model/SANA-WM_bidirectional-diffusers", torch_dtype=torch.bfloat16, ).to("cuda") out = pipe(image=img, prompt="...", action="w-80,jw-40,w-40", intrinsics=[fx, fy, cx, cy]) End-to-end smoke test (stage-1 + refiner + VAE decode) passes on H100. Co-Authored-By: Claude Opus 4.7 * feat(sana-wm): align pipeline with merged sana_video style; fix mp4 export transformer_sana_wm.py: * License header switched to the "HuggingFace Team and SANA-WM Authors" style used by merged sana_video. * Imports rewritten in stdlib -> third-party -> diffusers order; use diffusers `from ...utils import logging` instead of stdlib `logging`. * Fix 9 `Optional[X]` annotations written as `X or None` (Python's `or` short-circuits and silently returns `X`). * Fix two `assert (cond, msg)` tuple-asserts in PatchEmbedMS3D.forward that always pass (SyntaxWarning at import time). * Remove duplicate `__all__` declarations (the second silently overwrote the first). * Remove dead `reset_bn` (imports a nonexistent `packages.apps.utils`, would crash on call). * Remove the duplicate `logger = logging.getLogger(__name__)` further down in the file. transformer_sana_wm_kernels.py: * License header normalized; collapse three duplicate triton/torch import blocks into one. pipeline_sana_wm.py: * License header normalized. * `_decode_latents` now returns `(T, H, W, 3)` float in [0, 1], matching the diffusers convention used by `VideoProcessor`. Returning uint8 silently broke `export_to_video`: it does `frame * 255` assuming float input, so uint8 overflows to `(-x) mod 256` and inverts colors. * `__call__` converts to PIL/uint8 only when `output_type="pil"`. * Intrinsics argument now accepts (4,), (F, 4), (3, 3), and (F, 3, 3) forms (auto-extracts fx, fy, cx, cy from a 3x3 K) and auto-trims to `num_frames` when a longer-than-needed trajectory is passed. * Inline `retrieve_timesteps` with the standard `# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps` marker, matching merged sana_video. * Docstrings + EXAMPLE_DOC_STRING updated to reflect the new return type. pipeline_output.py: * Update `frames` field docstring to describe the new float [0, 1] return. refiner.py, cam_utils.py, scripts/sana_wm/convert_sana_wm_to_diffusers.py: * License headers normalized. Docs: * New `docs/source/en/api/pipelines/sana_wm.md` and `docs/source/en/api/models/sana_wm_transformer3d.md`, modeled on sana_video.md / sana_video_transformer3d.md, wired into `docs/source/en/_toctree.yml` under Models and Pipelines. 5s end-to-end smoke test (81 frames @ 16fps, 30 stage-1 steps + 3-step LTX-2 refiner) passes on 1x H100 80GB with `enable_model_cpu_offload`. Round-trip diff vs raw float frames is 2.06/255 mean (h264 lossy noise), confirming the export_to_video fix. * feat(sana-wm): port chunk-causal AR refiner mode (RefinerChunkRunner + KV cache hooks) The first cleanup pass only kept the legacy single-shot refiner path. That path is what the model was *not* trained on — its docstring even says "feeding the full sequence at once is out-of-distribution" — and its cost is O(T^2) attention over the full latent volume, which made longer videos unusable (~21 min per refiner step at 321 frames on an H100). Port the chunk-causal AR mode from the upstream reference so the refiner matches the training contract: * `refine_latents` now defaults to `block_size=3, kv_max_frames=11` (the canonical AR recipe). Pass `block_size=None` to fall back to the legacy single-shot path. * New `_refine_latents_ar` + `_RefinerChunkRunner` orchestrate the sliding window: pre-capture pre-RoPE sink K/V on `z_sana[:source_sink_frames]` at sigma=0, then for each `block_size`-frame chunk run a 3-step Euler with prefix `{sink_k_pre, sink_v, sink_pe, history_k, history_v}` and capture post-RoPE K/V to feed the next window. History is bounded to `kv_max_frames - source_sink_frames` so per-block compute is constant. * New `_predict_x0_active_block` runs the transformer on the active block only (Q from active, K/V from prefix+active). * New `_capture_block_kv` runs sigma=0 forward with a pre_rope/post_rope capture flag set on each `attn1`. * New `_forward_video_only_with_rope` takes a pre-built RoPE so each block can use absolute frame positions in the source video. * `_streaming_self_attention` extended with the `_kv_cache_capture`, `_tf_capture_kv`, `_tf_kv_prefix` hook contract that AR mode uses to inject and capture K/V on each block. * New helpers: `_build_rotary_emb_for_absolute_positions`, `_set_kv_prefix_on_blocks`, `_clear_kv_prefix_on_blocks`, `_set_capture_flag_on_blocks`, `_collect_captured_kv_from_blocks`. * `_encode_prompt` now also moves the Gemma-3 text encoder back to CPU after producing the embeds — otherwise it stays resident through the entire AR loop and gates how much GPU memory the refiner transformer has left. Module-level docstring updated to document both modes; existing single-shot path preserved verbatim. * feat(sana-wm): block-level checkpoint for AR refiner (resume after preemption) The AR refiner is expensive (~3-5 min per block) and the refinement loop ran end-to-end has no in-progress state to recover, so a SLURM preemption mid-refinement loses all progress. With the canonical ``block_size=3, kv_max_frames=11`` setup, refining a 50s video is 34 blocks of work that has to make it through without preemption on a backfill queue. Add per-block atomic checkpointing: * ``SanaWMLTX2Refiner.refine_latents(checkpoint_dir=Path)`` and ``_refine_latents_ar`` accept a directory. After each completed AR block, the AR loop writes ``checkpoint_dir/state.pt`` atomically (tmp + os.replace). * The payload is ``{block_idx_done, n_blocks, sink_size, block_size, output_shape, output, runner_state}``. ``runner_state`` is a CPU snapshot of the runner's ``_sink_kv_pre``, ``_history_kv_post``, ``_history_frames`` and ``torch.Generator`` state. * On entry, if ``state.pt`` exists with a compatible shape signature, the AR loop loads the persisted output tensor + runner state and resumes from ``block_idx_done + 1`` instead of recomputing from scratch. * ``SanaWMPipeline.__call__(refiner_checkpoint_dir=...)`` plumbs the directory through to the refiner. Checkpoint size: ~output_volume + sink_KV (~360MB for 50 layers) + rolling history KV (~3-4GB at full capacity) — saved once per block, total per-block save overhead ~10s on lustre. * test(sana-wm): add CPU unit tests + slow GPU integration stub * CPU unit tests for cam_utils helpers (action DSL → c2w, intrinsics rescale-for-crop, resize+center-crop, snap_num_frames 8k+1 rounding). * Public-surface registration tests (top-level diffusers symbols, SanaWMPipelineOutput dataclass shape, refiner signature has AR defaults + checkpoint_dir, pipeline __call__ accepts c2w/action/intrinsics/ refiner_checkpoint_dir). * @slow @require_torch_accelerator integration stub for an end-to-end I2V against the public checkpoint, currently @unittest.skip — wires up the nightly GPU path without exploding regular CI. SanaWMTransformer3DModel has hardcoded depth/hidden_size/num_heads inside its inner SanaMSVideoCamCtrl (not exposed through register_to_config), so the usual PipelineTesterMixin small-config fast tests aren't applicable without a transformer refactor (followup PR). * feat(sana-wm): make triton optional + auto-fallback to pure-PyTorch attention `transformer_sana_wm_kernels.py` previously did a hard `import triton` at the top of the file. That blocked importing the SANA-WM transformer on any environment without Triton (CPU-only, ROCm without Triton, older Triton, etc.), even though the model has pure-PyTorch attention classes for every `*Triton` variant. Make Triton optional and have the dispatcher transparently fall back: * Wrap `import triton` / `import triton.language as tl` in try/except. When unavailable, install a shim where `@triton.jit` is a no-op so the kernel function definitions still load (they just aren't compiled by Triton). Module-level `triton.X` / `tl.X` lookups return a self-shimming sentinel so signature parsing doesn't blow up either. * Add `is_triton_available()` + `_require_triton(entry_point)`. The four Triton-backed entry points called by the model (`fused_qk_inv_rms`, `fused_bigdn_func`, `cam_prep_func`, `cam_scan_bidi_chunkwise`) now raise a clear RuntimeError on a Triton-less host with a hint to use the pure-PyTorch attention variants — but the dispatcher does this automatically (see below) so users shouldn't ever see it. * Delete the leftover duplicate `import torch / triton / triton.language` block at line 262 (left over from the upstream port). * Register `BidirectionalGDNUCPESinglePathLiteLA` in `ATTENTION_BLOCKS` so the fallback chain can find it. * New `_resolve_attention_block(name, role)` walks the requested class's MRO at dispatch time. If Triton isn't usable AND the requested class name ends in `Triton`, route to the closest registered non-`Triton` ancestor (BidirectionalGDNUCPESinglePathLiteLABothTriton -> BidirectionalGDNUCPESinglePathLiteLA, etc.) and log a one-shot warning. * Rewire both `SanaVideoMSCamCtrlBlock` dispatch sites to use `_resolve_attention_block` for the GDN+UCPE camera branch and the main attention branch (the `BidirectionalSoftmaxUCPESinglePathLiteLA` branch doesn't use Triton at all so it stays hard-coded). Tests: * `test_kernels_module_imports_with_triton_hidden` — reloads the kernels module with `sys.modules['triton'] = None` and verifies the module imports, `is_triton_available()` is False, and the pure-PyTorch helpers remain callable. * `test_resolve_attention_block_cpu_fallback` — on a CPU-only host, the three `*Triton` attn types resolve to the correct non-Triton ancestor. * `test_triton_entry_point_raises_clean_error_without_triton` — verifies the `_require_triton` guard yields a RuntimeError that mentions Triton. * fix(sana-wm): make optional deps lazy + register transformer in __init__ Three CI checks were failing on the PR: 1. `check_code_quality` (43 ruff errors): mix of unused imports / import sorting / E731 lambdas (auto-fixable) plus a handful of F821 dead-code references inherited from the upstream research codebase (`xformers.*` inside `if _xformers_available:` blocks, an undefined `BlockHook` type annotation, two `x_sa`/`mlp_out` references in a block forward whose live assignment was already overridden by subclasses). Ran `ruff check --fix --unsafe-fixes` + `ruff format`, fixed the type annotation manually, and added targeted `# noqa: F821` markers on the conditionally unreachable lines. 2. `check_torch_dependencies`: `transformer_sana_wm.py` hard-imported `einops`, `fla`, `timm`, `termcolor`. The minimum-deps CI environment doesn't have them, and diffusers' lazy loader rewrites `ModuleNotFoundError` as `RuntimeError` so `test_pipeline_imports` blew up. Wrapped each of the four optional imports in a try/except shim — `rearrange`/ `ShortConvolution`/`DropPath`/`Attention_`/`Mlp` become placeholders that raise a clear `ImportError` on construction, `colored` falls back to plain text. Class bodies that subclass these still parse at module load, so `import diffusers.models.transformers.transformer_sana_wm` succeeds anywhere. Same treatment for the kernels file's `from einops import rearrange, repeat`. 3. `build_pr_documentation`: doc-builder imported `SanaWMTransformer3DModel` from `diffusers.models.transformers` (not the diffusers top level) and that subpackage's `__init__.py` was missing the entry. Added the import. * style(sana-wm): apply make style + fix-copies * `doc-builder style src/diffusers docs/source --max_len 119` rewraps docstrings in the six SANA-WM files (transformer, kernels, pipeline, refiner, output, cam_utils) to the repo-wide 119-column limit. No behaviour change — purely whitespace inside docstrings. * `make fix-copies` regenerates `dummy_pt_objects.py` and `dummy_torch_and_transformers_objects.py` to add `DummyObject` stubs for the three new public classes (`SanaWMTransformer3DModel`, `SanaWMPipeline`, `SanaWMLTX2Refiner`), so `from diffusers import …` gives the standard "missing backend" message on installs without torch / transformers. Verified: `make quality` passes (ruff check, ruff format check, doc-builder style check_only, check_doc_toc). Test suite still 15 passed / 1 skipped. * refactor(sana-wm): drop einops dependency, inline with torch ops Per @dg845's review on `transformer_sana_wm.py:44`. The 9 call sites all match well-known patterns that have one-liner torch equivalents: rearrange(s, "b d -> (b d)") -> s.reshape(b * d) rearrange(x, "(b d) d2 -> b (d d2)", ...) -> x.reshape(b, d * d2) rearrange(R, "b t h w i j -> b t h w j i") -> R.transpose(-1, -2) repeat(x, "b h w c -> b t h w c", t=T) -> x.unsqueeze(1).expand(-1, T, -1, -1, -1) repeat(x, "b t c -> b t h w c", h=H, w=W) -> x[:, :, None, None, :].expand(-1, -1, H, W, -1) repeat(x, "... -> b ...", b=B) -> x.unsqueeze(0).expand(B, *x.shape) repeat(x, "H W C -> B T H W C", B, T) -> x[None, None].expand(B, T, -1, -1, -1) repeat(x, "B H W C -> B T H W C", T) -> x.unsqueeze(1).expand(-1, T, -1, -1, -1) Each replacement is bit-identical to the einops original — verified against a fresh `einops` install on random tensors before swapping. The optional `from einops import ...` shim block is gone from both `transformer_sana_wm.py` and `transformer_sana_wm_kernels.py`. * refactor(sana-wm): remove dead code per @dg845's review Drop 23 unused symbols (1451 lines) from `transformer_sana_wm.py` that aren't reachable from the public SANA-WM checkpoint's `SanaMSVideoCamCtrl` -> `SanaVideoMSCamCtrlBlock` -> GDN/UCPE attention path. Each was verified to have zero call sites outside of its own definition (or only within other-deleted items). Classes: * `SanaMS`, `SanaMSBlock` — alternative `Sana` subclass + block, not used by the SANA-WM checkpoint (which goes through `SanaMSVideoCamCtrl`). * `ChunkCausalAttention`, `CachedCausalAttention`, `ChunkedLiteLAReLURope`, `LiteLAReLURope` — chunk-causal / cached attention variants and their common base; SANA-WM uses the bidi GDN path. `LiteLAReLURope` had only the three (now-deleted) subclasses referencing it. * `PAGCFGIdentitySelfAttnProcessorLiteLA`, `PAGIdentitySelfAttnProcessorLiteLA`, `SelfAttnProcessorLiteLA`, `SelfAttnProcessorLiteLAReLURope` — PAG processors; we don't expose PAG in the SANA-WM pipeline. * `ChunkGLUMBConvTemp`, `CachedGLUMBConvTemp`, `MBConvPreGLU` — alternative FFN/conv blocks; the checkpoint uses `GLUMBConvTemp`. * `MaskFinalLayer`, `DecoderLayer` — alternative final layers; the checkpoint uses `T2IFinalLayer`. * `LabelEmbedder`, `CaptionEmbedderDoubleBr` — alternative embedders; the checkpoint uses `CaptionEmbedder`. Helpers: * `set_grad_checkpoint`, `prepare_prompt_ar`, `resize_and_crop_tensor`, `generate_temporal_head_mask_mod`, `is_chunk_causal_request`, `get_chunk_index_from_config` — training-only or chunk-causal-mode helpers with no inference call sites. Verified after the diff: * `make quality` clean. * CPU test suite 15 passed / 1 skipped (slow GPU integration). * `import diffusers ; SanaWMPipeline / SanaWMTransformer3DModel / SanaWMLTX2Refiner` resolve identically. * fix(sana-wm): defer transformers import + document mask/return_dict Two PR-CI failures, both ours: * `check_torch_dependencies`: line 33 hard-imported `transformers`, which isn't present in the minimum-deps environment. Move the lone `AutoModelForCausalLM` use site inside `initialize_gemma_params` (a training-only helper). * `check_repository_consistency`: `SanaWMTransformer3DModel.forward`'s docstring was missing entries for `mask` and `return_dict`. Added. * refactor(sana-wm): address review feedback on SanaWMPipeline Per @dg845's review comments on pipeline_sana_wm.py: * Read the VAE spatial/temporal strides from the ``vae`` component (``vae_spatial_compression_ratio`` / ``vae_temporal_compression_ratio``) instead of hardcoding 32/8, mirroring LTX2Pipeline. * Inline the stage-1 sampling loop into ``__call__`` and factor the noise init into a standard ``prepare_latents`` method. * Use ``self.scheduler`` for the flow-matching Euler steps instead of constructing a new scheduler per call. * Drive the sampling loop with ``self.progress_bar`` (respects ``set_progress_bar_config``) instead of a bare tqdm. * Move input validation/normalization into a ``check_inputs`` method. * Move first-frame resize+center-crop, [-1, 1] normalization and the intrinsics rescale into a ``SanaWMImageProcessor(VaeImageProcessor)`` subclass (new ``image_processor.py``). * Add ``generator`` as a ``__call__`` argument (``seed`` kept as a convenience shortcut). * Post-process decoded latents with ``VideoProcessor.postprocess_video`` rather than a hand-rolled conversion; drop the ``.cpu()`` casts on the ``output_type="latent"`` path. * Move the pipeline README into the docs (docs/.../sana_wm.md); delete the in-package README. * refactor(sana-wm): make SanaWMLTX2Refiner a standalone DiffusionPipeline Per @dg845's review: the refiner has components (transformer, connectors, text encoder, tokenizer) and runs a denoising loop, so it fits better as a pipeline than a ModelMixin. * SanaWMLTX2Refiner now subclasses DiffusionPipeline and registers its components via ``register_modules`` (dropping the bespoke ``from_pretrained`` / ``save_pretrained``); standard load/save now handle the ``refiner/`` subfolder. * Add a ``FlowMatchEulerDiscreteScheduler`` component (shift=1.0) and drive the Euler steps through ``scheduler.step`` / ``scheduler.scale_noise`` (single-shot and per-AR-block) instead of hand-rolled updates. Numerically equivalent to the previous flow-matching update. * Rename the entry point ``refine_latents`` -> ``__call__``; add a ``device`` arg so the parent can hand it the execution device without a bulk move. * SanaWMPipeline: keep the refiner as an optional nested component; free the parent's GPU weights before running it (it manages its own sub-module placement) and bring the VAE back for decode. * Conversion script emits the new refiner layout (model_index.json + scheduler/ + tokenizer/); test asserts the refiner is a DiffusionPipeline with the canonical AR ``__call__`` defaults. Validated end-to-end on 1xH100 (stage-1 + refiner on the official demo, coherent video output). * docs(sana-wm): document the refiner __call__ `device` arg Fixes the check_repository_consistency failure — `SanaWMLTX2Refiner.__call__` gained a `device` parameter but its docstring wasn't updated. * refactor(sana-wm): clean up dead code + reuse shared utils in the DiT Addresses several of @dg845's transformer review comments (the safe, output-preserving subset — GPU smoke gives byte-identical output): * Reuse diffusers' shared `FP32LayerNorm` (models/normalization.py) and `get_1d_rotary_pos_embed` (models/embeddings.py); delete the local copies. * Remove dead inference code paths: - the `if self.diagonal_mask is not None:` flex-attention block (`diagonal_mask` is always `None`) + the now-unused `create_block_mask_cached` helper and `create_block_mask` import; - the `SANA_FSDP2_BLOCK_TIMING` block-timing/profiling scaffolding; - the `save_qkv` / `qkv_store_buffer` visualization hooks (never enabled at inference), at both the attention and model level. * Collapse `SanaVideoMSCamCtrlBlock.forward_frame_aware` into `forward` (the pipeline/refiner always pass >=3D timesteps, so the non-frame-aware branch was dead — it even referenced undefined locals). * Drop the 191-line `SanaMSVideoCamCtrl.load_state_dict` shape-remapping override — it's never reached by the shipped convert/inference flow (`nn.Module.load_state_dict` on the wrapper doesn't call it), and the release checkpoint already ships correctly-shaped weights. * Delete the unused `PatchEmbedMS` module. Net -503 lines; no numerics change (stage-1 + refiner smoke output mean identical to the pre-refactor run). * refactor(sana-wm): merge the 3 DiT classes into one SanaWMTransformer3DModel Per @dg845's review: `SanaWMTransformer3DModel` was a thin wrapper over `SanaMSVideoCamCtrl`, which in turn subclassed `Sana` and overwrote most of it. Fold all three into a single `SanaWMTransformer3DModel(ModelMixin, ConfigMixin)`: * The `@register_to_config` __init__ signature is unchanged (so config.json is identical); the body builds the modules directly on `self` instead of a nested `self._inner`. * Only the surviving `Sana.__init__` pieces are kept (t_embedder, cfg_embedder, attention_y_norm, config attrs, initialize_weights); the parts the subclass overwrote are gone. * `forward` takes the diffusers signature (hidden_states / timestep / encoder_hidden_states / encoder_attention_mask / return_dict) and returns `Transformer2DModelOutput`, folding in the old wrapper's arg-renaming. * Deleted now-unreachable code: `class Sana`, `class SanaMSVideoCamCtrl`, `class SanaBlock`, `class PatchEmbed` (the video model uses `PatchEmbedMS3D`), the `add_inner_prefix` helper, and the dead `sincos`/`flux_rope` pos-embed branches in forward (release uses `wan_rope`). Net -468 lines. This drops the `_inner.` state-dict prefix, so the conversion script no longer adds it. State-dict is otherwise identical: the merged model's `state_dict()` has exactly the same 871 param keys as before (verified), and a stage-1 + refiner GPU smoke on the public checkpoint gives byte-identical output (frame mean 0.5560, matching the pre-merge run). Addresses the class-merge review comment and removes the unused `SanaBlock`/`PatchEmbed` modules. * refactor(sana-wm): drop the no-op custom grad-checkpoint wrapper `auto_grad_checkpoint` gated on a `grad_checkpointing` attribute that was never set, so it always fell through to `module(*args, **kwargs)` — i.e. a no-op. Call the transformer blocks directly instead and delete the unused `auto_grad_checkpoint` / `checkpoint_sequential` helpers and the `torch.utils.checkpoint` import. (Inference is unchanged — the block never read the `use_reentrant` kwarg the wrapper passed. Full training-time gradient checkpointing via the standard `_gradient_checkpointing_func` would need the block's dynamic kwargs — camera_conditions / prope_fns / chunk_index — threaded through, so it's left as a follow-up for this inference-focused release.) * refactor(sana-wm): drop the fla-core and termcolor dependencies Per @dg845 / @sayakpaul: the SANA-WM transformer required `fla-core` to be constructible at all (and fla's `ShortConvolution` cannot even run on CPU — its dispatch does `torch.cpu.device(...)`, which doesn't exist). `ShortConvolution` was only ever used as a depthwise *causal* conv1d with `activation=None`, so it is replaced by a ~20-line self-contained PyTorch module with the same parameter layout (`weight` of shape `(C, 1, K)`, no bias) and the same `(output, cache)` return signature. No `fla-core` and no `kernels` dependency is needed now, and the model can be built and run on CPU. Verified on an H100 in bf16 at the real config (hidden 2240, kernel 4): `max|Δ|` vs fla is exactly `0.0` for every shape tested, and the state dict keeps the identical 871 keys. Also: * Use the existing `is_timm_available()` utility for the timm imports. The `else` branch keeps a placeholder class rather than raising at module scope, because several layers subclass these symbols and this module is imported eagerly by `diffusers.models` — a module-level raise would break plain `import diffusers` (and the `check_torch_dependencies` CI job) when timm isn't installed. The error is instead raised on construction. * Drop `termcolor` entirely: replace the `self.logger = print` + `colored(...)` pattern with the module-level `logger`, so there is no optional dependency left to gate. * Only attempt the optional null-embedding load when `null_embed_path` is actually set (it is unset for the public checkpoint, so this previously logged a spurious failure on every construction) and load it with `weights_only=True`. * fix(sana-wm): address review nits in the pipeline / refiner / cam utils * refiner: load the AR resume checkpoint with `weights_only=True`. The payload is only tensors / ints / tuples / dicts plus the generator's uint8 state, so it round-trips safely under the restricted unpickler. * refiner: fix `_tokens_per_frame` to divide by `patch_size_t` rather than multiply — `_pack_latents` emits `(T // patch_size_t) * (H // p) * (W // p)` tokens, so one latent frame contributes `(H // p) * (W // p) / patch_size_t`. No-op for LTX-2 (`patch_size_t=1`) but the history trimming would have kept too many tokens otherwise. * pipeline: replace `torch.cuda.empty_cache()` with diffusers' backend-agnostic `empty_device_cache(device.type)`. * pipeline: extract the offload probe into `_model_cpu_offload_active()`, matching the `hasattr(self, "_all_hooks") and len(self._all_hooks) > 0` idiom `DiffusionPipeline` uses internally, instead of a bare truthiness check on the attribute. * pipeline_output: add the missing Apache-2.0 license header. * cam_utils: drop the unnecessary `+ 1e-6` when normalizing the forward / right vectors — the branch is only taken when the norm is already > 0, so the epsilon just introduced a small systematic bias. * refactor(sana-wm): strip research-repo code from the DiT, fix model conventions From the `self-review` skill run against `.ai/models.md` / `.ai/AGENTS.md`. `transformer_sana_wm.py`: 6555 -> 4592 lines. Conventions: * `_no_split_modules` was `["blocks"]` — an attribute name, but accelerate matches on class name, so it matched nothing and `device_map="auto"` could split a block across devices and crash. Now `["SanaVideoMSCamCtrlBlock"]`. * Add `_repeated_blocks` (enables `compile_repeated_blocks()`) and `_skip_layerwise_casting_patterns`. `_keep_in_fp32_modules` is deliberately left unset with a note: the blocks apply the timestep modulation inline, so keeping `scale_shift_table` / `t_embedder` in fp32 upcasts the hidden states and feeds fp32 activations to bf16 weights (caught by a GPU smoke run). * Expose `num_layers` / `hidden_size` / `num_attention_heads` / `patch_size` through `register_to_config` instead of hardcoding the release architecture, so a tiny variant can be built for tests. Defaults are the released values, so `config.json` and the state dict are unchanged. * `torch.float64` -> `torch.float32` on the live RoPE paths (gotcha 5), `torch.empty` -> `torch.zeros` for parameter init (gotcha 6), and stop reading `self.proj.weight.dtype` to cast activations (gotcha 4). * `WanRotaryPosEmbed.freqs` is a non-persistent buffer instead of a plain attribute reassigned inside `forward` (which broke `.to()` and compile). Dead code (`AGENTS.md`: "delete training-time code paths, experimental flags, and ablation branches entirely — only keep the inference path"): * All weight-init / transfer-learning helpers — `from_pretrained` overwrites them, and they also printed ~20 lines of noise on every construction. * `CaptionEmbedder.initialize_gemma_params` (fetched `google/gemma-2b-it` out-of-band at runtime, and referenced an attribute that never exists), `token_drop`, and the training branch of its forward. * The cam-debug statistics apparatus, the `save_block_output` hooks (whose `get_block_output` was never defined), and `block_hook`. * Both unreachable recurrence variants, `_maybe_drop_cam_branch`, the xformers branches (`_xformers_available` was a literal `False`, defined twice, and `xformers` was never imported), 3 env-var escape hatches, 9 unused classes, 12 unused module-level helpers, and the sincos family — which also removes the NumPy import, satisfying the no-NumPy-in-forward rule. * Inline the `fp32_attention` mechanism: it was set on every submodule via `model.apply` and read at 18 sites, so it was folded to the shipped always-on behaviour rather than deleted. Ephemeral comments (commit SHAs from a private tree, references to files that don't exist in diffusers, FSDP2 rationale, stale docstrings) removed. State dict is unchanged: 871/871 keys match the released checkpoint. * fix(sana-wm): honour scheduler shift, use randn_tensor, drop cluster scaffolding From the `self-review` skill run against `.ai/pipelines.md` / `.ai/testing.md`. `flow_shift` was a dead knob. `__call__` set `self.scheduler.config.shift`, but `FlowMatchEulerDiscreteScheduler.set_timesteps` reads `self.shift` (i.e. `self._shift`), so the documented argument had no effect and every run used the checkpoint's `shift=9.8`. It also half-mutated a `FrozenDict` — the freeze guard checks a name-mangled attribute, so the assignment silently desynced `config.shift` from `config["shift"]` on a shared component. Removed the argument so the scheduler owns the shift (`pipelines.md` gotcha 3), which also drops a per-call mutation of a registered component (gotcha 7). Other pipeline fixes: * `torch.randn` -> `randn_tensor` in `prepare_latents` and the refiner (gotcha 10). `__call__` advertises `generator: Generator | list[Generator]`, but `torch.randn` raises on a generator list and on a CPU generator with a CUDA device, so that path could not work. * Delete the SLURM preemption/resume feature (`checkpoint_dir`, `_atomic_save_state`, `_capture_state` / `_restore_state`) and the single-shot refiner path the docstring itself called a debugging fallback — roughly 270 lines of research-cluster scaffolding. * `_empty_cuda_cache` (CUDA-only) -> `empty_device_cache`; `@torch.inference_mode()` -> `@torch.no_grad()` and removed from inner helpers the decorator already covers (gotcha 2). * Remove dead `_callback_tensor_inputs` (no `callback_on_step_end` exists), `_exclude_from_cpu_offload` (a no-op — the base class already skips non-`nn.Module` components), `_kv_max_frames`, and `latents.detach()`. Tests: drop the `@slow` integration stub — `testing.md` says integration and slow tests don't belong in the initial PR. `refiner.py`: 1286 -> 1017 lines. * refactor(sana-wm): drop the timm dependency Per @yiyixuxu's review. The three timm imports were all trivially replaceable: * `Mlp` — both local subclasses fully overrode `forward` and only inherited `fc1` / `act` / `drop1` / `fc2` / `drop2`, so it is now a self-contained module and the redundant `class Mlp(Mlp)` wrapper is gone. * `Attention` — `GDN` inherited only `qkv` and `proj` from it (it already defines its own `q_norm` / `k_norm`), so those two layers are declared inline and `GDN` subclasses `nn.Module`. * `DropPath` — `drop_path` was hardcoded to `0.0`, so this was always `nn.Identity()`. Removed along with the rest of the plumbing. No `dummy_timm_objects.py` is needed since the dependency is gone rather than optional, and `setup.py` is untouched (other models still use timm). State dict is unchanged (871/871 keys match the released checkpoint) and the end-to-end GPU smoke gives byte-identical output (frame mean 0.5554). * refactor(sana-wm): drop the Triton kernels and apply review feedback Per @yiyixuxu's inline review. Removes `transformer_sana_wm_kernels.py` entirely (3234 lines) — the Triton kernels will be published to `kernels-community` and wired back in a follow-up PR. * Delete the three `*Triton` attention classes. Each only overrode `forward` to call a fused kernel; their pure-PyTorch parents (`BidirectionalGDN`, `BidirectionalGDNUCPESinglePathLiteLA`) compute the same result. The Triton-availability probe and the MRO-walking fallback go with them. The released `config.json` still names the `*Triton` variants, so `ATTENTION_BLOCKS` maps those strings onto the pure-PyTorch classes. * Move the nine pure-PyTorch camera-math helpers that were reachable from the model into `transformer_sana_wm.py`; drop the rest with the kernels file (they were only reachable from the deleted Triton path). Other review items: * Use `get_activation()` from `..activations`; delete the local activation and normalization registries (`build_norm` was only ever called with `None`). * Delete `GLUMBConv` and make `GLUMBConvTemp` self-contained. `ConvLayer` is slimmed and renamed `SanaWMConvLayer`, but has to stay a module — the checkpoint keys are `mlp.inverted_conv.conv.weight`, so collapsing it to a bare `nn.Conv2d` would drop the `.conv` level. * Remove the closure-returning helpers the review flagged as hard to read and compile-hostile: `prepare_prope_fns` and friends now return plain tensors via `_prepare_ucpe_ray_transforms` + `_apply_ucpe_transform`, and the `_register_block` decorator factory is an explicit dict update. * Remove all seven `@torch.compile` decorators — users can reach for `compile_repeated_blocks()` instead. * Drop `partial` (pass `chunk_size` at the call site), `DWMlp`, the tuple helpers, `get_same_padding`, `modulate`, an unused `apply_rotary_emb`, and the validation of internal-only arguments. * Tests: drop `SanaWMTritonFallbackTests` along with the mechanism it covered. State dict is unchanged (871/871 keys) and the end-to-end GPU smoke passes with correct video output. * refactor(sana-wm): inline the modulation helper, drop the backward-only shim Two more items from @yiyixuxu's inline review: * Inline `t2i_modulate` at its four call sites. * Delete `_IdentityForwardContiguousBackward` / `_contiguous_backward`. It is the identity in forward and only exists to hand a contiguous gradient to the backward pass, so it is dead weight in an inference-only port. State dict unchanged (871/871 keys); GPU smoke output matches the previous run exactly (frame mean 0.5560). * Update src/diffusers/models/transformers/transformer_sana_wm.py Co-authored-by: YiYi Xu * Update src/diffusers/models/transformers/transformer_sana_wm.py Co-authored-by: YiYi Xu * style(sana-wm): drop the blank line left by removing the camera-branch init * refactor(sana-wm): address @dg845's review round Transformer: * Stop stashing patch-grid shape on `self` during `forward` (`self.f/h/w`) and thread it through as locals; `unpatchify` now takes it explicitly. * Replace the 7 `assert`s with `ValueError`s. * Drop `attn_drop`/`proj_drop` from `MultiHeadCrossAttention` (training-only, and `attn_drop` was never applied) plus a stale training-era banner comment. Pipeline / refiner: * Don't mutate the components handed to the pipeline. The VAE tiling + framewise settings move to the docs, `padding_side="right"` is passed per tokenizer call, and the `.eval()` calls are gone (no `self.training` branches remain, and `from_pretrained` already returns eval-mode modules). * Gate `cam_utils`' optional imports on `is_torchvision_available()` and a new `is_pi3_available()` helper. * Annotate `SanaWMLTX2Refiner.__init__` and inline `_refine_latents_ar` into its single caller. Conversion script: * Move to `scripts/` alongside the other Sana converters. * Raise on missing/unexpected keys instead of printing, so a bad mapping can't silently emit a broken transformer. State dict unchanged (871/871 keys). GPU smoke on the public checkpoint is byte-identical to the previous run (frame mean 0.5560), including with the VAE settings applied by the caller rather than the pipeline. * refactor(sana-wm): rename WanRotaryPosEmbed to SanaWMRotaryPosEmbed The name implied this was Wan's rotary embedding, but it isn't: the per-axis split is configurable through `fhw_dim`, and the frequencies stay complex in a single `freqs` buffer instead of being split into real cos/sin buffers. So it can't carry a `# Copied from`. Renamed, with a docstring recording why. The buffer is `persistent=False`, so the state dict is unchanged (871/871). * test(sana-wm): add transformer model tests, migrate pipeline tests to pytest * `tests/models/transformers/test_models_transformer_sana_wm.py` — generated with `utils/generate_model_tests.py` and filled in, following `test_models_transformer_sana_video.py`. The tiny config sets `softmax_every_n=2` so one block exercises the GDN camera branch and the other the softmax variant. Dummy inputs supply the conditioning the forward requires: `encoder_attention_mask`, `(B, F, 20)` camera conditions, and `chunk_plucker`. * `tests/pipelines/sana_wm/test_sana_wm.py` — rewritten in the pytest style of `tests/pipelines/sana_video/test_sana_video.py`: no `unittest`, bare asserts, module-level imports, and `parametrize` in place of loop-style cases (15 test functions become 31 cases). `AttentionTesterMixin` is skipped because the model calls `F.scaled_dot_product_attention` directly rather than going through a diffusers attention processor. * fix(sana-wm): make the transformer layerwise-casting and offload safe * `TimestepEmbedder.dtype` used `next(self.parameters()).dtype`, which reports the storage dtype under layerwise casting rather than the compute dtype. Use `get_parameter_dtype`, which is layerwise-casting aware. (The model-level `self.dtype` already routes through it via `ModelMixin`.) * Stop reading submodule `.weight`/`.bias` inside `forward`. Group-offload hooks fire on a module's `forward`, so reading its parameters directly leaves them offloaded and trips a device mismatch. The frame-gate and output-gate helpers now call their submodules, and the fused camera QKV projection becomes three `q/k/v` calls -- algebraically the same as one GEMM over the concatenated weights. * Don't mutate the tokenizer in `SanaWMLTX2Refiner._encode_prompt`; pass `padding_side="left"` per call, matching `SanaWMPipeline`. * Add a `torch.compile` test with `recompile_limit=2` -- the repeated block compiles once per attention variant. GPU smoke unchanged (frame mean 0.5560). * refactor(sana-wm): explicit arguments instead of **kwargs, generator instead of seed Replace the `**kwargs` bag threaded through model -> block -> attention with explicit keyword arguments. Only three runtime keys were ever read at the leaves (`frame_valid_mask`, `precomputed_gates`, `ucpe_ray_transforms`); the rest were forwarded and silently swallowed. `camera_embedding`, `chunk_index`, `chunk_index_global` and `chunk_split_strategy` turned out to be pure dead plumbing -- written into the per-block kwargs dicts and never read by any attention or MLP forward -- so they are gone. The pipeline also grows the standard diffusers arguments: * `prompt_embeds` / `prompt_attention_mask` / `negative_prompt_embeds` / `negative_prompt_attention_mask` on `encode_prompt` and `__call__`. * `seed` / `refiner_seed` are replaced by `generator` / `refiner_generator`. `generator=torch.Generator(device).manual_seed(42)` reproduces exactly what `seed=42` used to build, so results are unchanged. State dict unchanged (871/871). CPU old-vs-new equality on a tiny config over both attention variants is exact (`torch.equal`, max diff 0.0), and the GPU smoke on the public checkpoint still gives frame mean 0.5560. * refactor(sana-wm): drop the unreachable chunk-split strategies @yiyixuxu was right that the model only ever runs uniform chunking; my earlier reply defending the strategies was wrong. The two sites that actually chunk both call `normalize_chunk_index(None, T, chunk_size)` with three positional arguments, so `chunk_split_strategy` always took its `"uniform"` default. The configured `first_chunk_plus_one` reached the per-block kwargs dict and was then swallowed by the attention forwards' `**kwargs` without ever being read. Flattening that bag into explicit arguments is what surfaced it. Also note the `chunk_size` those call sites use is `chunk_gdn_chunk_size`, a different attribute from the `chunk_size` that was being threaded through. So `chunk_index_from_chunk_size`, `normalize_chunk_index`, `is_uniform_chunking` and `compute_chunk_sizes` are gone, the uniform boundaries are inlined at both call sites, and `chunk_split_strategy` is dropped from the model and block constructors. The released `config.json` still carries the key; loading is unaffected (`extract_init_dict` ignores it) but it logs an "not expected and will be ignored" warning, so it should come out of the checkpoint config on the next export. State dict unchanged (871/871); GPU smoke still frame mean 0.5560, which confirms those branches were never taken. * refactor(sana-wm): use the shared RMSNorm @yiyixuxu asked whether the shared `RMSNorm` could be used here and my earlier reply overstated the obstacles. Re-checking each one: * `scale_factor` is not a blocker. `attention_y_norm.weight` is one of the 871 checkpoint keys, so `from_pretrained` overwrites whatever the constructor initialised -- exactly the same reasoning that removed the other hand-written inits in this PR. It only ever affected from-scratch models. * `norm_dim` is not a blocker either; it is always the default `-1`. * The numerics do differ: ours ran the normalisation and the weight multiply in fp32, while the shared class computes only the variance in fp32 and does the scaling and weight multiply in the input/weight dtype. Measured rather than argued: the end-to-end GPU smoke on the public checkpoint moves from frame mean 0.5560 to 0.5561, and the decoded video is visually identical. That is well inside bf16 noise, so the local class is not worth keeping. The now-unused `y_norm_scale_factor` config argument goes too. As with `chunk_split_strategy`, the released `config.json` still carries it and will log an "not expected and will be ignored" warning until the checkpoint config is re-exported. * refactor(sana-wm): migrate MultiHeadCrossAttention to the diffusers attention API * Split the attention maths into a `SanaWMCrossAttnProcessor` and make `MultiHeadCrossAttention` an `AttentionModuleMixin`, so the processor can be swapped and the standard tooling applies. * Use `dispatch_attention_fn` instead of calling `F.scaled_dot_product_attention` directly, which lets the model pick up alternative attention backends. This also drops the two transposes, since the dispatcher takes `(batch, seq, heads, dim)`. * Build a boolean padding mask rather than an additive float one. The old `(1 - mask) * -10000.0` form isn't supported by the varlen backends. Projection names are unchanged (`q_linear`, `kv_linear`, `proj`, `q_norm`, `k_norm`), so the state dict stays at 871/871 keys. The GPU smoke on the public checkpoint is unchanged at frame mean 0.5561, so `-inf` masking makes no difference against the old `-10000.0` at this precision. * feat(sana-wm): give the refiner its own transformer model Per @yiyixuxu and @dg845: the stage-2 refiner no longer reaches into `LTX2VideoTransformer3DModel`. It used to drive AR refinement by setting attributes on LTX-2's submodules (`attn._tf_kv_prefix`, `_kv_cache_capture`, `_tf_capture_kv`), which meant a second `forward` implementation smuggled in from the pipeline. New `transformer_sana_wm_refiner.py`, following the `transformer_wan_vace.py` pattern of importing the reusable pieces rather than duplicating them: * `SanaWMRefinerKVLayerCache` / `SanaWMRefinerKVCache`, modelled on `Flux2KVCache`, holding the per-layer sink and history K/V. * `SanaWMLTX2RefinerTransformerBlock`, which subclasses `LTX2VideoTransformerBlock` and overrides only `forward`, so the submodule structure stays byte-identical. * `SanaWMLTX2RefinerTransformer3DModel`, taking `kv_cache` / `kv_cache_mode` explicitly through `forward`. `refiner.py` drops from 982 to 638 lines: the four `*_kv_prefix_on_blocks` / `*_captured_kv_*` helpers, `_forward_video_block`, `_streaming_self_attention` and `_forward_video_only_with_rope` are all gone, along with a dead `n_context_tokens > 0` branch that neither call site could reach. The state dict is unchanged at 3510 keys, identical to `LTX2VideoTransformer3DModel` at the same config, so only the class name in the checkpoint metadata has to change. CPU equivalence against the old path is exact across four AR blocks -- covering sink injection, history growth and history trimming -- with a negative control confirming the cache path is actually exercised. The GPU smoke on the public checkpoint is unchanged at frame mean 0.5561. * refactor(sana-wm): run the refiner as a separate pipeline Per @dg845, the two stages are no longer nested. `SanaWMPipeline` produces stage-1 latents and `SanaWMLTX2Refiner` consumes them: latents = SanaWMPipeline(...) video = SanaWMLTX2Refiner(...) The refiner takes an optional `vae` -- pass the base pipeline's so the weights are shared, the way SDXL shares components between base and refiner -- and decodes to video when given one. The manual device juggling in `SanaWMPipeline.__call__` is gone with the nesting that forced it; each pipeline now manages its own placement through the standard offload hooks. Un-nesting turned up a real bug: `torch_dtype` never reached the nested sub-pipeline, so `SanaWMPipeline.from_pretrained(..., torch_dtype=bfloat16)` silently ran the whole of stage 2 in float32 -- transformer, text encoder and connectors alike. That is where the "~87 GB of refiner weights" in the old comment came from. Loaded on its own the refiner honours the requested dtype, which halves its footprint; the GPU smoke moves from frame mean 0.5561 to 0.5548 and the decoded video is visually identical. The conversion script writes the refiner to its own output directory (`--dst-refiner`, default `-refiner`) rather than a `refiner/` subfolder. `DiffusionPipeline.from_pretrained` has no `subfolder` argument, so a nested folder silently loads the *base* pipeline's same-named component folders instead -- exactly the kind of wrong-weights failure that only surfaces because `connectors/` happens not to exist at the top level. * Apply style fixes * docs(sana-wm): document the refiner's output_type argument `check_forward_call_docstrings` flagged `SanaWMLTX2Refiner.__call__` as missing an entry for `output_type`, which the un-nesting added. * refactor(sana-wm): support only what the released checkpoint uses @yiyixuxu asked whether we can support just the released configuration and flatten what's left. We can — the checkpoint only ever instantiates two attention classes, so the other five in the chain were dead weight. * Collapse `GDN -> BidirectionalGDN`, `GDN -> _GDNUCPEBase`, `BidirectionalGDNUCPELiteLA -> ...PostUCPERenorm -> ...SinglePathLiteLA` into a self-contained `BidirectionalGDNUCPESinglePathLiteLA` plus `_SoftmaxUCPESinglePathLiteLA`. Inheritance depth goes 5 -> 1. * Blocks take the class, not a string: `SanaVideoMSCamCtrlBlock` gains an `attn_cls` argument, and the model picks per layer with a plain loop. That removes the `ATTENTION_BLOCKS` registry, `_resolve_attention_block`, `_GDN_TO_SOFTMAX_CAMCTRL`, `_inject_softmax_layers`, the `attn_type` / `camctrl_type` validation, and the unreachable no-camera branch. * `attn_type` and `camctrl_type` leave the config signature. As with `chunk_split_strategy` and `y_norm_scale_factor`, the released `config.json` still carries them and loading logs a benign "not expected" warning until the checkpoint is re-exported. * Also drop the zero-init on `GLUMBConvTemp.t_conv` (overwritten by `from_pretrained`) and inline the one-line `restore_shape` helper. The model docs lose a paragraph describing an `_inner.` prefix and an `add_inner_prefix` helper that no longer exist. 3908 -> 3153 lines. State dict unchanged at 871/871 keys with identical ordering, and 15 GDN + 5 softmax blocks as before. Old-vs-new CPU equality is exact (`torch.equal`, max diff 0.0) on both a tiny config and the released geometry, across a `softmax_every_n` sweep and on the backward pass, with negative controls confirming the comparison isn't vacuous. The GPU smoke is byte-identical too: refined latents mean 0.035024 / std 0.887467, frame mean 0.5548. * refactor(sana-wm): make the softmax attention class standalone Drops the last bit of inheritance between the two attention classes, per @yiyixuxu. `_SoftmaxUCPESinglePathLiteLA` now subclasses `nn.Module` directly. It turned out to be a simplification rather than duplication: of the parent's thirteen methods it only ever called three (`_apply_output_gate`, `_prepare_frame_valid_masks`, `_stabilize_cam_transforms`), and the other ten were dead for it. It is also always built with `conv_kernel_size=0`, so the whole short-convolution half of the inherited `__init__` -- `conv_q/k/v`, the `*_cam` variants, `kernel_func` -- was permanently `None` or unused and is now gone. State dict unchanged: 871/871 keys in the same order, and the softmax blocks still emit the same 24 keys each. Old-vs-new equality is exact (`torch.equal`, max abs diff 0.0) on two input configurations, with negative controls from 0.109 to 14.62 confirming the comparison isn't vacuous; the inlined `__init__` even consumes the RNG stream identically, so seeded from-scratch weights are bit-for-bit the same. GPU smoke unchanged at refined-latent mean 0.035024 / std 0.887467. * feat(sana-wm): export the refiner weights in bfloat16 The refiner's Gemma-3 text encoder was the one component still stored in float32, at 45.4 GB. Since the pipeline is loaded with `torch_dtype=torch.bfloat16`, those weights were cast on every load anyway, so storing them cast changes nothing and halves the download: text_encoder 45.4 GB -> 22.7 GB refiner total 84 GB -> 61 GB The VAE stays in float32, as the docs recommend. Verified rather than assumed: the GPU smoke is bit-identical (refined latents mean 0.035024 / std 0.887467, frame mean 0.5548), and it now runs in 124s instead of 384s because there is far less to load. * refactor(sana-wm): pipeline clean-up from @yiyixuxu's review `SanaWMLTX2Refiner`: * Dissolve `_RefinerChunkRunner` (191 lines) into `__call__`, so the AR loop reads top-to-bottom like `pipeline_wan_animate.py`, and swap the bare `tqdm` for `self.progress_bar(...)`. * Actually support `enable_model_cpu_offload()`. The docs already advertised it, but three things blocked it: the `device` default resolved through `next(self.transformer.parameters()).device`, which is `cpu` under offload; `_encode_prompt` called the inner Gemma backbone directly, bypassing the hook that `cpu_offload_with_hook` installs on the top-level module; and `_decode_latents` read `self.vae.device`. All three now go through `self._execution_device`, and the manual `.to(...)` juggling is gone. * TODO on the `LTX2TextConnectors` import pointing at huggingface/diffusers#14749. Two argument removals fall out: `progress=` (use `set_progress_bar_config(disable=True)`, which is what every other pipeline does) and `device=`. `SanaWMPipeline`: * `check_inputs` only validates now; the camera-trajectory normalisation moves to `prepare_camera_trajectory`. * Add the standard `guidance_scale` / `do_classifier_free_guidance` / `num_timesteps` / `current_timestep` / `interrupt` properties. * `return_dict=False` returned only `(frames,)`, silently dropping `c2w` and `latent`. It now returns the same three values either way -- the flag picks the container, not the contents. * Take a `PIL.Image`, not a path, and drop the `_model_cpu_offload_active` helper that had no callers left after the un-nesting. The Pi3X intrinsics estimator leaves `diffusers` entirely: it pulled in an extra dependency and fetched a second checkpoint without the user ever calling `from_pretrained`. It is packaged as a remote modular block hosted alongside the checkpoint, and the docs show `ModularPipeline.from_pretrained` instead. `is_pi3_available` goes with it. Equivalence for the refiner rewrite is exact across six configurations -- varying sink size, block size, KV window and seed so history growth and trimming both fire -- `torch.equal` with max abs diff 0.0, negative controls from 1.9e-2 to 5.55. GPU smoke on the public checkpoint is unchanged at refined-latent mean 0.035024 / std 0.887467. * refactor(sana-wm): model-side clean-up from @yiyixuxu's review * Split each attention class's main and camera branches into their own modules, so `forward` is a flat list of layer calls. The branch modules deliberately hold no parameters: if they owned `qkv` / `q_proj_cam` / ... PyTorch would prefix by attribute name and rename 390 of the 871 checkpoint keys. (`state_dict` duplicates a module aliased in two places rather than de-duplicating it, so aliasing isn't a way out either.) The compute is split; the ownership isn't. * `transformer_sana_wm_refiner.py` is self-contained -- nine `# Copied from` markers replace the `transformer_ltx2` import, at whole-function, whole-class and method-level granularity as each piece required. That costs +831 lines, since the copied block `__init__` drags in both attention processors. * The refiner's streaming self-attention is a real attention processor, `SanaWMLTX2RefinerAttnProcessor`, and the block calls `self.attn1(...)`. The processor swap happens in the model `__init__` after the blocks are built, because `attn1` is constructed inside the `# Copied from` body. * Inline the KV-cache mode strings, drop `plucker_proj`'s zero-init. State dicts unchanged: 871/871 and 3510/3510, ordered lists identical. Equivalence is exact -- max abs diff 0.0 -- against the real released checkpoint for the attention split, and across five refiner configurations driven through two AR blocks, with negative controls throughout. `check_copies` passes. GPU smoke unchanged at refined-latent mean 0.035024 / std 0.887467. * test(sana-wm): cover the GDN branch, which was silently dead in these tests `SanaWMTemporalShortConvolution` allocates its weight with `torch.zeros` and nothing initialises it any more -- the hand-written init came out earlier in this PR because `from_pretrained` overwrites it. That is true for real checkpoints, but a from-scratch model gets a zero `conv_k`, so the GDN key is identically zero and the whole branch returns zeros. Every test in this file builds a random model, so they were all passing while exercising none of the GDN maths. Two tests: one pins the zero-init as deliberate rather than accidental, the other fills the short convolutions with noise and asserts the output moves. Checked that the second one bites -- forcing the short convolution to return zeros makes it fail. * [SANA-WM] Address review: reuse diffusers embeddings, grad checkpointing, callbacks, offload fixes - Reuse Timesteps/TimestepEmbedding; drop unused cfg_embedder / y_embedding / model_max_length - Rename Mlp/PatchEmbed/CaptionEmbedder with SanaWM prefix - Support gradient checkpointing in SanaWMTransformer3DModel - Keep complex RoPE freqs out of buffers so Module.to(dtype) doesn't drop the imaginary part - Disable autocast around camera geometry (torch.cross) - Pipeline: callback_on_step_end, Literal output_type, _execution_device, maybe_free_model_hooks - Fix output docstring, bump copyright to 2026 - Port tests to new pytest mixins, add refiner transformer tests * [SANA-WM] Drop cam_utils and registration tests * [SANA-WM] Add return type annotations to forward methods --------- Co-authored-by: HaoyiZhu Co-authored-by: Claude Opus 4.7 Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: YiYi Xu Co-authored-by: github-actions[bot] Co-authored-by: Dhruv Nair Co-authored-by: Sayak Paul --- docs/source/en/_toctree.yml | 6 + .../models/sana_wm_refiner_transformer3d.md | 55 + .../en/api/models/sana_wm_transformer3d.md | 42 + docs/source/en/api/pipelines/sana_wm.md | 170 + scripts/convert_sana_wm_to_diffusers.py | 247 ++ src/diffusers/__init__.py | 10 + src/diffusers/models/__init__.py | 4 + src/diffusers/models/transformers/__init__.py | 2 + .../transformers/transformer_sana_wm.py | 3117 +++++++++++++++++ .../transformer_sana_wm_refiner.py | 1556 ++++++++ src/diffusers/pipelines/__init__.py | 10 + src/diffusers/pipelines/sana_wm/__init__.py | 49 + src/diffusers/pipelines/sana_wm/cam_utils.py | 321 ++ .../pipelines/sana_wm/image_processor.py | 67 + .../pipelines/sana_wm/pipeline_output.py | 44 + .../pipelines/sana_wm/pipeline_sana_wm.py | 659 ++++ src/diffusers/pipelines/sana_wm/refiner.py | 515 +++ src/diffusers/utils/dummy_pt_objects.py | 30 + .../dummy_torch_and_transformers_objects.py | 45 + .../test_models_transformer_sana_wm.py | 149 + ...test_models_transformer_sana_wm_refiner.py | 236 ++ tests/pipelines/sana_wm/__init__.py | 0 tests/pipelines/sana_wm/test_sana_wm.py | 250 ++ 23 files changed, 7584 insertions(+) create mode 100644 docs/source/en/api/models/sana_wm_refiner_transformer3d.md create mode 100644 docs/source/en/api/models/sana_wm_transformer3d.md create mode 100644 docs/source/en/api/pipelines/sana_wm.md create mode 100644 scripts/convert_sana_wm_to_diffusers.py create mode 100644 src/diffusers/models/transformers/transformer_sana_wm.py create mode 100644 src/diffusers/models/transformers/transformer_sana_wm_refiner.py create mode 100644 src/diffusers/pipelines/sana_wm/__init__.py create mode 100644 src/diffusers/pipelines/sana_wm/cam_utils.py create mode 100644 src/diffusers/pipelines/sana_wm/image_processor.py create mode 100644 src/diffusers/pipelines/sana_wm/pipeline_output.py create mode 100644 src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py create mode 100644 src/diffusers/pipelines/sana_wm/refiner.py create mode 100644 tests/models/transformers/test_models_transformer_sana_wm.py create mode 100644 tests/models/transformers/test_models_transformer_sana_wm_refiner.py create mode 100644 tests/pipelines/sana_wm/__init__.py create mode 100644 tests/pipelines/sana_wm/test_sana_wm.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 3abfb9052c..ab98a52430 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -411,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 @@ -653,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/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/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/models/__init__.py b/src/diffusers/models/__init__.py index 476b80f94d..cbe283b739 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -153,6 +153,8 @@ if is_torch_available(): _import_structure["transformers.transformer_qwenimage"] = ["QwenImageTransformer2DModel"] _import_structure["transformers.transformer_qwenimage21"] = ["QwenImage21Transformer2DModel"] _import_structure["transformers.transformer_sana_video"] = ["SanaVideoTransformer3DModel"] + _import_structure["transformers.transformer_sana_wm"] = ["SanaWMTransformer3DModel"] + _import_structure["transformers.transformer_sana_wm_refiner"] = ["SanaWMLTX2RefinerTransformer3DModel"] _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"] _import_structure["transformers.transformer_skyreels_v2"] = ["SkyReelsV2Transformer3DModel"] _import_structure["transformers.transformer_stable_audio3"] = ["StableAudio3DiTModel"] @@ -306,6 +308,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: QwenImageTransformer2DModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, + SanaWMLTX2RefinerTransformer3DModel, + SanaWMTransformer3DModel, SD3Transformer2DModel, SkyReelsV2Transformer3DModel, StableAudio3DiTModel, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index d29d49d168..079c4a37a2 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -64,6 +64,8 @@ if is_torch_available(): from .transformer_qwenimage import QwenImageTransformer2DModel from .transformer_qwenimage21 import QwenImage21Transformer2DModel from .transformer_sana_video import SanaVideoTransformer3DModel + from .transformer_sana_wm import SanaWMTransformer3DModel + from .transformer_sana_wm_refiner import SanaWMLTX2RefinerTransformer3DModel from .transformer_sd3 import SD3Transformer2DModel from .transformer_skyreels_v2 import SkyReelsV2Transformer3DModel from .transformer_stable_audio3 import StableAudio3DiTModel diff --git a/src/diffusers/models/transformers/transformer_sana_wm.py b/src/diffusers/models/transformers/transformer_sana_wm.py new file mode 100644 index 0000000000..7a38630973 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_sana_wm.py @@ -0,0 +1,3117 @@ +# 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. +# +# This file is modified from https://github.com/PixArt-alpha/PixArt-sigma + +from __future__ import annotations + +import math +from copy import deepcopy +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import logging +from ..activations import get_activation +from ..attention import AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..embeddings import TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import RMSNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class SanaWMMlp(nn.Module): + """Two-layer feed-forward block (`fc1` -> activation -> `fc2`).""" + + def __init__( + self, + in_features: int, + hidden_features: int | None = None, + out_features: int | None = None, + act_layer: type[nn.Module] = nn.GELU, + bias: bool = True, + drop: float = 0.0, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.drop1 = nn.Dropout(drop) + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop2 = nn.Dropout(drop) + + def forward(self, hidden_states: torch.Tensor, HW: tuple[int, int] | None = None) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.drop1(hidden_states) + hidden_states = self.fc2(hidden_states) + hidden_states = self.drop2(hidden_states) + return hidden_states + + +class SanaWMTemporalShortConvolution(nn.Module): + """Depthwise short convolution over the temporal axis, run in both directions. + + SANA-WM's GDN attention applies a short depthwise conv to Q/K/V before the linear-attention kernel. The reference + implementation used the *causal* `fla.modules.ShortConvolution` layer (with `activation=None`); this is a + self-contained PyTorch implementation -- so the model needs no `fla-core` dependency and can be built on any device + -- that runs the causal kernel forwards and backwards to obtain the non-causal filter the bidirectional model + needs. + + A causal depthwise Conv1d with kernel ``[w_0, w_1, ..., w_{k-1}]`` computes at time *t*: + + ``y_fwd[t] = w_0 * x[t-k+1] + ... + w_{k-1} * x[t]`` + + Running the same kernel on the time-flipped input and flipping back gives: + + ``y_bwd[t] = w_{k-1} * x[t] + ... + w_0 * x[t+k-1]`` + + Both passes include the current timestep ``x[t]`` with the center weight ``w_{k-1}``. To avoid double-counting one + copy of the center contribution is subtracted: + + ``y = y_fwd + y_bwd - w_{k-1} * x`` + + The result is a symmetric temporal filter where every position in the window ``[t-k+1, t+k-1]`` is counted exactly + once. + + Args: + hidden_size (`int`): Number of channels (the conv is depthwise, one group per channel). + kernel_size (`int`): Temporal kernel width. + bias (`bool`, defaults to `False`): Whether to add a per-channel bias. + """ + + def __init__(self, hidden_size: int, kernel_size: int, bias: bool = False) -> None: + super().__init__() + self.hidden_size = hidden_size + self.kernel_size = kernel_size + # Same parameter layout as the reference implementation: (C, 1, K). + self.weight = nn.Parameter(torch.zeros(hidden_size, 1, kernel_size)) + self.bias = nn.Parameter(torch.zeros(hidden_size)) if bias else None + + def _causal_conv(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Depthwise causal conv over `(batch, seq_len, hidden_size)` inputs.""" + seq_len = hidden_states.shape[1] + # Left-pad by (K - 1) and drop the tail so output[t] only sees inputs <= t. + hidden_states = F.conv1d( + hidden_states.transpose(1, 2), + self.weight.to(hidden_states.dtype), + None if self.bias is None else self.bias.to(hidden_states.dtype), + groups=self.hidden_size, + padding=self.kernel_size - 1, + )[..., :seq_len] + return hidden_states.transpose(1, 2) + + def forward(self, hidden_states: torch.Tensor, num_frames: int) -> torch.Tensor: + """Apply the bidirectional conv along the temporal axis, with the spatial axis merged into the batch. + + Args: + hidden_states (`torch.Tensor`): Input of shape `(batch, num_frames * spatial_size, hidden_size)`. + num_frames (`int`): Number of frames the sequence axis is split into. + + Returns: + `torch.Tensor`: Tensor of the same shape and dtype as `hidden_states`. + """ + batch_size, seq_len, channels = hidden_states.shape + spatial_size = seq_len // num_frames + dtype_in = hidden_states.dtype + + # (B, T * S, C) -> (B * S, T, C). The causal conv backward is not reliable on the non-contiguous + # strided layout this permutation produces, hence the explicit `contiguous()`. + hidden_states = ( + hidden_states.reshape(batch_size, num_frames, spatial_size, channels) + .permute(0, 2, 1, 3) + .contiguous() + .reshape(batch_size * spatial_size, num_frames, channels) + ) + + causal_forward = self._causal_conv(hidden_states) + causal_backward = self._causal_conv(hidden_states.flip(1)).flip(1) + + # Subtract the shared center tap (last weight of the causal kernel). Weight shape: (channels, 1, kernel_size), + # so the last element along dim=-1 is the weight applied to x[t]. + center_term = hidden_states * self.weight[:, 0, -1].unsqueeze(0).unsqueeze(0) + + hidden_states = causal_forward + causal_backward - center_term + if hidden_states.dtype != dtype_in: + hidden_states = hidden_states.to(dtype_in) + + return ( + hidden_states.reshape(batch_size, spatial_size, num_frames, channels) + .permute(0, 2, 1, 3) + .reshape(batch_size, seq_len, channels) + ) + + +# Safe element-count threshold for a single conv call: PyTorch's 2D conv kernels (both cuDNN and the ATEN fallback) +# use 32-bit indexing internally, so very large ``(batch * frames, channels, height, width)`` inputs (e.g. minute-scale +# video at default CFG) can overflow. Empirically a single call up to ~1B elements is safe; above that we split along +# the leading dim. Set so short videos stay on the original fused path (no chunking, no overhead). +_INT32_SAFE_CONV_ELEMENTS = 1 << 30 # 1,073,741,824 + + +class SanaWMConvLayer(nn.Module): + """2D convolution with an optional activation. + + Wraps the convolution in a ``conv`` submodule to keep the checkpoint's parameter names + (``mlp.inverted_conv.conv.weight``, ...) unchanged. + + Args: + in_dim (`int`): Input channels. + out_dim (`int`): Output channels. + kernel_size (`int`, defaults to 3): Spatial kernel size (odd, so ``same`` padding is exact). + groups (`int`, defaults to 1): Convolution groups. + use_bias (`bool`, defaults to `False`): Whether the convolution has a bias. + act (`str`, *optional*): Activation name resolved through + [`~models.activations.get_activation`], or `None` for no activation. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + kernel_size: int = 3, + groups: int = 1, + use_bias: bool = False, + act: Optional[str] = None, + ) -> None: + super().__init__() + self.conv = nn.Conv2d( + in_dim, + out_dim, + kernel_size=(kernel_size, kernel_size), + padding=kernel_size // 2, + groups=groups, + bias=use_bias, + ) + self.act = get_activation(act) if act is not None else None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv(hidden_states) + if self.act is not None: + hidden_states = self.act(hidden_states) + return hidden_states + + +class GLUMBConvTemp(nn.Module): + """SANA-WM feed-forward block: a gated inverted-bottleneck conv over space plus a residual temporal conv. + + Args: + in_features (`int`): Input channels. + hidden_features (`int`): Width of the inverted bottleneck (doubled internally for the GLU gate). + out_feature (`int`, *optional*): Output channels, defaults to `in_features`. + kernel_size (`int`, defaults to 3): Spatial kernel size of the depthwise convolution. + use_bias (`tuple[bool, bool, bool]`, defaults to `(False, False, False)`): Bias flag per convolution. + act (`tuple`, defaults to `("silu", "silu", None)`): Activation for the inverted conv, the GLU gate and the + point conv respectively; `None` means no activation. + t_kernel_size (`int`, defaults to 3): Temporal kernel size of the residual temporal convolution. + """ + + def __init__( + self, + in_features: int, + hidden_features: int, + out_feature: Optional[int] = None, + kernel_size: int = 3, + use_bias: Tuple[bool, bool, bool] = (False, False, False), + act: Tuple[Optional[str], Optional[str], Optional[str]] = ("silu", "silu", None), + t_kernel_size: int = 3, + ) -> None: + super().__init__() + out_feature = out_feature or in_features + + self.glu_act = get_activation(act[1]) + self.inverted_conv = SanaWMConvLayer( + in_features, hidden_features * 2, kernel_size=1, use_bias=use_bias[0], act=act[0] + ) + self.depth_conv = SanaWMConvLayer( + hidden_features * 2, + hidden_features * 2, + kernel_size=kernel_size, + groups=hidden_features * 2, + use_bias=use_bias[1], + act=None, + ) + self.point_conv = SanaWMConvLayer( + hidden_features, out_feature, kernel_size=1, use_bias=use_bias[2], act=act[2] + ) + self.t_conv = nn.Conv2d( + out_feature, + out_feature, + kernel_size=(t_kernel_size, 1), + padding=(t_kernel_size // 2, 0), + bias=False, + ) + + def forward(self, hidden_states: torch.Tensor, HW: Tuple[int, int, int]) -> torch.Tensor: + batch_size, seq_len, channels = hidden_states.shape + num_frames, height, width = HW + hidden_states = hidden_states.reshape(batch_size * num_frames, height, width, channels).permute(0, 3, 1, 2) + + # Split the leading dim so each conv launch stays under PyTorch's 32-bit indexing limit (no-op for short clips). + rows_per_call = max(1, _INT32_SAFE_CONV_ELEMENTS // (self.inverted_conv.conv.out_channels * height * width)) + spatial_chunks = [] + for start in range(0, hidden_states.shape[0], rows_per_call): + chunk = self.inverted_conv(hidden_states[start : start + rows_per_call]) + chunk = self.depth_conv(chunk) + value, gate = torch.chunk(chunk, 2, dim=1) + spatial_chunks.append(self.point_conv(value * self.glu_act(gate))) + hidden_states = spatial_chunks[0] if len(spatial_chunks) == 1 else torch.cat(spatial_chunks, dim=0) + + # Residual temporal aggregation over the frame axis. + hidden_states = hidden_states.view(batch_size, num_frames, channels, height * width).permute(0, 2, 1, 3) + hidden_states = hidden_states + self.t_conv(hidden_states) + return hidden_states.permute(0, 2, 3, 1).reshape(batch_size, seq_len, channels) + + +class SanaWMCrossAttnProcessor: + """Cross-attention from image queries to the text condition.""" + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "MultiHeadCrossAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size, _, channels = hidden_states.shape + + query = attn.q_linear(hidden_states) + key, value = attn.kv_linear(encoder_hidden_states).view(batch_size, -1, 2, channels).unbind(2) + + query = attn.q_norm(query).view(batch_size, -1, attn.heads, attn.head_dim) + key = attn.k_norm(key).view(batch_size, -1, attn.heads, attn.head_dim) + value = value.view(batch_size, -1, attn.heads, attn.head_dim) + + # A boolean mask (rather than an additive float one) keeps the varlen backends usable. + if attention_mask is not None and attention_mask.ndim == 2: + attention_mask = attention_mask.bool()[:, None, None, :] + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.reshape(batch_size, -1, channels).type_as(query) + return attn.proj(hidden_states) + + +class MultiHeadCrossAttention(torch.nn.Module, AttentionModuleMixin): + _default_processor_cls = SanaWMCrossAttnProcessor + _available_processors = [SanaWMCrossAttnProcessor] + + def __init__(self, d_model, num_heads, qk_norm=False, processor=None, **block_kwargs): + super().__init__() + if not (d_model % num_heads == 0): + raise ValueError("d_model must be divisible by num_heads") + + self.d_model = d_model + self.heads = num_heads + self.head_dim = d_model // num_heads + self.inner_dim = d_model + + self.q_linear = nn.Linear(d_model, d_model) + self.kv_linear = nn.Linear(d_model, d_model * 2) + self.proj = nn.Linear(d_model, d_model) + if qk_norm: + self.q_norm = RMSNorm(d_model, eps=1e-6) + self.k_norm = RMSNorm(d_model, eps=1e-6) + else: + self.q_norm = nn.Identity() + self.k_norm = nn.Identity() + + self.set_processor(processor if processor is not None else self._default_processor_cls()) + + def forward(self, x, cond, mask=None) -> 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/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/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/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/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.""" From 916a5791ed5f29d966200d2eee53f06862ec4434 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20Graci=C3=A1?= <145584850+Sergasgr@users.noreply.github.com> Date: Tue, 13 Oct 2026 05:14:36 +0200 Subject: [PATCH 12/13] [Z-Image] Apply the LoRA scale passed via joint_attention_kwargs (#15015) Fixes #15007 Co-authored-by: Sayak Paul --- src/diffusers/models/transformers/transformer_z_image.py | 6 ++++++ src/diffusers/pipelines/z_image/pipeline_z_image.py | 6 +++++- .../pipelines/z_image/pipeline_z_image_controlnet.py | 1 + .../z_image/pipeline_z_image_controlnet_inpaint.py | 1 + .../pipelines/z_image/pipeline_z_image_img2img.py | 1 + .../pipelines/z_image/pipeline_z_image_inpaint.py | 1 + src/diffusers/pipelines/z_image/pipeline_z_image_omni.py | 1 + tests/pipelines/z_image/test_z_image.py | 8 -------- 8 files changed, 16 insertions(+), 9 deletions(-) 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/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/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.""" From 08d69774f08f19ae903f97bc393a8abdeafc98aa Mon Sep 17 00:00:00 2001 From: Javier De Jesus Date: Wed, 14 Oct 2026 05:34:11 +0200 Subject: [PATCH 13/13] [utils] preserve method metadata in apply_forward_hook (#13803) * [utils] preserve method metadata in apply_forward_hook Add @functools.wraps to the inner wrapper so it inherits __name__, __doc__, __module__ and the signature of the wrapped method. Without this, methods decorated with @apply_forward_hook (e.g. AutoencoderKL's encode and decode) lose their docstring and signature, so doc-builder cannot render them through the [[autodoc]] directive and they show up as "wrapper" with a generic (*args, **kwargs) signature on the live API docs. The same affects VQModel, ConsistencyDecoderVAE and the other autoencoders that use the decorator. Add regression tests asserting that the wrapped method exposes the right metadata (name, docstring, signature including default values, __wrapped__) and that the underlying offload hook still fires through the wrapper. Fixes #6575 * Trim apply_forward_hook test to a single assertion --- src/diffusers/utils/accelerate_utils.py | 3 +++ tests/others/test_utils.py | 12 ++++++++++++ 2 files changed, 15 insertions(+) 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/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."""