mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
required-input-no-default
7015 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
yiyixuxu
|
2bca34340d | fix numbering | ||
|
yiyixuxu
|
c95198eace | flip the drop the required | ||
|
yiyixuxu
|
53ac3af32b |
require num_inference_steps in blocks that run after set_timesteps
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
yiyixuxu
|
949e55ffde |
drop the InputParam unit test
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
yiyixuxu
|
5142760fc7 |
reword the required-plus-default gotcha
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
yiyixuxu
|
f5c8ae5f4f |
regenerate auto docstrings
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
yiyixuxu
|
a73e96786d |
rebase follow-ups: new qwenimage21 site, keep MiniMax-H3 steps required, unit test and guide note
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
yiyixuxu
|
3be086bc2c |
drop the default of a required InputParam, with a warning
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> |
||
Aayush D C Dangi
|
784ae091a8 |
[LoRA] apply .alpha to Qwen-Image lora_A/lora_B LoRAs (#15016)
* [LoRA] apply .alpha in lora_A/lora_B Qwen-Image LoRAs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * [LoRA] remove unnecessary Qwen-Image alpha tests --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Eliott C.
|
29d199aba8 |
Use HF Trusted Publisher token in update_metadata workflow (#15096)
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Sayak Paul
|
652c80bf84 |
[ci] skip bf16 precision test on T4. (#15091)
skip bf16 precision test on T4. |
||
Christopher
|
7e4383c0a2 |
[LTX-2] Add motion_speed to condition the positional time base apart from the playback frame rate (#14967)
* [LTX-2] Add `conditioning_frame_rate` to decouple the positional time base from the playback frame rate LTX-2 places video tokens on a time axis in seconds (`pixel_frame / fps` in the RoPE coordinates), derived from `frame_rate`. Adapters such as the LTX-2.5 Slow-Motion-Control LoRA are conditioned on a different rate (`fps / speed`) than the video plays at. Add an optional `conditioning_frame_rate` to the LTX-2 classic pipelines and modular blocks: it drives the video, keyframe and reference coordinates and the transformer's `fps`, while the audio length, the duration head and the returned video keep following `frame_rate`. `None` resolves to `frame_rate`, leaving current outputs unchanged. * [LTX-2] Replace `conditioning_frame_rate` with `motion_speed` `frame_rate` keeps one meaning, the playback rate; `motion_speed` (default 1.0, must be positive) is the speed of the motion relative to real time, and the positional time axis is conditioned on `frame_rate / motion_speed`. * [LTX-2] Add an end-to-end `motion_speed` example to the LTX-2 doc Image-to-video with the LTX-2.5 Slow-Motion-Control LoRA at motion_speed 1.0 and 0.2 (same seed and caption), plus a note that two-stage generation must pass the same motion_speed to both stages. |
||
|
Sayak Paul
|
0c7d56d6f4 |
[ci] fix misnamed test steps in push_tests.yml (#15094)
* [ci] fix misnamed test steps in push_tests.yml * remove xformers. |
||
|
Sayak Paul
|
a68262b0e9 |
[ci] use python 3.11 when using metadata. (#15093)
use python 3.11 when using metadata. |
||
Muhammad Usman Mateen
|
3421fed4a1 |
Fix model imports when torchao is older than 0.15 (#15095)
* Add a test for importing models with torchao older than 0.15 * Only import FqnToConfig from torchao 0.15 on torchao_quantizer.py imported FqnToConfig whenever torchao was installed, but it only exists from torchao 0.15. With an older torchao, importing any model failed, including on Google Colab. * Remove the torchao import test |
||
|
Sayak Paul
|
1ddaf8f4e7 |
[tests] xfail some kandinsky6 tests (#14977)
* xfail some kandinsky6 tests. * remove compile tests |
||
Kieran Didi
|
5378870419 |
Add RFD4-Proteina (#15088)
* Add RFD4-Proteina Adds the RFD4-Proteina pipeline with its transformer, autoencoder and per-modality schedulers, a LoRA loader mixin, a checkpoint conversion script, tests and documentation. * Address review: pass tensors instead of dicts, name the pipeline example inputs prepare_latents returns a tuple and the sampling loop works on named tensors. The pipeline output has bb_ca, local_latents and clean_* fields instead of dicts, and the callback tensors are named. The pipeline properties move below __init__, the transformer return annotation names its tuple, LoRA moves from the loaders page to the pipeline docs, and the pipeline tests spell out their helpers. * Remove the constructor options that are constant for the released checkpoint The transformer always recycles the trunk, updates the pair representation with triangle multiplication (without bias) and aligns the chirality feature, as the released checkpoint does. The converter now checks that a checkpoint has exactly these settings instead of carrying them as options. * Spell out the scheduler test arguments and the docs examples, default the autoencoder to its released features The scheduler test helper takes its options by name, and the usage and batching examples pass the five required inputs by name. The autoencoder's default encoder features now match the released checkpoint, and the decoder moves its allowed-residue table to the token types' device. * Follow the repository's test and docs checks Restore the LoRA mixin on the loaders page, which utils/check_support_list.py requires for every LoRA loader mixin, and keep the mention in the pipeline docs. Compose the attention and torch.compile testers on the models, drop the model-level LoRA tests that the first pull request should not carry, and make the pipeline's backbone-noising step public as a lifecycle method. * Keep the pipeline test helpers at the top of the test class * Drop the training-only detach calls and make the noise helpers static methods * Remove num_inference_steps and skip the unused encode for inverse folding The number of steps is fixed by the released recipe and is the length of recycle_depths. Inverse folding starts both modalities at time zero, so its latents are noise and the encoded input was only consuming random draws. --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
muhammadali-afridi
|
de988b7ee7 |
Add Bria fibo-2 (#14991)
* Add Bria fibo-2: BriaFibo2Pipeline and BriaFibo2Transformer2DModel
fibo-2 is Bria's 8.1B flow-matching DiT, built on the Z-Image block and
conditioned on Qwen3-VL hidden states two ways: a Perceiver resampler turns
them into gist tokens that share the stream with the image tokens, and five
blocks add gated cross-attention to their own bundle of Qwen3-VL layers. It
uses the Flux 2 VAE.
One pipeline class covers both of the model's jobs and both checkpoints:
text-to-image generation, instruction editing of up to five images with an
optional mask, the turbo checkpoint (4 steps, no guidance) and the guided one.
- BriaFibo2Transformer2DModel, with images to edit passed as context latents
- BriaFibo2Pipeline, including:
- captions cleaned and serialized the way training did, for generation and
edits, in every caption form fibo-2 was trained with: a full caption, a
short description, or a short description with its text_render list
- images of the caption's text, which Qwen3-VL sees beside the prompt:
drawn the way training drew them, by default (`render_text_images`), or
passed as `text_images`
- masked edits, the region greyed out through a binary mask as in training
- a conversion script from the original checkpoint layout
- model and pipeline tests, and docs
Verified against Bria's reference implementation on an A100 and an H100.
* Apply RoPE with explicit casting instead of an autocast block
The rotation already runs in float32 and casts back to the input's dtype, and
none of its ops are autocast ops, so the autocast block changed nothing.
* Mark the full-graph compile tests as expected failures instead of skipping them
Both fail with torch._dynamo.exc.Unsupported, on the torch.device context in
which RopeEmbedder, copied from Z-Image, builds its frequency table. Strict, so
they show up once Dynamo can trace it.
* Test the pipeline through its public call only
Removes the tests of private prompt and text-image helpers. The caption
cleaning test now checks that a caption renders like its trained form, and
the text images test draws its strips with the public draw_text_image.
* Annotate the transformer's tuple return as tuple[torch.Tensor]
* Make _get_qwen_prompt_embeds a staticmethod that takes the text encoder and processor
The chat template, the sequence length and the Qwen3-VL layers become
arguments, defaulting to module constants; encode_prompt passes the
pipeline's own. The edit and text prompt helpers become classmethods.
Embeddings are unchanged.
* Point to the release checkpoints, briaai/Fibo-2.0-dev and briaai/Fibo-2.0-turbo
The docs and the example load Fibo-2.0-dev and show the turbo checkpoint's
4-step edits without guidance. Text-image fonts download from Fibo-2.0-dev,
and the conversion script's usage converts a local original checkpoint.
* Parse prompts in one method, take fonts as an argument, and size and mask images in BriaFibo2ImageProcessor
- `_parse_prompt` replaces `_clean_caption`, `_normalize_caption`, `_split_edit_prompt`, `_get_generation_prompt` and
`_get_caption`. `_text_render_entries`, `_get_text_prompt` and `_draw_text_images` are inlined.
- The pipeline no longer fetches fonts: `__call__` takes `fonts`, and the docs show how to get the ones training
used from briaai/Fibo-2.0-dev. A text that needs them warns without them.
- `BriaFibo2ImageProcessor` holds the trained sizes, the cover-and-crop and the masking. The masking stays before the
crop, at the image's own size, where training greyed the images.
The prompts and images Qwen3-VL sees, the VAE's inputs and the outputs are byte-identical to before.
* Prepare the images to edit and the mask in BriaFibo2ImageProcessor.preprocess, as Flux3's image processor does
* Point to the turbo checkpoint's new name, briaai/Fibo-2.0-edit-turbo
* Show a working example for each use in the docs, and call the model Fibo 2.0, as its model cards do
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com>
|
||
|
Sayak Paul
|
a5eb4c1364 | Fix PyArrow compatibility in CPU CI (#15082) | ||
Alexander Ivanov
|
aea7b70be0 |
[feat] add KVAE-Audio autoencoder (#14384)
* [feat] add KVAE-Audio autoencoder * [fix] non-contiguous tensor into in_proj in AutoencoderKLKVAEAudio * [chore] sync KVAE-Audio with upstream test and annotation changes * [refactor] dedicated attention class and processor for AutoencoderKLKVAEAudio * [refactor] remove the runtime sample_rate argument from AutoencoderKLKVAEAudio * [refactor] always apply attention in AutoencoderKLKVAEAudio, drop use_attn * [refactor] drop the derived latent_dim default from AutoencoderKLKVAEAudio * [docs] fix encoder_rates docstring in AutoencoderKLKVAEAudio * [refactor] rename Snake1d to KVAEAudioSnake1d in AutoencoderKLKVAEAudio * [tests] skip layerwise casting training for AutoencoderKLKVAEAudio |
||
Akshan Krithick
|
6184688f71 |
move unpack out of the flux, krea2 and ltx decoders into the core denoise groups (#14737)
* add a modular test that pins the canonical latents form each family leaves in the state * move unpack out of the flux, krea2 and ltx decoders into the core denoise groups * pop the unpack step from a copy of the blocks in the deprecation tests * let families that decode inside the denoise loop opt out of the latents-form test * document the canonical latents form in the modular reference * [docs] canonical latents form: name normalize, patchify and pack separately Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * require geometry on the unpack steps and test the deprecated decode path directly * skip the canonical latents test explicitly for ltx 2.5 too * remove the packed latents path in 0.47.0, five releases after 0.42 * update the latents_output_name comment for the explicit skips * move the packed test latents to the test device --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
Sayak Paul
|
e04f0c5f24 |
post-release 0.41.0 (#15062)
* post-release 0.41.0 * up |
||
|
YiYi Xu
|
4df4a5b357 |
FLUX.3: match the released BFL reference image sizing (#15073)
* FLUX.3: match the released BFL reference image sizing A single reference is kept whole: downscaled so its padded canvas fits under 4 MP, reflect-padded to the token grid (`resize_mode="fill"` of the FLUX.3 image processor), and cropped back on output. Several references are fitted under 1 MP each. The default output size follows the first reference. --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
YiYi Xu
|
b800f828f3 |
[docs] FLUX.3: modular-only notes, lazy loading, callbacks (#15069)
* [docs] FLUX.3: modular-only notes, lazy loading, callbacks Add a tips callout on what differs from DiffusionPipeline (lazy loading, outputs, guider, offloading, editing blocks), move the prompting advice into the upsampling section with the pop-the-step example, and add a per-step callbacks section with a block that takes callback_on_step_end. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * [docs] FLUX.3: callback block walks all denoising loops Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * [docs] modular: add "Coming from DiffusionPipeline" to the quickstart and move per-step callbacks to the callbacks guide Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * [docs] FLUX.3: show the upsampled caption; clarify the index note in the quickstart Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * [docs] modular: make "Coming from DiffusionPipeline" its own page Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
YiYi Xu
|
d42ebf9c6e |
FLUX.3 prompt upsampler block (#15058)
* adding FLUX.3 prompt upsampler block |
||
Stephen Batifol
|
af4271c3e1 |
[docs] FLUX.3: prompt upsampler required, ours or any LLM (#15059)
* docs(flux3): the prompt upsampler is required, ours or any LLM * docs(flux3): every example goes through upsample() * docs(flux3): strongly advise the prompt upsampler * docs(flux3): the model is FLUX 3 Image Klein |
||
|
Sayak Paul
|
7ecb1e1b4f |
Flux3 image (#15046)
* Bump the actions group with 17 updates (#12) Bumps the actions group with 17 updates: | Package | From | To | | --- | --- | --- | | [actions/checkout](https://github.com/actions/checkout) | `6` | `7` | | [actions/upload-artifact](https://github.com/actions/upload-artifact) | `4.6.2` | `7.0.1` | | [docker/setup-buildx-action](https://github.com/docker/setup-buildx-action) | `3.12.0` | `4.2.0` | | [docker/login-action](https://github.com/docker/login-action) | `3.7.0` | `4.6.0` | | [docker/build-push-action](https://github.com/docker/build-push-action) | `6.19.2` | `7.3.0` | | [huggingface/doc-builder/.github/workflows/build_main_documentation.yml](https://github.com/huggingface/doc-builder) | `e60a538eea9817ab312196d0d233604b01697265` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` | | [actions/setup-python](https://github.com/actions/setup-python) | `5.6.0` | `7.0.0` | | [huggingface/doc-builder/.github/workflows/build_pr_documentation.yml](https://github.com/huggingface/doc-builder) | `e60a538eea9817ab312196d0d233604b01697265` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` | | [anthropics/claude-code-action](https://github.com/anthropics/claude-code-action) | `1.0.155` | `1.0.193` | | [huggingface/security-workflows/.github/workflows/codeql-reusable.yml](https://github.com/huggingface/security-workflows) | `1.1.0` | `1.2.0` | | [actions/download-artifact](https://github.com/actions/download-artifact) | `4.3.0` | `8.0.1` | | [actions/labeler](https://github.com/actions/labeler) | `5.0.0` | `7.0.0` | | [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) | `1.13.0` | `1.14.2` | | [huggingface/tailscale-action](https://github.com/huggingface/tailscale-action) | `7d53c9737e53934c30290b5524d1c9b4a7c98c8a` | `07d36e4c3b68707bd3fe070dbb96f03f993a6efd` | | [trufflesecurity/trufflehog](https://github.com/trufflesecurity/trufflehog) | `3.94.2` | `3.97.0` | | [crate-ci/typos](https://github.com/crate-ci/typos) | `1.42.1` | `1.49.0` | | [huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml](https://github.com/huggingface/doc-builder) | `9ad2de8582b56c017cb530c1165116d40433f1c6` | `af49b8dc3859c3d80c7220c8c0ba146431132e49` | Updates `actions/checkout` from 6 to 7 - [Release notes](https://github.com/actions/checkout/releases) - [Commits](https://github.com/actions/checkout/compare/v6...v7) Updates `actions/upload-artifact` from 4.6.2 to 7.0.1 - [Release notes](https://github.com/actions/upload-artifact/releases) - [Commits](https://github.com/actions/upload-artifact/compare/v4.6.2...v7.0.1) Updates `docker/setup-buildx-action` from 3.12.0 to 4.2.0 - [Release notes](https://github.com/docker/setup-buildx-action/releases) - [Commits]( |
||
|
Akshan Krithick
|
62cde2b011 |
add modular pipeline for qwen-image 2.1 (#14976)
* add modular pipeline for qwen-image 2.1 * split out a turbo blockset and drop the guider from the text encoders * Update tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21_turbo.py * Update tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
Jingya HUANG
|
4113984327 |
refactor: drop the Neuron RoPE path in Qwen Image (#15042)
Neuron now runs complex tensors natively, so Qwen Image uses the default complex RoPE on every device. This removes `apply_rotary_emb_qwen_neuron`, `ROPE_PER_DEVICE` and the `torch.angle` branch in `_get_device_freqs`, added in #13718. Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Jingya HUANG
|
ee849197ce |
[core] Tensor parallelism for Qwen-Image-2.1 (#14865)
* feat: qwen image 21 tp plan * test: vallidated on neuron * fix: lru_cache_unless_export leak * fix: use complex RoPE on TPU for Qwen Image 2.1 TorchTPU reports TPU tensors as "tpu" and supports complex dtypes, so only Neuron needs the angle-based RoPE path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * refactor: inline the Neuron device check for Qwen Image 2.1 RoPE * review: cache Qwen Image 2.1 RoPE device freqs with lru_cache * review: drop redundant torch.device * fix: use torch.accelerator.device_count() when resolving the parallel device * refactor: drop the Neuron branch in Qwen Image 2.1 block ids * restore: neuron repeat_interleave lowering error fix * review: add TensorParallelNeuronTesterMixin for Neuron TP tests * review: revert Qwen Image 2.1 RoPE freqs cache to a per-instance dict * review: drop the Neuron RoPE path in Qwen Image 2.1 Neuron now runs complex tensors natively, so Qwen Image 2.1 uses the default complex RoPE on every device. This removes `apply_rotary_emb_qwen_neuron`, `ROPE_PER_DEVICE` and the per-device freqs cache, and keeps `self.freqs = [freq.to(device) ...]` as on main. --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
efe46a5b17 |
ci: move to python 3.11. (#15030)
* move to python 3.11. * update setup.py along with others. * fix(ci): pin GitHub Actions to commit SHAs (#15030) (#15033) fix(ci): pin actions in workflows flagged on #15030 Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com> * up * fix(ci): pin GitHub Actions to commit SHAs (#15030) (#15036) fix(ci): pin actions in workflows flagged on #15030 Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com> * up --------- Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com> |
||
Dhruv Nair
|
ea981b8b7d |
[CI] Clean up test_modeling_common.py (#14988)
clean up |
||
Pauline Bailly-Masson
|
40f63690da |
🔒 Pin GitHub Actions to commit SHAs (#15028)
Pin GitHub Actions to commit SHAs Replace mutable tags and branch refs in workflow files with full commit SHAs, keeping the original version as a trailing comment. Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> |
||
Vinh H. Pham
|
7c6eb78306 |
[hooks] Support Flux2 in MagCache (#14756)
* [hooks] Support Flux2 in MagCache MagCache could not be enabled on `Flux2Transformer2DModel` (Flux2 and Flux2 Klein): 1. `Flux2TransformerBlock` and `Flux2SingleTransformerBlock` were not registered in `TransformerBlockRegistry`, so `apply_mag_cache` raised `ValueError: Model class ... not registered`. 2. Once registered, the tail hook raised a shape error. The head hook records the image-only input of `transformer_blocks[0]`, but the Flux2 model concatenates the text and image streams before the single-stream blocks and calls them with `encoder_hidden_states=None`, so the last single block returns the fused `[text; image]` sequence. The residual `tail_output - head_input` could not be formed. Register the two Flux2 blocks (the single block with `return_encoder_hidden_states_index=None`, matching how the model calls it) and add `_align_to_hidden_states`, which keeps only the trailing image tokens of the fused tensor when computing the residual. The cached residual is therefore image-only, which is what the reference MagCache4FLUX implementation does, and it matches the head input on skipped steps. This also replaces a dead "fallback" branch in the tail hook that subtracted mismatched shapes anyway. Tests: `MagCacheTesterMixin` for `Flux2Transformer2DModel`, `Flux2Pipeline` and `Flux2KleinPipeline`. FirstBlockCache has the same head/tail structure and still fails on Flux2; left for a follow-up. * [hooks] Slice the fused tail output in the existing mismatch branch Implement the Flux2 residual fix inside the tail hook's sequence-length mismatch branch instead of a separate `_align_to_hidden_states` helper. The surrounding structure stays as it was: the `diff == 0` case still covers a batch-size mismatch via broadcasting, and only the mismatched-length case changes to subtract the trailing image tokens. --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
Wang, Yi
|
839af3a043 |
Update FlashAttention Hub integration and guard CUDA-only tests (#15024)
* Update FlashAttention Hub integration and guard CUDA-only tests Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * fmt Signed-off-by: Wang, Yi <yi.a.wang@intel.com> --------- Signed-off-by: Wang, Yi <yi.a.wang@intel.com> |
||
Javier De Jesus
|
08d69774f0 |
[utils] preserve method metadata in apply_forward_hook (#13803)
* [utils] preserve method metadata in apply_forward_hook Add @functools.wraps to the inner wrapper so it inherits __name__, __doc__, __module__ and the signature of the wrapped method. Without this, methods decorated with @apply_forward_hook (e.g. AutoencoderKL's encode and decode) lose their docstring and signature, so doc-builder cannot render them through the [[autodoc]] directive and they show up as "wrapper" with a generic (*args, **kwargs) signature on the live API docs. The same affects VQModel, ConsistencyDecoderVAE and the other autoencoders that use the decorator. Add regression tests asserting that the wrapped method exposes the right metadata (name, docstring, signature including default values, __wrapped__) and that the underlying offload hook still fires through the wrapper. Fixes #6575 * Trim apply_forward_hook test to a single assertion |
||
Sergio Graciá
|
916a5791ed |
[Z-Image] Apply the LoRA scale passed via joint_attention_kwargs (#15015)
Fixes #15007 Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Junsong Chen
|
1628cb1c33 |
Add SANA-WM camera-controlled image-to-video pipeline (#13881)
* feat(sana-wm): add diffusers-style SANA-WM camera-controlled I2V pipeline
Adds the public SANA-WM bidirectional camera-controlled image-to-video
model as a first-class diffusers pipeline + transformer. Layout mirrors
``sana_video``: the model lives under ``src/diffusers/models/transformers/``
as a near-single-file (kernels split off so the ``@triton.jit`` decorators
don't drown the model body); the pipeline lives under
``src/diffusers/pipelines/sana_wm/``.
Files added:
src/diffusers/models/transformers/
├── transformer_sana_wm.py # SanaWMTransformer3DModel + blocks + helpers
└── transformer_sana_wm_kernels.py # fused Triton kernels + camera math
src/diffusers/pipelines/sana_wm/
├── __init__.py
├── pipeline_sana_wm.py
├── pipeline_output.py
├── refiner.py
└── cam_utils.py
Pipeline architecture:
* Stage 1: 1600M ``SanaWMTransformer3DModel`` DiT with bidirectional
GDN-Triton linear attention + UCPE camera-control branch, LTX-style
flow-matching Euler scheduler with per-token timesteps.
* Stage 2: LTX-2 sink-bidirectional Euler refiner (3 distilled sigma
steps, reuses diffusers' ``LTX2VideoTransformer3DModel`` +
``LTX2TextConnectors`` + Gemma-3 text encoder).
* Decode through the LTX-2 VAE (``AutoencoderKLLTX2Video``).
One-line usage:
pipe = SanaWMPipeline.from_pretrained(
"Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
torch_dtype=torch.bfloat16,
).to("cuda")
out = pipe(image=img, prompt="...", action="w-80,jw-40,w-40",
intrinsics=[fx, fy, cx, cy])
End-to-end smoke test (stage-1 + refiner + VAE decode) passes on H100.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* feat(sana-wm): align pipeline with merged sana_video style; fix mp4 export
transformer_sana_wm.py:
* License header switched to the "HuggingFace Team and SANA-WM Authors"
style used by merged sana_video.
* Imports rewritten in stdlib -> third-party -> diffusers order; use
diffusers `from ...utils import logging` instead of stdlib `logging`.
* Fix 9 `Optional[X]` annotations written as `X or None` (Python's `or`
short-circuits and silently returns `X`).
* Fix two `assert (cond, msg)` tuple-asserts in PatchEmbedMS3D.forward
that always pass (SyntaxWarning at import time).
* Remove duplicate `__all__` declarations (the second silently overwrote
the first).
* Remove dead `reset_bn` (imports a nonexistent `packages.apps.utils`,
would crash on call).
* Remove the duplicate `logger = logging.getLogger(__name__)` further
down in the file.
transformer_sana_wm_kernels.py:
* License header normalized; collapse three duplicate triton/torch import
blocks into one.
pipeline_sana_wm.py:
* License header normalized.
* `_decode_latents` now returns `(T, H, W, 3)` float in [0, 1], matching
the diffusers convention used by `VideoProcessor`. Returning uint8
silently broke `export_to_video`: it does `frame * 255` assuming float
input, so uint8 overflows to `(-x) mod 256` and inverts colors.
* `__call__` converts to PIL/uint8 only when `output_type="pil"`.
* Intrinsics argument now accepts (4,), (F, 4), (3, 3), and (F, 3, 3)
forms (auto-extracts fx, fy, cx, cy from a 3x3 K) and auto-trims to
`num_frames` when a longer-than-needed trajectory is passed.
* Inline `retrieve_timesteps` with the standard `# Copied from
diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps`
marker, matching merged sana_video.
* Docstrings + EXAMPLE_DOC_STRING updated to reflect the new return type.
pipeline_output.py:
* Update `frames` field docstring to describe the new float [0, 1] return.
refiner.py, cam_utils.py, scripts/sana_wm/convert_sana_wm_to_diffusers.py:
* License headers normalized.
Docs:
* New `docs/source/en/api/pipelines/sana_wm.md` and
`docs/source/en/api/models/sana_wm_transformer3d.md`, modeled on
sana_video.md / sana_video_transformer3d.md, wired into
`docs/source/en/_toctree.yml` under Models and Pipelines.
5s end-to-end smoke test (81 frames @ 16fps, 30 stage-1 steps + 3-step
LTX-2 refiner) passes on 1x H100 80GB with `enable_model_cpu_offload`.
Round-trip diff vs raw float frames is 2.06/255 mean (h264 lossy noise),
confirming the export_to_video fix.
* feat(sana-wm): port chunk-causal AR refiner mode (RefinerChunkRunner + KV cache hooks)
The first cleanup pass only kept the legacy single-shot refiner path. That
path is what the model was *not* trained on — its docstring even says
"feeding the full sequence at once is out-of-distribution" — and its cost
is O(T^2) attention over the full latent volume, which made longer videos
unusable (~21 min per refiner step at 321 frames on an H100).
Port the chunk-causal AR mode from the upstream reference so the refiner
matches the training contract:
* `refine_latents` now defaults to `block_size=3, kv_max_frames=11`
(the canonical AR recipe). Pass `block_size=None` to fall back to the
legacy single-shot path.
* New `_refine_latents_ar` + `_RefinerChunkRunner` orchestrate the sliding
window: pre-capture pre-RoPE sink K/V on `z_sana[:source_sink_frames]`
at sigma=0, then for each `block_size`-frame chunk run a 3-step Euler
with prefix `{sink_k_pre, sink_v, sink_pe, history_k, history_v}` and
capture post-RoPE K/V to feed the next window. History is bounded to
`kv_max_frames - source_sink_frames` so per-block compute is constant.
* New `_predict_x0_active_block` runs the transformer on the active block
only (Q from active, K/V from prefix+active).
* New `_capture_block_kv` runs sigma=0 forward with a pre_rope/post_rope
capture flag set on each `attn1`.
* New `_forward_video_only_with_rope` takes a pre-built RoPE so each block
can use absolute frame positions in the source video.
* `_streaming_self_attention` extended with the `_kv_cache_capture`,
`_tf_capture_kv`, `_tf_kv_prefix` hook contract that AR mode uses to
inject and capture K/V on each block.
* New helpers: `_build_rotary_emb_for_absolute_positions`,
`_set_kv_prefix_on_blocks`, `_clear_kv_prefix_on_blocks`,
`_set_capture_flag_on_blocks`, `_collect_captured_kv_from_blocks`.
* `_encode_prompt` now also moves the Gemma-3 text encoder back to CPU
after producing the embeds — otherwise it stays resident through the
entire AR loop and gates how much GPU memory the refiner transformer
has left.
Module-level docstring updated to document both modes; existing
single-shot path preserved verbatim.
* feat(sana-wm): block-level checkpoint for AR refiner (resume after preemption)
The AR refiner is expensive (~3-5 min per block) and the refinement loop
ran end-to-end has no in-progress state to recover, so a SLURM preemption
mid-refinement loses all progress. With the canonical
``block_size=3, kv_max_frames=11`` setup, refining a 50s video is 34
blocks of work that has to make it through without preemption on a
backfill queue.
Add per-block atomic checkpointing:
* ``SanaWMLTX2Refiner.refine_latents(checkpoint_dir=Path)`` and
``_refine_latents_ar`` accept a directory. After each completed AR
block, the AR loop writes ``checkpoint_dir/state.pt`` atomically
(tmp + os.replace).
* The payload is ``{block_idx_done, n_blocks, sink_size, block_size,
output_shape, output, runner_state}``. ``runner_state`` is a CPU snapshot
of the runner's ``_sink_kv_pre``, ``_history_kv_post``,
``_history_frames`` and ``torch.Generator`` state.
* On entry, if ``state.pt`` exists with a compatible shape signature, the
AR loop loads the persisted output tensor + runner state and resumes
from ``block_idx_done + 1`` instead of recomputing from scratch.
* ``SanaWMPipeline.__call__(refiner_checkpoint_dir=...)`` plumbs the
directory through to the refiner.
Checkpoint size: ~output_volume + sink_KV (~360MB for 50 layers) +
rolling history KV (~3-4GB at full capacity) — saved once per block,
total per-block save overhead ~10s on lustre.
* test(sana-wm): add CPU unit tests + slow GPU integration stub
* CPU unit tests for cam_utils helpers (action DSL → c2w, intrinsics
rescale-for-crop, resize+center-crop, snap_num_frames 8k+1 rounding).
* Public-surface registration tests (top-level diffusers symbols,
SanaWMPipelineOutput dataclass shape, refiner signature has AR defaults
+ checkpoint_dir, pipeline __call__ accepts c2w/action/intrinsics/
refiner_checkpoint_dir).
* @slow @require_torch_accelerator integration stub for an end-to-end I2V
against the public checkpoint, currently @unittest.skip — wires up the
nightly GPU path without exploding regular CI.
SanaWMTransformer3DModel has hardcoded depth/hidden_size/num_heads inside
its inner SanaMSVideoCamCtrl (not exposed through register_to_config), so
the usual PipelineTesterMixin small-config fast tests aren't applicable
without a transformer refactor (followup PR).
* feat(sana-wm): make triton optional + auto-fallback to pure-PyTorch attention
`transformer_sana_wm_kernels.py` previously did a hard `import triton`
at the top of the file. That blocked importing the SANA-WM transformer
on any environment without Triton (CPU-only, ROCm without Triton,
older Triton, etc.), even though the model has pure-PyTorch attention
classes for every `*Triton` variant.
Make Triton optional and have the dispatcher transparently fall back:
* Wrap `import triton` / `import triton.language as tl` in try/except.
When unavailable, install a shim where `@triton.jit` is a no-op so
the kernel function definitions still load (they just aren't compiled
by Triton). Module-level `triton.X` / `tl.X` lookups return a
self-shimming sentinel so signature parsing doesn't blow up either.
* Add `is_triton_available()` + `_require_triton(entry_point)`. The four
Triton-backed entry points called by the model (`fused_qk_inv_rms`,
`fused_bigdn_func`, `cam_prep_func`, `cam_scan_bidi_chunkwise`) now
raise a clear RuntimeError on a Triton-less host with a hint to use
the pure-PyTorch attention variants — but the dispatcher does this
automatically (see below) so users shouldn't ever see it.
* Delete the leftover duplicate `import torch / triton / triton.language`
block at line 262 (left over from the upstream port).
* Register `BidirectionalGDNUCPESinglePathLiteLA` in `ATTENTION_BLOCKS`
so the fallback chain can find it.
* New `_resolve_attention_block(name, role)` walks the requested class's
MRO at dispatch time. If Triton isn't usable AND the requested class
name ends in `Triton`, route to the closest registered non-`Triton`
ancestor (BidirectionalGDNUCPESinglePathLiteLABothTriton ->
BidirectionalGDNUCPESinglePathLiteLA, etc.) and log a one-shot warning.
* Rewire both `SanaVideoMSCamCtrlBlock` dispatch sites to use
`_resolve_attention_block` for the GDN+UCPE camera branch and the main
attention branch (the `BidirectionalSoftmaxUCPESinglePathLiteLA` branch
doesn't use Triton at all so it stays hard-coded).
Tests:
* `test_kernels_module_imports_with_triton_hidden` — reloads the kernels
module with `sys.modules['triton'] = None` and verifies the module
imports, `is_triton_available()` is False, and the pure-PyTorch helpers
remain callable.
* `test_resolve_attention_block_cpu_fallback` — on a CPU-only host, the
three `*Triton` attn types resolve to the correct non-Triton ancestor.
* `test_triton_entry_point_raises_clean_error_without_triton` — verifies
the `_require_triton` guard yields a RuntimeError that mentions Triton.
* fix(sana-wm): make optional deps lazy + register transformer in __init__
Three CI checks were failing on the PR:
1. `check_code_quality` (43 ruff errors): mix of unused imports / import
sorting / E731 lambdas (auto-fixable) plus a handful of F821 dead-code
references inherited from the upstream research codebase (`xformers.*`
inside `if _xformers_available:` blocks, an undefined `BlockHook` type
annotation, two `x_sa`/`mlp_out` references in a block forward whose
live assignment was already overridden by subclasses). Ran `ruff check
--fix --unsafe-fixes` + `ruff format`, fixed the type annotation
manually, and added targeted `# noqa: F821` markers on the conditionally
unreachable lines.
2. `check_torch_dependencies`: `transformer_sana_wm.py` hard-imported
`einops`, `fla`, `timm`, `termcolor`. The minimum-deps CI environment
doesn't have them, and diffusers' lazy loader rewrites `ModuleNotFoundError`
as `RuntimeError` so `test_pipeline_imports` blew up. Wrapped each of
the four optional imports in a try/except shim — `rearrange`/
`ShortConvolution`/`DropPath`/`Attention_`/`Mlp` become placeholders
that raise a clear `ImportError` on construction, `colored` falls back
to plain text. Class bodies that subclass these still parse at module
load, so `import diffusers.models.transformers.transformer_sana_wm`
succeeds anywhere. Same treatment for the kernels file's
`from einops import rearrange, repeat`.
3. `build_pr_documentation`: doc-builder imported `SanaWMTransformer3DModel`
from `diffusers.models.transformers` (not the diffusers top level) and
that subpackage's `__init__.py` was missing the entry. Added the import.
* style(sana-wm): apply make style + fix-copies
* `doc-builder style src/diffusers docs/source --max_len 119` rewraps
docstrings in the six SANA-WM files (transformer, kernels, pipeline,
refiner, output, cam_utils) to the repo-wide 119-column limit. No
behaviour change — purely whitespace inside docstrings.
* `make fix-copies` regenerates `dummy_pt_objects.py` and
`dummy_torch_and_transformers_objects.py` to add `DummyObject` stubs
for the three new public classes (`SanaWMTransformer3DModel`,
`SanaWMPipeline`, `SanaWMLTX2Refiner`), so `from diffusers import …`
gives the standard "missing backend" message on installs without
torch / transformers.
Verified: `make quality` passes (ruff check, ruff format check,
doc-builder style check_only, check_doc_toc). Test suite still
15 passed / 1 skipped.
* refactor(sana-wm): drop einops dependency, inline with torch ops
Per @dg845's review on `transformer_sana_wm.py:44`. The 9 call sites all
match well-known patterns that have one-liner torch equivalents:
rearrange(s, "b d -> (b d)") -> s.reshape(b * d)
rearrange(x, "(b d) d2 -> b (d d2)", ...) -> x.reshape(b, d * d2)
rearrange(R, "b t h w i j -> b t h w j i") -> R.transpose(-1, -2)
repeat(x, "b h w c -> b t h w c", t=T) -> x.unsqueeze(1).expand(-1, T, -1, -1, -1)
repeat(x, "b t c -> b t h w c", h=H, w=W) -> x[:, :, None, None, :].expand(-1, -1, H, W, -1)
repeat(x, "... -> b ...", b=B) -> x.unsqueeze(0).expand(B, *x.shape)
repeat(x, "H W C -> B T H W C", B, T) -> x[None, None].expand(B, T, -1, -1, -1)
repeat(x, "B H W C -> B T H W C", T) -> x.unsqueeze(1).expand(-1, T, -1, -1, -1)
Each replacement is bit-identical to the einops original — verified
against a fresh `einops` install on random tensors before swapping. The
optional `from einops import ...` shim block is gone from both
`transformer_sana_wm.py` and `transformer_sana_wm_kernels.py`.
* refactor(sana-wm): remove dead code per @dg845's review
Drop 23 unused symbols (1451 lines) from `transformer_sana_wm.py` that
aren't reachable from the public SANA-WM checkpoint's
`SanaMSVideoCamCtrl` -> `SanaVideoMSCamCtrlBlock` -> GDN/UCPE attention
path. Each was verified to have zero call sites outside of its own
definition (or only within other-deleted items).
Classes:
* `SanaMS`, `SanaMSBlock` — alternative `Sana` subclass + block, not
used by the SANA-WM checkpoint (which goes through `SanaMSVideoCamCtrl`).
* `ChunkCausalAttention`, `CachedCausalAttention`,
`ChunkedLiteLAReLURope`, `LiteLAReLURope` — chunk-causal / cached
attention variants and their common base; SANA-WM uses the bidi GDN
path. `LiteLAReLURope` had only the three (now-deleted) subclasses
referencing it.
* `PAGCFGIdentitySelfAttnProcessorLiteLA`,
`PAGIdentitySelfAttnProcessorLiteLA`, `SelfAttnProcessorLiteLA`,
`SelfAttnProcessorLiteLAReLURope` — PAG processors; we don't expose
PAG in the SANA-WM pipeline.
* `ChunkGLUMBConvTemp`, `CachedGLUMBConvTemp`, `MBConvPreGLU` —
alternative FFN/conv blocks; the checkpoint uses `GLUMBConvTemp`.
* `MaskFinalLayer`, `DecoderLayer` — alternative final layers; the
checkpoint uses `T2IFinalLayer`.
* `LabelEmbedder`, `CaptionEmbedderDoubleBr` — alternative embedders;
the checkpoint uses `CaptionEmbedder`.
Helpers:
* `set_grad_checkpoint`, `prepare_prompt_ar`, `resize_and_crop_tensor`,
`generate_temporal_head_mask_mod`, `is_chunk_causal_request`,
`get_chunk_index_from_config` — training-only or
chunk-causal-mode helpers with no inference call sites.
Verified after the diff:
* `make quality` clean.
* CPU test suite 15 passed / 1 skipped (slow GPU integration).
* `import diffusers ; SanaWMPipeline / SanaWMTransformer3DModel /
SanaWMLTX2Refiner` resolve identically.
* fix(sana-wm): defer transformers import + document mask/return_dict
Two PR-CI failures, both ours:
* `check_torch_dependencies`: line 33 hard-imported `transformers`, which
isn't present in the minimum-deps environment. Move the lone
`AutoModelForCausalLM` use site inside `initialize_gemma_params` (a
training-only helper).
* `check_repository_consistency`: `SanaWMTransformer3DModel.forward`'s
docstring was missing entries for `mask` and `return_dict`. Added.
* refactor(sana-wm): address review feedback on SanaWMPipeline
Per @dg845's review comments on pipeline_sana_wm.py:
* Read the VAE spatial/temporal strides from the ``vae`` component
(``vae_spatial_compression_ratio`` / ``vae_temporal_compression_ratio``)
instead of hardcoding 32/8, mirroring LTX2Pipeline.
* Inline the stage-1 sampling loop into ``__call__`` and factor the noise
init into a standard ``prepare_latents`` method.
* Use ``self.scheduler`` for the flow-matching Euler steps instead of
constructing a new scheduler per call.
* Drive the sampling loop with ``self.progress_bar`` (respects
``set_progress_bar_config``) instead of a bare tqdm.
* Move input validation/normalization into a ``check_inputs`` method.
* Move first-frame resize+center-crop, [-1, 1] normalization and the
intrinsics rescale into a ``SanaWMImageProcessor(VaeImageProcessor)``
subclass (new ``image_processor.py``).
* Add ``generator`` as a ``__call__`` argument (``seed`` kept as a
convenience shortcut).
* Post-process decoded latents with ``VideoProcessor.postprocess_video``
rather than a hand-rolled conversion; drop the ``.cpu()`` casts on the
``output_type="latent"`` path.
* Move the pipeline README into the docs (docs/.../sana_wm.md); delete
the in-package README.
* refactor(sana-wm): make SanaWMLTX2Refiner a standalone DiffusionPipeline
Per @dg845's review: the refiner has components (transformer, connectors,
text encoder, tokenizer) and runs a denoising loop, so it fits better as a
pipeline than a ModelMixin.
* SanaWMLTX2Refiner now subclasses DiffusionPipeline and registers its
components via ``register_modules`` (dropping the bespoke
``from_pretrained`` / ``save_pretrained``); standard load/save now handle
the ``refiner/`` subfolder.
* Add a ``FlowMatchEulerDiscreteScheduler`` component (shift=1.0) and drive
the Euler steps through ``scheduler.step`` / ``scheduler.scale_noise``
(single-shot and per-AR-block) instead of hand-rolled updates. Numerically
equivalent to the previous flow-matching update.
* Rename the entry point ``refine_latents`` -> ``__call__``; add a ``device``
arg so the parent can hand it the execution device without a bulk move.
* SanaWMPipeline: keep the refiner as an optional nested component; free the
parent's GPU weights before running it (it manages its own sub-module
placement) and bring the VAE back for decode.
* Conversion script emits the new refiner layout (model_index.json +
scheduler/ + tokenizer/); test asserts the refiner is a DiffusionPipeline
with the canonical AR ``__call__`` defaults.
Validated end-to-end on 1xH100 (stage-1 + refiner on the official demo,
coherent video output).
* docs(sana-wm): document the refiner __call__ `device` arg
Fixes the check_repository_consistency failure — `SanaWMLTX2Refiner.__call__`
gained a `device` parameter but its docstring wasn't updated.
* refactor(sana-wm): clean up dead code + reuse shared utils in the DiT
Addresses several of @dg845's transformer review comments (the safe,
output-preserving subset — GPU smoke gives byte-identical output):
* Reuse diffusers' shared `FP32LayerNorm` (models/normalization.py) and
`get_1d_rotary_pos_embed` (models/embeddings.py); delete the local copies.
* Remove dead inference code paths:
- the `if self.diagonal_mask is not None:` flex-attention block
(`diagonal_mask` is always `None`) + the now-unused
`create_block_mask_cached` helper and `create_block_mask` import;
- the `SANA_FSDP2_BLOCK_TIMING` block-timing/profiling scaffolding;
- the `save_qkv` / `qkv_store_buffer` visualization hooks (never
enabled at inference), at both the attention and model level.
* Collapse `SanaVideoMSCamCtrlBlock.forward_frame_aware` into `forward`
(the pipeline/refiner always pass >=3D timesteps, so the non-frame-aware
branch was dead — it even referenced undefined locals).
* Drop the 191-line `SanaMSVideoCamCtrl.load_state_dict` shape-remapping
override — it's never reached by the shipped convert/inference flow
(`nn.Module.load_state_dict` on the wrapper doesn't call it), and the
release checkpoint already ships correctly-shaped weights.
* Delete the unused `PatchEmbedMS` module.
Net -503 lines; no numerics change (stage-1 + refiner smoke output mean
identical to the pre-refactor run).
* refactor(sana-wm): merge the 3 DiT classes into one SanaWMTransformer3DModel
Per @dg845's review: `SanaWMTransformer3DModel` was a thin wrapper over
`SanaMSVideoCamCtrl`, which in turn subclassed `Sana` and overwrote most of
it. Fold all three into a single `SanaWMTransformer3DModel(ModelMixin,
ConfigMixin)`:
* The `@register_to_config` __init__ signature is unchanged (so config.json
is identical); the body builds the modules directly on `self` instead of
a nested `self._inner`.
* Only the surviving `Sana.__init__` pieces are kept (t_embedder,
cfg_embedder, attention_y_norm, config attrs, initialize_weights); the
parts the subclass overwrote are gone.
* `forward` takes the diffusers signature (hidden_states / timestep /
encoder_hidden_states / encoder_attention_mask / return_dict) and returns
`Transformer2DModelOutput`, folding in the old wrapper's arg-renaming.
* Deleted now-unreachable code: `class Sana`, `class SanaMSVideoCamCtrl`,
`class SanaBlock`, `class PatchEmbed` (the video model uses
`PatchEmbedMS3D`), the `add_inner_prefix` helper, and the dead
`sincos`/`flux_rope` pos-embed branches in forward (release uses
`wan_rope`). Net -468 lines.
This drops the `_inner.` state-dict prefix, so the conversion script no
longer adds it. State-dict is otherwise identical: the merged model's
`state_dict()` has exactly the same 871 param keys as before (verified),
and a stage-1 + refiner GPU smoke on the public checkpoint gives
byte-identical output (frame mean 0.5560, matching the pre-merge run).
Addresses the class-merge review comment and removes the unused
`SanaBlock`/`PatchEmbed` modules.
* refactor(sana-wm): drop the no-op custom grad-checkpoint wrapper
`auto_grad_checkpoint` gated on a `grad_checkpointing` attribute that was
never set, so it always fell through to `module(*args, **kwargs)` — i.e. a
no-op. Call the transformer blocks directly instead and delete the unused
`auto_grad_checkpoint` / `checkpoint_sequential` helpers and the
`torch.utils.checkpoint` import.
(Inference is unchanged — the block never read the `use_reentrant` kwarg
the wrapper passed. Full training-time gradient checkpointing via the
standard `_gradient_checkpointing_func` would need the block's dynamic
kwargs — camera_conditions / prope_fns / chunk_index — threaded through,
so it's left as a follow-up for this inference-focused release.)
* refactor(sana-wm): drop the fla-core and termcolor dependencies
Per @dg845 / @sayakpaul: the SANA-WM transformer required `fla-core` to be
constructible at all (and fla's `ShortConvolution` cannot even run on CPU —
its dispatch does `torch.cpu.device(...)`, which doesn't exist).
`ShortConvolution` was only ever used as a depthwise *causal* conv1d with
`activation=None`, so it is replaced by a ~20-line self-contained PyTorch
module with the same parameter layout (`weight` of shape `(C, 1, K)`, no
bias) and the same `(output, cache)` return signature. No `fla-core` and no
`kernels` dependency is needed now, and the model can be built and run on
CPU. Verified on an H100 in bf16 at the real config (hidden 2240, kernel 4):
`max|Δ|` vs fla is exactly `0.0` for every shape tested, and the state dict
keeps the identical 871 keys.
Also:
* Use the existing `is_timm_available()` utility for the timm imports. The
`else` branch keeps a placeholder class rather than raising at module
scope, because several layers subclass these symbols and this module is
imported eagerly by `diffusers.models` — a module-level raise would break
plain `import diffusers` (and the `check_torch_dependencies` CI job) when
timm isn't installed. The error is instead raised on construction.
* Drop `termcolor` entirely: replace the `self.logger = print` +
`colored(...)` pattern with the module-level `logger`, so there is no
optional dependency left to gate.
* Only attempt the optional null-embedding load when `null_embed_path` is
actually set (it is unset for the public checkpoint, so this previously
logged a spurious failure on every construction) and load it with
`weights_only=True`.
* fix(sana-wm): address review nits in the pipeline / refiner / cam utils
* refiner: load the AR resume checkpoint with `weights_only=True`. The
payload is only tensors / ints / tuples / dicts plus the generator's uint8
state, so it round-trips safely under the restricted unpickler.
* refiner: fix `_tokens_per_frame` to divide by `patch_size_t` rather than
multiply — `_pack_latents` emits `(T // patch_size_t) * (H // p) * (W // p)`
tokens, so one latent frame contributes `(H // p) * (W // p) / patch_size_t`.
No-op for LTX-2 (`patch_size_t=1`) but the history trimming would have kept
too many tokens otherwise.
* pipeline: replace `torch.cuda.empty_cache()` with diffusers'
backend-agnostic `empty_device_cache(device.type)`.
* pipeline: extract the offload probe into `_model_cpu_offload_active()`,
matching the `hasattr(self, "_all_hooks") and len(self._all_hooks) > 0`
idiom `DiffusionPipeline` uses internally, instead of a bare truthiness
check on the attribute.
* pipeline_output: add the missing Apache-2.0 license header.
* cam_utils: drop the unnecessary `+ 1e-6` when normalizing the forward /
right vectors — the branch is only taken when the norm is already > 0, so
the epsilon just introduced a small systematic bias.
* refactor(sana-wm): strip research-repo code from the DiT, fix model conventions
From the `self-review` skill run against `.ai/models.md` / `.ai/AGENTS.md`.
`transformer_sana_wm.py`: 6555 -> 4592 lines.
Conventions:
* `_no_split_modules` was `["blocks"]` — an attribute name, but accelerate
matches on class name, so it matched nothing and `device_map="auto"` could
split a block across devices and crash. Now `["SanaVideoMSCamCtrlBlock"]`.
* Add `_repeated_blocks` (enables `compile_repeated_blocks()`) and
`_skip_layerwise_casting_patterns`. `_keep_in_fp32_modules` is deliberately
left unset with a note: the blocks apply the timestep modulation inline, so
keeping `scale_shift_table` / `t_embedder` in fp32 upcasts the hidden states
and feeds fp32 activations to bf16 weights (caught by a GPU smoke run).
* Expose `num_layers` / `hidden_size` / `num_attention_heads` / `patch_size`
through `register_to_config` instead of hardcoding the release architecture,
so a tiny variant can be built for tests. Defaults are the released values,
so `config.json` and the state dict are unchanged.
* `torch.float64` -> `torch.float32` on the live RoPE paths (gotcha 5),
`torch.empty` -> `torch.zeros` for parameter init (gotcha 6), and stop
reading `self.proj.weight.dtype` to cast activations (gotcha 4).
* `WanRotaryPosEmbed.freqs` is a non-persistent buffer instead of a plain
attribute reassigned inside `forward` (which broke `.to()` and compile).
Dead code (`AGENTS.md`: "delete training-time code paths, experimental flags,
and ablation branches entirely — only keep the inference path"):
* All weight-init / transfer-learning helpers — `from_pretrained` overwrites
them, and they also printed ~20 lines of noise on every construction.
* `CaptionEmbedder.initialize_gemma_params` (fetched `google/gemma-2b-it`
out-of-band at runtime, and referenced an attribute that never exists),
`token_drop`, and the training branch of its forward.
* The cam-debug statistics apparatus, the `save_block_output` hooks (whose
`get_block_output` was never defined), and `block_hook`.
* Both unreachable recurrence variants, `_maybe_drop_cam_branch`, the xformers
branches (`_xformers_available` was a literal `False`, defined twice, and
`xformers` was never imported), 3 env-var escape hatches, 9 unused classes,
12 unused module-level helpers, and the sincos family — which also removes
the NumPy import, satisfying the no-NumPy-in-forward rule.
* Inline the `fp32_attention` mechanism: it was set on every submodule via
`model.apply` and read at 18 sites, so it was folded to the shipped
always-on behaviour rather than deleted.
Ephemeral comments (commit SHAs from a private tree, references to files that
don't exist in diffusers, FSDP2 rationale, stale docstrings) removed.
State dict is unchanged: 871/871 keys match the released checkpoint.
* fix(sana-wm): honour scheduler shift, use randn_tensor, drop cluster scaffolding
From the `self-review` skill run against `.ai/pipelines.md` / `.ai/testing.md`.
`flow_shift` was a dead knob. `__call__` set `self.scheduler.config.shift`,
but `FlowMatchEulerDiscreteScheduler.set_timesteps` reads `self.shift` (i.e.
`self._shift`), so the documented argument had no effect and every run used
the checkpoint's `shift=9.8`. It also half-mutated a `FrozenDict` — the freeze
guard checks a name-mangled attribute, so the assignment silently desynced
`config.shift` from `config["shift"]` on a shared component. Removed the
argument so the scheduler owns the shift (`pipelines.md` gotcha 3), which also
drops a per-call mutation of a registered component (gotcha 7).
Other pipeline fixes:
* `torch.randn` -> `randn_tensor` in `prepare_latents` and the refiner
(gotcha 10). `__call__` advertises `generator: Generator | list[Generator]`,
but `torch.randn` raises on a generator list and on a CPU generator with a
CUDA device, so that path could not work.
* Delete the SLURM preemption/resume feature (`checkpoint_dir`,
`_atomic_save_state`, `_capture_state` / `_restore_state`) and the
single-shot refiner path the docstring itself called a debugging fallback —
roughly 270 lines of research-cluster scaffolding.
* `_empty_cuda_cache` (CUDA-only) -> `empty_device_cache`;
`@torch.inference_mode()` -> `@torch.no_grad()` and removed from inner
helpers the decorator already covers (gotcha 2).
* Remove dead `_callback_tensor_inputs` (no `callback_on_step_end` exists),
`_exclude_from_cpu_offload` (a no-op — the base class already skips
non-`nn.Module` components), `_kv_max_frames`, and `latents.detach()`.
Tests: drop the `@slow` integration stub — `testing.md` says integration and
slow tests don't belong in the initial PR.
`refiner.py`: 1286 -> 1017 lines.
* refactor(sana-wm): drop the timm dependency
Per @yiyixuxu's review. The three timm imports were all trivially replaceable:
* `Mlp` — both local subclasses fully overrode `forward` and only inherited
`fc1` / `act` / `drop1` / `fc2` / `drop2`, so it is now a self-contained
module and the redundant `class Mlp(Mlp)` wrapper is gone.
* `Attention` — `GDN` inherited only `qkv` and `proj` from it (it already
defines its own `q_norm` / `k_norm`), so those two layers are declared
inline and `GDN` subclasses `nn.Module`.
* `DropPath` — `drop_path` was hardcoded to `0.0`, so this was always
`nn.Identity()`. Removed along with the rest of the plumbing.
No `dummy_timm_objects.py` is needed since the dependency is gone rather than
optional, and `setup.py` is untouched (other models still use timm).
State dict is unchanged (871/871 keys match the released checkpoint) and the
end-to-end GPU smoke gives byte-identical output (frame mean 0.5554).
* refactor(sana-wm): drop the Triton kernels and apply review feedback
Per @yiyixuxu's inline review. Removes `transformer_sana_wm_kernels.py`
entirely (3234 lines) — the Triton kernels will be published to
`kernels-community` and wired back in a follow-up PR.
* Delete the three `*Triton` attention classes. Each only overrode `forward`
to call a fused kernel; their pure-PyTorch parents (`BidirectionalGDN`,
`BidirectionalGDNUCPESinglePathLiteLA`) compute the same result. The
Triton-availability probe and the MRO-walking fallback go with them.
The released `config.json` still names the `*Triton` variants, so
`ATTENTION_BLOCKS` maps those strings onto the pure-PyTorch classes.
* Move the nine pure-PyTorch camera-math helpers that were reachable from
the model into `transformer_sana_wm.py`; drop the rest with the kernels
file (they were only reachable from the deleted Triton path).
Other review items:
* Use `get_activation()` from `..activations`; delete the local activation
and normalization registries (`build_norm` was only ever called with
`None`).
* Delete `GLUMBConv` and make `GLUMBConvTemp` self-contained. `ConvLayer` is
slimmed and renamed `SanaWMConvLayer`, but has to stay a module — the
checkpoint keys are `mlp.inverted_conv.conv.weight`, so collapsing it to a
bare `nn.Conv2d` would drop the `.conv` level.
* Remove the closure-returning helpers the review flagged as hard to read
and compile-hostile: `prepare_prope_fns` and friends now return plain
tensors via `_prepare_ucpe_ray_transforms` + `_apply_ucpe_transform`, and
the `_register_block` decorator factory is an explicit dict update.
* Remove all seven `@torch.compile` decorators — users can reach for
`compile_repeated_blocks()` instead.
* Drop `partial` (pass `chunk_size` at the call site), `DWMlp`, the tuple
helpers, `get_same_padding`, `modulate`, an unused `apply_rotary_emb`,
and the validation of internal-only arguments.
* Tests: drop `SanaWMTritonFallbackTests` along with the mechanism it covered.
State dict is unchanged (871/871 keys) and the end-to-end GPU smoke passes
with correct video output.
* refactor(sana-wm): inline the modulation helper, drop the backward-only shim
Two more items from @yiyixuxu's inline review:
* Inline `t2i_modulate` at its four call sites.
* Delete `_IdentityForwardContiguousBackward` / `_contiguous_backward`. It is
the identity in forward and only exists to hand a contiguous gradient to
the backward pass, so it is dead weight in an inference-only port.
State dict unchanged (871/871 keys); GPU smoke output matches the previous
run exactly (frame mean 0.5560).
* Update src/diffusers/models/transformers/transformer_sana_wm.py
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* Update src/diffusers/models/transformers/transformer_sana_wm.py
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* style(sana-wm): drop the blank line left by removing the camera-branch init
* refactor(sana-wm): address @dg845's review round
Transformer:
* Stop stashing patch-grid shape on `self` during `forward` (`self.f/h/w`) and
thread it through as locals; `unpatchify` now takes it explicitly.
* Replace the 7 `assert`s with `ValueError`s.
* Drop `attn_drop`/`proj_drop` from `MultiHeadCrossAttention` (training-only,
and `attn_drop` was never applied) plus a stale training-era banner comment.
Pipeline / refiner:
* Don't mutate the components handed to the pipeline. The VAE tiling +
framewise settings move to the docs, `padding_side="right"` is passed per
tokenizer call, and the `.eval()` calls are gone (no `self.training`
branches remain, and `from_pretrained` already returns eval-mode modules).
* Gate `cam_utils`' optional imports on `is_torchvision_available()` and a new
`is_pi3_available()` helper.
* Annotate `SanaWMLTX2Refiner.__init__` and inline `_refine_latents_ar` into
its single caller.
Conversion script:
* Move to `scripts/` alongside the other Sana converters.
* Raise on missing/unexpected keys instead of printing, so a bad mapping can't
silently emit a broken transformer.
State dict unchanged (871/871 keys). GPU smoke on the public checkpoint is
byte-identical to the previous run (frame mean 0.5560), including with the VAE
settings applied by the caller rather than the pipeline.
* refactor(sana-wm): rename WanRotaryPosEmbed to SanaWMRotaryPosEmbed
The name implied this was Wan's rotary embedding, but it isn't: the per-axis
split is configurable through `fhw_dim`, and the frequencies stay complex in a
single `freqs` buffer instead of being split into real cos/sin buffers. So it
can't carry a `# Copied from`. Renamed, with a docstring recording why.
The buffer is `persistent=False`, so the state dict is unchanged (871/871).
* test(sana-wm): add transformer model tests, migrate pipeline tests to pytest
* `tests/models/transformers/test_models_transformer_sana_wm.py` — generated
with `utils/generate_model_tests.py` and filled in, following
`test_models_transformer_sana_video.py`. The tiny config sets
`softmax_every_n=2` so one block exercises the GDN camera branch and the
other the softmax variant. Dummy inputs supply the conditioning the forward
requires: `encoder_attention_mask`, `(B, F, 20)` camera conditions, and
`chunk_plucker`.
* `tests/pipelines/sana_wm/test_sana_wm.py` — rewritten in the pytest style of
`tests/pipelines/sana_video/test_sana_video.py`: no `unittest`, bare
asserts, module-level imports, and `parametrize` in place of loop-style
cases (15 test functions become 31 cases).
`AttentionTesterMixin` is skipped because the model calls
`F.scaled_dot_product_attention` directly rather than going through a
diffusers attention processor.
* fix(sana-wm): make the transformer layerwise-casting and offload safe
* `TimestepEmbedder.dtype` used `next(self.parameters()).dtype`, which reports
the storage dtype under layerwise casting rather than the compute dtype. Use
`get_parameter_dtype`, which is layerwise-casting aware. (The model-level
`self.dtype` already routes through it via `ModelMixin`.)
* Stop reading submodule `.weight`/`.bias` inside `forward`. Group-offload
hooks fire on a module's `forward`, so reading its parameters directly
leaves them offloaded and trips a device mismatch. The frame-gate and
output-gate helpers now call their submodules, and the fused camera QKV
projection becomes three `q/k/v` calls -- algebraically the same as one
GEMM over the concatenated weights.
* Don't mutate the tokenizer in `SanaWMLTX2Refiner._encode_prompt`; pass
`padding_side="left"` per call, matching `SanaWMPipeline`.
* Add a `torch.compile` test with `recompile_limit=2` -- the repeated block
compiles once per attention variant.
GPU smoke unchanged (frame mean 0.5560).
* refactor(sana-wm): explicit arguments instead of **kwargs, generator instead of seed
Replace the `**kwargs` bag threaded through model -> block -> attention with
explicit keyword arguments. Only three runtime keys were ever read at the
leaves (`frame_valid_mask`, `precomputed_gates`, `ucpe_ray_transforms`); the
rest were forwarded and silently swallowed. `camera_embedding`,
`chunk_index`, `chunk_index_global` and `chunk_split_strategy` turned out to
be pure dead plumbing -- written into the per-block kwargs dicts and never
read by any attention or MLP forward -- so they are gone.
The pipeline also grows the standard diffusers arguments:
* `prompt_embeds` / `prompt_attention_mask` / `negative_prompt_embeds` /
`negative_prompt_attention_mask` on `encode_prompt` and `__call__`.
* `seed` / `refiner_seed` are replaced by `generator` / `refiner_generator`.
`generator=torch.Generator(device).manual_seed(42)` reproduces exactly what
`seed=42` used to build, so results are unchanged.
State dict unchanged (871/871). CPU old-vs-new equality on a tiny config over
both attention variants is exact (`torch.equal`, max diff 0.0), and the GPU
smoke on the public checkpoint still gives frame mean 0.5560.
* refactor(sana-wm): drop the unreachable chunk-split strategies
@yiyixuxu was right that the model only ever runs uniform chunking; my earlier
reply defending the strategies was wrong. The two sites that actually chunk
both call `normalize_chunk_index(None, T, chunk_size)` with three positional
arguments, so `chunk_split_strategy` always took its `"uniform"` default. The
configured `first_chunk_plus_one` reached the per-block kwargs dict and was
then swallowed by the attention forwards' `**kwargs` without ever being read.
Flattening that bag into explicit arguments is what surfaced it.
Also note the `chunk_size` those call sites use is `chunk_gdn_chunk_size`, a
different attribute from the `chunk_size` that was being threaded through.
So `chunk_index_from_chunk_size`, `normalize_chunk_index`,
`is_uniform_chunking` and `compute_chunk_sizes` are gone, the uniform
boundaries are inlined at both call sites, and `chunk_split_strategy` is
dropped from the model and block constructors.
The released `config.json` still carries the key; loading is unaffected
(`extract_init_dict` ignores it) but it logs an "not expected and will be
ignored" warning, so it should come out of the checkpoint config on the next
export.
State dict unchanged (871/871); GPU smoke still frame mean 0.5560, which
confirms those branches were never taken.
* refactor(sana-wm): use the shared RMSNorm
@yiyixuxu asked whether the shared `RMSNorm` could be used here and my earlier
reply overstated the obstacles. Re-checking each one:
* `scale_factor` is not a blocker. `attention_y_norm.weight` is one of the 871
checkpoint keys, so `from_pretrained` overwrites whatever the constructor
initialised -- exactly the same reasoning that removed the other
hand-written inits in this PR. It only ever affected from-scratch models.
* `norm_dim` is not a blocker either; it is always the default `-1`.
* The numerics do differ: ours ran the normalisation and the weight multiply
in fp32, while the shared class computes only the variance in fp32 and does
the scaling and weight multiply in the input/weight dtype.
Measured rather than argued: the end-to-end GPU smoke on the public
checkpoint moves from frame mean 0.5560 to 0.5561, and the decoded video is
visually identical. That is well inside bf16 noise, so the local class is not
worth keeping.
The now-unused `y_norm_scale_factor` config argument goes too. As with
`chunk_split_strategy`, the released `config.json` still carries it and will
log an "not expected and will be ignored" warning until the checkpoint config
is re-exported.
* refactor(sana-wm): migrate MultiHeadCrossAttention to the diffusers attention API
* Split the attention maths into a `SanaWMCrossAttnProcessor` and make
`MultiHeadCrossAttention` an `AttentionModuleMixin`, so the processor can be
swapped and the standard tooling applies.
* Use `dispatch_attention_fn` instead of calling
`F.scaled_dot_product_attention` directly, which lets the model pick up
alternative attention backends. This also drops the two transposes, since
the dispatcher takes `(batch, seq, heads, dim)`.
* Build a boolean padding mask rather than an additive float one. The old
`(1 - mask) * -10000.0` form isn't supported by the varlen backends.
Projection names are unchanged (`q_linear`, `kv_linear`, `proj`, `q_norm`,
`k_norm`), so the state dict stays at 871/871 keys. The GPU smoke on the
public checkpoint is unchanged at frame mean 0.5561, so `-inf` masking makes
no difference against the old `-10000.0` at this precision.
* feat(sana-wm): give the refiner its own transformer model
Per @yiyixuxu and @dg845: the stage-2 refiner no longer reaches into
`LTX2VideoTransformer3DModel`. It used to drive AR refinement by setting
attributes on LTX-2's submodules (`attn._tf_kv_prefix`, `_kv_cache_capture`,
`_tf_capture_kv`), which meant a second `forward` implementation smuggled in
from the pipeline.
New `transformer_sana_wm_refiner.py`, following the `transformer_wan_vace.py`
pattern of importing the reusable pieces rather than duplicating them:
* `SanaWMRefinerKVLayerCache` / `SanaWMRefinerKVCache`, modelled on
`Flux2KVCache`, holding the per-layer sink and history K/V.
* `SanaWMLTX2RefinerTransformerBlock`, which subclasses
`LTX2VideoTransformerBlock` and overrides only `forward`, so the submodule
structure stays byte-identical.
* `SanaWMLTX2RefinerTransformer3DModel`, taking `kv_cache` / `kv_cache_mode`
explicitly through `forward`.
`refiner.py` drops from 982 to 638 lines: the four `*_kv_prefix_on_blocks` /
`*_captured_kv_*` helpers, `_forward_video_block`, `_streaming_self_attention`
and `_forward_video_only_with_rope` are all gone, along with a dead
`n_context_tokens > 0` branch that neither call site could reach.
The state dict is unchanged at 3510 keys, identical to
`LTX2VideoTransformer3DModel` at the same config, so only the class name in
the checkpoint metadata has to change. CPU equivalence against the old path
is exact across four AR blocks -- covering sink injection, history growth and
history trimming -- with a negative control confirming the cache path is
actually exercised. The GPU smoke on the public checkpoint is unchanged at
frame mean 0.5561.
* refactor(sana-wm): run the refiner as a separate pipeline
Per @dg845, the two stages are no longer nested. `SanaWMPipeline` produces
stage-1 latents and `SanaWMLTX2Refiner` consumes them:
latents = SanaWMPipeline(...)
video = SanaWMLTX2Refiner(...)
The refiner takes an optional `vae` -- pass the base pipeline's so the weights
are shared, the way SDXL shares components between base and refiner -- and
decodes to video when given one. The manual device juggling in
`SanaWMPipeline.__call__` is gone with the nesting that forced it; each
pipeline now manages its own placement through the standard offload hooks.
Un-nesting turned up a real bug: `torch_dtype` never reached the nested
sub-pipeline, so `SanaWMPipeline.from_pretrained(..., torch_dtype=bfloat16)`
silently ran the whole of stage 2 in float32 -- transformer, text encoder and
connectors alike. That is where the "~87 GB of refiner weights" in the old
comment came from. Loaded on its own the refiner honours the requested dtype,
which halves its footprint; the GPU smoke moves from frame mean 0.5561 to
0.5548 and the decoded video is visually identical.
The conversion script writes the refiner to its own output directory
(`--dst-refiner`, default `<dst>-refiner`) rather than a `refiner/`
subfolder. `DiffusionPipeline.from_pretrained` has no `subfolder` argument, so
a nested folder silently loads the *base* pipeline's same-named component
folders instead -- exactly the kind of wrong-weights failure that only
surfaces because `connectors/` happens not to exist at the top level.
* Apply style fixes
* docs(sana-wm): document the refiner's output_type argument
`check_forward_call_docstrings` flagged `SanaWMLTX2Refiner.__call__` as
missing an entry for `output_type`, which the un-nesting added.
* refactor(sana-wm): support only what the released checkpoint uses
@yiyixuxu asked whether we can support just the released configuration and
flatten what's left. We can — the checkpoint only ever instantiates two
attention classes, so the other five in the chain were dead weight.
* Collapse `GDN -> BidirectionalGDN`, `GDN -> _GDNUCPEBase`,
`BidirectionalGDNUCPELiteLA -> ...PostUCPERenorm -> ...SinglePathLiteLA`
into a self-contained `BidirectionalGDNUCPESinglePathLiteLA` plus
`_SoftmaxUCPESinglePathLiteLA`. Inheritance depth goes 5 -> 1.
* Blocks take the class, not a string: `SanaVideoMSCamCtrlBlock` gains an
`attn_cls` argument, and the model picks per layer with a plain loop. That
removes the `ATTENTION_BLOCKS` registry, `_resolve_attention_block`,
`_GDN_TO_SOFTMAX_CAMCTRL`, `_inject_softmax_layers`, the `attn_type` /
`camctrl_type` validation, and the unreachable no-camera branch.
* `attn_type` and `camctrl_type` leave the config signature. As with
`chunk_split_strategy` and `y_norm_scale_factor`, the released
`config.json` still carries them and loading logs a benign "not expected"
warning until the checkpoint is re-exported.
* Also drop the zero-init on `GLUMBConvTemp.t_conv` (overwritten by
`from_pretrained`) and inline the one-line `restore_shape` helper. The
model docs lose a paragraph describing an `_inner.` prefix and an
`add_inner_prefix` helper that no longer exist.
3908 -> 3153 lines. State dict unchanged at 871/871 keys with identical
ordering, and 15 GDN + 5 softmax blocks as before. Old-vs-new CPU equality is
exact (`torch.equal`, max diff 0.0) on both a tiny config and the released
geometry, across a `softmax_every_n` sweep and on the backward pass, with
negative controls confirming the comparison isn't vacuous. The GPU smoke is
byte-identical too: refined latents mean 0.035024 / std 0.887467, frame mean
0.5548.
* refactor(sana-wm): make the softmax attention class standalone
Drops the last bit of inheritance between the two attention classes, per
@yiyixuxu. `_SoftmaxUCPESinglePathLiteLA` now subclasses `nn.Module`
directly.
It turned out to be a simplification rather than duplication: of the parent's
thirteen methods it only ever called three (`_apply_output_gate`,
`_prepare_frame_valid_masks`, `_stabilize_cam_transforms`), and the other ten
were dead for it. It is also always built with `conv_kernel_size=0`, so the
whole short-convolution half of the inherited `__init__` -- `conv_q/k/v`,
the `*_cam` variants, `kernel_func` -- was permanently `None` or unused and
is now gone.
State dict unchanged: 871/871 keys in the same order, and the softmax blocks
still emit the same 24 keys each. Old-vs-new equality is exact
(`torch.equal`, max abs diff 0.0) on two input configurations, with negative
controls from 0.109 to 14.62 confirming the comparison isn't vacuous; the
inlined `__init__` even consumes the RNG stream identically, so seeded
from-scratch weights are bit-for-bit the same. GPU smoke unchanged at
refined-latent mean 0.035024 / std 0.887467.
* feat(sana-wm): export the refiner weights in bfloat16
The refiner's Gemma-3 text encoder was the one component still stored in
float32, at 45.4 GB. Since the pipeline is loaded with
`torch_dtype=torch.bfloat16`, those weights were cast on every load anyway,
so storing them cast changes nothing and halves the download:
text_encoder 45.4 GB -> 22.7 GB
refiner total 84 GB -> 61 GB
The VAE stays in float32, as the docs recommend. Verified rather than
assumed: the GPU smoke is bit-identical (refined latents mean 0.035024 / std
0.887467, frame mean 0.5548), and it now runs in 124s instead of 384s because
there is far less to load.
* refactor(sana-wm): pipeline clean-up from @yiyixuxu's review
`SanaWMLTX2Refiner`:
* Dissolve `_RefinerChunkRunner` (191 lines) into `__call__`, so the AR loop
reads top-to-bottom like `pipeline_wan_animate.py`, and swap the bare
`tqdm` for `self.progress_bar(...)`.
* Actually support `enable_model_cpu_offload()`. The docs already advertised
it, but three things blocked it: the `device` default resolved through
`next(self.transformer.parameters()).device`, which is `cpu` under offload;
`_encode_prompt` called the inner Gemma backbone directly, bypassing the
hook that `cpu_offload_with_hook` installs on the top-level module; and
`_decode_latents` read `self.vae.device`. All three now go through
`self._execution_device`, and the manual `.to(...)` juggling is gone.
* TODO on the `LTX2TextConnectors` import pointing at
huggingface/diffusers#14749.
Two argument removals fall out: `progress=` (use
`set_progress_bar_config(disable=True)`, which is what every other pipeline
does) and `device=`.
`SanaWMPipeline`:
* `check_inputs` only validates now; the camera-trajectory normalisation
moves to `prepare_camera_trajectory`.
* Add the standard `guidance_scale` / `do_classifier_free_guidance` /
`num_timesteps` / `current_timestep` / `interrupt` properties.
* `return_dict=False` returned only `(frames,)`, silently dropping `c2w` and
`latent`. It now returns the same three values either way -- the flag picks
the container, not the contents.
* Take a `PIL.Image`, not a path, and drop the `_model_cpu_offload_active`
helper that had no callers left after the un-nesting.
The Pi3X intrinsics estimator leaves `diffusers` entirely: it pulled in an
extra dependency and fetched a second checkpoint without the user ever
calling `from_pretrained`. It is packaged as a remote modular block hosted
alongside the checkpoint, and the docs show `ModularPipeline.from_pretrained`
instead. `is_pi3_available` goes with it.
Equivalence for the refiner rewrite is exact across six configurations --
varying sink size, block size, KV window and seed so history growth and
trimming both fire -- `torch.equal` with max abs diff 0.0, negative controls
from 1.9e-2 to 5.55. GPU smoke on the public checkpoint is unchanged at
refined-latent mean 0.035024 / std 0.887467.
* refactor(sana-wm): model-side clean-up from @yiyixuxu's review
* Split each attention class's main and camera branches into their own
modules, so `forward` is a flat list of layer calls. The branch modules
deliberately hold no parameters: if they owned `qkv` / `q_proj_cam` /
... PyTorch would prefix by attribute name and rename 390 of the 871
checkpoint keys. (`state_dict` duplicates a module aliased in two places
rather than de-duplicating it, so aliasing isn't a way out either.) The
compute is split; the ownership isn't.
* `transformer_sana_wm_refiner.py` is self-contained -- nine `# Copied from`
markers replace the `transformer_ltx2` import, at whole-function,
whole-class and method-level granularity as each piece required. That costs
+831 lines, since the copied block `__init__` drags in both attention
processors.
* The refiner's streaming self-attention is a real attention processor,
`SanaWMLTX2RefinerAttnProcessor`, and the block calls `self.attn1(...)`. The
processor swap happens in the model `__init__` after the blocks are built,
because `attn1` is constructed inside the `# Copied from` body.
* Inline the KV-cache mode strings, drop `plucker_proj`'s zero-init.
State dicts unchanged: 871/871 and 3510/3510, ordered lists identical.
Equivalence is exact -- max abs diff 0.0 -- against the real released
checkpoint for the attention split, and across five refiner configurations
driven through two AR blocks, with negative controls throughout.
`check_copies` passes. GPU smoke unchanged at refined-latent mean 0.035024 /
std 0.887467.
* test(sana-wm): cover the GDN branch, which was silently dead in these tests
`SanaWMTemporalShortConvolution` allocates its weight with `torch.zeros` and
nothing initialises it any more -- the hand-written init came out earlier in
this PR because `from_pretrained` overwrites it. That is true for real
checkpoints, but a from-scratch model gets a zero `conv_k`, so the GDN key is
identically zero and the whole branch returns zeros. Every test in this file
builds a random model, so they were all passing while exercising none of the
GDN maths.
Two tests: one pins the zero-init as deliberate rather than accidental, the
other fills the short convolutions with noise and asserts the output moves.
Checked that the second one bites -- forcing the short convolution to return
zeros makes it fail.
* [SANA-WM] Address review: reuse diffusers embeddings, grad checkpointing, callbacks, offload fixes
- Reuse Timesteps/TimestepEmbedding; drop unused cfg_embedder / y_embedding / model_max_length
- Rename Mlp/PatchEmbed/CaptionEmbedder with SanaWM prefix
- Support gradient checkpointing in SanaWMTransformer3DModel
- Keep complex RoPE freqs out of buffers so Module.to(dtype) doesn't drop the imaginary part
- Disable autocast around camera geometry (torch.cross)
- Pipeline: callback_on_step_end, Literal output_type, _execution_device, maybe_free_model_hooks
- Fix output docstring, bump copyright to 2026
- Port tests to new pytest mixins, add refiner transformer tests
* [SANA-WM] Drop cam_utils and registration tests
* [SANA-WM] Add return type annotations to forward methods
---------
Co-authored-by: HaoyiZhu <hyizhu1108@gmail.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
|
||
|
Wang, Yi
|
1e004eb6ec |
xpu supports TP as well (#15013)
xpu support TP as well Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> |
||
|
Wang, Yi
|
883d51ef02 |
Make Stable Diffusion device-map comparisons deterministic (#15010)
Description: The balanced device-map tests compare images from single-device and multi-device inference. On XPU, repeated fp16 inference can produce different results even with the same seed, causing the pixel comparison to fail. Enable PyTorch deterministic algorithms for the two image-comparison tests and restore the previous setting afterward. This keeps other tests unaffected. Validation: Both affected tests passed on a multi-XPU runner (RUN_NIGHTLY=1 RUN_SLOW=1). Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Lucain
|
b1f6cde2ac |
Bump huggingface_hub upper bound to <3.0 (diffusers) (#14863)
Bump huggingface_hub upper bound to <3.0 Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: célina <hanouticelina@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Lucas Ruan
|
e1a6c67dc7 | Fix Anima text encoder exports and standalone core denoising (#14943) | ||
|
YiYi Xu
|
85c9fa4cae |
[cli] detect composite block classes in custom_blocks (#15001)
`diffusers-cli custom_blocks` only recognized classes whose base was spelled exactly `ModularPipelineBlocks`, so a `SequentialPipelineBlocks` or `AutoPipelineBlocks` subclass (the usual thing to publish) was reported as "could not be retrieved". Accept the composite subclasses as well, and match on the class name so dotted bases like `diffusers.modular_pipelines.SequentialPipelineBlocks` work too. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> |
||
|
Jingya HUANG
|
1d5d056ecf |
[TPU] TorchTPU backend integration - eager / torch.compile / tp (#14039)
* feat: add torchtpu * feat:draft TorchTPU support * fix: wan overflow issue + compile mode error on sdxl * doc: enhance with TorchTPU doc * doc: enhance with TorchTPU doc * style: remove unused imports flagged by ruff Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs: remove Debug Eager and Fused Eager sections from tpu.md Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat: add TP for TPU * style: fix import sorting in TPU test scripts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: ruff format TPU test scripts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: style * fix: test for native 4 devices * fix: propagate TPU device fixes to Flux/Flux2/Wan-family copies; register tpu.md; drop redundant execution_device check - Propagate the text_encoder.device-based fix (introduced for TPU CPU-offload support) from FluxPipeline/Flux2KleinPipeline/WanPipeline into their `# Copied from` copies (flux/*, flux2_klein_inpaint, visualcloze, anyflow, chronoedit, lucy_edit, skyreels_v2/*). SDXL-family copies of StableDiffusionXLPipeline.encode_prompt are intentionally left untouched; they'll be handled in a follow-up PR that fixes device placement for every pipeline component (not just text encoders). - Register docs/source/en/optimization/tpu.md in _toctree.yml (was breaking the docs build: "not present in the table of contents"). - Remove the redundant "prefer non-CPU, non-meta component" loop from DiffusionPipeline._execution_device: PR #14383 already fixed this in DiffusionPipeline.device, which _execution_device falls back to. Verified on TPU hardware that _execution_device still resolves correctly for a split-placement pipeline after the removal. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * tests: cleanup * fix: fix compile mode * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * test: remove flux2 e2e test * doc: apply suggestions * doc: apply suggestions * review: remove monkey patch * review: revert neuron-specific changes in the tests * review: apply suggestions * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: add tpu to tp doc * review: remove unnecessary for tpu * test: assert every _tp_plan parameter is sharded after a TP load Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * removal: delete redundant tp shard helpers * removal: drop TPU workarounds no longer needed after recent main changes * doc: keep the FLUX.2-dev text encoder off a single TPU chip in the TP example * doc: encode the prompt on CPU in the TPU TP example so it runs end to end * test: tighten TPU TP tolerance and simplify the TPU TP worker * doc: use enable_model_cpu_offload() in the TPU eager example * doc: shard both FLUX.2-dev text encoder and transformer in the TPU TP example * doc: run the text encoders on TPU in the compiled example * test: drop the redundant TPU sync in the TP worker * test: trim comments in the TPU TP tests * test: run the TPU TP test with mp.spawn like the CUDA one * test: pin the TPU TP test to 4 chips * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: use DistributedConfig and dtype, and drop the torch_tpu import in the TPU examples * test: skip the TPU TP test for models without a _tp_plan * doc: correct the warning of recompilation * fix: clone the SD3 timestep so torch.compile doesn't recompile every step, thus we can use it for the tpu.md compile mode snippet * doc: update advice for compile mode when one chip doesn't fit * test: pass the TPU TP tolerances as test arguments * review: apply suggestions to the doc * review: view compilation patch for tpu-only --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: JingyaHuang <huang_jingya@outlok.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Christopher
|
78bf2c6266 |
[LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) (#14163)
* [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) Adds loading of LoKr (LyCORIS Kronecker product) adapters: - `load_lora_adapter` detects `lokr_` keys and injects a peft `LoKrConfig`, inferred from the tensor shapes via `_create_lokr_config` (decompose factor, per-module rank/alpha patterns). - State dict conversions for the formats in the wild: ai-toolkit Z-Image (dotted diffusers paths under `diffusion_model.`), ai-toolkit BFL Flux2 (fused qkv), LyCORIS underscore format, and bare dotted diffusers paths. - BFL fused-QKV LoKr cannot be split exactly into separate Q/K/V Kronecker factors, so `Flux2LoraLoaderMixin.load_lora_weights` fuses the model's QKV projections and maps the adapter 1:1 (exact). - Alpha follows the LyCORIS convention: scaling applies only to rank-decomposed factors and is baked into the weights at conversion. Fixes #13221 * [LoRA] LoKr review follow-ups: adapter-agnostic error message, refuse fused-QKV load over an unfused adapter - "Invalid adapter checkpoint. We currently support LoRA and LoKr." replaces the message that still said "LoRA checkpoint" and described the substring check. - Loading a fused-QKV LoKr checkpoint now refuses when an adapter is already injected on to_q/to_k/to_v or add_{q,k,v}_proj: fuse_qkv_projections() would replace those modules and orphan it. Covered by a test. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * [LoRA] rename the LoKr tests to test_lokr_conversion.py The file exercises checkpoint conversion and loading for LoKr adapters, not LoRA. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [LoRA] LoKr review follow-ups: scope to Z-Image/Flux2, model-level fused QKV, LoKrTesterMixin - Only the Z-Image and Flux2 loader mixins accept LoKr checkpoints, the two that convert them; the other mixins are back to their LoRA-only check. - The fused-QKV handling moves to `_maybe_fuse_qkv_projections_for_lokr` in peft_utils, called from `load_lora_adapter`, so the model fuses its projections whichever entry point loads the adapter. The quantized-model error explains why. - `_bake_lokr_alpha` becomes `_bake_lokr_alpha_` as it edits the dict in place. - The BFL Flux2 converter no longer accepts expanded diffusers block names under `diffusion_model.`: no LoKr checkpoint uses that layout; diffusers-named checkpoints have no prefix and take the generic converter. - The LyCORIS converter names the keys it does not recognize. - Tests move from tests/lora to a `LoKrTesterMixin` in tests/models/testing_utils, with the checkpoint formats tested on the Flux2 and Z-Image model test classes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [LoRA] mark ZImageLoraLoaderMixin.load_lora_weights as copied from Flux2 --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
d7d0b9c2c7 |
[ci] consolidate labeling on PRs. (#14911)
consolidate labeling on PRs. Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
Patrick Ribbsaeter
|
d961a388fd |
fix(image_processor): maintain 3-tuple return contract in InpaintProcessor.preprocess when mask is None (#14481)
fix(image_processor): preserve 3-tuple return contract when mask is None (#14470) Ensure InpaintProcessor.preprocess returns (image, None, postprocessing_kwargs) when mask is None, preserving the expected 3-tuple return contract and signature consistency across all preprocess branches. Consolidate unit tests into a parameterized test covering with-mask, without-mask, and crop configurations. |
||
Subodh
|
e0415c41f8 |
Fix active adapter detection for text encoders (#14961)
* Fix active adapter detection for text encoders Signed-off-by: Subodh Bansode <subodh171011@gmail.com> * test: move text encoder adapter regression test Signed-off-by: Subodh Bansode <subodh171011@gmail.com> --------- Signed-off-by: Subodh Bansode <subodh171011@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Steven Liu
|
122b1e11fd |
[docs] Community methods (#14973)
* freeu * cachedit |
||
|
Steven Liu
|
c6df88a511 |
[docs] Jobs guide (#14926)
hf jobs Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |