mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-24 20:59:27 +03:00
linked-issue-self-opened
1390 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
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>
|
||
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 |
||
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> |
||
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> |
||
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>
|
||
|
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> |
||
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> |
||
|
Sayak Paul
|
bc5e3bd9c5 |
Kandinsky6 (#14949)
* implement Kandinsky 6 - add TI2VA and SR pipelines - additonally implement PiflowScheduler and MMAudioVAE * update docs for Kandinsky 6 * remove redundant bigvgan code and manual checkpoint loading inside model code * add missing _no_split_modules * use attention processor in SR * fix unguarded dependencies - removed: einops, pydantic - lazy: torchvision, av, librosa * remove dead manual weight loading code; flattent MMAudioVAE structure * move attention processors to transformer files * remove duplicate code in K6 latent_upscaler * address 1st batch of comments * patch transformer k6 * patch SR * patch pipeline style * replace TimestepEmbedding and fix batching * refactor ti2va dit and pipeline * refactor piflow scheduler * refactor SR * patch sr transformer autocast * major refactor * untrack notebooks directory Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * patch mmaudio VAE * patch mmaudio VAE * fix inits * fix flex attn * refactoring * remove __all__ * separate MMAudio VAE and Vocoder * fix naming * refactor checkpointing, magcache, transformer, sr vae for K6 * refactor * address apply_scale_shift_norm problem * update docs * apply style chanes * address review feedback on transformer, SR VAE, and flex attention Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * address comments * fix critical bugs * Revert checkpoint-breaking module renames from |
||
naykun
|
da1d3829cf |
Configure fixed sampling sigmas on QwenImage21Pipeline (#14950)
* feat(qwenimage21): configure sampling sigmas on the pipeline * docs(qwenimage21): clarify default and runtime sampling sigmas |
||
|
Sayak Paul
|
cff9dafbbf |
[modular] Echo team joy future academy jd echo by @Gelercatty (#14941)
* Add Echo modular pipeline * Address Echo modular pipeline review feedback * Inline Echo input list normalization * Align Echo latent boundaries and batch expansion * Polish Echo documentation per review * Load official Echo checkpoint and separate block tests * consolidate testing. * formatting --------- Co-authored-by: Yanwen Ma <gelercat117@gmail.com> |
||
Dhruv Nair
|
578c9b2c66 |
Add single file support for Krea 2 (#14914)
update |
||
|
Dhruv Nair
|
acbabca385 |
Consolidate torch device backend dispatch (#14792)
* introduce TorchDeviceBackend * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Dhruv Nair
|
a91ff967f5 |
[Modular] Avoid downloading weights when loading from an existing local path (#14797)
* avoid redownload from local path * update * update * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Jingya HUANG
|
863092ffec |
[core] Shard tensor-parallel checkpoints on load and save (#14544)
* [core] Shard tensor-parallel checkpoints on load and save Stream each rank's slice of a tensor-parallel checkpoint straight off disk instead of materializing the full checkpoint on every rank and resharding it afterwards, and gather the shards back on save. - `from_pretrained(..., parallel_config=TensorParallelConfig(...))` resolves the shard specs on the still-meta model, then slices each safetensors tensor before the dtype cast, so host memory peaks at ~1/tp_degree of the checkpoint. - `save_pretrained` all-gathers the DTensors into an ordinary checkpoint, or writes a distributed checkpoint with `dcp=True` so no full tensor is ever formed. The writing `tp_degree` is recorded, since a packed weight's stored layout is interleaved by it. - Factor the plan interpretation out of the Neuron pre-shard path into shared `TPShardSpec` / `resolve_tp_shard_specs` / `_local_shard` / `_hooks_only_styles` helpers, so both backends and both the load and save paths shard identically. * Raise when tensor parallelism is combined with quantization, offloading or LoRA Addresses the remaining two items of the review on #13718: tensor parallelism was rejected alongside quantization and `device_map` only on the `from_pretrained` streaming path, while `enable_parallelism` — which the quantization error message itself recommended — accepted a quantized, offloaded or adapter-injected model and sharded it anyway. - Add `_check_tp_model_state`, called from `apply_tensor_parallel`, the one chokepoint every TP entry point funnels through. It rejects a model that is quantized, group-offloaded, placed by accelerate (`device_map` or CPU offload), or has PEFT layers injected. Placed before the device-type check so the reported reason is the useful one. - Guard the reverse order too: `enable_group_offload`, the two pipeline CPU-offload methods, and `load_lora_adapter` now refuse a tensor-parallel model. - `save_pretrained` refuses a quantized tensor-parallel model. Previously the `dcp=True` branch returned before the quantizer's serialization step, writing shards with no quantization metadata and no error. - The DCP load guard checked the `quantization_config` kwarg only, so a pre-quantized checkpoint directory loaded silently; check the config's own entry too, and add the missing `_tp_plan` check that otherwise surfaced as a raw `AttributeError`. - Correct the `from_pretrained` message and the doc sentence that pointed at `enable_parallelism` as a way to shard a quantized model. The new tests are the first tensor-parallel tests that need neither an accelerator nor more than one rank: every case asserts a raise before any collective, so they run single-process on gloo. * review: remove tp save/ dcp related * review: add a new helper . * review: moved the TP checks into enable_parallelism * review: move the warning under caution block * review: add loading time numbers * review: move tp size dividende check * review: delete test * review: consolidate tp checks * review: consolidate the tp checks, all in tensor_parallel.py * review: improve conditionals * review: debug info * review: check weights format before instantiating the model * review: restore duplicated _find_mismatched_keys since irrelevant * revert: delete duplicated helper * review: update doc for loading benchmark of flux2-dev * review: improve the raising check * review: add an args section * Test tensor-parallel loading from sharded checkpoints --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
yzhautouskay
|
78405ec9a1 |
[Cosmos3] Fix Transfer SeaCache artifacts with control CFG (#14897)
* Fix Cosmos 3 Transfer SeaCache indicators * Clarify SeaCache indicator guidance; tests refactor --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
56aa3cd669 |
[docs] document release philosophy (#14904)
* document release philosophy * Update docs/source/en/conceptual/contribution.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/conceptual/contribution.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/conceptual/contribution.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * clarify how folks can reach out to coordinate a release. --------- Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> |
||
|
Sayak Paul
|
fef717ffb0 |
[core] propagate sage attention updates. (#14584)
* propagate sage attention updates. * add sage blackwell. * up * style. * address feedback * style. |
||
|
Sayak Paul
|
4ac08e9408 |
[diffusers-cli] add a dockerfile for diffusers cli hf jobs (#14550)
* add a dockerfile for diffusers cli hf jobs * up * up * use diffusers from pypi * fix(ci): harden GitHub Actions workflows (#14550) (#14851) fix(ci): harden workflow files flagged on #14550 Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com> * Apply batched suggestions from code review Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * lower bound for tokenizers. --------- Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
|
Steven Liu
|
4295ee3ec5 |
[docs] Adapters (#14842)
* lora * legacy adapters * ip adapter * controlnet * feedback |
||
|
Sayak Paul
|
0121a91f9d |
test and docs for xet dedup based file sharing. (#14835)
* test and docs for xet dedup based file sharing. * up[ |
||
|
Dhruv Nair
|
f2edf6ea30 |
Add single file support for Qwen 2.1 (#14836)
update |
||
|
Steven Liu
|
7263f3317f |
[docs] Using diffusion pipelines (#14805)
* refactor * more refactor * improve diagram, fix link * legacy ckpts |
||
Aritra Roy Gosthipaty
|
71f7c96a47 |
[docs] typo with hyperlink (#14826)
typo with link |
||
|
yzhautouskay
|
344d6e7300 |
Add SeaCache support for Cosmos3 pipelines (#14663)
* Add optional FP32 sampling state for Cosmos3 * Add cache execution contexts to Cosmos3 pipelines * Add SeaCache support for Cosmos3 * Cap consecutive SeaCache reuses * Align SeaCache defaults with Cosmos3 inference * Enable SeaCache by default for Cosmos3 * Document SeaCache for Cosmos3 * Use FP32 sampling state by default for Cosmos3 * Fix SeaCache reconfiguration, transfer state isolation, and FP32 sampling * Remove SeaCache stats and ablation instrumentation * Move SeaCache to transformer-level and disable by default * Align SeaCache and tests with model-level cache patterns * Make Cosmos3 sampling state always FP32 * Add pipe definition to the docs snippet * Remove redundant hook registry cache resets * Use raw vision latents for SeaCache indicators * Make Cosmos3 SeaCache safe for regional compilation * Fix SeaCache with Cosmos3 model parallelism * Support single-stream models in SeaCache * Add Wan T2V support to SeaCache * pass denoising step info to cache hooks via cache_context Replaces SeaCache's pipeline-state callbacks: the pipeline attaches the loop coordinates (step_index, sigma, num_inference_steps) to the denoiser call with `cache_context(name, ...)`, carried by a frozen CacheContext dataclass and stored on each stateful hook's StateManager (the registry no longer holds context). Co-Authored-By: Dhruv Nair <dhruv.nair@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Fix Wan cache metadata and Cosmos modular cleanup Pass denoising step metadata to Wan T2V conditional and unconditional cache contexts. Restore stateful cache reset after Cosmos modular pipeline calls, restore the required Literal import, and update regression tests and documentation. * Refactor Cosmos3 SeaCache decoder hook * Apply CI formatting updates * Fix distilled system prompt fallback * Preserve Cosmos3 fp32 sampling state --------- Co-authored-by: yiyixuxu <yixu310@gmail.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
naykun
|
6256aa7666 |
Add Qwen-Image 2.1 (#14804)
* feat: add Qwen-Image 2.1 pipeline with block-causal attention and KV cache New classes: - QwenImage21Transformer2DModel: single-stream transformer with block-causal attention and t=0 modulation for the text and condition-image prefix - AutoencoderKLQwenImage21: 64-channel VAE (z_dim=64, decoder_base_dim=144) - QwenImage21Pipeline: text-to-image and image-conditioned generation Attention: - Block-causal: the joint text/image sequence is causal while each image block (condition and target) stays internally bidirectional. Built as a compiled flex_attention BlockMask, which keeps the score matrix block-sparse and makes 2048x2048 feasible. - flex_attention is optional. Without it the mask is approximated by a two-pass prefill (prefix causally, then the target image over the cached prefix). Exact for text and for the target image, approximate for condition images. KV cache: - The text and condition-image prefix is modulated from t=0, so its activations do not change across denoising steps and its keys and values are cached after the first step. Later steps only recompute the target image's tokens. Also: separate text-to-image and image-conditioned prompt templates with image-pad token downsampling, and plain classifier-free guidance. Includes model tests, docs, and full registration. * fix style * refactor: address PR review feedback for Qwen-Image 2.1 - Remove `causal_block` config flag (always on for released checkpoint) - Split attention into QwenImage21FlexAttnProcessor and QwenImage21SDPAAttnProcessor - Replace two-pass approximate SDPA prefill with exact multi-pass prefill (each image block gets bidirectional attention, text gets causal mask) - Refactor KV cache to QwenImage21KVCache/QwenImage21KVLayerCache classes with explicit kv_cache_mode="extract"/"cached"/"extend" - Lazy-compile flex_attention on first use (fixes OOM on uncompiled path) - Fix edit pipeline: remove broken _downsample_image_pad_tokens, add mm_token_type_ids for transformers 5.x, auto-convert RGB to RGBA - Use @apply_lora_scale decorator, extract _IMG_TOKENS_PER_SLOT constant - Add # Copied from markers for retrieve_latents and _encode_vae_image - Fix mutable default feat_idx=[0] in all 8 VAE forward methods - Update docs: add usage snippet, remove stale causal_block references - Delete examples/qwenimage21/ (snippet moved to docs) * fix: KV cache pinned the whole prefill sequence at batch size 1 The "extract" branch stored the prefix as `key[:, cache_write_slice].contiguous()`. At batch size 1 that slice already counts as contiguous, because PyTorch ignores size-1 dimensions in the check, so `contiguous()` returned the same view and the cache pinned the full prefill K/V for every step of the denoising loop: 8.0 GiB of resident memory at 2048x2048 across 32 layers. At batch size 2 and above the slice is not contiguous, the copy happens, and the leak disappears, so no test caught it. Store a `clone()` instead, and assert in the tests that the cached prefix owns its storage. Measured at 2048x2048 on one H100, bf16, batch 1, 20 steps: peak memory for a full pipeline call drops from 64.5 GiB to 56.5 GiB. * move the block-causal segmentation into the SDPA processor Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WpU4nuugzCw8T1c4tL6N2A * refactor: let the processor pick the block-causal path Builds on the previous commit by @yiyixuxu, which moved the segmentation into the processor. - Drop the internal `torch.compile(flex_attention)`. The convention is that the caller compiles, so `QwenImage21FlexAttnProcessor` goes through `dispatch_attention_fn(..., backend="flex")` and warns once when it finds an uncompiled `flex_attention` — that falls back to a dense fp32 score matrix and runs out of memory at high resolution. - `_attention_backend` is `None` on the flex processor. It is what `set_attention_backend()` sets and only applies to the cached decode steps; the prefill needs the flex kernel for its `BlockMask` and is not configurable. The processor raises from `__init__` when flex_attention is unavailable. - Pad the sequence axis directly instead of transposing around `F.pad`, so the padded tensors stay contiguous, which the compiled flex kernel requires. - Rename `QwenImage21SDPAAttnProcessor` to `QwenImage21AttnProcessor` and make it the default. Once compiled, flex is 1.5% faster end to end at 2048x2048 (31.8s vs 32.3s over 20 steps) and 3.8% faster with two condition images, in exchange for 37s of compilation; uncompiled it cannot render 2048x2048 at all. A default that only works when the caller compiles is the wrong trade, so flex is documented as the opt-in path instead. - Derive the prefix segment boundaries once per forward rather than once per layer. They only depend on `image_ids` and `prefix_len`, so the per-layer version repeated the same `tolist()` device sync 32 times. `forward` passes down whichever representation the installed processors read, and builds neither for a processor that does not need it. - Replace `test_non_flex_backend_rejected_when_causal`, which no longer describes the intended behaviour, with a check that `set_attention_backend` only affects decode. Measured on one H100, bf16, batch 1, 20 steps at 2048x2048: the default path runs out of the box in 32.7s at 56.5 GiB peak, and `set_attn_processor(QwenImage21AttnProcessor())` now works at all — it used to hand the flex `BlockMask` to a non-flex kernel and raise. * fix: VAE class defaults did not describe the released model `scale_factor_spatial` was 8 while the encoder applies four spatial downsamples: `encode` takes a 1024x1024 image to a (1, 64, 1, 64, 64) latent, so the ratio is 16. Every tile-to-latent conversion divides by it, so tiling silently produced a wrong-shaped latent — a 2048x2048 encode came out as 168x168 instead of 128x128. The in/out channel defaults were 3 while this VAE takes four channels, so the class could not be instantiated from its own defaults. Add the model test file that was missing, covering the ratio against the architecture and the shape of a tiled encode. Tile values are not compared: each tile starts the causal convolution feature cache fresh, which is a property of the tiling implementation rather than of these defaults. * address review feedback: copies, pipeline, docs VAE: - Mark the classes that are byte-identical to their Wan originals with `# Copied from` (`DupUp3D`, `WanUpsample`, `WanRMS_norm`, `WanAttentionBlock`, `patchify`, `unpatchify`). Adopting the upstream `RMS_norm.forward` in the process also picks up a fix we had missed: it normalizes in fp32 for fp16/bf16/fp8 inputs. - Drop the `non_linearity` argument, which was always "silu", from the blocks that take it, and validate `AvgDown3D`'s channel divisibility with a `ValueError` before the fields are assigned rather than with an `assert` after. Pipeline: - Take `calculate_dimensions` verbatim from the edit pipeline so it can carry a `# Copied from`, which is what fixes `check_repository_consistency`: the marker on `_encode_vae_image` was one blank line out of sync. - Move the `QwenImage21KVCache` import to the top of the module. Docs: - Apply @stevhliu's suggestions. `models.autoencoders.vae.AutoencoderKLOutput` does not exist, so that autodoc reference was broken; the module is `autoencoder_kl`. - Describe the two attention processors instead of the old "with and without the flex backend" split, and add the snippet for opting into flex, which has to be compiled. `make style` also reflowed a few docstrings from the previous commit. * address review feedback: validate before doing work The `kv_cache` checks in the transformer's `forward` ran after the input projections, the joint sequence build, the rotary embeddings and the modulation, so a bad `kv_cache_mode` was only reported once that work had been done. They now run first. Following the same point through the pipeline turned up a check that could never fire: `check_inputs` warns when `height` and `width` are not divisible by `vae_scale_factor * 2`, but the rounding happened before the call, so the values it saw were always divisible. It now runs before the rounding, and `output_resolution=1000` warns and yields 992x992. * fix: num_images_per_prompt > 1 raised in the pipeline `encode_prompt` expands `prompt_embeds` and its mask to `batch_size * num_images_per_prompt`, but returns `image_pad_mask` unexpanded, and the target slots appended to that mask were sized from `latents`, which is expanded. Any call with `num_images_per_prompt > 1` therefore died in the concatenation: RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 1 but got size 2 for tensor number 1 in the list. Size the slots from each mask's own batch instead. The transformer reads the layout from row 0 because samples share it, so the mask does not need expanding. Verified end to end: `num_images_per_prompt=2`, a list of two prompts, both together, and each of those with classifier-free guidance and with a condition image. * Use the recommended sampling defaults: 40 steps, no guidance Qwen-Image 2.1 is meant to be sampled in 40 steps without classifier-free guidance, so `num_inference_steps` defaults to 40 and `true_cfg_scale` to 1.0. The other QwenImage pipelines default to 50 and 4.0, which is why this differs from its siblings. It also removes a warning from every default call: `true_cfg_scale=4.0` with no negative prompt took the "guidance is not enabled" branch. Passing a `negative_prompt` without raising `true_cfg_scale` still warns, which is the case worth warning about. The docs and the example docstring rely on the defaults now instead of passing a step count, and the docs state the recommendation. * docs: complete the forward and __call__ docstrings `utils/check_forward_call_docstrings.py` on main checks that every argument in a forward/__call__ signature has a docstring entry and that a non-None return type has a Returns section. It landed after this branch's base, so it only started running here once the copy check stopped failing ahead of it. Add the missing entries: `sample_posterior` and `generator` plus Returns on the VAE's forward, `attention_kwargs` and `return_dict` plus Returns on the transformer's, and the four embedding arguments and `callback_on_step_end_tensor_inputs` on the pipeline's __call__. * add pipeline tests for qwenimage 2.1 (#6) * add more copied froms (#5) * docs: describe multiple condition images, and inline the flex warning The pipeline page now has a section on passing several condition images, which is where @sayakpaul asked for it, and the flex section emphasises that the processor wants a compiled model. `_warn_if_flex_attention_is_uncompiled()` is inlined at its only call site, as requested. A class-level flag keeps it to one warning per process: the default processor is constructed per attention module, so 32 instances would otherwise each warn, and the logger has no `warning_once`. * fix-copies: drop the AvgDown3D marker `AvgDown3D` validates its channel divisibility with a `ValueError` before assigning its fields, where Wan still asserts after, so the copy is not consistent and `check_copies` fails on it. The other ten markers from #5 are fine and stay. * fix: feed the transformer the pre-norm text hidden state The transformer was trained on the last decoder layer's output of the text encoder, before the encoder's final RMSNorm. Up to transformers 4.x that is what `hidden_states[-1]` holds. From transformers 5.0 the output capturing ties that entry to `last_hidden_state`, so it comes back normalized instead, and the transformer reads something a third of the way off — visible first as garbled text in the rendered image. Neutralize the final norm for the encoder call with a forward hook that returns the module's input, so `hidden_states[-1]` is the layer output on either version. Nothing else is touched: the weights stay untouched, which matters because they are on `meta` under offloading, and no version check is needed. On the released checkpoint the prompt embeddings now match the pre-norm value exactly and the rendered image is pixel-identical to it, where before the whole image shifted by 5.35/255 on average. * fix: match the checkpoint's text conditioning, and repair the unexercised paths Aligned with the text encoder the checkpoint was trained with: - The image marker is `<image1>`, `<image2>`, … as in training, not `Picture 1: `. The two tokenize to different lengths (4 tokens against 5), so the conditioning the transformer read was a sequence that never occurred in training, and the rendered image moves 4.02/255 on average. With the marker corrected the output is pixel-identical to the training template. This also retires the `random.choice` over four spellings of that word, which ran on the global RNG and left image-conditioned generation irreproducible from `generator`. - Condition images reach the vision encoder with their alpha composited over white, as in training. The VAE still reads all four channels. - The processor pads on the left, as in training. The joint sequence is re-padded on the right either way, so this only changes the positions the encoder itself sees for a batch of prompts of different lengths. - An empty prompt becomes a space. Qwen has no bos token, so the encoder would otherwise have nothing to read. Prompt embeddings: - The 2D prompt mask was repeated with `repeat(1, n, 1)`, which prepends an axis and tiles the rows where the 3D embeddings interleave theirs. With more than one prompt and more than one image per prompt, each sample was denoised against another prompt's padding. - Supplying `prompt_embeds` raised: `image_pad_mask` only comes from `encode_prompt`'s own encoding, and the target slots appended to it were sized from `latents`. It is synthesized for text embeddings now, and required when the embeddings cover condition images. Supplying embeddings without a mask raised too. - `has_neg_prompt` no longer requires `negative_prompt_embeds_mask`, so a caller who passes `encode_prompt`'s own output back in keeps guidance. An unpadded prompt returns `None` for the mask, the pattern `pipeline_qwenimage.py` uses, since a mask that carries no information costs the backends that reject one. Condition images: - A tensor or ndarray `image` raised on `image.size`. They are normalized to PIL up front; a latents tensor is rejected with a message, because the text encoder has to see the image, and a per-prompt nested list with another, because one flat set applies to the whole batch. That also removes the half-wired path where a latents tensor was silently dropped from both the prompt and the latents. - A list of prompts with a condition image raised a bare `StopIteration`: every prompt's template repeats the placeholders, but the processor was handed one set of images. Denoising and validation: - Interrupting on the first step used to `continue`, skipping the step that prefills the KV cache and leaving the next one to decode from an empty one. It breaks out now. - `kv_cache_mode` without a `kv_cache` is rejected instead of failing later on a shape. - `prepare_latents` checks the generator list before spending a VAE encode per image. - The VAE's image convolution names its limitation instead of asserting: it folds the single frame away and has no temporal context, so it cannot take a feature cache. Removed: `_downsample_image_pad_tokens`, `_max_length`, the `_drop_idx_ti2i` alias, and the KV cache's `is_populated` and `clear()`, none of which anything reached. Training collapses each run of `<|image_pad|>` to one token; skipping that is equivalent here, since those positions are overwritten by the VAE latents either way. The text encoder hook now points at huggingface/transformers#48087, which lets the config untie `hidden_states[-1]` from 5.18 and will make the hook unnecessary. * fix callback test (#2) * Add a TODO at the text encoder hook It can be replaced with `tie_last_hidden_states=False` in the text encoder's config once huggingface/transformers#48087 ships in a stable release. --------- Co-authored-by: yiyixuxu <yixu310@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
陈志谦
|
7221eef457 |
docs: pool eight small doc fixes into one PR (#14757)
* docs: fix latent -> latents in LTX latent upsampler docstrings __call__ takes 'latents' (the docstring said 'latent', and its own description says 'Input latents to normalize'). * docs: fix 'promt' typo in kandinsky5 default prompt templates The default system template sent to the Qwen text encoder misspelled 'prompt engineer' as 'promt engineer' (4 pipelines). Note: this changes the token sequence of the default template, so default outputs will differ numerically from before / from the upstream Kandinsky5 reference. If exact parity with the reference repo is preferred, the reference itself would need the same fix. * docs: fix docstring parameter names that do not match signatures - flux controlnet encode_prompt documented a clip_skip parameter that does not exist (copy-paste from SD-style docs; flux has no CLIP layer-skipping) - hunyuan_video1_5 pipelines documented num_images_per_prompt while the parameter is num_videos_per_prompt - kandinsky5 _encode_prompt_qwen/_encode_prompt_clip documented num_videos_per_prompt, which neither method accepts * docs(zh): fix English anchors pointing at translated Chinese headings speed-memory-optims.md linked ./memory#model-offloading and #group-offloading while the zh memory page translates those headings (模型卸载 / 组卸载); using_peft_for_inference.md linked #hotswapping while the zh heading is 热切换. Repointed at the Chinese slugs. * docs: fix PAG anchor and the StableDiffusionPipelineSafe row in the SD overview - pag.md linked #pagappliedlayers; the heading '### pag_applied_layers' keeps its underscores (#pag_applied_layers) - the SD overview table linked ./stable_diffusion_safe, a doc page that no longer exists anywhere under docs/source; the row now names the deprecated pipeline as plain text instead of a dead link * docs: fix dead Masked Im2Im TOC anchors in community pipelines README The TOC rows linked #stable-diffusion-masked-im2im / #stable-diffusion-xl-masked-im2im; the actual headings are 'Masked Im2Im Stable Diffusion Pipeline' and 'Masked Im2Im Stable Diffusion Pipeline XL' (lines 3184/3219). * docs(examples): drop duplicated 'with' in T5/Gemma sequence-length help strings 16 example scripts' --max_sequence_length help read 'to use with with the T5 text encoder' (or the Gemma variant); surfaced in --help output. * fix(links): keep the pag_applied_layers anchor in doc-builder form doc-builder slugifies headings by dropping underscores, so the heading '### pag_applied_layers' anchors as #pagappliedlayers — the form the check-links CI expects. GitHub-style underscores broke it. Signed-off-by: simpleqt <89645338+simpleqt@users.noreply.github.com> * fix(kandinsky5): recompute prompt_template_encode_start_idx after the promt fix Correcting 'promt' -> 'prompt' in the system templates shifts each template's token count by one (verified with the Qwen2.5-VL tokenizer: 41->40, 55->54, 129->128, 129->128). Without this, the hardcoded start indices slice the encoded prompt one token off. --------- Signed-off-by: simpleqt <89645338+simpleqt@users.noreply.github.com> |
||
|
Steven Liu
|
37ce5add81 |
[docs] Refactor get started section (#14783)
* docs * feedback * feedback |
||
Yiming Zhao
|
759164b7ad |
[Cosmos3] Mixed W8A8/W8A16 denoising for ModelOpt FP8 checkpoints (#14664)
* Add opt-in mixed W8A8/W8A16 denoising for Cosmos3 ModelOpt FP8 checkpoints. Keep native ModelOpt GEMM on middle steps and dequant-linear W8A16 on the first/last steps so CFG cond/uncond share one precision per scheduler step. * Read Cosmos3 mixed-precision schedule from the checkpoint policy. Enable first/last W8A16 only when transformer/config.json declares diffusion_step_policy, so distilled FP8 stays native W8A8 instead of inheriting a hardcoded 3+3 window. * Harden Cosmos3 mixed-precision loading and document official Hub fp8 schedules. Read the checkpoint runtime policy from on-disk transformer/config.json when the live ModelOpt config omits it, fail closed on incomplete policies, and allow FP32 activations on the W8A16 path. * Address review: simplify FP8 mixed docs, no-op non-ModelOpt backends, drop unit tests. W8A8/W8A16 is explained without listing call-site overrides; TorchAO and other quantizers keep their native forwards; focused tests are removed until Hub fp8 usage is clearer. * Document FP8 mixed default vs none trade-offs without ModelOpt restore boilerplate. Keep serialized restore in the ModelOpt guide and spell out that all-W8A8 is faster but can flicker on multi-step video. * Apply style fixes * Regenerate Cosmos modular auto docstrings for mixed-precision inputs. --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
Steven Liu
|
b27b69c1d6 |
[docs] Improve navigation (#14537)
* docs * rearrange * cross links * titles |
||
|
YiYi Xu
|
ebe93ee5f7 |
agent doc: add a requirement for PR communications (#14709)
* modular guide: helpers take explicit arguments, not `components` Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS * agent guide: human-in-the-loop rules, helpers take explicit arguments Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS * clarify human-in-the-loop scope to communication Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS * contribution guide + agent guide: PR communication must be concise and human-readable Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS * Apply batched suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Apply suggestion from @yiyixuxu --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> |
||
|
YiYi Xu
|
7643c48266 |
ModularPipeline.save_pretrained: change the default overwrite_modular_index to be True (#14659)
* flipt default for save_pretrained overwrite_modular_index to be True * update doc * make style Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS * Update docs/source/en/modular_diffusers/modular_pipeline.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/modular_diffusers/modular_pipeline.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/modular_diffusers/modular_pipeline.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/modular_diffusers/modular_pipeline.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * address review comments Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> |
||
Alex Artyomov
|
82e309a2b1 |
Add LTX-2.5 DFR pipeline (keyframe slots, spatial detailing, tiled temporal rounds) (#14567)
* Add LTX-2.5 DFR pipeline with generated keyframe slots Ports DFRPipeline from the Lightricks reference. Stage 1 generates video plus single-pixel-frame keyframe slots at a fraction of the requested resolution on a VAE-aligned segment grid; both are spatially latent-upsampled and stage 2 re-denoises at twice that resolution with the slots re-attached and an optional spatial detailing IC-LoRA active for that stage only. Optional temporal x2/x4 refine rounds tile the canvas at shared keyframes and densify with ancestral Euler. With spatial_upscalings=2 a full-resolution detailing epilogue follows the rounds. The transformer already stored keyframes_abs_pos_embedding for load/save; this wires it into the forward through a new video_keyframes_mask argument, which only a keyframes-aware pipeline passes, so other pipelines are unaffected. The epilogue denoises the whole canvas in one loop and tiles the transformer call inside it, so every Euler step steps a canvas whose tiles have already agreed on their overlaps. Spatial tiles blend under a trapezoidal mask, since neither side of a height or width border holds a known answer. Temporal tiles are cut on the keyframe seams the last refine round stitched on: both windows reproduce a shared keyframe there, so the later one drops its run-up under a rectangular mask rather than averaging it. Conditionings are attached once on the whole canvas and filtered per tile at the token level, and a keyframe two windows share is one token they both read. The epilogue is handed its keyframes rather than asked to generate them. Each carry plane is decoded on its own -- the VAE is causal, so a stacked decode would bleed neighbours -- then Lanczos-stretched x2 in RGB and encoded again at the output resolution, and pinned fully clean. Only the video latent is spatially upsampled. Conditioning fps is snapped to 60 above 30 rather than merely capped there, at every stage. RoPE time is pixel_frame / fps, and the transformer is trained around 24/25/30 and 60; a temporal round taking 24 fps to 48 lands between those, and it shows as stutter at the latent borders. Playback fps is unchanged, so 24 fps with one round still ships 48 fps. A condition's index is read on the canvas num_frames asks for, and the moment it names is carried onto each refine round's longer canvas by scaling its pixel position by 2**round. The scaled position does not generally land on a latent boundary, so it travels as a pixel index rather than through the public latent index; a keyframe conditioning is appended as extra tokens instead of being spliced into the base grid, so it does not need to. height and width must be divisible by 2**spatial_upscalings times the VAE's spatial compression ratio, which makes 4K 3840x2176 rather than 3840x2160. That rule is checked ahead of the looser one every LTX-2 pipeline applies, so the error names the divisor a DFR caller actually has to satisfy. Four details are easy to get wrong, and each is covered by a test after showing up as a visible seam at a tile handover: - The ancestral step injects noise into every token, so the conditioning blend has to be re-applied afterwards. Skipping it lets the strength-0.95 anchor keyframes erode over the schedule, and those anchors are the only thing pinning adjacent tiles onto the same content. - Velocity is converted to x0 with each token's own noise level, not the scalar schedule sigma: a token held at strength s sits at (1 - s) * sigma. - Each tile draws its ancestral noise from a generator seeded seed + 1000 * round + tile, kept separate from the main generator so the draws do not consume state the next tile's initial noising reads. - Two tiles invent the slot that falls in the later one's dropped lead-in. The stitch keeps the earlier tile's frames there, so the earlier tile's copy is the one the canvas holds, and the one the next round must anchor on. * Split the DFR pipeline into composable stages Addresses the review on #14567: use the same compose pattern as the other LTX two-stage pipelines, get ancestral Euler from an existing scheduler, and put the temporal rounds in their own pipeline so the schedule is not switched mid-call. `LTX2DFRPipeline.__call__` is now one denoise pass at `height` x `width`. Callers compose stage 1, `LTX2LatentUpsamplePipeline`, stage 2 and each temporal round, and a documented recipe is the copy-paste 1080p path. The recipe knobs (`spatial_upscalings`, `temporal_upscalings`, `detailing_lora_adapter_name`) and the required upsampler components are gone; `height`/`width` are this pass, not the final output. `ancestral_euler_step` is replaced by `LTXEulerAncestralRFScheduler.step` plus a re-application of the conditioning blend, which ancestral noise would otherwise erode on the strength-0.95 seam anchors. The new `LTX2DFRTemporalRefinePipeline` owns that scheduler and one round; stage 1, stage 2 and the epilogue stay on `FlowMatchEulerDiscreteScheduler`. It refuses any other scheduler rather than silently taking a deterministic step and returning a softer canvas. Pack/unpack, `prepare_latents`, `denoise` and `encode_conditions` move to `LTX2DFRCoreMixin` so neither pipeline subclasses the other. Public latents are raw on both sides of every boundary, `output_type="latent"` returns the untrimmed canvas so a slot on the pad is not dropped, and `trim_canvas` does the trim before decode. `LTX2DFRPipelineOutput` carries `keyframes` and `keyframe_positions`, which cannot be re-derived after a round. Verified against the pre-split implementation: bit-exact on dummy components in fp32 and bf16, with and without the IC-LoRA reference, over one and two rounds; and within one bf16 ulp on the real checkpoint, where the only difference is that the split normalizes upsampled latents in fp32 rather than bf16. Also in this pass: - `__call__` takes `video_tiles` (the `epilogue_tiles` layout) instead of a resolved token plan. Resolving one needs the RoPE coordinates that only exist once `prepare_latents` has run, so a caller could not build the plan at all. - `rebuild_epilogue_keyframes` is public and takes and returns raw latents. The composed epilogue needs it, so it was public API in practice while named private, and its normalized return forced callers into `_denormalize_latents`. - Drop prompt enhancement and `num_videos_per_prompt` from the temporal pipeline. Enhancement belongs to stage 1 -- re-running it would denoise the canvas under a different prompt than the one that generated it -- and the batch is set by the incoming latent canvas, so `num_videos_per_prompt > 1` only ever raised. - Fix the docs recipe: `requested_frames` counted latent frames where `trim_canvas` wants pixel frames, truncating a 241-frame render to 25; the three pipelines share components, so place them together instead of offloading one and leaving `temporal_latent_upsampler` off the device; and the detailing IC-LoRA is applied at 0.5, the strength the reference hardcodes. - Pass `crf=0` on the conditions in the refine-round test. The default CRF sends the image through H.264 re-compression, which needs PyAV, so the test failed on any environment without it while testing nothing about re-compression. * Drop LTX2DFRCoreMixin in favor of # Copied from Each DFR pipeline now owns pack/unpack, prepare_latents, and denoise so the classes stay self-contained. Shared helpers copy from LTX2Pipeline / LTX2ConditionPipeline, and DFR-specific methods copy from LTX2DFRPipeline. dfr_core.py keeps only the constants and canvas helpers. * inline _finalize_output Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * remove dfr_core.py Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * make return_dict=False return a tuple of 4, consistent with diffusers convention; update docs and pipeline output accordingly Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * remove the dfr_layout test suite Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
|
Steven Liu
|
7a8c9da2a9 |
[CI] Repair broken links (#14627)
fix Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
daad99a04a |
[tests] remove quanto tests. (#14254)
* remove quanto tests. * deprecate quanto and remove tests |
||
dg845
|
c969cf2f7e |
Move LTX2Guidance Guider to guiders/ Folder (#14558)
* Move LTX2Guidance from /modular_pipelines/ltx2/ to guiders/ to allow LTX-2 guiders to be correctly loaded from checkpoint * Add regression test for LTX-2 guiders declared as from_pretrained components Covers the round trip this move fixes: a checkpoint that declares `guider` and `audio_guider` in `modular_model_index.json` must record a resolvable (library, class_name) pair and reload its non-default guidance scales. The existing `test_modular_index_consistency` skips components without a `pretrained_model_name_or_path`, so `from_config` guiders were uncovered. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Migrate LTX2Guidance docs to guider docs after move * Improve LTX2Guidance docs and add arg docstrings --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> |
||
kaixuanliu
|
b9bc22bd57 |
make examples in docs device agnostic (#14414)
* docs: annotate hardcoded cuda device with alternatives in English docs * docs: annotate hardcoded device_map="cuda" with alternatives * update Signed-off-by: kaixuanliu <kaixuan.liu@intel.com> --------- Signed-off-by: kaixuanliu <kaixuan.liu@intel.com> |
||
Shivam Shrirao
|
06e0f2a81c |
Fibo Edit: multi-reference conditioning and batching (#14566)
* Fibo Edit: multi-reference conditioning, batching, timesteps * docs: simplify Fibo Edit page, fix EXAMPLE_DOC_STRING to the VLM prompt flow * docs: point Fibo Edit at the Gemini prompt-to-JSON block and the 1.5 base/turbo ids briaai/FIBO-edit-prompt-to-JSON is retired in favor of briaai/FIBO-edit-gemini-prompt-to-JSON, which handles multiple reference images and masks. Also corrects the checkpoint ids to their canonical briaai/Fibo-Edit-1.5-base and briaai/Fibo-Edit-1.5-turbo spellings. * apply review feedback on reference handling and the attention mask - collapse the mask branch to `if not attention_mask.all():` in both Fibo pipelines - encode references via `self.image_processor.preprocess(...)`, dropping the manual numpy normalization and the `_vae_safe_size` helper - normalize `image` into a reference list once in `__call__`, pass it to `check_inputs` * make style * make style |
||
|
Sayak Paul
|
9222edcce3 |
Deprecate 'fp16_safetensors' command in CLI documentation (#14583)
Add deprecation notice for 'fp16_safetensors' command. |
||
|
dg845
|
9841a24724 |
[docs] LTX-2.5 Modular Pipeline Docs (#14450)
* LTX-2.5 modular pipeline init docs * Update LTX-2.5 modular docs now that the modular pipeline automatic downloads the Gemma 4 prompt enhancer and processor * fix some modular docs typos * fix more typos * Update docs/source/en/api/pipelines/ltx2.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Remove unused transformers imports from LTX-2.5 modular examples --------- Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> |
||
iridescentWen
|
2f7e0154a9 |
docs: fix 14 typos in docs, comments and docstrings (#14541)
docs: fix typos in docs, comments and docstrings
Prose, comments and docstrings only. No executable line changed.
docs/:
- optimization/memory.md: maxmium -> maximum; heigh -> height (in the
channels_last layout tuple, which read "(batch size, heigh, width,
channels)")
- api/pipelines/bria_fibo.md: proffesional -> professional
- api/pipelines/wan.md: involed -> involved
src/ comments and docstrings:
- guiders/magnitude_aware_guidance.py: supression -> suppression
- models/attention.py: spliting -> splitting
- models/attention_dispatch.py: abritrary -> arbitrary (x2)
- models/transformers/transformer_ltx2.py: timstamps -> timestamps (x2)
- pipelines/ltx2/pipeline_ltx2{,_image2video}.py: corrct -> correct
- pipelines/deprecated/unidiffuser/modeling_uvit.py: tbe -> the
- pipelines/visualcloze/pipeline_visualcloze_{combined,generation}.py:
"whe they are lists" -> "when they are lists" in a user-facing
ValueError message
Deliberately not touched:
- pipeline_kandinsky*.py "promt"/"scren" live inside prompt_template,
which is fed to the text encoder verbatim. Changing them would change
generation output.
- ANE (Apple Neural Engine), MoT (Mixture-of-Transformers): initialisms.
- "racoon": appears in prompt examples whose reference images were
generated with that exact prompt.
- "Shuting Wang": a paper author's name.
|
||
|
Steven Liu
|
0eff7f0664 |
[docs] Clean up links (#14535)
* docs * more cleaning |
||
|
Dhruv Nair
|
3681e65996 |
Restructure skill files and reference guides for CLI based installation. (#14454)
* restructure guides as skill references * reorganize AGENTS.md and references * clean up * clean up * clean up * fix AGENTS.md references and add check_ai to Makefile * clean up * update paths * clean up |