AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity

Merge branch 'main' into remove-lora-utils-tests

Browse Source
This commit is contained in:
Sayak Paul
2026-10-14 10:48:42 +05:30
committed by GitHub
parent aee9bdd078 08d69774f0
commit c5feabd80c
66 changed files with 8790 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
.ai/skills/custom-blocks/SKILL.md
View File

@@ -57,7 +57,8 @@ diffusers-cli custom_blocks [--block_module_name <file.py>] [--block_class_name
### What it does
1. **AST scan**: parses `<file>` 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 <file>: [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 `<ClassName>()` 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__`.

49
.github/workflows/pr_labeler.yml vendored
View File

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

8
docs/source/en/_toctree.yml
View File

@@ -114,6 +114,8 @@
title: Intel Gaudi
- local: optimization/neuron
title: AWS Neuron
- local: optimization/tpu
title: TPU
title: Hardware-specific acceleration
- isExpanded: false
sections:
@@ -409,6 +411,10 @@
title: SanaTransformer2DModel
- local: api/models/sana_video_transformer3d
title: SanaVideoTransformer3DModel
- local: api/models/sana_wm_refiner_transformer3d
title: SanaWMLTX2RefinerTransformer3DModel
- local: api/models/sana_wm_transformer3d
title: SanaWMTransformer3DModel
- local: api/models/sd3_transformer2d
title: SD3Transformer2DModel
- local: api/models/skyreels_v2_transformer_3d
@@ -651,6 +657,8 @@
title: Sana Sprint
- local: api/pipelines/sana_video
title: Sana Video
- local: api/pipelines/sana_wm
title: SANA-WM
- local: api/pipelines/shap_e
title: Shap-E
- local: api/pipelines/stable_cascade

55
docs/source/en/api/models/sana_wm_refiner_transformer3d.md Normal file
View File

@@ -0,0 +1,55 @@
<!-- 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. -->
# 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

42
docs/source/en/api/models/sana_wm_transformer3d.md Normal file
View File

@@ -0,0 +1,42 @@
<!-- 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. -->
# 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

170
docs/source/en/api/pipelines/sana_wm.md Normal file
View File

@@ -0,0 +1,170 @@
<!-- 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. -->
# 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

157
docs/source/en/optimization/tpu.md Normal file
View File

@@ -0,0 +1,157 @@
<!--Copyright 2026 The HuggingFace Team. 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.
-->
# 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.

5
docs/source/en/using-diffusers/cli.md
View File

@@ -290,8 +290,9 @@ hf sandbox kill <id>
## `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

247
scripts/convert_sana_wm_to_diffusers.py Normal file
View File

@@ -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)``:
<output_dir>/
├── 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: `<dst>-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 <src>/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()

2
setup.py
View File

@@ -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",

10
src/diffusers/__init__.py
View File

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

14
src/diffusers/commands/custom_blocks.py
View File

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

2
src/diffusers/dependency_versions_table.py
View File

@@ -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",

2
src/diffusers/hooks/tensor_parallel.py
View File

@@ -22,7 +22,7 @@ from ..utils import get_logger, is_peft_available
logger = get_logger(__name__) # pylint: disable=invalid-name
_SUPPORTED_TP_DEVICES = ("cuda", "neuron")
_SUPPORTED_TP_DEVICES = ("cuda", "neuron", "tpu", "xpu")
class PackedColwiseParallel:

11
src/diffusers/image_processor.py
View File

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

2
src/diffusers/loaders/lora_base.py
View File

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

154
src/diffusers/loaders/lora_conversion_utils.py
View File

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

32
src/diffusers/loaders/lora_pipeline.py
View File

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

116
src/diffusers/loaders/peft.py
View File

@@ -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
# <Unsafe code
# We can be sure that the following works as it just sets attention processors, lora layers and puts all in the same dtype
@@ -304,13 +318,13 @@ class PeftAdapterMixin:
try:
if hotswap:
state_dict = map_state_dict_for_hotswap(state_dict)
check_hotswap_configs_compatible(self.peft_config[adapter_name], lora_config)
check_hotswap_configs_compatible(self.peft_config[adapter_name], adapter_config)
try:
hotswap_adapter_from_state_dict(
model=self,
state_dict=state_dict,
adapter_name=adapter_name,
config=lora_config,
config=adapter_config,
)
except Exception as e:
logger.error(f"Hotswapping {adapter_name} was unsuccessful with the following error: \n{e}")
@@ -320,7 +334,7 @@ class PeftAdapterMixin:
incompatible_keys = None
else:
inject_adapter_in_model(
lora_config, self, adapter_name=adapter_name, state_dict=state_dict, **peft_kwargs
adapter_config, self, adapter_name=adapter_name, state_dict=state_dict, **peft_kwargs
)
incompatible_keys = set_peft_model_state_dict(self, state_dict, adapter_name, **peft_kwargs)
@@ -331,7 +345,7 @@ class PeftAdapterMixin:
# - before the model is compiled and the 2nd adapter is being hotswapped in
# Therefore, it needs to be called here
prepare_model_for_compiled_hotswap(
self, config=lora_config, **self._prepare_lora_hotswap_kwargs
self, config=adapter_config, **self._prepare_lora_hotswap_kwargs
)
# We only want to call prepare_model_for_compiled_hotswap once
self._prepare_lora_hotswap_kwargs = None

4
src/diffusers/models/__init__.py
View File

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

2
src/diffusers/models/_modeling_parallel.py
View File

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

2
src/diffusers/models/transformers/__init__.py
View File

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

3117
src/diffusers/models/transformers/transformer_sana_wm.py Normal file
View File

File diff suppressed because it is too large Load Diff

1556
src/diffusers/models/transformers/transformer_sana_wm_refiner.py Normal file
View File

File diff suppressed because it is too large Load Diff

6
src/diffusers/models/transformers/transformer_z_image.py
View File

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

8
src/diffusers/modular_pipelines/anima/encoders.py
View File

@@ -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.",
),

2
src/diffusers/modular_pipelines/anima/modular_pipeline.py
View File

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

10
src/diffusers/pipelines/__init__.py
View File

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

49
src/diffusers/pipelines/sana_wm/__init__.py Normal file
View File

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

321
src/diffusers/pipelines/sana_wm/cam_utils.py Normal file
View File

@@ -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>-<duration>'.")
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)

67
src/diffusers/pipelines/sana_wm/image_processor.py Normal file
View File

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

44
src/diffusers/pipelines/sana_wm/pipeline_output.py Normal file
View File

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

659
src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py Normal file
View File

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

515
src/diffusers/pipelines/sana_wm/refiner.py Normal file
View File

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

5
src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
View File

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

6
src/diffusers/pipelines/z_image/pipeline_z_image.py
View File

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

1
src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py
View File

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

1
src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py
View File

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

1
src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py
View File

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

1
src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py
View File

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

1
src/diffusers/pipelines/z_image/pipeline_z_image_omni.py
View File

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

1
src/diffusers/utils/__init__.py
View File

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

3
src/diffusers/utils/accelerate_utils.py
View File

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

30
src/diffusers/utils/dummy_pt_objects.py
View File

@@ -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"]

45
src/diffusers/utils/dummy_torch_and_transformers_objects.py
View File

@@ -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"]

5
src/diffusers/utils/import_utils.py
View File

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

124
src/diffusers/utils/peft_utils.py
View File

@@ -368,6 +368,126 @@ def check_peft_version(min_version: str) -> None:
)
def _maybe_fuse_qkv_projections_for_lokr(model, state_dict) -> None:
"""
Fuse the model's QKV projections when a peft-format LoKr state dict targets fused `to_qkv` / `to_added_qkv`
projections that the model does not have yet.
BFL-format Flux2 LoKr checkpoints apply LoKr to the fused QKV projections. Unlike a LoRA delta, a Kronecker product
delta over the fused projection cannot be split exactly into separate Q/K/V factors, so the model's projections are
fused instead and the adapter maps 1:1.
"""
fused_targets = {
module
for module in (k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k)
if module.rsplit(".", 1)[-1] in ("to_qkv", "to_added_qkv")
}
named_modules = dict(model.named_modules())
if all(module in named_modules for module in fused_targets) or not hasattr(model, "fuse_qkv_projections"):
return
if getattr(model, "is_quantized", False):
raise ValueError(
"This LoKr checkpoint targets fused QKV projections. Fusing concatenates the Q/K/V weights into a new "
"`nn.Linear`, which is not possible with quantized weights. Please load the transformer without "
"quantization."
)
# Fusing replaces to_q/to_k/to_v (and the add_*_proj) with a single projection, which would orphan any adapter
# already injected on the unfused ones.
from peft.tuners.tuners_utils import BaseTunerLayer
unfused_projections = {"to_q", "to_k", "to_v", "add_q_proj", "add_k_proj", "add_v_proj"}
adapted = [
name
for name, module in named_modules.items()
if isinstance(module, BaseTunerLayer) and name.rsplit(".", 1)[-1] in unfused_projections
]
if adapted:
raise ValueError(
"This LoKr checkpoint targets fused QKV projections, but an adapter is already loaded on the unfused "
f"projections (e.g. `{adapted[0]}`). Unload it with `unload_lora_weights()` before loading this checkpoint."
)
logger.info("The LoKr checkpoint targets fused QKV projections; calling `fuse_qkv_projections()` on the model.")
model.fuse_qkv_projections()
def _create_lokr_config(state_dict, metadata):
"""
Create a `LoKrConfig` from a peft-format LoKr state dict (keys like `{module}.lokr_w1`).
Without metadata, the config is inferred from the tensor shapes. The checkpoint alpha is expected to be already
baked into the weights by the state dict conversion, so `alpha` is set equal to the rank (runtime scaling 1.0).
peft re-derives each module's Kronecker factorization from `decompose_factor` and only creates rank-decomposed
factors when the rank is small compared to the factorized dimensions, so for modules whose checkpoint factors are
full matrices the rank is set to `max(lokr_w2.shape)` to make peft create full matrices as well.
"""
from peft import LoKrConfig
from peft.tuners.lokr.layer import factorization
if metadata is not None:
try:
return LoKrConfig(**metadata)
except TypeError as e:
raise TypeError("`LoKrConfig` class could not be instantiated.") from e
modules = sorted({k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k})
# Reconstruct each module's factorized dimensions, (out_l, out_k) x (in_m, in_n), from the checkpoint. A factor
# is either stored as a full matrix (`lokr_w1`) or rank-decomposed into `lokr_w1_a @ lokr_w1_b` (same for w2).
factorizations = {}
rank_dict = {}
for module in modules:
w1, w1_a, w1_b = (state_dict.get(f"{module}.lokr_w1{s}") for s in ("", "_a", "_b"))
w2, w2_a, w2_b = (state_dict.get(f"{module}.lokr_w2{s}") for s in ("", "_a", "_b"))
out_l, in_m = w1.shape if w1 is not None else (w1_a.shape[0], w1_b.shape[1])
out_k, in_n = w2.shape if w2 is not None else (w2_a.shape[0], w2_b.shape[1])
factorizations[module] = ((out_l, out_k), (in_m, in_n))
if w2_a is not None:
rank_dict[module] = w2_a.shape[1]
elif w1_a is not None:
rank_dict[module] = w1_a.shape[1]
else:
rank_dict[module] = max(w2.shape)
# Find the `decompose_factor` under which peft reproduces the checkpoint factorizations. A fixed factor shows up
# as the left dimension of the modules it divides (modules it does not divide fall back to a near-square
# factorization, like with factor -1), so every observed left dimension is a candidate.
left_dims = {dims[0][0] for dims in factorizations.values()}
decompose_factor = None
for candidate in sorted(left_dims) + [-1]:
if all(
factorization(out_l * out_k, candidate) == (out_l, out_k)
and factorization(in_m * in_n, candidate) == (in_m, in_n)
for (out_l, out_k), (in_m, in_n) in factorizations.values()
):
decompose_factor = candidate
break
if decompose_factor is None:
raise ValueError(
"Could not infer a `decompose_factor` that reproduces the Kronecker factorizations of this LoKr "
"state dict. Please open an issue: https://github.com/huggingface/diffusers/issues/new"
)
r = collections.Counter(rank_dict.values()).most_common(1)[0][0]
rank_pattern = {k: v for k, v in rank_dict.items() if v != r}
lokr_config_kwargs = {
"r": r,
"alpha": r,
"rank_pattern": rank_pattern,
"alpha_pattern": dict(rank_pattern),
"target_modules": modules,
"decompose_both": any(".lokr_w1_a" in k for k in state_dict),
"decompose_factor": decompose_factor,
}
try:
return LoKrConfig(**lokr_config_kwargs)
except TypeError as e:
raise TypeError("`LoKrConfig` class could not be instantiated.") from e
def _create_lora_config(
state_dict, network_alphas, metadata, rank_pattern_dict, is_unet=True, model_state_dict=None, adapter_name=None
):
@@ -397,7 +517,7 @@ def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name: str) -> None
# Check only for unexpected keys.
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
if unexpected_keys:
lora_unexpected_keys = [k for k in unexpected_keys if ".lora_" in k]
lora_unexpected_keys = [k for k in unexpected_keys if ".lora_" in k or "lokr_" in k]
if lora_unexpected_keys:
warn_msg = (
f"Loading adapter weights from state_dict led to unexpected keys found in the model:"
@@ -407,7 +527,7 @@ def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name: str) -> None
# Filter missing keys specific to the current adapter.
missing_keys = getattr(incompatible_keys, "missing_keys", None)
if missing_keys:
lora_missing_keys = [k for k in missing_keys if ".lora_" in k and adapter_name in k]
lora_missing_keys = [k for k in missing_keys if (".lora_" in k or "lokr_" in k) and adapter_name in k]
if lora_missing_keys:
warn_msg += (
f"Loading adapter weights from state_dict led to missing keys in the model:"

4
tests/models/testing_utils/__init__.py
View File

@@ -17,12 +17,14 @@ from .cache import (
from .common import BaseModelTesterConfig, ModelTesterMixin
from .compile import TorchCompileTesterMixin
from .ip_adapter import IPAdapterTesterMixin
from .lokr import LoKrTesterMixin
from .lora import LoraHotSwappingForModelTesterMixin, LoraTesterMixin
from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin
from .parallelism import (
ContextParallelAttentionBackendsTesterMixin,
ContextParallelTesterMixin,
TensorParallelTesterMixin,
TensorParallelTPUTesterMixin,
)
from .quantization import (
AutoRoundCompileTesterMixin,
@@ -67,6 +69,7 @@ __all__ = [
"ContextParallelTesterMixin",
"ContextParallelAttentionBackendsTesterMixin",
"TensorParallelTesterMixin",
"TensorParallelTPUTesterMixin",
"CPUOffloadTesterMixin",
"FasterCacheConfigMixin",
"FasterCacheTesterMixin",
@@ -80,6 +83,7 @@ __all__ = [
"GroupOffloadTesterMixin",
"IPAdapterTesterMixin",
"LayerwiseCastingTesterMixin",
"LoKrTesterMixin",
"LoraHotSwappingForModelTesterMixin",
"LoraTesterMixin",
"MemoryTesterMixin",

139
tests/models/testing_utils/lokr.py Normal file
View File

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

135
tests/models/testing_utils/parallelism.py
View File

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

9
tests/models/transformers/test_models_transformer_flux.py
View File

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

123
tests/models/transformers/test_models_transformer_flux2.py
View File

@@ -17,9 +17,11 @@ import os
import subprocess
import sys
import pytest
import torch
from diffusers import Flux2Transformer2DModel
from diffusers.loaders.lora_pipeline import Flux2LoraLoaderMixin
from diffusers.models.transformers.transformer_flux2 import (
Flux2KVAttnProcessor,
Flux2KVCache,
@@ -36,17 +38,20 @@ from ..testing_utils import (
ContextParallelTesterMixin,
GGUFCompileTesterMixin,
GGUFTesterMixin,
LoKrTesterMixin,
LoraHotSwappingForModelTesterMixin,
LoraTesterMixin,
MemoryTesterMixin,
ModelTesterMixin,
SingleFileTesterMixin,
TensorParallelTesterMixin,
TensorParallelTPUTesterMixin,
TorchAoCompileTesterMixin,
TorchAoTesterMixin,
TorchCompileTesterMixin,
TrainingTesterMixin,
)
from ..testing_utils.lokr import check_lokr_deltas, make_lokr_factors
enable_full_determinism()
@@ -176,6 +181,14 @@ def make_neuron_tp_spec():
return Flux2Transformer2DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu")
class TestFlux2TransformerTensorParallelTPU(Flux2TransformerTesterConfig, TensorParallelTPUTesterMixin):
"""Tensor Parallel inference test for Flux2 Transformer on TPU."""
def get_init_dict(self):
# One head per chip.
return {**super().get_init_dict(), "num_attention_heads": self.tp_world_size}
@is_tensor_parallel
@require_torch_neuron
class TestFlux2TransformerTensorParallelNeuron:
@@ -202,6 +215,116 @@ class TestFlux2TransformerLoRA(Flux2TransformerTesterConfig, LoraTesterMixin):
"""LoRA adapter tests for Flux2 Transformer."""
class TestFlux2TransformerLoKr(Flux2TransformerTesterConfig, LoKrTesterMixin):
"""LoKr adapter tests for Flux2 Transformer, including the Flux2 LoKr checkpoint formats."""
# ai-toolkit stores a placeholder alpha for full-matrix factors, where LoKr applies no scaling.
placeholder_alpha = torch.tensor(9999220736.0)
def get_bfl_qkv_state_dict(self, model):
"""A BFL-format LoKr state dict on the fused QKV projections of the first double block."""
to_q = model.transformer_blocks[0].attn.to_q
state_dict, expected_deltas = {}, {}
for bfl_path, diffusers_path in [
("double_blocks.0.img_attn.qkv", "transformer_blocks.0.attn.to_qkv"),
("double_blocks.0.txt_attn.qkv", "transformer_blocks.0.attn.to_added_qkv"),
]:
factors, expected_deltas[diffusers_path] = make_lokr_factors(3 * to_q.out_features, to_q.in_features)
state_dict.update({f"diffusion_model.{bfl_path}.{k}": v for k, v in factors.items()})
state_dict[f"diffusion_model.{bfl_path}.alpha"] = self.placeholder_alpha
return state_dict, expected_deltas
@torch.no_grad()
def test_lokr_bfl_checkpoint(self):
# BFL checkpoints (e.g. ai-toolkit) apply LoKr to the fused QKV projections. A Kronecker product delta cannot
# be split exactly into Q/K/V, so loading fuses the model's projections and maps the adapter 1:1.
torch.manual_seed(0)
model = self.model_class(**self.get_init_dict()).eval().to(torch_device)
state_dict, expected_deltas = self.get_bfl_qkv_state_dict(model)
for bfl_path, diffusers_path in [
("single_blocks.0.linear1", "single_transformer_blocks.0.attn.to_qkv_mlp_proj"),
("double_blocks.0.img_attn.proj", "transformer_blocks.0.attn.to_out.0"),
("double_blocks.0.img_mlp.0", "transformer_blocks.0.ff.linear_in"),
]:
linear = model.get_submodule(diffusers_path)
factors, expected_deltas[diffusers_path] = make_lokr_factors(linear.out_features, linear.in_features)
state_dict.update({f"diffusion_model.{bfl_path}.{k}": v for k, v in factors.items()})
state_dict[f"diffusion_model.{bfl_path}.alpha"] = self.placeholder_alpha
converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict)
model.load_lora_adapter(converted, prefix="transformer", adapter_name="default")
assert model.transformer_blocks[0].attn.fused_projections
check_lokr_deltas(model, expected_deltas)
def test_lokr_fused_qkv_checkpoint_refuses_when_unfused_projections_are_adapted(self):
# Fusing would replace to_q and orphan the adapter already injected there.
from peft import LoraConfig
model = self.model_class(**self.get_init_dict()).eval().to(torch_device)
model.add_adapter(LoraConfig(r=2, target_modules=["to_q"]), adapter_name="lora")
state_dict, _ = self.get_bfl_qkv_state_dict(model)
converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict)
with pytest.raises(ValueError, match="already loaded on the unfused projections"):
model.load_lora_adapter(converted, prefix="transformer", adapter_name="lokr")
assert not model.transformer_blocks[0].attn.fused_projections
@torch.no_grad()
def test_lokr_lycoris_checkpoint(self):
# LyCORIS wraps the diffusers model and encodes module paths with underscores under a `lycoris_` prefix.
torch.manual_seed(0)
model = self.model_class(**self.get_init_dict()).eval().to(torch_device)
state_dict, expected_deltas = {}, {}
for diffusers_path in [
"single_transformer_blocks.0.attn.to_qkv_mlp_proj",
"transformer_blocks.0.attn.to_q",
"transformer_blocks.0.attn.to_out.0",
"transformer_blocks.0.ff.linear_in",
]:
linear = model.get_submodule(diffusers_path)
factors, expected_deltas[diffusers_path] = make_lokr_factors(linear.out_features, linear.in_features)
lycoris_path = "lycoris_" + diffusers_path.replace(".", "_")
state_dict.update({f"{lycoris_path}.{k}": v for k, v in factors.items()})
state_dict[f"{lycoris_path}.alpha"] = torch.tensor(16.0)
converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict)
model.load_lora_adapter(converted, prefix="transformer", adapter_name="default")
check_lokr_deltas(model, expected_deltas)
def test_lokr_lycoris_checkpoint_with_unknown_keys_raises(self):
state_dict = {
"lycoris_transformer_blocks_0_attn_to_q.lokr_w1": torch.randn(4, 4),
"lycoris_transformer_blocks_0_attn_norm_q.lokr_w1": torch.randn(4, 4),
}
with pytest.raises(ValueError, match="lycoris_transformer_blocks_0_attn_norm_q.lokr_w1"):
Flux2LoraLoaderMixin.lora_state_dict(state_dict)
@torch.no_grad()
def test_lokr_diffusers_names_checkpoint(self):
# Checkpoints that store the diffusers module paths directly, with alpha keys and no prefix (e.g. SimpleTuner,
# `bghira/flux2-klein-9b-distillation-lokr`). Alpha scales the rank-decomposed factors only.
torch.manual_seed(0)
model = self.model_class(**self.get_init_dict()).eval().to(torch_device)
rank, alpha = 2, 1.0
state_dict, expected_deltas = {}, {}
for diffusers_path, factor_rank in [
("single_transformer_blocks.0.attn.to_out", None),
("transformer_blocks.0.attn.to_k", rank),
]:
linear = model.get_submodule(diffusers_path)
factors, delta = make_lokr_factors(linear.out_features, linear.in_features, rank=factor_rank)
state_dict.update({f"{diffusers_path}.{k}": v for k, v in factors.items()})
state_dict[f"{diffusers_path}.alpha"] = torch.tensor(alpha)
expected_deltas[diffusers_path] = delta if factor_rank is None else (alpha / rank) * delta
converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict)
model.load_lora_adapter(converted, prefix="transformer", adapter_name="default")
check_lokr_deltas(model, expected_deltas)
class TestFlux2TransformerLoRAHotSwap(Flux2TransformerTesterConfig, LoraHotSwappingForModelTesterMixin):
"""LoRA hot-swapping tests for Flux2 Transformer."""

13
tests/models/transformers/test_models_transformer_qwenimage.py
View File

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

149
tests/models/transformers/test_models_transformer_sana_wm.py Normal file
View File

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

236
tests/models/transformers/test_models_transformer_sana_wm_refiner.py Normal file
View File

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

33
tests/models/transformers/test_models_transformer_z_image.py
View File

@@ -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."""

51
tests/modular_pipelines/anima/test_modular_pipeline_anima.py
View File

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

28
tests/others/test_cli_commands.py
View File

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

38
tests/others/test_image_processor.py
View File

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

12
tests/others/test_utils.py
View File

@@ -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."""

0
tests/pipelines/sana_wm/__init__.py Normal file
View File

250
tests/pipelines/sana_wm/test_sana_wm.py Normal file
View File

@@ -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."""

14
tests/pipelines/stable_diffusion/test_stable_diffusion.py
View File

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

12
tests/pipelines/testing_utils/lora.py
View File

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

8
tests/pipelines/z_image/test_z_image.py
View File

@@ -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."""

9
tests/testing_utils.py
View File

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