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

1733 Commits

Author SHA1 Message Date
Gong Junmin
1a8a17b71b Add ACE-Step pipeline for text-to-music generation (#13095)
* Add ACE-Step pipeline for text-to-music generation

Rebased on origin/main from the original pr-13095 branch (3 commits squashed).

- AceStepDiTModel: Diffusion Transformer with RoPE, GQA, sliding window,
  AdaLN timestep conditioning, and cross-attention.
- AceStepConditionEncoder: fuses text / lyric / timbre into a single
  cross-attention sequence.
- AceStepPipeline: text2music / cover / repaint / extract / lego / complete.
- Conversion script for the original checkpoint layout.
- Docs + tests.

* Fix ACE-Step pipeline audio quality and auto-detect turbo/base/sft variants

The PR's original inference produced low-quality audio on turbo because the
pipeline (a) mangled the SFT prompt format, (b) applied classifier-free guidance
with the wrong unconditional embedding (empty-string encoded vs. the learned
`null_condition_emb`), and (c) hardcoded turbo defaults even when loading a
base/SFT checkpoint.

Changes:

* Converter preserves `null_condition_emb` (stored under the condition encoder)
  and propagates `is_turbo`/`model_version` into the transformer config so the
  pipeline can route per-variant defaults.
* `AceStepConditionEncoder` registers `null_condition_emb` as a learned
  parameter matching the original module.
* Pipeline auto-detects variant via `is_turbo`/`model_version` and picks
  defaults that match `acestep/inference.py`:
    * turbo:  steps=8,  shift=3.0, guidance_scale=1.0 (no CFG)
    * base/SFT: steps=27, shift=1.0, guidance_scale=7.0
* Base/SFT timestep schedule uses the linear+shift transform from
  `acestep/models/base/modeling_acestep_v15_base.py`; turbo still uses the
  hardcoded 8-step `SHIFT_TIMESTEPS` table.
* CFG reuses the learned `null_condition_emb` and batches the
  conditional+unconditional forwards into a single transformer call.
* `SFT_GEN_PROMPT` matches the newline layout in `acestep/constants.py` so the
  text encoder sees the same prompt distribution it was trained on.

DiT parity vs. the original ACE-Step 1.5 turbo DiT is bit-identical
(max_abs=0.0 in fp32 eager/SDPA across 4 seed/shape cases) — see
scripts/dit_parity_test.py.

* Add ACE-Step parity test scripts

Two developer-facing parity harnesses live under scripts/:

* dit_parity_test.py — loads the same converted turbo weights into the
  original AceStepDiTModel and the diffusers AceStepDiTModel, drives
  identical (hidden_states, timestep, timestep_r, encoder_hidden_states,
  context_latents) inputs, and asserts max-abs-diff ≤ 1e-5 in fp32
  eager/SDPA. Currently passes bit-identical (max_abs=0) across four
  shape/seed cases including batched + odd-length paths.

* audio_parity_jieyue.py — full end-to-end audio parity. Given the same
  JSON example, runs both the original ACE-Step 1.5 pipeline and the
  diffusers AceStepPipeline at matched seed/precision (bf16 + FA2 by
  default) and saves side-by-side .wav files for listening verification.
  Supports text2music / cover / repaint × turbo / base / sft via a
  --matrix mode that writes 18 wavs named
  {variant}_{task}_{official,diffusers}.wav.

* Route SFT parity to acestep-v15-sft checkpoint

On jieyue the release tree has a dedicated SFT checkpoint at
checkpoints/acestep-v15-sft with its own modeling_acestep_v15_base.py
shipped under acestep/models/sft/. Point the SFT row of the parity matrix
at that checkpoint / module so we're testing the actual SFT weights, not
the plain base ones.

* audio_parity_jieyue: fix doubled 'acestep-' in cache path; --converted-root flag

Previously the converted-pipeline cache dir was
`/tmp/acestep-<variant>-diffusers` but <variant> already starts with
"acestep-", giving `/tmp/acestep-acestep-v15-turbo-diffusers`. Drop the
prefix.

On jieyue the overlay rootfs (including /tmp) only has a few GB free; a
full turbo conversion needs ~5 GB per variant. Add --converted-root (env
ACESTEP_CONVERTED_ROOT) so the cache can live on vepfs.

* audio_parity_jieyue: two-phase matrix bootstraps cover/repaint from text2music

The ACE-Step release bundle on jieyue doesn't ship sample .wav/.mp3
files, so matrix mode had no default --src-audio and would skip
cover/repaint entirely. Run text2music first for every variant, then
reuse the TURBO official text2music output as the shared source for the
cover/repaint rows. Users can still override with --src-audio.

* audio_parity_jieyue: seed the diffusers generator on the pipeline device

The ORIGINAL ACE-Step pipeline seeds on the execution device
(`torch.Generator(device=device).manual_seed(seed)`), i.e. the CUDA RNG
stream when running on GPU. Previously the parity harness seeded the
diffusers side with a CPU generator, so even though the seed integer
matched, the two sides drew different noise from the outset and the
final outputs were essentially uncorrelated. Use the execution-device
generator on both sides for a fair comparison.

* Fix ACE-Step pipeline: switch to APG guidance + peak normalization

Two issues found after the first jieyue audio parity run:

1. The original base/SFT pipeline uses APG (Adaptive Projected Guidance,
   acestep/models/common/apg_guidance.py) with a stateful momentum
   buffer and norm/projection steps — NOT vanilla CFG. Using vanilla CFG
   produced uncorrelated outputs vs. the reference (pearson ~0.0 on
   20 s samples); this PR ports `_apg_forward` + `_APGMomentumBuffer`
   and plugs them into the denoising loop when `guidance_scale > 1`.
   Momentum is instantiated once per pipeline call (persists across
   denoising steps) to match the reference semantics.

2. The post-VAE "anti-clipping normalization" in this pipeline was
   `audio /= std * 5` with a `std<1 -> std=1` guard. The original
   post-processing in
   acestep/core/generation/handler/generate_music_decode.py is simple
   peak normalization: `if audio.abs().max() > 1: audio /= peak`. The
   std-based proxy both (a) let clips with peak < 1 leak through
   unchanged (over-quiet) and (b) failed to bring clipping peaks to
   exactly 1 in a bunch of base/SFT cases (observed max=1.000, std=0.200
   repeatedly in the first parity run). Switch to peak normalization on
   both sides.

Tested via scripts/audio_parity_jieyue.py on A800; re-run pending to
confirm the base/SFT correlation improvements.

* Fix ACE-Step chunk mask values to match the original pipeline

The DiT receives `context_latents = concat(src_latents, chunk_mask)` on the
channel dim, and was trained with chunk_mask values drawn from the three
sentinels documented in acestep/inference.py:

  2.0 -> model-decided (default for text2music / cover / full-generation)
  1.0 -> keep this latent frame from src_latents (repaint preserved region)
  0.0 -> explicitly repaint this frame (only inside the repaint window)

Previously _build_chunk_mask returned all-1.0 for text2music (and cover /
lego), and an inverted 0/1 mask for repaint (1 inside the window, 0 outside).
Either case puts context_latents out of distribution. Switch text2music /
cover to the 2.0 sentinel and flip the repaint mask so it's 1.0 outside /
0.0 inside. Update the repaint src_latents zero-out to multiply by the new
mask (was `1 - chunk_mask`) so the zero region still lines up with the
repaint window.

* Add direct invoker for ACE-Step generate_music (ground truth)

Our earlier audio_parity_jieyue.py reconstructs the original pipeline by
calling AceStepConditionGenerationModel.generate_audio() directly, which
silently skips a lot of the real handler plumbing (conditioning masks,
silence-latent tiling, cover/repaint pre-processing, etc.). That made the
'official' wavs we saved sound wrong — flat, drone-like, not music.

This new script calls acestep.inference.generate_music end-to-end through
the real AceStepHandler, with LM + CoT explicitly disabled so we still have
a deterministic comparison. Use it to generate the ground-truth 'official'
wav for a given JSON example, then separately run the diffusers pipeline
with the same inputs and diff the two.

* run_official_generate_music: call initialize_service to bind a DiT variant

AceStepHandler() is a shell — you have to call handler.initialize_service(
project_root=..., config_path=..., device=..., use_flash_attention=..., ...)
before generate_music will work. Mirror what cli.py does at the equivalent
spot (around cli.py:1400).

* Fix silence-reference for ACE-Step timbre encoder

The root cause for the flat / drone-like outputs I was seeing (including
in my 'official' reconstruction): when no reference_audio is provided the
pipeline was feeding literal zeros to the timbre encoder. The real
handler feeds a slice of the learned `silence_latent` tensor.

The handler also transposes silence_latent on load (see
acestep/core/generation/handler/init_service_loader.py:214:
  self.silence_latent = torch.load(...).transpose(1, 2)
) converting [1, 64, 15000] -> [1, 15000, 64] so that
`silence_latent[:, :750, :]` yields the expected [1, 750, 64] shape.

Changes:

* Converter: load silence_latent.pt, transpose to [1, T, C], bake into
  the condition_encoder safetensors under key `silence_latent`.
  (Also keeps the raw .pt file at the pipeline root for debugging.)
* AceStepConditionEncoder: register `silence_latent` as a persistent
  buffer so from_pretrained loads it alongside the trained weights.
* Pipeline: when reference_audio is None, slice
  `condition_encoder.silence_latent[:, :timbre_fix_frame, :]` and
  broadcast across the batch instead of zeros. Emits a loud warning
  (and falls back to zeros) if the buffer is all-zero — that means the
  checkpoint was produced by an older converter and should be rebuilt.
* audio_parity_jieyue.py: the reference path now matches the handler's
  silence-latent slicing.

Without this fix, every variant/task combo produced drone-like audio
even when my numeric DiT-forward parity claimed they were identical.

* Fix three more ACE-Step pipeline bugs I found by dumping real inputs

Instrumented the live generate_audio call in the real ACE-Step handler and
observed the exact tensors it sees — my diffusers pipeline was wrong in
three independent ways:

1. src_latents for text2music should be silence_latent tiled to
   latent_length, NOT zeros. The handler fills no-target cases from
   silence_latent_tiled (observed std=0.96). Zeros are OOD for the DiT
   context_latents concat and produce drone-like outputs.

2. chunk_mask values cap at 1.0 (not 2.0). The handler starts with a
   bool tensor (True inside the generate span, False outside); the
   chunk_mask_modes=auto -> 2.0 override does NOT take effect because
   the underlying tensor is bool, so setting entry = 2.0 casts to True.
   After the later .to(dtype) float cast, the DiT sees 1.0/0.0 — exactly
   what I observed in the captured tensor (unique values = [True]).

3. Default shift is 1.0 for ALL variants, including turbo. I was
   defaulting turbo to shift=3.0 which picks a different SHIFT_TIMESTEPS
   table (the 8-step schedule is keyed by shift, not variant).

Also:
* Added _silence_latent_tiled() helper that slices / tiles the learned
  silence_latent (now loaded as a buffer on the condition encoder) to
  the requested latent length.
* Repaint path now substitutes silence_latent (not raw zeros) inside
  the repaint window — matches conditioning_masks.py.
* audio_parity_jieyue.py mirrors the same src/chunk/shift choices on
  its 'original' leg for apples-to-apples parity once the buggy
  reconstruction is removed from the picture.

* Add peak+loudness post-normalization to AceStepPipeline

The real pipeline normalizes audio in two stages (see
acestep/audio_utils.py:72 normalize_audio + generate_music_decode.py):
  1. if peak > 1: audio /= peak  (anti-clip)
  2. audio *= target_amp / peak   (target_amp = 10 ** (-1/20) ~ 0.891)

Step 2 is loudness normalization to -1 dBFS. Without it diffusers outputs
had peak=1.0 vs the real 0.891 — same music content (pearson was ~0.86
already), just 1.12x louder. Add step 2 after the existing anti-clip step.

* Match acestep/inference.py inference_steps=8 for ALL variants

GenerationParams.inference_steps default is 8 — turbo AND base/SFT. I had
base/SFT defaulting to 27 here, so every base/SFT parity run was comparing
a 27-step diffusers trajectory against an 8-step real trajectory. Different
number of denoising steps means different audio even at fixed seed.

This likely explains the lower base/SFT correlation in my earlier jieyue
runs (turbo was 0.86, base/SFT were 0.32-0.34). Aligning step counts
should bring base/SFT closer to turbo parity.

* Address PR #13095 review: rename classes + reuse diffusers primitives

Response to dg845's PR comments batch 1+2. DiT parity harness still bit-identical
(max_abs=0 on fp32 / SDPA across 4 shape cases).

Transformer file:
* Rename AceStepDiTModel -> AceStepTransformer1DModel (alias kept).
* Rename AceStepDiTLayer -> AceStepTransformerBlock (alias kept).
* Inherit AttentionMixin + CacheMixin on the DiT model.
* Swap in diffusers.models.normalization.RMSNorm for the hand-rolled
  AceStepRMSNorm (weight-key-compatible).
* Swap the hand-rolled rotary embedding + apply_rotary for diffusers'
  get_1d_rotary_pos_embed + apply_rotary_emb (use_real_unbind_dim=-2 to
  match the cat-half convention ACE-Step inherits from Qwen3).
* Use get_timestep_embedding with flip_sin_to_cos=True — keeps the
  (cos, sin) ordering of the original sinusoidal. State-dict-compatible.
* Drop max_position_embeddings arg from DiT config (RoPE computes freqs
  per call based on seq_len); converter drops it.
* Gradient-checkpoint call now takes just the layer module (matches the
  Flux2 idiom).

Pipeline modeling file (pipelines/ace_step/modeling_ace_step.py):
* Moved _pack_sequences + AceStepEncoderLayer here — they aren't used
  by the DiT, so they shouldn't live in the transformer file.
* AceStepLyricEncoder + AceStepTimbreEncoder set
  _supports_gradient_checkpointing = True and wrap encoder-layer calls
  through the checkpointing func when enabled.
* Use diffusers RMSNorm + the RoPE helper from the transformer file
  (shared single implementation).

Converter (scripts/convert_ace_step_to_diffusers.py):
* model_index.json now carries AceStepTransformer1DModel.
* Drop max_position_embeddings / use_sliding_window from the emitted
  configs.

No numerical regressions: scripts/dit_parity_test.py PASSES with
max_abs=0.0 on fp32/SDPA across short, long, batched, and
padding-path shape variants.

* Address PR #13095 review: pipeline polish + converter HF-hub support

Response to dg845 review comments on the pipeline side. DiT parity still
bit-identical (max_abs=0 across 4 shape cases).

Pipeline (pipelines/ace_step/pipeline_ace_step.py):
* Add `sample_rate` + `latents_per_second` properties sourced from the
  VAE config so the pipeline no longer hardcodes 48000 / 25 / 1920.
  Propagates through prepare_latents, chunk_mask window math, and the
  audio-duration round-trip.
* Add `do_classifier_free_guidance` property (matches LTX2 et al.).
* Add `check_inputs(...)` called from `__call__` before allocating noise.
  Validates prompt type, lyrics type, task_type, step count, guidance
  scale, shift, cfg interval bounds and repaint window ordering.
* Add `callback_on_step_end` + `callback_on_step_end_tensor_inputs` —
  the modern callback form. The legacy `callback` / `callback_steps`
  pair is kept for back-compat. Setting `pipe._interrupt = True` inside
  the callback stops the loop early.
* Expose `encode_audio(audio)` as a public helper that wraps the tiled
  VAE encode + (B, T, D) transpose the pipeline performs internally.

Converter (scripts/convert_ace_step_to_diffusers.py):
* Accept a Hugging Face Hub repo id for `--checkpoint_dir`; resolves it
  via `huggingface_hub.snapshot_download` when the argument isn't a
  local path.

Exports:
* Register `AceStepTransformer1DModel` in the top-level __init__,
  models/__init__, models/transformers/__init__, and dummy_pt_objects so
  `from diffusers import AceStepTransformer1DModel` works and the
  pipeline loader resolves the new class name from model_index.json.

Deferred for a follow-up (commented inline in the PR): full
`Attention + AttnProcessor + dispatch_attention_fn` refactor and
`FlowMatchEulerDiscreteScheduler` migration — both would benefit from a
dedicated parity re-run and review.

* Fix stale ACE-Step 1.0-era docs / class names in the 1.5 integration

Docs and docstrings still carried a mix of 1.0 paper title, non-existent
`ACE-Step/ACE-Step-v1-5-turbo` hub id, `shift=3.0` turbo default, and
the old `AceStepDiTModel` class name. Cleaned up to match the actual
1.5 release:

* pipelines/ace_step.md: correct citation title ("ACE-Step 1.5: Pushing
  the Boundaries of Open-Source Music Generation"), correct repo
  (`ace-step/ACE-Step-1.5`), new variants table with real HF ids
  (`Ace-Step1.5` / `acestep-v15-base` / `acestep-v15-sft`) and their
  per-variant step/CFG defaults, drop the wrong `shift=3.0` tip.
* models/ace_step_transformer.md: page renamed to
  `AceStepTransformer1DModel` with a short 1.5-specific description;
  `AceStepDiTModel` noted as a backwards-compat alias.
* pipeline_ace_step.py: import, docstring, `Args`, and `__init__`
  annotation reference `AceStepTransformer1DModel`; example model id
  now `ACE-Step/Ace-Step1.5`; `_variant_defaults` docstring and the
  `__call__` variant-fallback comment no longer claim `shift=3.0` /
  `27 steps` — real defaults are 8 steps / shift=1.0 across all
  variants, guidance=1.0 (turbo) vs 7.0 (base+sft).

* Address PR #13095 review: VAE tiling on AutoencoderOobleck + Timesteps class

Two more deferred review threads from dg845 addressed:

* Move tiled encode/decode onto AutoencoderOobleck
  (https://github.com/huggingface/diffusers/pull/13095#discussion_r2785513647).
  AutoencoderOobleck now carries `use_tiling` + `tile_sample_min_length` /
  `tile_sample_overlap` / `tile_latent_min_length` / `tile_latent_overlap`
  attributes and private `_tiled_encode` / `_tiled_decode` methods; the
  existing `encode` / `_decode` dispatch to them when tiling is enabled and
  the input exceeds the threshold. `AutoencoderMixin.enable_tiling()` is
  already inherited.

  AceStepPipeline's private `_tiled_encode` / `_tiled_decode` and the
  `use_tiled_decode` `__call__` arg are gone; `__init__` now calls
  `self.vae.enable_tiling()` so the long-audio memory behaviour is preserved
  by default. Users can opt out with `pipe.vae.disable_tiling()`.

  Note: the VAE-side tiling concatenates encoder features (h) and samples
  the posterior once, instead of the old per-tile `.sample()` calls. This
  is the standard diffusers pattern; numerically differs only in the
  structure of the noise across tile boundaries.

* Use the Timesteps nn.Module for the sinusoid
  (https://github.com/huggingface/diffusers/pull/13095#discussion_r2785420234).
  `AceStepTimestepEmbedding` wraps `Timesteps(in_channels, flip_sin_to_cos=
  True, downscale_freq_shift=0)` instead of calling `get_timestep_embedding`
  directly — reviewer asked for the Module form.

* Address PR #13095 review: refactor AceStepAttention to Attention + AttnProcessor

Splits the monolithic AceStepAttention into the diffusers standard
Attention + AttnProcessor layout:
  - AceStepAttention (torch.nn.Module, AttentionModuleMixin) holds the
    to_q/to_k/to_v/to_out projections and norm_q/norm_k RMSNorms.
  - AceStepAttnProcessor2_0 runs the attention dispatch through
    dispatch_attention_fn so users can pick flash / sage / native backends
    via model.set_attention_backend(...) or the attention_backend context
    manager.

GQA (Q has 16 heads / K,V have 8) is preserved by passing enable_gqa=True
to dispatch_attention_fn instead of repeat_interleave; fusion is disabled
(_supports_qkv_fusion = False) because Q and K,V have different output
sizes.

The converter is updated to rename the six attention sub-keys
(q_proj -> to_q, k_proj -> to_k, v_proj -> to_v, o_proj -> to_out.0,
q_norm -> norm_q, k_norm -> norm_k) on both the DiT decoder path and the
condition encoder path, since AceStepLyricEncoder / AceStepTimbreEncoder
share the same AceStepAttention class.

Addresses review comments r2785433213 and r2785450463.

* Address PR #13095 review: migrate to FlowMatchEulerDiscreteScheduler

Replace the hand-rolled flow-matching Euler loop with
`FlowMatchEulerDiscreteScheduler`. ACE-Step still computes its own shifted /
turbo sigma schedule via `_get_timestep_schedule`, but now passes it to
`scheduler.set_timesteps(sigmas=...)` and delegates the ODE step to
`scheduler.step()`. The scheduler is configured with `num_train_timesteps=1`
and `shift=1.0` so `scheduler.timesteps` stays in `[0, 1]` (the convention the
DiT was trained on) and the scheduler doesn't re-shift already-shifted sigmas.

The scheduler's appended terminal `sigma=0` reproduces the old loop's
final-step "project to x0" case exactly: `prev = x + (0 - t_curr) * v`.

Parity on jieyue (seed=42, bf16 + flash-attn, turbo text2music, 8 steps):
  waveform Pearson = 0.999999
  spectral Pearson = 1.000000
  max |diff|       = 2.5e-3  (fp32 step-math vs previous bf16 step-math)

fp32 Euler-loop A/B against the hand-rolled path: max |diff| = 3.6e-7.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Address PR #13095 review: move DiT tests + drop stale test kwargs

- Move the DiT transformer tests out of the pipeline test file into a new
  tests/models/transformers/test_models_transformer_ace_step.py that follows
  the standard BaseModelTesterConfig + ModelTesterMixin scaffold (matches
  test_models_transformer_longcat_audio_dit.py).
- Drop `max_position_embeddings` from the remaining AceStepDiTModel and
  AceStepConditionEncoder test fixtures — neither constructor accepts that
  argument anymore.
- Drop `use_sliding_window` from the same fixtures — also no longer a
  constructor argument (the actual `sliding_window` int kwarg is kept).
- Wire `FlowMatchEulerDiscreteScheduler(num_train_timesteps=1, shift=1.0)`
  into `get_dummy_components()` now that the pipeline requires it.

Resolves https://github.com/huggingface/diffusers/pull/13095#discussion_r3115653554,
r3115664850, r3115673059, r3115676580, r3115680700.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Address PR #13095 review from dg845 (2026-04-23)

Fixes 5 review threads + style:

1. Converter now builds `AceStepPipeline` in memory and calls
   `save_pretrained`. Previously the hand-written `model_index.json` was
   missing the `scheduler` entry — fresh converter output couldn't be loaded
   by `AceStepPipeline.from_pretrained` (r3127767785). This also makes the
   converter robust to future `__init__` signature changes.

2. `latent_length` uses `math.ceil(...)` instead of `int(...)` so non-integer
   products (e.g. `latents_per_second=2.0, audio_duration=0.4 → 0.8`) round up
   to `1` instead of truncating to `0` and crashing shape checks (r3127790939).

3. Add `_callback_tensor_inputs = ["latents"]` on `AceStepPipeline` so the
   standard diffusers callback tests pick up the right tensor (r3127795954).

4. `AceStepConditionEncoder.silence_latent` no longer hard-codes the channel
   dim to 64. The placeholder buffer now uses the `timbre_hidden_dim`
   constructor argument, so smaller test configs with `timbre_hidden_dim != 64`
   load without shape errors (r3127812932).

5. Revert `self.vae.enable_tiling()` from `AceStepPipeline.__init__`. Users can
   call `pipe.vae.enable_tiling()` themselves for long-form generation; that
   matches the opt-in convention used by the rest of diffusers (r3127777296).

6. `ruff check --fix` + `ruff format` over all ACE-Step sources (the style fix
   dg845 asked for via `@bot /style`).

Also: converter now accepts sharded `model.safetensors.index.json` layouts
alongside the single-file `model.safetensors`, so the 5B XL turbo variant
converts without a pre-processing step.

Parity on jieyue (seed=42, bf16 + flash-attn, turbo text2music 160s, fresh
converter output loaded via `from_pretrained`):
  waveform Pearson  = 0.999954
  spectral Pearson  = 0.999977
  max |a-b| bf16    = 4.3e-02  (dominated by the VAE tiling default flip)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Address PR #13095 review from yiyixuxu (2026-04-23)

Code-level (22 threads):

1. Delete 3 dev/parity scripts (`scripts/audio_parity_jieyue.py`,
   `scripts/dit_parity_test.py`, `scripts/run_official_generate_music.py`)
   that shouldn't have been committed.
2. Rename `AutoencoderOobleck._encode_one` → `_encode` to match the convention
   used by other diffusers VAEs.
3. Delete the hard-coded `SHIFT_TIMESTEPS` / `VALID_SHIFTS` table in
   `pipeline_ace_step.py`: the per-shift turbo schedules are recovered
   exactly by `linspace(1, 0, N+1)[:-1]` plus the flow-match shift formula
   that the non-turbo branch already uses, so a single code path covers both.
4. Drop the backwards-compat `AceStepDiTModel` / `AceStepDiTLayer` aliases
   and every reference (top-level `__init__`, `models/__init__`,
   `transformers/__init__`, dummy objects, tests, docs toctree, model card).
   `AceStepTransformer1DModel` is the only exported name now.
5. Remove the unused `attention_mask` / `encoder_attention_mask` args from
   `AceStepTransformer1DModel.forward`; the model rebuilds its masks from
   the sequence shape and never consumed them.
6. In the DiT forward and both encoders, pass `None` instead of an all-zero
   `full_attn_mask` / `encoder_4d_mask` to non-sliding attention layers — SDPA
   dispatches to a faster kernel when the mask is None.
7. Inline the shared `_run_encoder_layers` helper directly into
   `AceStepLyricEncoder.forward` / `AceStepTimbreEncoder.forward` so layer
   calls are visible at the forward boundary (diffusers style).
8. Move `is_turbo` / `sample_rate` / `latents_per_second` from `@property`s
   that re-read module configs each call to cached attributes populated in
   `__init__` (Flux2-style), with a default-ACE-Step fallback when
   `self.vae` is offloaded. Drop the now-unused `SAMPLE_RATE = 48000`
   module-level constant and the three property definitions.
9. Warn + coerce `guidance_scale` to 1.0 on turbo (guidance-distilled)
   checkpoints, following `pipeline_flux2_klein`. Prevents over-guided
   audio when users forward their base/sft CFG settings to a turbo pipe.
10. Remove the `logger.warning(...)` paths that triggered on
    `silence_latent` missing/zero — those only fired for author-side
    unconverted checkpoints and tests; end users always load converted
    weights where the buffer is baked in.
11. Drop the redundant `with torch.no_grad():` wrappers inside
    `encode_prompt` — the pipeline's `__call__` runs under `torch.no_grad`
    already.
12. Strip "reviewer comment on PR #13095" attribution comments from three
    docstrings (here and everywhere).

Parity on jieyue (seed=42, bf16 + flash-attn, XL turbo 160s text2music):
  waveform Pearson = 0.9747
  spectral Pearson = 0.9895

The shift comes from full-attention layers switching `attn_mask=0_tensor` →
`attn_mask=None`, which dispatches to a different SDPA kernel on bf16. The
two outputs are algebraically equivalent for fp32 eager; on bf16+FA the
delta is dominated by kernel-level ULPs, well within the sampler-noise
band (ear-check on the 160s example confirms no audible regression).

Still open — AudioTokenizer/Detokenizer (deferred) + APG guider follow-up
(dims differ from `diffusers.guiders.adaptive_projected_guidance`, not a
drop-in; worth a separate PR).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* Address ACE-Step audio token and APG review

* Fix ACE-Step docs CI

* Address ACE-Step pipeline cleanup review

* Fix ACE-Step flash attention sliding windows

* Add ACE-Step callback properties

* Address ACE-Step final review comments

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-04-30 18:30:44 -10:00
Akshan Krithick
b231a6a896 Add Ernie-Image modular pipeline (#13498)
* Add Ernie-Image modular pipeline

* Address review

* Fix alphabetical ordering and generator type_hint

* Address review

* Add height,width as outputs of prompt enhancer

---------

Co-authored-by: YiYi Xu <yixu310@gmail.com>
2026-04-27 09:25:19 -10:00
Sayak Paul
f7fd76adcd [attention backends] fix ring CP for flash and flash 3 (#13182)
* tests: add cp backend and attention backend tests.

* up

* up

* up

* fix ring for flash and flash_3

* generate.

* Apply suggestions from code review

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

* up

* up

---------

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
2026-04-24 20:35:08 -03:00
hlky
a37f6f8394 Improve trust_remote_code (#13448)
* Robust trust check for custom_pipeline parameter of DiffusionPipeline.from_pretrained method

* test_custom_components_from_local_dir

* Apply style fixes

* fix

* Update src/diffusers/utils/dynamic_modules_utils.py

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

* Adjust tests and allow community pipeline

* DIFFUSERS_DISABLE_REMOTE_CODE

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
2026-04-23 11:07:43 +05:30
Sayak Paul
7c88e5fe21 [tests] fix group offloading with disk tests (#13491)
fix group offloading with disk tests
2026-04-22 11:41:37 -03:00
Dhruv Nair
50987b1290 [CI] Fix BnB tests (#13481)
* update

* update

* update
2026-04-22 13:19:32 +05:30
kaixuanliu
3d30b7d9d2 Fix non-deterministic T5 outputs in HiDream pipeline tests (#13534)
avoid dropout of t5 model in hidream-image tests

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-04-21 21:58:40 +05:30
Ricardo-M-L
62bfa5a23e fix(freeu): run FFT in float32 for float16 inputs to avoid ComplexHalf (#13503)
* fix(freeu): run FFT in float32 for float16 inputs to avoid ComplexHalf

`fourier_filter` already upcasts `bfloat16` inputs to `float32` before
calling `torch.fft.fftn`, because PyTorch's FFT does not support bf16.
The same is true for `float16`: depending on the PyTorch version,
`fftn` either

- produces the experimental `torch.complex32` (ComplexHalf) dtype and
  emits a `UserWarning: ComplexHalf support is experimental`, or
- raises `RuntimeError: Unsupported dtype Half` outright.

Both paths were reachable from FreeU with half-precision models
(e.g. `sd-turbo` + `fp16` + `enable_freeu`) as reported in #12504.

Extend the existing upcast branch to cover `float16` too. The function
already downcasts the result back to `x_in.dtype` at the end, so the
externally observable dtype is unchanged.

Closes #12504.

* Address review: generalize upcast to non-float32 + fix ruff F821

- Apply @sayakpaul's suggestion: use `elif x.dtype != torch.float32:`
  so any non-float32 dtype (bf16, fp16, and future half-precision
  dtypes) is upcast to float32 before the FFT.
- Drop the `"torch.Tensor"` return annotation on the test helper
  that triggered ruff F821 in CI (torch is imported inside the
  method body, not at module scope).
2026-04-21 08:45:10 -03:00
Aditya Borate
7448258505 Add FLUX.2 Klein Inpaint Pipeline (#13050)
* Add Flux2KleinInpaintPipeline

* Fixed mask channel mismatch and a bit of cleaning

* Added tests and minor refactors

* Added support for reference images for inpainting

* Style fixes

* Fixed the example docstring

* Corrected mask latent preparation for correct dimensional alignment

* replace masked_image_latents context with clean_source_latents, fix mask spatial alignment and remove unused VAE encoding

* Fix T-coordinate collision for conditioning

* Changed the default strength from 0.6 to 0.8

* Added reference image test and updated the frozenset

* Validated ref image, latent passing support and fixed ref image preprocessing

* Refined preprocessing with 1MP resolution cap and timestep tracking

* Updated typing, improved validation and changed the example docstring

* Style fixes

* Fixed batch inference discrepancy and addressed review comments

* Fixed a typo

Co-authored-by: Álvaro Somoza <asomoza@users.noreply.github.com>

* Apply suggestion from @asomoza

Co-authored-by: Álvaro Somoza <asomoza@users.noreply.github.com>

* Reused encoded latents and fix channel check consistency

* fixed pre-encoded latent preprocessing for source and ref images

* Apply style fixes

* Updated the docstring with the shape requirements

* Apply style fixes

* Fixed copies
2026-04-17 17:34:36 -04:00
Baihao You
8d30d05de1 fix(qwen): fix CFG failing when passing neg prompt embeds with none mask (#13379)
* fix(qwen): fix CFG failing when passing neg prompt embeds with none mask

* fix(qwen): safely handle missing embeds masks in edit and inpaint pipelines

* test(qwen): add tests for true cfg scale without neg prompt mask

* fix(qwen): correct comments for copied functions in controlnet and inpaint pipelines

* fix(qwen): add warnings for missing prompt and negative prompt masks in pipelines

* test(qwen): use torch_device and clarify dummy inputs in cfg mask tests

* fix(qwen): address Claude PR review feedback

* fix(qwen): fix warning message based on reviewer suggestion

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-04-17 13:33:58 +05:30
Akshan Krithick
b3889ea478 Add modular pipeline for HunyuanVideo 1.5 (#13389)
* Add modular pipeline support for HunyuanVideo 1.5

* Fix I2V latent/cond spatial dimension mismatch

* Fix guidance_scale default to 7.5 matching ClassifierFreeGuidance

* Fix tokenizer type: use Qwen2TokenizerFast to match model

* Fix system message string formatting to match standard pipeline

* Rewrite HunyuanVideo 1.5 modular: use standard pipeline methods directly

* Remove I2V exports (T2V only for now)

* Fix encoder: use static methods directly instead of encode_prompt

* Inline all standard pipeline methods, remove runtime dependency

* Add HunyuanVideo 1.5 image-to-video modular blocks

* Fix missing FrozenDict import in before_denoise.py

* auto-generated docstrings via #auto_docstring

* Fix ruff lint and format issues

* use InputParam/OutputParam templates and fix ruff

* Address LTX review feedback here like add AutoBlocks, refactor I2V latents, lift encoders

* Add workflow map, workflow tests, auto docstrings, export only AutoBlocks

* Address Claude CI review

* Address claude CI review 2

---------

Co-authored-by: YiYi Xu <yixu310@gmail.com>
2026-04-16 10:47:51 -10:00
Lancer
c41a3c3ed8 [Feat] Adds LongCat-AudioDiT pipeline (#13390)
* Add LongCat-AudioDiT pipeline

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* upd

* Apply suggestions from code review

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* Apply style fixes

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* Apply style fixes

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

* Apply style fixes

* upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

---------

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
2026-04-15 00:47:38 -07:00
Sayak Paul
f65f135f64 [tests] xfail clip related issues. (#13454)
xfail clip related issues./
2026-04-14 22:33:25 +05:30
Sayak Paul
273b445426 [tests] fix deprecated attention processor testing. (#13469)
fix deprecated attention processor testing.
2026-04-14 22:05:39 +05:30
hlky
5063aa5566 FlashPack (#12700)
* FlashPack

* setup

* save_pretrained

* dtype is property

* destination_path

* logging

* pipeline

* ruff

* flashpack_kwargs

* download

* Fix docstring

* Apply suggestions from code review

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

* tests

* ignore_cleanup_errors

* -load_flashpack_checkpoint

* Apply style fixes

---------

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
2026-04-13 20:42:17 +05:30
HsiaWinter
dc8d903217 Add ernie image (#13432)
* Add ERNIE-Image

* Update doc

* Update doc

* Change from Custom-Attention to Diffusers Style Attention

* Change from Custom-Attention to Diffusers Style Attention

* 兼容SGLang

* 优化PE模块的加载与offload策略

* 更新Doc文件与config配置相关内容

* Fix官方反馈的内容

* 根据官方建议优化代码

* Update code

* update

* update

* Apply style fixes

* update

* update

* Apply style fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
2026-04-10 17:06:31 -10:00
Akshan Krithick
5a9a941a89 [modular] Add LTX Video modular pipeline (#13378)
* Add modular pipeline support for LTX Video

* Fix guidance_scale passthrough to guider

* Add LTX modular pipeline tests

* Add LTX image-to-video modular pipeline

* Fix i2v VAE dtype mismatch

* Add cache_context to denoiser for CFG parity

* Address review feedback

* Generate auto docstrings for LTX assembled blocks

* Fix ruff lint and format issues

* use InputParam/OutputParam templates and ruff check

* address all review

* Lift LTXVaeEncoderStep out of I2V core denoise into its own auto block

* unused declarations

* Fix check_copies: sync retrieve_timesteps and drop unsupported Copied from tags

* Address review: remove unused code, add LTXAutoBlocks, refactor I2V latents flow

* removed LTXBlocks,LTXImage2VideoBlocks

* Update test to use LTXAutoBlocks

* workflow map and auto docstring

* Add LTXVideoPachifier, workflow map

---------

Co-authored-by: YiYi Xu <yixu310@gmail.com>
2026-04-10 16:53:32 -10:00
Sayak Paul
896fec351b [tests] tighten dependency testing. (#13332)
* tighten dependency testing.

* invoke dependency testing temporarily.

* f
2026-04-10 18:12:12 +05:30
Dhruv Nair
a2583e55ff [CI] Add GLM Image Transformer Model Tests (#13344)
* update

* update

* update

* update
2026-04-07 16:28:05 +05:30
kaixuanliu
10ba0be991 Fix IndexError in HunyuanVideo I2V pipeline (#13244)
* add fallback logic for Hunyuan pipeline to make it compatible with
latest transformers

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

* use the last <|end_header_id|> token position + 1 as the assistant section marker

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

* fix format

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

* update variant name

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

---------

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
2026-04-06 10:33:34 -10:00
Sayak Paul
c39fba2ac4 [tests] fix autoencoderdc tests (#13424)
* fix autoencoderdc tests

* up
2026-04-06 21:05:20 +02:00
andrewor14
24b4c259fb Remove references to torchao's AffineQuantizedTensor (#13405)
**Summary:** TorchAO recently deprecated AffineQuantizedTensor
and related classes (https://github.com/pytorch/ao/issues/2752).
These will be removed in the next release. We should remove
references of these classes in diffusers before then.

**Test Plan:**
python -m pytest -s -v tests/quantization/torchao/test_torchao.py

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-04-06 20:41:26 +02:00
Dhruv Nair
ee3c352315 [CI] Hunyuan Transformer Tests Refactor (#13342)
* update

* update

* update

* update

* update

* update

* update
2026-04-06 20:16:20 +05:30
Sayak Paul
357b681890 [tests] refactor autoencoderdc tests (#13369)
* refactor autoencoderdc tests

* fix

* propagate new changes.
2026-04-06 11:10:21 +02:00
Dhruv Nair
065e36937a [CI] Refactor Cosmos Transformer Tests (#13335)
update

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-04-06 10:05:37 +05:30
sippycoder
447e571ada NucleusMoE-Image (#13317)
* adding NucleusMoE-Image model

* update system prompt

* Add text kv caching

* Class/function name changes

* add missing imports

* add RoPE credits

* Update src/diffusers/models/transformers/transformer_nucleusmoe_image.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/models/transformers/transformer_nucleusmoe_image.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/models/transformers/transformer_nucleusmoe_image.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/models/transformers/transformer_nucleusmoe_image.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* update defaults

* Update src/diffusers/pipelines/nucleusmoe_image/pipeline_nucleusmoe_image.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* review updates

* fix the tests

* clean up

* update apply_text_kv_cache

* SwiGLUExperts addition

* fuse SwiGLUExperts up and gate proj

* Update src/diffusers/hooks/text_kv_cache.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/hooks/text_kv_cache.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/hooks/text_kv_cache.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/hooks/text_kv_cache.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/models/transformers/transformer_nucleusmoe_image.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/models/transformers/transformer_nucleusmoe_image.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* _SharedCacheKey -> TextKVCacheState

* Apply style fixes

* Run python utils/check_copies.py --fix_and_overwrite
python utils/check_dummies.py --fix_and_overwrite

* Apply style fixes

* run `make fix-copies`

* fix import

* refactor text KV cache to be managed by StateManager

---------

Co-authored-by: Murali Nandan Nagarapu <nmn@withnucleus.ai>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
2026-04-03 02:01:13 -07:00
Sayak Paul
5adc544b79 [tests] refactor wan autoencoder tests (#13371)
* refactor wan autoencoder tests

* up

* address dhruv's feedback.
2026-04-03 07:36:40 +02:00
Steven Liu
e365d749a1 [docs] deprecate pipelines (#13157)
* deprecate

* fix

* fix

* fix

* fix

* remove deprecated .md files

* update links

* fix
2026-04-01 10:16:23 -07:00
Andrew Ross
b9353819a4 corrects single file path validation logic (#13363)
* corrects single file path validation logic

* Update tests/modular_pipelines/test_modular_pipelines_common.py

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

---------

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
2026-04-01 20:38:42 +05:30
YangKai0616
b88e60bd1b Fix: ensure consistent dtype and eval mode in pipeline save/load tests (#13339)
* Fix: ensure consistent dtype and eval mode in pipeline save/load tests

* Modify according to the comments

* Update according to the comments

* Update comment

* Code quality

* cast buffers to torch.float16

* conflict

* Fix

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-03-31 14:21:28 +05:30
Cheung Ka Wai
e1e7d58a4a Fix Ulysses SP backward with SDPA (#13328)
* add UT for backward

* fix SDPA attention backward
2026-03-30 15:15:27 +05:30
Sayak Paul
10ec3040a2 [ci] move to assert instead of self.Assert* (#13366)
move to assert instead of self.Assert*
2026-03-30 11:09:14 +05:30
Howard Zhang
f2be8bd6b3 change minimum version guard for torchao to 0.15.0 (#13355) 2026-03-28 09:11:51 +05:30
Howard Zhang
1fe2125802 remove str option for quantization config in torchao (#13291)
* remove str option for quantization config in torchao

* Apply style fixes

* minor fixes

* Added AOBaseConfig docs to torchao.md

* minor fixes for removing str option torchao

* minor change to add back int and uint check

* minor fixes

* minor fixes to tests

* Update tests/quantization/torchao/test_torchao.py

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Update docs/source/en/quantization/torchao.md

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Update tests/quantization/torchao/test_torchao.py

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* version=2 update to test_torchao.py

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-03-27 08:52:37 +05:30
Sayak Paul
85ffcf1db2 [tests] Tests for conditional pipeline blocks (#13247)
* implement test suite for conditional blocks.

* remove

* another fix.

* Revert "another fix."

This reverts commit ab07b603ab.
2026-03-26 08:48:16 +05:30
Kashif Rasul
5d207e756e [Discrete Diffusion] Add LLaDA2 pipeline (#13226)
* feat: add LLaDA2 and BlockRefinement pipelines for discrete text diffusion

Add support for LLaDA2/LLaDA2.1 discrete diffusion text generation:
- BlockRefinementPipeline: block-wise iterative refinement with confidence-based
  token commitment, supporting editing threshold for LLaDA2.1 models
- LLaDA2Pipeline: convenience wrapper with LLaDA2-specific defaults
- DiscreteDiffusionPipelineMixin: shared SAR sampling utilities (top-k, top-p,
  temperature) and prompt/prefix helpers
- compute_confidence_aware_loss: CAP-style training loss
- Examples: sampling scripts for LLaDA2 and block refinement, training scripts
  with Qwen causal LM
- Docs and tests included

* feat: add BlockRefinementScheduler for commit-by-confidence scheduling

Extract the confidence-based token commit logic from BlockRefinementPipeline
into a dedicated BlockRefinementScheduler, following diffusers conventions.

The scheduler owns:
- Transfer schedule computation (get_num_transfer_tokens)
- Timestep management (set_timesteps)
- Step logic: confidence-based mask-filling and optional token editing

The pipeline now delegates scheduling to self.scheduler.step() and accepts
a scheduler parameter in __init__.

* test: add unit tests for BlockRefinementScheduler

12 tests covering set_timesteps, get_num_transfer_tokens, step logic
(confidence-based commits, threshold behavior, editing, prompt masking,
batched inputs, tuple output).

* docs: add toctree entries and standalone scheduler doc page

- Add BlockRefinement and LLaDA2 to docs sidebar navigation
- Add BlockRefinementScheduler to schedulers sidebar navigation
- Move scheduler autodoc to its own page under api/schedulers/

* feat: add --revision flag and fix dtype deprecation in sample_llada2.py

- Add --revision argument for loading model revisions from the Hub
- Replace deprecated torch_dtype with dtype for transformers 5.x compat

* fix: use 1/0 attention mask instead of 0/-inf for LLaDA2 compat

LLaDA2 models expect a boolean-style (1/0) attention mask, not an
additive (0/-inf) mask. The model internally converts to additive,
so passing 0/-inf caused double-masking and gibberish output.

* refactor: consolidate training scripts into single train_block_refinement.py

- Remove toy train_block_refinement_cap.py (self-contained demo with tiny model)
- Rename train_block_refinement_qwen_cap.py to train_block_refinement.py
  (already works with any causal LM via AutoModelForCausalLM)
- Fix torch_dtype deprecation and update README with correct script names

* fix formatting

* docs: improve LLaDA2 and BlockRefinement documentation

- Add usage examples with real model IDs and working code
- Add recommended parameters table for LLaDA2.1 quality/speed modes
- Note that editing is LLaDA2.1-only (not for LLaDA2.0 models)
- Remove misleading config defaults section from BlockRefinement docs

* feat: set LLaDA2Pipeline defaults to recommended model parameters

- threshold: 0.95 -> 0.7 (quality mode)
- max_post_steps: 0 -> 16 (recommended for LLaDA2.1, harmless for 2.0)
- eos_early_stop: False -> True (stop at EOS token)

block_length=32, steps=32, temperature=0.0 were already correct.
editing_threshold remains None (users enable for LLaDA2.1 models).

* feat: default editing_threshold=0.5 for LLaDA2.1 quality mode

LLaDA2.1 is the current generation. Users with LLaDA2.0 models can
disable editing by passing editing_threshold=None.

* fix: align sampling utilities with official LLaDA2 implementation

- top_p filtering: add shift-right to preserve at least one token above
  threshold (matches official code line 1210)
- temperature ordering: apply scaling before top-k/top-p filtering so
  filtering operates on scaled logits (matches official code lines 1232-1235)
- greedy branch: return argmax directly when temperature=0 without
  filtering (matches official code lines 1226-1230)

* refactor: remove duplicate prompt encoding, reuse mixin's _prepare_input_ids

LLaDA2Pipeline._prepare_prompt_ids was a near-copy of
DiscreteDiffusionPipelineMixin._prepare_input_ids. Remove the duplicate
and call the mixin method directly. Also simplify _extract_input_ids
since we always pass return_dict=True.

* formatting

* fix: replace deprecated torch_dtype with dtype in examples and docstrings

- Update EXAMPLE_DOC_STRING to use dtype= and LLaDA2.1-mini model ID
- Fix sample_block_refinement.py to use dtype=

* remove BlockRefinementPipeline

* cleanup

* fix readme

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: YiYi Xu <yixu310@gmail.com>

* removed DiscreteDiffusionPipelineMixin

* add support for 2d masks for flash attn

* Update src/diffusers/training_utils.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/training_utils.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* fix issues from review

* added tests

* formatting

* add check_eos_finished to scheduler

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/schedulers/scheduling_block_refinement.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/schedulers/scheduling_block_refinement.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* fix renaming issues and types

* remove duplicate check

* Update docs/source/en/api/pipelines/llada2.md

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

* Update src/diffusers/pipelines/llada2/pipeline_llada2.py

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

---------

Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
2026-03-25 16:17:50 +05:30
Sayak Paul
e358ddcce6 fix to device and to dtype tests. (#13323) 2026-03-25 11:47:02 +05:30
Sayak Paul
f1fd515257 [tests] fix lora logging tests for models. (#13318)
* fix lora logging tests for models.

* make style
2026-03-24 15:48:03 +05:30
Cheung Ka Wai
afdda57f61 Fix the attention mask in ulysses SP for QwenImage (#13278)
* fix mask in SP

* change the modification to qwen specific

* drop xfail since qwen-image mask is fixed

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-03-24 02:12:50 -07:00
Sayak Paul
6350a7690a [chore] properly deprecate src.diffusers.utils.testing_utils. (#13314)
properly deprecate src.diffusers.utils.testing_utils.
2026-03-24 10:54:35 +05:30
Cheung Ka Wai
9d4c9dcf21 change QwenImageTransformer UT to batch inputs (#13312)
* UT expands to batch inputs

* update according to suggestion

* update according to suggestion 2

* fix CI

* update according to suggestion 3

* clean line
2026-03-24 08:56:40 +05:30
ddavidchick
ef309a1bb0 Add KVAE 1.0 (#13033)
* add kvae2d

* add kvae3d video

* add docs for kvae2d and kvae3d video

* style fixes

* fix kvae3d docs

* fix normalzation

* fix kvae video for code style

* fix kvae video

* kvae minor fixes

* add gradient ckpting for kvaes

* get rid of inplace ops kvae video

* add tests for KVAEs

* kvae2d normalization style change

* kvaes fix style

* update dummy_pt_objects test for kvaes

---------

Co-authored-by: YiYi Xu <yixu310@gmail.com>
2026-03-23 12:56:49 -10:00
Dhruv Nair
52558b45d8 [CI] Flux2 Model Test Refactor (#13071)
* update

* update

* update

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-03-23 16:56:08 +05:30
Sayak Paul
c02c17c6ee [tests] test load_components in modular (#13245)
* test load_components.

* fix

* fix

* u[

* up
2026-03-21 09:41:48 +05:30
Sayak Paul
522b523e40 [ci] hoping to fix is_flaky with wanvace. (#13294)
* hoping to fix is_flaky with wanvace.

* revert changes in src/diffusers/utils/testing_utils.py and propagate them to tests/testing_utils.py.

* up
2026-03-20 16:02:16 +05:30
Dhruv Nair
32b4cfc81c [Modular] Test for catching dtype and device issues with AutoModel type hints (#13287)
* update

* update

* update
2026-03-20 10:36:03 +05:30
dg845
072d15ee42 Add Support for LTX-2.3 Models (#13217)
* Initial implementation of perturbed attn processor for LTX 2.3

* Update DiT block for LTX 2.3 + add self_attention_mask

* Add flag to control using perturbed attn processor for now

* Add support for new video upsampling blocks used by LTX-2.3

* Support LTX-2.3 Big-VGAN V2-style vocoder

* Initial implementation of LTX-2.3 vocoder with bandwidth extender

* Initial support for LTX-2.3 per-modality feature extractor

* Refactor so that text connectors own all text encoder hidden_states normalization logic

* Fix some bugs for inference

* Fix LTX-2.X DiT block forward pass

* Support prompt timestep embeds and prompt cross attn modulation

* Add LTX-2.3 configs to conversion script

* Support converting LTX-2.3 DiT checkpoints

* Support converting LTX-2.3 Video VAE checkpoints

* Support converting LTX-2.3 Vocoder with bandwidth extender

* Support converting LTX-2.3 text connectors

* Don't convert any upsamplers for now

* Support self attention mask for LTX2Pipeline

* Fix some inference bugs

* Support self attn mask and sigmas for LTX-2.3 I2V, Cond pipelines

* Support STG and modality isolation guidance for LTX-2.3

* make style and make quality

* Make audio guidance values default to video values by default

* Update to LTX-2.3 style guidance rescaling

* Support cross timesteps for LTX-2.3 cross attention modulation

* Fix RMS norm bug for LTX-2.3 text connectors

* Perform guidance rescale in sample (x0) space following original code

* Support LTX-2.3 Latent Spatial Upsampler model

* Support LTX-2.3 distilled LoRA

* Support LTX-2.3 Distilled checkpoint

* Support LTX-2.3 prompt enhancement

* Make LTX-2.X processor non-required so that tests pass

* Fix test_components_function tests for LTX2 T2V and I2V

* Fix LTX-2.3 Video VAE configuration bug causing pixel jitter

* Apply suggestions from code review

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Refactor LTX-2.X Video VAE upsampler block init logic

* Refactor LTX-2.X guidance rescaling to use rescale_noise_cfg

* Use generator initial seed to control prompt enhancement if available

* Remove self attention mask logic as it is not used in any current pipelines

* Commit fixes suggested by claude code (guidance in sample (x0) space, denormalize after timestep conditioning)

* Use constant shift following original code

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-03-19 14:58:29 -07:00
kaixuanliu
8e4b5607ed skip invalid test case for helios pipeline (#13218)
* skip invalid test case for helio pipeline

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

* update skip reason

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

---------

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>
2026-03-17 20:58:35 -10:00
Dhruv Nair
11a3284cee [CI] Qwen Image Model Test Refactor (#13069)
* update

* update

* update

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
2026-03-17 16:44:04 +05:30
Sayak Paul
16e7067647 [tests] fix llava kwargs in the hunyuan tests (#13275)
fix llava kwargs in the hunyuan tests
2026-03-17 10:11:47 +05:30