mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
b800f828f3951c4d206e154411bcbfa693cd322e
263 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
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]( |
||
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>
|
||
|
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> |
||
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> |
||
Buffett Liu
|
360bef8074 |
Feat -- Stable Audio 3 (#14119)
* Stable Audio 3 First Version Commit, require model dict loading on GPUs * Fixing bugs for sa3 VAE * Make quality fix * Fix the unit test failure * Fix SA3 pipeline noise: learned text-padding embedding + no cross-attn mask The diffusers Stable Audio 3 pipeline produced noise instead of music because the cross-attention conditioning was built differently from the reference: - The reference text conditioner uses padding_mode="learned": padded text positions (~245 of 256 for a short prompt) are filled with a trained `padding_embedding`, and the DiT attends to *all* positions (its cross-attention mask is intentionally disabled). Our pipeline instead zeroed padded positions and masked them out, wiping ~95% of the conditioning signal. Changes: - transformer_stable_audio3: add learned `prompt_padding_embedding`; in forward, replace masked text positions with it (in cond_token_dim space, before to_cond_embed) then attend to the full context, matching the reference. - pipeline_stable_audio_3: stop zeroing padded positions (the DiT now handles them); default `silence_padding_duration` to 6.0 (reference headroom default). - convert_..._to_diffusers: convert `conditioner.conditioners.prompt.padding_embedding` into the DiT. - scheduling_ping_pong: pin schedule endpoints (sigmas[0]=1.0, sigmas[-1]=0.0) to match the reference LogSNRShift endpoint preservation. - run_..._inference: coerce float16->float32 on CPU. - tests: update transformer state-dict expectations (523 tensors, add prompt_padding_embedding) and ping-pong sigma endpoint assertions. Verified against the reference with identical noise + identical conditioning: full 8-step ping-pong trajectories agree to 4e-6 (final latent). Removed the dev-only parity scripts (verify_*_dit_parity, verify_*_vae_parity) from the tree. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Style fix * Autopipeline available * Adding document description for testing examples. * Solving fix for serge review bot suggestion * Design review comments fix * CI/CD Fix for ping pong scheduler md file * Fix doc build: remove autodoc for nonexistent StableAudio3EulerSchedulerOutput * Fix doc build: qualify non-exported autodoc references with module paths * Fix details in scheduler and some nits * suggestion fix * fix-copies: propagate prepare_latents dtype fix to audio2audio/inpaint pipelines --------- Co-authored-by: Buffett Liu <jengliu@linkedin.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> |
||
apolinário
|
2da7040be1 |
Add MiniMax Music 3 (#14456)
* Add MiniMax Music 3 transformer, pipeline sub-models, and conversion script
* Add MiniMax Music 3 pipeline, tests, and docs
* Wire example docstring, point diffusers-format references at MiniMaxAI/MiniMax-Music-3
* Use MiniMaxAI/MiniMax-Music3 as the official diffusers weights repo id
* Apply doc-builder style
* Use dtype in examples, demonstrate the structured prompt format
* Use soundfile in examples like the other audio pipelines
* Convert to a modular pipeline and move models under src/diffusers/models
Per review: the conditioner moves to models/condition_embedders, the RVQ depth
decoder and vocoder to models/transformers and models/autoencoders, and the
standard pipeline is replaced by MiniMaxMusic3ModularPipeline (anima-style
blocks, helios-style chunk loop, guider-abstracted CFG with a zeros
unconditional branch). Index configs now reference all components under the
diffusers library. Also applies the transformer review suggestions
(Transformer2DModelOutput, inlined rotary forward, explicit processor default).
* Point modular index loading specs at the Hub repo in the converter
* Apply doc-builder style to the modular blocks docstrings
* Address review-bot comments: overlap comment accuracy, single pipeline-level progress bar, zero-frame duration guard, import ordering, docstring cross-reference
* Support CPU offloading in the AR step: trigger offload hooks before submodule calls (minimax_h3 workaround), clear error when the AR models cannot co-reside
* Docs: PR-install tip and memory-usage section with measured offloading numbers
* Address review comments on the block assembly
Take the individual models rather than `components` in `_generate_depth_codes` and
`_embed_audio_frame`; `num_codebooks` / `audio_vocab_size` come off the depth decoder's
config, so the two pipeline properties that existed only for them are gone.
Assemble the blockset as three top-level children, each owning models and runnable on its
own — generate the semantics once and denoise it with different settings:
semantic_generator MiniMaxMusic3SemanticGenerationStep [tokenize, generate]
denoise MiniMaxMusic3CoreDenoiseStep [prepare_chunks, denoise]
decode MiniMaxMusic3VocoderDecodeStep
`MiniMaxMusic3TextEncoderStep` is renamed `MiniMaxMusic3TokenizeStep` (it only tokenizes)
and the autoregressive leaf `MiniMaxMusic3AutoregressiveStep`, freeing its old name for
the sequential.
Point the tests at `hf-internal-testing/tiny-minimax-music3-modular-pipe`; the previous
fixture repo's component specs pointed at a local path, so nothing loaded off that machine.
Drop the PR-branch install tip from the docs.
20s and 60s generations are bit-identical to before these changes.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---------
Co-authored-by: yiyixuxu <yixu310@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
|
||
|
Sayak Paul
|
7564fb016d |
Ltx 2.5 (#14447)
* Add LTX-2.4 support to the LTX2 transformer, pipelines, and conversion script LTX-2.4 reuses the existing LTX2 model and pipeline classes. The delta over 2.3 is a small set of additive config flags plus a Gemma 3 -> Gemma 4 text encoder swap, so no new model or pipeline classes are introduced. - transformer_ltx2.py: add `ff_bias`/`audio_ff_bias` (2.4's video FFN drops its bias) and `use_prompt_adaln_single` (toggles timestep-dependent prompt cross-attention modulation; when off, cross-attention K/V becomes timestep-independent and cacheable across denoising steps for a given prompt). Both default to their 2.3 behavior. The flag is read back from `self.config` rather than mirrored onto an instance attribute, and all three new constructor args are documented. - pipeline_ltx2.py / pipeline_ltx2_image2video.py: add an optional `prompt_enhancer` component to both pipelines (previously T2V-only). LTX-2.4's fine-tuned text encoder is conditioning-only, so enhancement uses a separate off-the-shelf google/gemma-4-E2B-it checkpoint with its own message format and decoding recipe -- unlike LTX-2.0/2.3, where one checkpoint serves both roles. `enhance_prompt()` resolves the format and `.generate` kwargs from whichever model is active, and frees a dedicated enhancer from GPU memory right after use (guarded so it does not interfere with accelerate's offload hooks), mirroring the existing text_encoder handling. Both `__call__`s gain `enable_prompt_enhancement: bool | None = None`, which resolves to `True` when a dedicated `prompt_enhancer` is configured (LTX-2.4) or when `system_prompt` was passed explicitly (matching prior LTX-2.0/2.3 behavior exactly), and `False` otherwise. Explicit `False` disables enhancement even on 2.4. When enabled with no `system_prompt` on a 2.4 pipeline, the matching default system prompt is injected. `max_new_tokens`/`seed` keep their literal defaults (512/10) unconditionally: greedy decoding consumes no randomness, so `seed` is inert for the dedicated-enhancer case, and 512 tokens comfortably covers the target caption length. No public API or default changes were needed for LTX-2.0/2.3. Validation of the enhancement arguments happens in `check_inputs`, so an unsatisfiable request fails before the prompt is encoded rather than partway through generation. Text encoder, tokenizer, processor and enhancer type hints name the concrete Gemma classes they accept. - utils.py: add a `PromptEnhancementConfig` dataclass plus `GEMMA3_PROMPT_ENHANCEMENT_CONFIG`/`GEMMA4_PROMPT_ENHANCEMENT_CONFIG` as the single source of truth for each model's message prefix and `.generate` kwargs, shared by both pipelines. Add the validated "capstyle_plus" LTX2_4_T2V/I2V_DEFAULT_SYSTEM_PROMPT strings, marked `docstyle-ignore` so `doc-builder style` cannot re-wrap them -- the prompts must stay byte-for-byte identical to the reference, newlines included. - convert_ltx2_to_diffusers.py: add a "2.4" branch to all five get_ltx2_*_config functions. Transformer/VAE/vocoder configs are structurally identical to 2.3 (verified against the checkpoint's own safetensors metadata); only `ff_bias=False` differs. Connector `caption_channels`/`text_proj_in_factor` are now derived from the live Gemma text config rather than hardcoded, since 2.4's text encoder is not pinned to a single checkpoint the way Gemma-3-12B is for 2.0/2.3. Swap `Gemma3ForConditionalGeneration`/`Gemma3Processor` for `AutoModelForImageTextToText`/`AutoProcessor` so one path covers Gemma 3 and Gemma 4. Raise a clear error when --version 2.4 is requested without pointing --text_encoder_model_id at a Gemma 4 (gemma4_unified) checkpoint, and add --prompt_enhancer_model_id, required whenever --version 2.4 is combined with --add_processor, since falling back to --text_encoder_model_id (correct for 2.0/2.3) would pair 2.4 with the wrong enhancement model. Also fixes a pre-existing vocoder class-selection bug that only checked for "2.3", and `processor` never being passed into the --full_pipeline LTX2Pipeline(...) construction. - docs/source/en/api/pipelines/ltx2.md: add an "LTX-2.4" section covering what carries over from 2.3 unchanged (guidance recommendations, aside from a different STG block index) and what does not (a single-stage checkpoint only, with no two-stage or distilled workflow yet), plus the corrected prompt-enhancement recipe and its enabled-by-default behavior for both LTX2Pipeline and LTX2ImageToVideoPipeline. * Fix resolution-dependent timestep shift being a no-op in all LTX2 pipelines Every LTX2Pipeline variant's `mu = calculate_shift(...)` call passed the scheduler's `max_image_seq_len` config value as the `image_seq_len` argument instead of the current generation's actual packed sequence length. Since calculate_shift's formula is a line through (base_seq_len, base_shift) and (max_seq_len, max_shift), passing image_seq_len == max_seq_len always evaluates to exactly max_shift -- so `mu` was pinned to a constant (2.05 for the LTX-2.4 scheduler config) regardless of height/width/num_frames, even though `use_dynamic_shifting: true` is set specifically to make this resolution-dependent. Found while benchmarking LTX-2.4 diffusers output against the reference pipeline with bit-identical starting noise: the reference computes this shift from the real video token count (ltx_core's LTX2Scheduler.execute), which diverges substantially from the constant diffusers was using at any resolution other than exactly the checkpoint's max-anchor token count. For a 768x512, 121-frame video (6144 tokens), the reference lands on mu ~= 2.78 versus diffusers' constant 2.05. Fixed by passing the current call's actual packed video latent length (`latents.shape[1]`) as `image_seq_len`, matching the pattern already used correctly in pipeline_flux.py (the source this was copied from). Applies to all 5 LTX2 pipeline variants, each with their own independent `__call__` (not linked via `# Copied from` for this method). * Add LTX-2.4 duration head for automatic num_frames prediction LTX-2.4 checkpoints ship a small regression head (~1.9M params) that predicts the natural duration of the shot implied by a caption, from the same text connector output the transformer is conditioned on. With it converted, a caller can let the model choose the video length instead of picking `num_frames`. - duration_head.py: `LTX2DurationHead` (a `ModelMixin` optional pipeline component) plus the `LTX2AutoDuration` request object. Modality-specific projections map the video and audio connector streams into a shared pooler dim, learnable modality embeddings tag them, one learnable query cross-attends the concatenation, and a small MLP regresses a log-duration. `forward` returns seconds as a tensor; `predict_num_frames` clamps to bounds and snaps to the VAE's causal temporal grid. The attention pooler uses explicit to_q/to_k/to_v/to_out with `dispatch_attention_fn` rather than `torch.nn.MultiheadAttention`, which both reference implementations use because that is the layout the checkpoint ships. `nn.MultiheadAttention` is documented in diffusers as breaking `enable_sequential_cpu_offload`; the split form also gets backend dispatch. Two details are load-bearing for numerical parity: the GELU must be tanh-approximated (the exact GELU gives different numbers against the JAX-trained head), and the clamp must precede the grid snap (a clamped frame count is not necessarily grid-aligned). Where narrow bounds convert to a frame window containing no grid point -- at 24 fps [1.0s, 1.02s] rounds to [24, 24], and 24 is not 8k + 1 -- the nearest grid point is used and a warning logged, rather than refusing to generate over a rounding artifact. The output MLP's config argument is `mlp_hidden_dim`, not the reference's `mlp_hidden`: the submodule keeps the checkpoint's `mlp_hidden` name, and `ModelMixin.__getattr__` resolves config keys ahead of submodules, so the two colliding would shadow the `nn.Linear` with an `int`. - pipeline_ltx2.py / pipeline_ltx2_image2video.py: `num_frames` becomes `int | LTX2AutoDuration | None`. Omitting it auto-predicts when the checkpoint ships a head and keeps the legacy 121 otherwise, mirroring the reference -- whose CLI also defaults to auto-prediction -- and matching the resolution these pipelines already do for `enable_prompt_enhancement`, which likewise switches on the presence of an optional LTX-2.4 component. Pre-2.4 pipelines have no head, so nothing changes for them. The prediction runs immediately after the connectors and before `num_frames` is first read. Only the positive half of the CFG-concatenated batch is used, and rows past the first are `num_videos_per_prompt` duplicates. Auto-duration is rejected from `check_inputs` -- before prompt enhancement and encoding, so a bad request costs nothing -- when there is no head, and when more than one prompt is supplied: the batch carries a single temporal dimension, so prompts with different natural lengths cannot share one frame count. Batched prompts with an explicit integer `num_frames` are unaffected. - convert_ltx2_to_diffusers.py: `--duration_head` (with `--duration_head_prefix`) alongside the other per-component flags. The head's keys sit at the checkpoint top level rather than under the DiT prefix, so the existing prefix helper handles them. The checkpoint's fused `in_proj_weight` is split into separate q/k/v, and dimensions are read back from weight shapes since checkpoint metadata carries no duration_head config; only `num_pooler_heads` is not recoverable and is fixed at 4. Pre-2.4 checkpoints yield no keys and are skipped rather than failing. Verified against ltx_core on the real 2.4 checkpoint: fed the reference's recorded connector tokens, the converted head predicts 12.9375s (video only), 2.671875s (audio only) and 3.515625s (both) -- bit-identical to the reference -- and the same frame counts through the clamp/snap path at 24/25/30/8 fps. The snapping arithmetic matches the reference `seconds_to_clamped_num_frames` across 260 combinations of duration, frame rate and bounds. * Add LTX-2 modular video and audio decoder blocks Add the first blocks of the LTX-2 modular pipeline under modular_pipelines/ltx2/decoders.py: - LTX2VaeDecoderStep: unpacks and decodes video latents (or returns latents for output_type="latent"), applying the optional decode-time noise on normalized latents before denormalizing, matching the standard LTX2Pipeline decode stage. - LTX2AudioDecoderStep: unpacks and decodes audio latents into a waveform via the audio VAE and vocoder in a single block. Pack/unpack/denormalize helpers are redefined at module level rather than imported, since modular blocks must not import from diffusers.pipelines.* (the vocoder class is imported from the pipelines path for now, flagged for relocation to models/). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Add LTX-2 modular pipeline (t2v + i2v) and parity harnesses Add the LTX-2 modular pipeline package under modular_pipelines/ltx2/, covering the joint video+audio text-to-video and image-to-video workflows for LTX-2.4: - encoders.py: dedicated Gemma-4 prompt enhancer (t2v/i2v), Gemma text encoder, text connectors, and the i2v image VAE encoder. - before_denoise.py: text-input expansion, flow-match timesteps (with a deep-copied audio scheduler), video/audio latent prep, and RoPE coords. - denoise.py: the joint video+audio denoise loop with manual guidance (CFG + spatio-temporal + modality-isolation), shared across t2v/i2v. - decoders.py: video VAE decode and audio VAE + vocoder decode. - modular_blocks_ltx2.py: LTX2Blocks (t2v), LTX2ImageToVideoBlocks (i2v), and LTX2AutoBlocks (both, default), plus the auto/conditional wrappers. - modular_pipeline.py: LTX2ModularPipeline with the compression-ratio and patch-size properties the blocks read. Wire up lazy imports and register the pipeline (top-level diffusers exports, modular_pipelines __init__, MODULAR_PIPELINE_MAPPING, and dummy objects). Also add integrations/ (temporary, for-visibility) parity harnesses that compare the modular t2v/i2v blocksets against the standard LTX-2 pipelines by sharing the same loaded components. This directory is meant to be removed before the final integration. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Add script to create tiny LTX-2.4 test pipeline ckpt (under assumption that it inherits all configs from LTX-2.3) * make fix-copies * Run modular connector once on CFG-concatenated batch for bitwise parity Run the LTX-2 text connector once on the CFG-concatenated `[uncond, cond]` batch (as `LTX2Pipeline` does) instead of once per branch, then split the outputs back into uncond/cond. The connector is applied per batch element, so both forms are mathematically equivalent, but its GEMM/attention kernels round identically for a given row only at batch >= 2; running the branches separately diverged from the standard pipeline by ~1e-6 at `num_videos_per_prompt=1`. The modular path is now bitwise-identical to the standard pipeline at any batch size. Also extend the T2V/I2V parity harnesses with a `--num_videos_per_prompt` argument and a `--check_tensor_stats` flag for per-output min/mean/std/max. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Use tiny LTX2VocoderWithBWE in the LTX-2.4 test checkpoint Swap the placeholder LTX2Vocoder for a scaled-down LTX2VocoderWithBWE that mirrors LTX-2.3's vocoder (snakebeta + antialiasing, no final activation/bias, 16kHz -> 48kHz bandwidth extension) while keeping the same in/out channel shapes. Dimensions are reduced but the shape invariants the two-stage forward requires are preserved: in_channels = audio_vae.output_channels * mel_bins, bwe_in_channels = out_channels * num_mel_channels, filter_length == window_length, and prod(bwe_upsample_factors) == (output_sr // input_sr) * hop_length so the BWE residual and the resampled stage-1 skip line up. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Match LTX-2.4 transformer flags in the test checkpoint Only `ff_bias=False` is a transformer-level delta from LTX-2.3 per the authoritative LTX-2.4 config in `scripts/convert_ltx2_to_diffusers.py`; `audio_ff_bias` and `use_prompt_adaln_single` keep their `True` defaults. The test checkpoint was incorrectly overriding both to `False`, so drop those two overrides. T2V/I2V modular-vs-standard parity remains bitwise. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Add CPU offload flag to parity scripts to be able to load full checkpoints * Support dynamic timestep shift and duration head in LTX-2 modular pipeline Track two recent additions to the standard LTX-2.4 pipelines in the modular blocks (t2v + i2v): - Resolution-aware timestep shift: LTX2SetTimestepsStep now computes `mu` from the actual packed video sequence length (derived from height/width/num_frames and the transformer patch sizes) instead of a constant, matching the standard pipeline's `latents.shape[1]`-based shift. Uses the compute-from-dims approach (like LTX-1 / Flux2), so no block reordering is needed. - Optional duration head: add LTX2DurationStep, which predicts a concrete `num_frames` from the connector text conditioning via the `duration_head` component when `num_frames` is an `LTX2AutoDuration` request, and re-emits it. Wrapped in the LTX2AutoDurationStep conditional (skipped for an integer `num_frames`) and wired into all three blocksets after the connector step, so `num_frames` is resolved before the shift and latent prep run. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Add duration predictor support for tiny checkpoint and parity scripts * Refactor LTX-2 modular guidance into video/audio guiders Replace the manual multi-term guidance in `LTX2LoopDenoiser` with a `LTX2Guidance` guider (in `guider.py`), instantiated once as the video `guider` and once as the `audio_guider`. Each combines CFG + spatio-temporal guidance (STG) + modality-isolation via the delta formulation in x0 space; the denoiser owns a `plan_guidance_passes` union plan across the two guiders, runs each transformer pass, converts velocity->x0, and delegates the per-modality combine to the guiders. Guidance scales are now guider config, not `__call__` kwargs. Parity harnesses updated to configure the guiders accordingly. Parity investigation (not yet resolved): - The refactor runs every guidance pass as its own single-batch transformer forward, whereas the standard `LTX2Pipeline` batches the cond+uncond CFG pair into one forward and runs STG/modality-isolation as separate single-batch passes. STG and modality already match (single-batch in both). - The change is mathematically equivalent, not a logic bug: in fp32 the denoised latents match to ~8e-6 mean abs diff (sparse outliers up to ~3.5e-4), and disabling STG does not move the diff. - But GPU matmul is not batch-invariant, so cond computed alone differs from cond computed inside a batch-of-2. Negligible in fp32 (~1e-6/op); ~1e-2/op in bf16, where amplification by the CFG delta and accumulation over sampler steps drives the modular vs. standard bf16 latents to ~10% mean-relative divergence. - Net: numerical, but the modular pipeline does NOT reproduce the standard pipeline bitwise in bf16 (the real inference dtype). Restoring parity would require re-batching the cond+uncond pair into a single forward to match the reference execution, keeping STG/modality single-batch. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Batch cond+uncond CFG forward to mirror standard LTX-2 pipeline Restore fp32 bitwise parity with `LTX2Pipeline` in the modular denoiser by running the cond+uncond CFG pair as a single batched transformer forward (`torch.cat([latents] * 2)` + `.chunk(2)`), keeping STG and modality-isolation as separate single-batch conditional forwards -- matching the reference op-for-op (batch sizes, repeated coords, cache-context names). `plan_guidance_passes` now emits forward-groups (`identifiers` / `conditioning` aligned lists + `flags` + `cache_context`) instead of one entry per pass; the denoiser runs each group once and chunks the CFG forward back into its `[uncond, cond]` identifiers. The `LTX2Guidance` combine is unchanged -- only how the four x0 tensors are obtained changed. Parity results: - fp32: bitwise (0.0 max abs diff), verified at full-checkpoint scale including under CPU offload. This is the authoritative parity gate. - bf16: the previous single-batch-per-pass design diverged ~10% mean-relative from the standard pipeline (GPU matmul is not batch-invariant: cond alone vs. cond in a batch-of-2). Batching the CFG pair removes that gap; a smaller ~1% (tiny) / ~5% (full) bf16 gap remains. It is not a logic difference (fp32 is bitwise across scale and offload); it is a bf16-kernel effect -- coarser mantissa amplifying non-associative accumulation order, plus bf16 using different kernels than fp32 (tensor-core GEMM algorithm selection, fused attention). bf16 is therefore a close-but-not-bitwise check, not a gate. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Drive LTX-2 modular guidance through the guider API (per-pass inputs) Replace the batched-CFG denoiser with one that runs each guidance pass as its own single-batch transformer forward, driven end-to-end through the standard guider API. `LTX2Guidance.prepare_inputs(guider_inputs)` builds one identifier- tagged batch per active pass from a dict whose values are 4-tuples indexed by pass [cond, uncond, stg, modality]; the per-pass model flags (`spatio_temporal_guidance_blocks`, `isolate_modalities`) ride in those tuples alongside the encoder inputs, so a pass fully describes its own forward. The denoiser unions both guiders' passes by identifier, runs each once (storing video+audio x0 on the batch), and combines each modality via its guider's `forward`/`__call__`, filtered to that guider's active passes so the batch count matches `num_conditions`. Removes the bespoke `plan_guidance_passes` union helper and the empty-dict `prepare_inputs_from_block_state` call: the plan is now expressed as the `guider_inputs` tuples + `active_predictions()`, so guidance logic lives behind the guider API rather than in the denoiser. Parity trade vs. the previous batched-CFG design: - Batched CFG matched the standard pipeline op-for-op and was fp32-bitwise. Running every pass single-batch is mathematically equivalent but, since GPU matmul is not batch-invariant, cond computed alone differs from cond inside a batch-of-2: ~1e-4 mean-relative in fp32 on a full checkpoint (sparse outliers up to ~3.5e-4 max), ~10% mean-relative in bf16. - This is numerical, not a logic difference, and fp32-within-tolerance (not bitwise) is the modular-ecosystem norm. The trade buys end-to-end guider-API usage (swappable within the LTX-2 guidance family, per-pass flags carried the same way as encoder inputs) at the cost of the bitwise guarantee. Parity harnesses: gate on magnitude-aware stats (mean abs diff relative to mean magnitude, plus a loose max-abs ceiling) instead of assert_close's near-bitwise fp32 defaults, which no single-batch design can clear. fp32 is the authoritative gate (1e-3/1e-3); bf16 is a loose sanity check (0.15/0.5). `--atol`/`--rtol` become `--mean_rel_tol`/`--max_abs_tol`; motivation documented in-file. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Drive LTX-2 modular guidance via prepare_inputs_from_block_state Move `LTX2LoopDenoiser` onto the standard Wan/Z-Image guider idiom: the denoiser now owns a `guider_input_fields` map (transformer arg -> per-pass block-state attribute names, indexed [cond, uncond, stg, modality]) and calls `guider.prepare_inputs_from_block_state(block_state, guider_input_fields)` instead of hand-building a literal `guider_inputs` dict in `__call__`. This lifts the cond/uncond/stg/modality field mapping to a construction-time arg (swappable per workflow) and resolves the connector_*->encoder_hidden_states name mismatch via the map keys. The two per-pass model flags (`spatio_temporal_guidance_blocks`, `isolate_modalities`) are pass-identity constants, not block-state conditioning, so they can't ride the name-referenced field map; the denoiser sets them on each batch by identifier after preparation, via a `pass_flags.get(identifier, (None, False))` lookup. The plain-conditional default keeps this correct for any guider that emits a subset of passes (e.g. a swapped-in `ClassifierFreeGuidance` -> just pred_cond/pred_uncond gets no STG and no modality isolation). `LTX2Guidance` gains `prepare_inputs_from_block_state` (names, via the base `_prepare_batch_from_block_state` helper); `prepare_inputs` (literals) is retained so both halves of the guider data-prep API are implemented. Pass structure and numerics are unchanged (still four single-batch forwards), so the fp32/bf16 parity story is untouched; user-confirmed parity OK. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Route LTX-2 batch-invariant denoiser kwargs via denoiser_input_fields Tag the three upstream-produced, batch-invariant transformer kwargs (`audio_num_frames` from the audio-latents step; `video_coords` / `audio_coords` from the coords step) with `kwargs_type="denoiser_input_fields"`, and have `LTX2LoopDenoiser` collect them from `block_state.denoiser_input_fields` filtered against the transformer's forward signature (à la qwenimage/cosmos3) instead of listing them as explicit inputs. This drops three explicit `InputParam`s in favor of one `denoiser_input_fields` template input. Per-pass conditioning (cond/uncond/stg/modality) stays on the guider field map -- the tag only delivers a flat dict, so it can't do the cond/uncond split or the connector_*->encoder_hidden_states rename. The locally-computed latent dims (num_frames/height/width/fps) are still supplied in-denoiser: they aren't upstream outputs and their names would clash with the pixel-space values in state. Parity unchanged (validated T2V + multi-frame I2V): the tagged values reach the transformer identically; the first-forward capture shows them bitwise-equal. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Fix LTX-2 parity harnesses: guidance overrides reach both pipelines The CLI guidance flags (`--guidance_scale`, etc.) only fed the modular guiders; the standard pipeline was always called with the hardcoded GUIDANCE dict. So `--guidance_scale 1.0` disabled CFG on the modular side only, comparing standard-with-full-CFG (batch-2) against modular-no-CFG (batch-1) -- an apples-to-oranges run that manifested as a huge (but spurious) I2V mismatch. Add `_resolve_guidance(args)` (CLI overrides on top of GUIDANCE) and drive BOTH the standard call and `_make_guiders` from the one resolved dict, in both the t2v and i2v harnesses. With CFG correctly disabled on both sides, multi-frame I2V is bitwise-ish (~5e-6 mean-rel), confirming the default-guidance ~8e-3 divergence is the documented cond/uncond batch-invariance (amplified by I2V's per-token masked timestep + clean anchor frame), not a logic bug. Also: - i2v: loosen the fp32 gate to (2e-2, 1.5e-1) to fit that amplified-but-numerical multi-frame divergence, with a comment pointing at the CFG-off run as the tight bug-catching gate. - i2v: add `--debug_forward`, which diffs the first transformer forward (inputs + outputs) between the two runs via a forward hook -- the diagnostic that pinpointed the batch-shape mismatch above. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Document denoiser_input_fields standalone-run caveat for LTX-2 denoiser `audio_num_frames` / `video_coords` / `audio_coords` reach `LTX2LoopDenoiser` via the `denoiser_input_fields` tag, not as named inputs. Note in the docstring that a standalone run (without the upstream tagging blocks) must pass them through `denoiser_input_fields={...}`; plain named kwargs are silently ignored (modular.md's kwargs_type standalone gotcha). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Address LTX-2 modular self-review: fix dead enhancer trigger + docstrings Blocking fix: `enable_prompt_enhancement` was declared only as a `block_trigger_inputs` entry / `select_block` param, never as an `InputParam`, so it was not an accepted pipeline input -- `pipe(enable_prompt_enhancement=True)` was dropped as "unexpected" and the prompt enhancer could never run. Declare it on both enhancer sub-blocks so it reaches `select_block` (mirrors how the `image` trigger is declared). Also from the self-review: - Regenerate the modular auto-docstrings: the guidance knobs (guidance_scale, stg_scale, ...) that moved onto the guider no longer show as block inputs, and the enhancer trigger now appears. - Add descriptions to the 16 `InputParam`s that rendered as "TODO: Add description." in the generated docstrings (conversion checklist requires none). - Drop the defensive `getattr(tokenizer, "padding_side", "left")` (a declared tokenizer always has it; gotcha #7). - Rewrite two ephemeral comments (the connector-output "reconcile when denoise.py is written" NOTE, now resolved; the guider "batched-CFG variant in git history" pointer) into standing rationale. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Fix stale I2V parity-harness --help text for the loosened fp32 gate The `--mean_rel_tol` / `--max_abs_tol` help strings still advertised the old fp32 defaults (1e-3, 1e-3); the I2V gate is now (2e-2, 1.5e-1). Update the help to match the actual DTYPE_TOLERANCES so `--help` doesn't misreport the gate. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * make style and make quality * Add the LTX-2.4 diffusion VAE decoder with a native-parity harness * Add a converter for the LTX-2.4 diffusion VAE decoder * Initialise use_slicing and document the diffusion VAE decisions * Record real-weight parity against the reference decoder * Add model-level tests for the LTX-2.4 diffusion VAE decoder The generated test file needed three decisions rather than fill-ins: the class had no forward() for the mixins to call, the decoder denoises so forward has to take a generator for any output comparison to mean anything, and MemoryTesterMixin.test_group_offloading reused one inputs_dict across four forwards without re-seeding it. The last is fixed in the mixin, reusing the helper test_group_offloading_with_disk already had for the same reason, now module-level and reading the signature off the class (offloading replaces model.forward with a *args wrapper). * Fold the diffusion VAE converter into convert_ltx2_to_diffusers.py --diffusion_vae now sits beside --vae instead of shipping a second script. It also drops the standalone script's --encoder-vae: the encoder is in the same checkpoint and goes through the conv VAE's own rename rules, with its config pinned in get_ltx2_diffusion_video_vae_config. Output verified bitwise identical to the standalone script's on the rc2 checkpoint, 491/491 tensors and the same config.json. * Make the diffusion decoder work through LTX2Pipeline The documented usage -- LTX2Pipeline.from_pretrained(repo, vae=diffusion_vae) -- had never been run: every real-weight run so far called vae.decode() directly. Through the pipeline it raised on vae.config.timestep_conditioning, which this decoder has no reason to carry, and would have bound the positional decode timestep to the decoder's generator. The decode step now branches on the decoder type, skipping the decode_timestep pre-noising and passing the generator instead, which also makes decoding seed-reproducible; two full runs from one seed agree bitwise. The conv path's own statements are unchanged. * Make the diffusion decoder work through the modular decode step Same fix as the standard pipeline, in LTX2VaeDecoderStep: skip the decode_timestep pre-noising and pass the generator instead of a timestep. Verified on rc2 weights by injecting the decoder into LTX2Blocks().init_pipeline() with update_components, the route the modular parity harness already uses -- two runs at 320x448x17 from one seed agree bitwise, so no modular_model_index.json repo was needed to run this. * Note that the pipeline decode block predates the LTX-2.4 integration PR * Address diffusers self-review on the diffusion VAE decoder Four things the project's own rules catch: - models.md forbids unconditional torch.float64 in a model (MPS/NPU/Neuron cannot run it). The RoPE frequency base now goes through maybe_adjust_dtype_for_device, the helper flux/flux2/wan use for exactly this. Measured cost of the downcast where it applies: 1.5e-08 on the frequencies, 1e-06 on the angles. - rope_dim_split was threaded through five classes but was never a config entry, so it was always None and always fell back to the default split. Removed, and the default-split helper inlined into its one caller. - FromOriginalModelMixin was declared with no SINGLE_FILE_LOADABLE_CLASSES entry, so from_single_file raised rather than working. Dropped; single-file support can come with its own mapping later. - set_attention_backend() overwrites the processor's _attention_backend, which handed the neighborhood BlockMask to backends that cannot read it. Guarded like AnyFlowCausalAttnProcessor does. Also runs make fix-copies, which was missing the dummy object for the new class. * Add the API doc page for AutoencoderKLLTX2VideoDiffusionDecoder * Mark the parity harness transient and drop the notes doc from the diff Matches the header the other integrations/ harnesses on this branch already carry, and states that the file imports the native reference package. The notes doc goes to the PR description instead: the self-review skill is explicit that it never ships in the diff. * Accept the diffusion decoder in the image-to-video pipeline too Same branch as the text-to-video pipeline: skip the decode_timestep pre-noising and pass the generator instead of a timestep. The remaining LTX2 pipelines that share this decode block (condition, hdr_lora, ic_lora, latent_upsample) are deliberately left for later, so the surface stays the two pipelines this PR actually exercises. * Reorganize LTX-2 modular blocks into coarser component-aligned phases Regroup the top-level LTX-2 block assemblies so each direct child of LTX2AutoBlocks is a coherent, component-aligned phase, mirroring Flux 2's [text_encoder, vae_encoder, denoise, decode] shape: - New LTX2TextConditioningStep wraps [text_encoder, text_input, connectors] as one "text_encoder" phase (LTX-2's text conditioning is genuinely a three-step model chain; text_input must precede connectors for batch-invariant parity, so the whole chain lives here rather than pushing text_input into denoise). - New LTX2DecoderStep wraps [video_decode, audio_decode] as one "decode" phase. - duration stays a top-level child: it is component-aligned (duration_head) and a ConditionalPipelineBlocks that self-skips for concrete num_frames, so it is a cleanly optional phase alongside prompt_enhancer / vae_encoder -- keeping it out of the (mandatory) denoise block. Top-level children become: t2v: [text_encoder, duration, denoise, decode] i2v: [text_encoder, duration, vae_encoder, denoise, decode] auto: [prompt_enhancer, text_encoder, duration, vae_encoder, denoise, decode] Also reorder the class definitions to follow this flow, keeping the core-denoise trio (LTX2CoreDenoiseStep, LTX2Image2VideoCoreDenoiseStep, LTX2AutoCoreDenoiseStep) contiguous instead of split by LTX2AutoVaeEncoderStep. Pure re-grouping: no block bodies change and intermediate outputs still propagate, so t2v/i2v parity is bit-identical to before (user-confirmed at 17 frames). Docstrings regenerated. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Enable LTX-2 prompt enhancement by default when the enhancer is loaded Mirror the standard pipeline's tri-state prompt-enhancement default (pipeline_ltx2.py:559-560), where `enable_prompt_enhancement` defaults to auto-on when a `prompt_enhancer` component is present (LTX-2.4) and off otherwise, and wire the enhancer into the t2v/i2v assemblies. - Both enhancer sub-blocks: default `enable_prompt_enhancement=None` (tri-state). Resolve `None` against component presence inside `__call__` -- select_block cannot see components, so the gate dispatches for None/True and the sub-block decides. `None` + no enhancer -> pass the prompt through unchanged (the Cosmos3 optional-component pattern); explicit `True` + no enhancer -> a clear ValueError instead of the previous cryptic `NoneType` crash from within the enhance helper. - LTX2AutoPromptEnhancerStep.select_block: skip only on an explicit `False`. - Add LTX2AutoPromptEnhancerStep to LTX2Blocks and LTX2ImageToVideoBlocks so all three assemblies share the prompt_enhancer-first shape (previously only LTX2AutoBlocks had it; passing the flag to the others was silently ignored). - Parity harnesses: force `enable_prompt_enhancement=False` on the modular run too (the modular sets now carry an enhancer), and fix the now-stale "contains no enhancer block" docstrings. Only the dedicated-enhancer (LTX-2.4) path is wired; the standard pipeline's text_encoder-as-enhancer fallback for LTX-2.0/2.3 remains a follow-up. Parity unchanged (both pipelines run with enhancement disabled in the harnesses). Docstrings regenerated. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Rename LTX-2.4 product identifiers to LTX-2.5. Update conversion --version, model ids, system-prompt constant names, and docs/comments to the public 2.5 name. Behavior is unchanged in this commit. * Add image-conditioning H.264 CRF matching Lightricks training. Resolve CRF 18 for LTX-2.5 (Gemma 4) and 33 for earlier generations, and apply the PyAV re-compress path for I2V and single-frame Condition/IC keyframes. * Align LTX-2 pipelines with the Lightricks reference behavior. Extract shared prompt-enhancement and check_inputs mixins; enable Condition duration+enhancement and IC enhancement (HDR remains encode-only); match enhancement recipe and LTX-2.3/2.5 sampling defaults. * Update LTX-2 docs for 2.5 naming, defaults, and opt-in enhancement. Document LTX_2_3/2.5 sampling defaults, STG block 28, and enable_prompt_enhancement=True as the opt-in path. * Replace LTX2AutoDuration with min_seconds/max_seconds call args. Drop the dataclass wrapper and expose duration bounds directly on the T2V, I2V, and Condition pipelines. * Align LTX-2 modular blocks with the merged standard-pipeline changes The standard LTX-2 pipelines moved to shared `LTX2PromptEnhancementMixin` / `LTX2CheckInputsMixin`, dropped `LTX2AutoDuration`, adopted the Lightricks reference defaults, and gained H.264 CRF re-compression of image conditionings. Port all of that to the modular blocks. - Prompt enhancement: re-derive `_enhance_prompt` from the new mixin -- the `user prompt:` / `User Raw Input Prompt:` templates, 896px long-side image prep, left-padding to a multiple of 8 for Flash Attention, `clean_response` on the decoded output, and `max_new_tokens=None` resolving to the Gemma-4 budget of 600. - Revert `enable_prompt_enhancement` to a plain `False` opt-in, matching the reference pipelines (the tri-state auto-on default is gone upstream). The friendly error for `True` with no enhancer loaded is kept. - Auto-duration: trigger on an omitted `num_frames` instead of an `LTX2AutoDuration` sentinel, with `min_seconds`/`max_seconds` as call args. `LTX2DurationStep` raises when no `duration_head` is loaded rather than silently falling back to 121 frames. - Every `num_frames` InputParam default becomes `None`: input defaults are seeded into the pipeline state before any block runs, so a default of 121 would leave the auto-duration branch permanently unreachable. - Defaults sweep: `num_inference_steps` 40 -> 30, `use_cross_timestep` False -> True, and the video/audio guider ComponentSpec configs updated to the reference guidance stack. - `LTX2VaeEncoderStep` gains `image_crf` and re-compresses the conditioning image before preprocess, declaring `text_encoder` only to resolve the model default CRF. - Rename LTX-2.4 -> LTX-2.5 throughout, following the upstream product rename. Parity harnesses updated for the new duration arguments, plus an `--image_crf` flag on the i2v harness defaulting to 0 so the comparison isolates block logic from the PyAV codec round-trip. Verified against a real checkpoint: t2v, i2v at `--image_crf 0`, and i2v at `--image_crf -1` all pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Evaluate the decoder's SwiGLU in token tiles Decode-only peak VRAM at 121 frames and 512x768 was 24.6 GiB, and 17.0 GiB of that was one transient inside the first stage-5 block: `w_gate(x)` and `w_up(x)` are both hidden-width and their product makes a third, so the whole video's worth of hidden activations was live three times over. At that size each one is 2,973,696 tokens x 1024 x 2 bytes = 5.67 GiB. The MLP is pointwise across tokens, so tiling it is exact -- it changes only how many hidden-width elements exist at once. Measured on a full-size decode: peak 24.8 -> 19.4 GiB (-22%), wall clock 2.79 -> 2.84 s (+1.6%), output bit-identical (max abs diff 0.000e+00). That also puts the diffusion decoder below the conv decoder, which needs 23.8 GiB at the same size. Tile size 16384 tokens is the reference decoder's own default (`DEFAULT_SWIGLU_TILE_SIZE`). Bounding a token count rather than a number of tiles keeps peak memory independent of resolution. This is the first of the two transients the reference's default decode mode avoids; the other is the full-volume QKV projection plus RoPE's fp32 upcast, which it addresses by chunking attention along width. That one is not bit-exact (it reassociates NATTEN's accumulation), so it is left for a follow-up rather than changed here. * Declare _no_split_modules on the diffusion decoder Both block types close over a residual add that combines outputs from different children, which is the case models.md says has to stay co-located, so neither can be split across devices. Unskips three memory tests that were gated on the attribute: the file goes from 35 passed / 7 skipped to 38 passed / 4 skipped, and the three (test_cpu_offload, test_disk_offload_with_safetensors, test_disk_offload_without_safetensors) pass. * Match the encoder defaults to the published LTX-2.4 config block_out_channels was (256, 512, 1024, 2048) and layers_per_block was (4, 6, 6, 2, 2); both were inherited from the 2.0 conv VAE. The published vae_diffusion/config.json has (256, 512, 1024, 1024) and (4, 6, 4, 2, 2). Real loads were unaffected because the config sets these explicitly, but the defaults described a model that does not exist. All 27 __init__ defaults that appear in the published config now match it. * Cover the token-tiled SwiGLU path in the tests The dummy video is 9x48x48, so its stage-5 grid is 5184 tokens against a 16384-token tile size: every other test in the file takes the untiled branch and the tiling loop ran nowhere in CI. This shrinks the tile size to force ~41 tiles and requires torch.equal against the untiled evaluation, since the MLP is pointwise across tokens and tiling must be exact rather than close. Checked the test is not vacuous: sabotaging the tiled branch alone (silu on w_up instead of w_gate) fails it at 8.83e-01. * Let the parity harness run against the LTX-2.5 reference too The harness only spoke LTX-2.4's decoder API and died on `forward_pre_diffusion` against 2.5, so it could not check the release the port is now targeted at. Three shims absorb the differences: * 2.5 selects a decode pathway by swapping `__class__` on every block, so an as-constructed decoder has no `forward_combined`. Install the COMBINED pathway explicitly, uncompiled, since that is the pathway this port implements. * stages 1-4 split into `forward_stages_1_to_3` + `forward_stage_4`, so the chunked pathway can defer stage 4. * the trailing NATTEN ghost pad moved out to the caller while its crop stayed behind `forward_stage_4(pad_trailing=True)`; the two have to be paired or the context loses real frames (3 of 17 survive at this config, which is how the mismatch first showed up). Same numbers against both references -- context 5.367e-04, step 7.337e-04 -- which is the check that the shims are semantically right and not just shape-compatible. * Add LTX-2.X modular condition blocks Ports LTX2ConditionPipeline to the modular pipeline system, including the LTX-2.5 additions (per-condition CRF, conditions-sourced prompt enhancement, auto-duration, the Lightricks reference defaults). New blocks: - encoders.py: LTX2ConditionEncoderStep (resize/center-crop, single-frame H.264 re-compression at the model CRF, temporal trim, VAE encode) and LTX2ConditionPromptEnhancerStep (grounds the rewrite in the first PIL frame found in `conditions`, text-only otherwise). - before_denoise.py: LTX2ConditionPrepareLatentsStep (first-frame overwrite + keyframe token append with their own RoPE coords, emitting conditioning_mask / clean_latents / appended_coords / base_token_count), LTX2ConditionSetTimestepsStep, LTX2ConditionPrepareAudioLatentsStep, LTX2ConditionPrepareCoordsStep, and the _prepare_keyframe_coords helper. - denoise.py: LTX2ConditionLoopBeforeDenoiser (mask-scaled per-token video timestep), LTX2ConditionLoopAfterDenoiser (x0 blend against the clean condition latents, then a full-sequence scheduler step), LTX2ConditionDenoiseStep. - decoders.py: LTX2TrimConditionTokensStep. - modular_blocks_ltx2.py: LTX2ConditionAutoPromptEnhancerStep, LTX2ConditionCoreDenoiseStep, LTX2ConditionDecoderStep, LTX2ConditionBlocks. LTX2LoopDenoiser is reused unchanged: the two LTX2Guidance guiders already express the whole CFG + STG + modality-isolation stack the condition pipeline uses, and the transformer takes the longer packed sequence with the base num_frames/height/width exactly as the standard pipeline does. The standard pipeline's CFG batch duplication (conditioning mask, coords) has no modular counterpart, since each guidance pass is its own single-batch forward. Block ordering: prepare-latents runs *before* set-timesteps, unlike the text-to-video and image-to-video core denoise steps. The resolution-aware shift `mu` is computed from the packed latent sequence length, which here includes the appended keyframe tokens, so the latents must exist first. This mirrors LTX2ConditionPipeline (its section 4 precedes section 5). The text-to-video and image-to-video blocksets are unaffected: with no appended tokens the grid-derived sequence length already equals latents.shape[1]. Two subtleties the parity harness surfaced: - The condition pipeline samples audio noise directly in packed shape [B, L, C * M], while the text-to-video and image-to-video pipelines sample unpacked [B, C, L, M] and pack afterwards. Both draw the same number of values from the generator but lay them out differently, so reusing the shared LTX2PrepareAudioLatentsStep silently desynchronizes the audio noise and, through the joint attention, the video too (1.4 mean-rel on audio, 0.13 on video before the fix). Hence a standalone condition audio block. - `noise_scale` means different things across workflows: 0.0 for text-to-video/image-to-video, `None -> sigmas[0] or 1.0` for conditions. A blockset keeps the first non-None default across its blocks, so a 0.0 would have shadowed the condition resolution and left the latents unnoised. Both condition-path blocks declare None; prepare-latents resolves once and writes the value back for the audio step. Keyframe tokens, masks and coords are expanded to the generation batch (the pattern the in-context pipeline already uses for reference tokens). The standard pipeline builds them at batch 1 and cats them onto a batch-B sequence, so it raises for num_videos_per_prompt > 1 with any condition at latent index > 0; the modular blocks handle that case. Verification: integrations/ltx2_condition_parity.py (transient, removed before merge) compares LTX2ConditionBlocks against LTX2ConditionPipeline on shared component objects. Eight cases on a tiny checkpoint -- first-frame only, keyframe only, both plus a multi-frame video condition, two keyframes, no conditions, --crf -1, --predict_duration, num_videos_per_prompt=2 -- all pass at 3e-7 to 1.3e-6 mean-rel. On a full checkpoint at 768x512x17 with two conditions through the CRF path: video 8.7e-04 mean-rel / 1.1e-02 max-abs, audio 1.5e-04 / 2.1e-03, in line with the image-to-video figures at the same frame count. make style, make quality and make fix-copies are clean; auto-docstrings regenerated with no TODO placeholders. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Add LTX-2.X modular in-context (IC-LoRA) blocks Ports LTX2InContextPipeline to the modular pipeline system on top of the condition blocks, and fixes two bugs in the standard in-context pipeline that the parity harness surfaced. New blocks: - encoders.py: LTX2ReferenceEncoderStep (preprocess, VAE-encode and pack each reference video, compute the coords that map its tokens into the target space, and build the per-token cross-attention strengths), plus the _downsample_mask_to_latent helper. - before_denoise.py: LTX2InContextPrepareLatentsStep (frame conditions as in the condition step, then the reference tokens appended last) and LTX2BuildVideoSelfAttentionMaskStep. - modular_blocks_ltx2.py: LTX2AutoReferenceEncoderStep, LTX2AutoBuildVideoSelfAttentionMaskStep, LTX2InContextCoreDenoiseStep, LTX2InContextBlocks. LTX2ConditionDenoiseStep is reused unchanged: with the blend gate fixed below, reference tokens are pinned by exactly the same x0 blend as frame conditions, matching the reference implementation where VideoConditionByReferenceLatent and VideoConditionByKeyframeIndex produce the same (denoise_mask, clean_latent) pair and post_process_latent runs unconditionally. The condition encoder, timesteps, audio-latents, coords and decoder blocks are shared as well. video_self_attention_mask reaches the transformer through the denoiser_input_fields tag, so LTX2LoopDenoiser needed no change. `reference_conditions` is optional, matching LTX2InContextPipeline: IC-LoRAs that carry their behavior in the adapter weights (camera control, style) take no reference video, which is the shape of the standard pipeline's own docstring example. The reference encoder and the attention-mask step are both skipped when it is absent. There is no duration step (the in-context pipeline ships without a duration_head), so LTX2ConditionEncoderStep now raises a concise error if `num_frames` is still None by the time the conditions are encoded. Per-reference token counts come from the encoder rather than an equal split of the total, so references of differing lengths (a reference video shorter than num_frames) get the right strengths. Standard-pipeline fixes in pipeline_ltx2_ic_lora.py: - The x0 blend was gated on `has_conditions` alone, so with reference conditions and no frame conditions the reference tokens were seeded clean and given a zeroed timestep but never re-pinned. Their x0 prediction error, divided by sigma each step, then accumulated and drifted them away from the encoded reference. Now gated on `has_conditions or num_ref_tokens > 0`, in line with the surrounding gates and the reference implementation. - video_self_attention_mask was rebound in the loop rather than expanded into a per-pass local. The main pass expanded it to the CFG-doubled batch, after which the STG pass's expand() to `latents.shape[0]` raised "RuntimeError: The expanded size of the tensor (1) must match the existing size (2)". Any run combining an attention mask with STG or modality guidance -- i.e. the default guidance settings -- died on the first step. Also removed the unused has_appended_tokens local and replaced the "# - TODO" placeholder above the attention-mask construction with the block structure, which the reference confirms. Verification: integrations/ltx2_in_context_parity.py (transient, removed before merge). Twelve cases on a tiny checkpoint -- reference only, reference at partial strength, scalar attention strength, pixel-space attention mask, downscale factor 2, references plus first-frame and keyframe conditions, two references, the CRF path, and four reference-free variants -- all pass at 3.4e-7 to 1.3e-6 mean-rel. The seven-case condition suite still passes. On a full checkpoint at 768x512x17 with a reference plus two conditions: video 4.4e-05 mean-rel / 5.6e-04 max-abs, audio 3.2e-05 / 2.4e-04. Known limitation, not addressed here: transformer_ltx2.py converts the multiplicative [0, 1] self-attention mask to an additive bias linearly, as (1 - mask) * -10000, whereas the reference uses log space (bias = log(mask), finfo.min at zero). Intermediate strengths therefore saturate to "fully masked", making conditioning_attention_strength effectively binary and spatially varying masks inert beyond their zero/non-zero pattern. Parity is unaffected -- both pipelines feed the same mask into the same transformer -- but the feature does not currently work as documented. make style, make quality and make fix-copies are clean; auto-docstrings regenerated with no TODO placeholders. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Make the diffusion decoder a decode pipeline instead of a VAE Addresses the review: a full diffusion model should not be an AutoencoderKL that hides its denoising behind `.decode()`, and it should not have to be threaded into every LTX-2 pipeline as a `vae`. AutoencoderKLLTX2VideoDiffusionDecoder becomes LTX2VideoDiffusionDecoderModel, a plain ModelMixin with no encoder. Encoding stays with AutoencoderKLLTX2Video, whose latent space this consumes unchanged, so latents remain interchangeable between the two decoders. Dropping the encoder also removes the converter's awkwardness of sourcing encoder weights and the real block_out_channels from the conv `vae/` folder: ten encoder-only config keys are gone. LTX2VideoDiffusionDecodePipeline drives it, shaped like LTX2LatentUpsamplePipeline: run any LTX-2 pipeline with output_type="latent", then decode. `vae` is an optional component, since it is only ever consulted for the latent statistics, and the decoder carries its own `latents_mean` / `latents_std` buffers to fall back on. That means a decode-only workflow never loads a second autoencoder. The `vae=` path is removed accordingly: the isinstance branches in LTX2Pipeline, LTX2ImageToVideoPipeline and the modular decode step are gone, along with the fixtures that built a diffusion decoder to pass as `vae`. The model's slicing/tiling test class went with them, since enable_tiling was an AutoencoderMixin concern and this is no longer an autoencoder. Tests: 37 passed / 3 skipped on the model, 83 passed / 2 skipped across the t2v and i2v pipelines. Still to come: mapping the 2.5 checkpoint to a blockset that selects the diffusion decode block, and the KandinskyCombinedPipeline pattern for running the two-step flow by default in the standard API. * Add a modular decode block for the LTX-2 diffusion decoder The previous commit removed the `isinstance(vae, ...)` branch that let the diffusion decoder stand in as the modular `vae`, which was the right call twice over: it is what the review objected to, and modular.md says a new case gets a new block rather than a branch inside an existing one. But removing it without a replacement left modular with no route to the diffusion decoder at all. LTX2DiffusionVaeDecoderStep is that replacement. It declares the decoder as its own `diffusion_decoder` component rather than borrowing the `vae` slot, and passes a generator instead of a decode timestep, since this decoder draws the noise it denoises. `LTX2VaeDecoderStep` is untouched, so the convolutional path is unaffected. What is still missing is the wiring: a blockset using this block, and an `_ltx2_map_fn` so a checkpoint that ships the diffusion decoder selects it automatically instead of the user choosing. That lands in modular_blocks_ltx2.py, which #12 is reorganizing, so it needs sequencing against that PR rather than racing it. * Test the diffusion decode pipeline, and give it a video-only output The decode pipeline had no coverage: it was never actually run. Writing the tests immediately turned up a bug, which is that it returned `LTX2PipelineOutput`, whose `audio` field is required. This pipeline decodes video only, so it now returns a video-only `LTX2VideoDecodeOutput` instead of passing a meaningless `audio=None`. Four tests: decode with no `vae` (the fallback to the decoder's own latent statistics), decode with a `vae` supplied (its statistics take precedence, checked by making the two disagree rather than asserting a shape), reproducibility under a seeded generator plus divergence under a different seed, and `denormalize=False` actually skipping the statistics. The decoder's stat buffers are set to non-trivial values in the fixture so a run that silently skipped denormalization could not pass by coincidence. * Fix the decode example: output_type="latent" is already denormalized The example in the docs and the pipeline docstring both showed `output_type="latent"` feeding straight into the decode pipeline on its defaults. That path denormalizes twice: `LTX2Pipeline` applies the latent statistics before returning latents (pipeline_ltx2.py), and the decode pipeline applies them again because `denormalize` defaults to True. Every channel ends up scaled by its own std a second time, so e.g. channel 0 (std 0.238) comes out roughly 4x too small. Passing `denormalize=False` in the example and saying so in both docstrings. Caught by running the flow end to end on real rc4 weights, which is also where the two decoders agree at 0.9978 once the statistics are applied only once. Whether the *default* should flip, or `output_type="latent"` should stop denormalizing instead, is a contract question across both pipelines and is still open. * Fix IC-LoRA video self-attention mask grouping The video self-attention mask deviated from the reference implementation (`ltx_core.conditioning.mask_utils.build_attention_mask`) in two ways. Both predate the modular blocks and applied to `LTX2InContextPipeline` as well; the modular block had ported the same construction. 1. Keyframe conditions cross-attended with reference tokens. `num_noisy_tokens` was derived as `latents.shape[1] - num_ref_tokens`, which folds the appended keyframe tokens into the noisy group, so they received the reference cross mask. The reference takes `num_noisy_tokens` from the target latent shape -- generated-video tokens only -- and tracks the extras offset separately, so keyframe <-> reference is 0.0. Split the two quantities apart: `_build_video_self_attention_mask` now takes `num_prefix_tokens` alongside `num_noisy_tokens`. 2. Multiple reference conditions fully attended to each other. All references were passed as a single attention group with a 1.0 self-block. The reference wraps each `LTX2ReferenceCondition` in its own `ConditioningItemAttentionStrengthWrapper`, building the mask once per item, so reference_i <-> reference_j is 0.0. The call site now splits `ref_cross_mask` into one group per reference before building the mask. Threading per-reference token counts out of `_encode_reference_conditions` also removes the equal-split assumption in `prepare_latents` and `prepare_reference_latents`, which assigned the wrong per-token strengths when references encoded to different token counts (reachable when a reference video is shorter than `num_frames`). API impact ---------- `prepare_latents` returns a 7-tuple instead of a 6-tuple, adding `ref_token_counts` (`list[int]`) as the new final element. Elements 1-6 keep the types they have on `main`; in particular `ref_cross_mask` remains the concatenated `[1, num_ref_tokens]` tensor. Callers that unpack the tuple fail immediately with a `ValueError` rather than silently receiving a changed type. `_encode_reference_conditions` (private) likewise gains `reference_token_counts` as a fourth return element. `prepare_reference_latents` keeps its documented 4-tuple unchanged. Verification ------------ Correctness against the reference. Both the standard helper and the modular block were checked for exact tensor equality (`torch.equal`) against `build_attention_mask`, driven the way `ic_lora.py` drives it -- one `ConditioningItemAttentionStrengthWrapper` per reference, so the reference builds iteratively with `existing_mask` carried forward. Five layouts, all matching, exercised through the same concatenate-then-split path the call site uses: - references only, 1 reference - keyframes + 1 reference (exercises fix 1) - keyframes + 2 references (exercises fixes 1 and 2) - 2 references of unequal length (exercises fix 2 and the token counts) - no keyframes, 3 references (exercises fix 2) Standard/modular parity on a real checkpoint, 512x512x17, fp32, via `integrations/ltx2_in_context_parity.py`. Video latents (1, 128, 3, 16, 16), audio latents (1, 8, 18, 16); gates are mean-rel 2e-2 and max-abs 1.5e-1: --reference 1.0 9 --condition 2 0.5 9 --conditioning_attention_strength 0.5 video: max abs 4.822e-05, mean rel 7.851e-06 audio: max abs 1.875e-04, mean rel 1.773e-05 --reference 1.0 9 --reference 0.8 9 --conditioning_attention_strength 0.5 video: max abs 1.953e-04, mean rel 9.101e-05 audio: max abs 9.251e-05, mean rel 1.026e-05 The first case is the layout where `num_prefix_tokens != num_noisy_tokens`; a wrong prefix/base split would have placed the group blocks at the wrong offsets with valid shapes, surfacing as a parity mismatch. The returned video shape also unpacks from exactly 768 base tokens (3 x 16 x 16), confirming `base_token_count` is the quantity now fed as `num_noisy_tokens`. Scope of the parity runs: both sides share this construction, so parity shows they did not desynchronize, not that either matches the reference -- that rests on the equality test above. Both references in the second case are 9 frames and so encode to equal token counts, meaning the equal-split and per-count paths coincide; the token-count change is not discriminated by these runs. A reference pair of differing lengths (e.g. `--reference 1.0 9 --reference 0.8 17`, giving 512 vs 768 tokens) would separate them, and would have failed parity before this commit since the modular block already used per-reference counts while the standard pipeline split evenly. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Add LTX25AutoBlocks and LTX25ModularPipeline Following the review: LTX-2.5 is a new checkpoint, so it gets its own blockset and pipeline rather than a config flag inside the LTX-2 ones. LTX25AutoBlocks subclasses LTX2AutoBlocks and overrides exactly one entry, `video_decode`, to LTX2DiffusionVaeDecoderStep. Every other leaf block is reused, so the two blocksets have identical block_names and cannot drift apart. That makes the diffusion decoder the default for LTX-2.5, which matches what the checkpoint ships natively, without the user having to pass a decoder at all. Going the other way stays a one-line swap rather than a second official blockset: blocks.sub_blocks["video_decode"] = LTX2VaeDecoderStep() LTX25ModularPipeline carries `default_blocks_name = "LTX25AutoBlocks"` and a new "ltx2.5" key in MODULAR_PIPELINE_MAPPING, so a repo shipping `modular_model_index.json` routes straight to the new blockset. The condition and in-context workflows belong in this blockset too, but those blocks live in the modular condition/IC PR, so folding them in belongs there rather than as a competing change here. * Drop the encoder half in the diffusion decoder converter The converter still assembled the whole VAE checkpoint, encoder included, and handed it to a strict `load_state_dict`. That was correct while the class was an autoencoder; after it became decoder-only the ~84 `encoder.*` keys have nowhere to land and the load raises. Nothing caught it because no test converts anything and every other check loads pre-converted weights, where unexpected keys are only warned about. Now the encoder is discarded rather than remapped, keeping the `decoder.` half and the per-channel statistics that become `latents_mean` / `latents_std`. `strict=True` stays, since it is what guards against a botched rename rule. Verified by actually running it: pulled the 395 `vae.*` tensors out of the native rc4 checkpoint by byte range, converted, and diffed against the already-published folder. 407/407 parameters, max abs difference 0.000e+00, so the fix changes nothing about the weights it produces. Also added a regression test, since the absence of any converter coverage is why this slipped. * Share the neighborhood attention mask and parameterise the AdaLN chunk count Two of the resolved review threads. Both are internal to the decoder modules and are unaffected by the pipeline-vs-model question still open on the thread above them. Build the FlexAttention mask once per stage instead of once per block. The mask depends only on the grid and the kernel, both fixed within a stage, so 24 builds collapse to 5; measured at 768x512x121 that is 582ms to 132ms. build_block_mask returns None for NATTEN, which encodes the window in its kernel. That is not merely an optimisation: create_block_mask materialises O(N^2) uncompiled, so a 69x64x96 stage would ask for 167 GiB and stage 5 for 8.2 TiB, and the mask must never be built for a path that would not read it. Make the AdaLN chunk count a constructor argument, num_chunks on the shared AdaLN and num_mod_params on the diffusion block, with the decoder passing shared_adaln.num_chunks through so the two cannot drift. num_mod_params is already the name LTX2AdaLayerNormSingle uses. The block's forward still destructures exactly seven chunks and reads four of them, which is the reference's shape. Decoding the test fixture is bitwise identical to before (max abs diff 0.000e+00). 39 passed, 4 skipped. * Export the LTX-2.5 modular classes from the top-level namespace `ModularPipeline.from_pretrained` resolves a repo's `_class_name` with `getattr(importlib.import_module("diffusers"), class_name)`, so a class that is only reachable from `diffusers.modular_pipelines.ltx2` cannot be selected by a checkpoint. `LTX25ModularPipeline` and `LTX25AutoBlocks` were exported there but not from `diffusers` itself, which left the whole checkpoint-to-blockset mapping dead: loading a repo whose index names them raised `AttributeError: module diffusers has no attribute LTX25ModularPipeline`. Exported alongside `LTX2ModularPipeline` / `LTX2AutoBlocks`, which is how the loader already finds those. Verified against a real repo: `ModularPipeline.from_pretrained(...)` now resolves LTX25ModularPipeline -> LTX25AutoBlocks -> LTX2DiffusionVaeDecoderStep, with the decoder supplied by the repo index and nothing passed by hand. Found by loading an actual checkpoint rather than constructing the blocks directly; calling `LTX25AutoBlocks().init_pipeline(repo)` goes blocks-to-pipeline and never exercises the lookup. * Name the diffusion decoder LTX-2.5 and save it to diffusion_decoder/ Two naming fixes, one of them load-bearing. The converter wrote the decoder to a `vae_diffusion/` subfolder, a leftover from when the class was an autoencoder. The folder name matters now: `DiffusionPipeline.from_pretrained` resolves each component from the subfolder named after it, so `LTX2VideoDiffusionDecodePipeline` (component name `diffusion_decoder`, same as the modular `ComponentSpec`) can never load itself out of a repo whose folder is called `vae_diffusion`. The converter now saves to `diffusion_decoder/`, matching the layout the LTX-2.5 weight repos already use, and the docs example loads from that name. The docstrings also called it the LTX-2.4 decoder. The diffusion decoder ships with LTX-2.5; the converter handles those checkpoints under the `--version 2.4` architecture family, which the version guard's message now spells out instead of misattributing the feature. Verified by re-running the converter on the rc4 distilled checkpoint after the change: the output lands in `diffusion_decoder/` and is bit-identical (sha256) to the pre-change conversion, so this renames and redocuments without touching what is produced. The converter regression test passes. * Move the LTX-2 text-input expansion into the core denoise step `LTX2TextInputStep` expanded the prompt embeddings to the effective batch size inside `LTX2TextConditioningStep`, so running the text-conditioning stage on its own baked `num_videos_per_prompt` into its outputs. Feeding them to `LTX2CoreDenoiseStep` with a different `num_videos_per_prompt` meant re-running the text stage or reshaping its outputs by hand. `LTX2TextInputStep` now expands the connector conditioning (the tensors the denoiser actually reads through `guider_input_fields`) instead of the raw Gemma embeds, and runs as the first block of both core denoise steps. The raw `prompt_embeds` expansion is dropped -- it only existed to feed the connector at the expanded batch. `batch_size` / `dtype` move up to `LTX2TextEncoderStep`, since `LTX2DurationStep` needs `batch_size` before the denoise stage. Two behavior notes: - The connectors now run at one row per prompt rather than `num_videos_per_prompt` rows. They are applied per batch element, so this is mathematically identical (verified to ~3.6e-7, float32 rounding) and N times cheaper, but not bitwise identical to `LTX2Pipeline` when `num_videos_per_prompt > 1`. - Expansion uses `repeat_interleave` for the masks as well as the embeds. The previous code (inherited from the standard pipeline) interleaved the embeds but tiled the masks, which misaligned them whenever both `batch_size > 1` and `num_videos_per_prompt > 1`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Describe the connector inputs on LTX2TextInputStep `LTX2TextInputStep` declares the connector conditioning tensors and now runs first in the denoise sequence, so its declarations -- not `LTX2LoopDenoiser`'s -- are the ones the docstring generator picks up. They carried no `description`, which left `TODO: Add description.` placeholders on the three core denoise blocks. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Build the condition and in-context workflows from LTX2AutoBlocks `LTX2AutoBlocks` covered text-to-video and image-to-video only; the condition and in-context (IC-LoRA) workflows were reachable solely through their own blocksets. Fold them into the auto blockset so all four are selected from the inputs, the way t2v and i2v already were. - Merge `LTX2ConditionAutoPromptEnhancerStep` into `LTX2AutoPromptEnhancerStep`, which now picks the condition, image or text enhancer off `conditions` / `image`. - Add `LTX2AutoConditionEncoderStep`, triggered by `conditions` or `reference_conditions` so a reference-only in-context run still gets the empty per-condition lists `LTX2InContextPrepareLatentsStep` requires. - Extend `LTX2AutoCoreDenoiseStep` to four branches and add `LTX2AutoDecoderStep`, which trims the appended keyframe/reference tokens whenever `base_token_count` is present. Two fixes this exposed: - `noise_scale` means an interpolation factor (t2v/i2v, 0.0) in one workflow and an initial noise level (condition/in-context, `None` -> `sigmas[0]` or 1.0) in another. A blockset keeps the first non-`None` default across its blocks, so sharing one blockset published 0.0 and silently zeroed the noise for condition workflows that omitted the argument. `LTX2PrepareLatentsStep` and `LTX2PrepareAudioLatentsStep` now declare `None` and resolve to 0.0 in the block, mirroring what the condition steps already do. - `LTX2AutoBuildVideoSelfAttentionMaskStep.select_block` compared `conditioning_attention_strength` against 1.0, but `get_execution_blocks` passes unset trigger inputs as an explicit `None`, so resolving any in-context workflow raised `TypeError`. The parity harnesses passed `noise_scale=None` explicitly, which took the argument-present path and hid the shadowing bug; `--omit_noise_scale` drops the key entirely so the default is exercised. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Run the text-input step in the condition and in-context denoise steps `LTX2TextInputStep` is the only block that expands the connector text conditioning by `num_videos_per_prompt`, and it was missing from `LTX2ConditionCoreDenoiseStep` and `LTX2InContextCoreDenoiseStep`. The prepare-latents/audio/coords steps all build their tensors at `batch_size * num_videos_per_prompt`, so the video latents arrived at the transformer at batch B*N while `encoder_hidden_states` was still at B — a shape error in cross-attention for any `num_videos_per_prompt > 1`. The step goes at the head of both sequences, before prepare-latents, so the prepare-latents-before-set-timesteps ordering these workflows need (`mu` is read off the full packed sequence length) is untouched. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Broadcast keyframe and reference conditioning across the generation batch Frame conditions are VAE-encoded at batch 1, but `latents`, `conditioning_mask` and `clean_latents` are built at `batch_size * num_videos_per_prompt`. Appending a keyframe condition (latent index > 0) therefore failed in `prepare_latents` with "Sizes of tensors must match except in dimension 1" whenever the batch was greater than 1. First-frame conditions were unaffected because they are applied by indexed assignment, which broadcasts. Expand the keyframe tokens, mask and coords to `batch_size` before the concatenation, and do the same for the IC-LoRA reference coords, which feed the same `appended_coords` block that is later concatenated onto batch-`batch_size` `video_coords`. This is what the modular LTX-2 blocks already do, and matches the `[B, 3, ...]` shape both docstrings promise. Also picks up a pending `make style` docstring rewrap in the same file. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat: use natten for processing. (#15) * feat: use natten for processing. * address dhruv's comments. * remove unneeded test * remove integration file * make vae properties accessible without loading it. (#17) * Address review: replace LTX2 mixins with Copied-from pipeline methods. Inline check_inputs/enhance_prompt on LTX2Pipeline and sync I2V/Condition/InContext via # Copied from, and update docs for shared 2.5 defaults. (cherry picked from commit c183923cdd9c951f9694f9f6cd903714330cff03) * Point the modular enhance-prompt comment at LTX2Pipeline `LTX2PromptEnhancementMixin` was inlined into the pipelines in c183923cd, so the mirrored-from reference no longer names an existing class. * Remove the integrations/ dev scripts The LTX-2 parity harnesses and the tiny-checkpoint builder were local development scaffolding and are not part of the shipped library. * moular tests (#19) * Keep LTX-2.5.1 keyframes_abs_pos_embedding through convert. (#20) Distilled 2.5.1 monoliths ship this (1, inner_dim) buffer with use_keyframes_abs_pos_embedding=True; store it on the transformer for load/save without wiring it into the regular distilled forward yet. * remove test_ltx2_keyframes_abs_pos_embedding_load_and_forward * Update the LTX-2.5 docs: two stage generation, distilled default, repo id (#21) * feat: implement tiling for the diffusion decoder. (#22) * add a note on tiling (#23) * fix documentation. --------- Co-authored-by: aartyomov <aartyomov@lightricks.com> Co-authored-by: Daniel Gu <daniel.gu@huggingface.co> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: linoytsaban <linoy.tsaban@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: DN6 <dhruv.nair@gmail.com> Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com> |
||
|
apolinário
|
f53d552036 |
Add MiniMax-H3 (#14355)
* Add MiniMax-H3 * Keep the MiniMax-H3 VAEs in float32 under torch_dtype casts * Make MiniMax-H3 modular only * Assert the VAE float32 pin as a positive contract * Point the modular tests at the hf-internal-testing tiny repo * Document performance recipes per hardware class * Slim hardware loading snippets and PR install note * Use pinnable int8 config and freeze quantized components * Scope load time quantization and disable low cpu mem usage under streamed offload * Minimax h3 follow up (review & refactor) (#14371) review & refactor * Regenerate the modular auto docstrings and restyle the H3 docstrings `make modular-autodoctrings` and `make quality` both failed on CI. The nine `# auto_docstring` blocks in `modular_blocks_minimax_h3.py` were stale, and the docstrings of eleven files had not been through `doc-builder style`. Nothing but docstrings and comments changes here. One trap worth recording: `utils/modular_auto_docstring.py` shells out to `ruff` and degrades silently when it is not on `PATH` ("Warning: tool not found ... Skipping formatting"), which makes it rewrite 42 unrelated pipelines instead of the one that is stale. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Fix the H3 fast tests against the refactored state contract `condition_latents` is the VAE encoder block's own output — a list of one latent per condition — and the after-denoise step unpacks `latents` and drops the conditioning rows, so neither is meaningful once a request has come back. The end-to-end tests no longer reach for them; a new standalone test runs `MiniMaxH3KeyframeVaeEncoderStep` on its own, which is where the list is real. Three other tests were failing for reasons of their own: - `test_check_inputs_references` was missing its `@parametrize` and errored at collection. Added four cases that hit real validation. - The duration ceiling had drifted between the two blocks that check it. `before_denoise` warned and reassigned before validating, so it reported `got 362` — a count the caller never passed, right after warning it had rounded their 346. It now validates first, like `before_encoder` already did, and both messages read `got 346 (rounded up to 362)`. - `height` without `width` raised `TypeError: unsupported operand type(s) for %: 'NoneType' and 'int'` on `t2va`, and died inside the resize on `fl2va`. Both blocks now raise the same error `MiniMaxH3Ref2VASetupStep` already raised. 58 passed -> 71 passed, 5 failed, no errors. What is left: `output_type="latent"` has no opt-out in the video decode block, `reference_image_short_edge` is a constructor argument so it does not survive save/reload, an image reference is PIL-only where the test expects three layouts, and `test_float16_inference`, which fails on the base branch too. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Stop claiming `num_inference_steps` is optional in the H3 tests The block declares it `InputParam.template("num_inference_steps", required=True)` and `expected_workflow_defaults` lists it under `required_inputs` for all three workflows, so `optional_params` said the opposite in the same file. It was inherited from the mixin's default set and never examined. Nothing enforces the requirement today: `_check_inputs` substitutes the declared default before it tests `required`, and the template carries `default=50`, so omitting the input runs 50 steps rather than raising. That is not specific to H3 — 14 call sites on main are in the same position — and it is tracked in #14388. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Fix the two-card recipe, and stop accepting `output_type="latent"` The two-card section did not run, on this branch or on the base one. A pipeline resolves a single execution device for all of its components, so the documented shape — one pipeline whose conditioner is `device_map`ped to the second card and whose denoiser is moved to the first — built the rotary positions on `cuda:1` and handed them to a transformer on `cuda:0`: transformer_minimax_h3.py, in Rope.forward RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! Reproduced against `abc5e9bf7` too, so it is not something the refactor broke; the section went in as a plausible recipe that was never executed. The split is therefore between pipelines rather than inside one: pop the text encoder block into a conditioner of its own, give each pipeline a `ComponentsManager` pinned to a card, and pass the state from one call to the other. The managers' hooks are what align `prompt_embeds` onto the denoiser's card; placing the components by hand instead works too, but then that one tensor has to be moved explicitly. Verified end to end against the tiny fixture. `output_type="latent"` is separate. It is not an output format — a request that wants latents runs a pipeline without the decode blocks — and it used to run the whole VAE decode before dying inside `postprocess_video` with "latent does not exist", which does not say what to do instead. The video decode block now rejects anything it cannot postprocess, before decoding. `test_output_type` covers the three formats it does accept, and the rejection is a `test_check_inputs` case. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Skip the workflow tests whose tiny repositories cannot serve them `test_from_pretrained_workflow`, `test_load_components_workflow` and `test_unload_components` build a pipeline from `pretrained_model_name_or_path` and compare it against the test class's own blocks. Three fixtures cannot answer that, in three different ways, and none of them is about the pipelines' code: - `tiny-anima-modular-pipe` has no `modular_model_index.json`. Anima assembles its dummy components in `get_pipeline` and never loads from a repository, so the path it declares had never been exercised. - `tiny-flux2-klein-modular` names `Flux2KleinBaseAutoBlocks` while its `_class_name` and `is_distilled` both say distilled. `from_pretrained` honours `_blocks_class_name`, so it builds base blocks, which declare a `guider` the distilled ones do not. - `tiny-qwenimage-edit-modular` names `QwenImageModularPipeline`, so `from_pretrained` falls back to `QwenImageAutoBlocks`, which declares no `image_conditioned` workflow. Skipped with a TODO each, since the fixes are Hub-side. The klein and qwenimage ones are a single key in `modular_model_index.json`; Anima needs a repository published. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Declare `reference_image_short_edge` as pipeline config `MiniMaxH3Ref2VASetupStep` resolves three pieces of released-checkpoint geometry — the canvas short edge, the canvas area cap, and the separate short edge an image reference is encoded at — and declared only the first two as `ConfigSpec`s. The third was a constructor argument, so it did not survive a save and reload: the reloaded pipeline rebuilt the block from the repository, took the 2048 default instead of the 64 the tests configure, resized reference images at a different resolution, and landed 0.398 away from the original output. All three now sit together, which is also what lets the test fixture stop reaching into the block tree to swap in a configured copy — the shrunken geometry is one `update_components` call next to the canvas rule it belongs with, and `test_reference_image_geometry` exercises the released default rather than constructing the block with an explicit 2048. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Accept a reference image in the three layouts a video reference takes `MiniMaxH3VideoReference.frames` takes a list of images, a channels-last array or a channels-first tensor; `MiniMaxH3ImageReference.image` took a PIL image only, and the two array layouts failed on `entry.image.size` — an element count on numpy, a bound method on torch — with `TypeError: 'int' object is not subscriptable`. Passing an array through as-is would have been worse than the crash. The resize would silently drop to `F.interpolate`'s nearest-neighbour where a PIL image gets LANCZOS, and the VAE encode does `np.array(image).permute(2, 0, 1)`, which fixes the axes of an image and scrambles those of a channels-first tensor. So the layouts are normalized onto a PIL image before any of that, through the processor component the block already declares: `pt_to_numpy` for the channel axis and `numpy_to_pil` for the array. Neither converts dtype, and `numpy_to_pil` scales by 255, so `uint8` is normalized on the original object first — after `pt_to_numpy` a `uint8` tensor is floats over `[0, 255]` and a later dtype check would miss it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Read the model's constants off the pipeline rather than the module `before_denoise.py` imported `MINIMAX_H3_AUDIO_CHANNELS`, `MINIMAX_H3_AUDIO_TAG` and `MINIMAX_H3_VIDEO_TAG` and read them inside its layout builders, even though two of the three were already exposed as pipeline properties and used that way everywhere else. They are arguments now, passed from `components.*` at the two `__call__` sites; `audio_tag` needed the property `text_tag` and `video_tag` already had, which is presumably why it was the one still reaching for the global. Same treatment where a helper cannot see a pipeline: `resolve_canvas_size` and `audio_latent_num_frames` take the constants as default arguments instead of closing over them, and `_normalize_video_condition` takes the rate it resamples onto, so `before_encoder.py` imports none of them at all. What is left of the constants is their definitions, the properties that expose them, and default argument values. Two helpers with a single caller each are inlined at it: the two rotary-time sums, which are a parity contract — the reference sums the same series pairwise in one place and sequentially in the other, and the orders differ in the last ulp from 16 latent frames onwards. Both inlined sums are bit-identical to the functions they replace across 1..399 latent frames, and the two still disagree on 373 of those counts. `_frame_position_grid` stays, and `build_packed_sequence` now calls it rather than repeating its five lines. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Reject a reference video too short for the conditioner to read The conditioner reads a reference video at 2 fps and Qwen3-VL merges the sampled frames in groups of two, so anything under 13 frames at 24 fps samples down to a single frame and dies inside transformers: image_processing_glm4v.py ValueError: t:1 must be larger than temporal_factor:2 which names neither the reference nor the rate that produced it, and only fires on the versions carrying that check — it passed locally and failed on CI. `_sample_video_condition_frames` now rejects it where the sampling happens, with the bound derived from the rates rather than hardcoded: "must run at least 13 frames at 24 fps (0.54 seconds), got 4". `test_reference_media_layouts` was building exactly such a reference, four frames, which is what surfaced this. It uses a second of video now, like `test_reference_combinations` already did. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Drop the padded-layout attention mask `forward` built a boolean attention mask whenever a row carried a negative modality tag, reproducing the reference implementation's `cu_seqlens = [0, used, S]` split of the padding tail it adds for FlashAttention. Nothing in this port ever produces such a row: both packers partition `[0, sequence_length)` and write only the text, audio and video tags, so the branch was unreachable and the mask was never built. The `clamp(min=0)` that kept a `-1` from indexing the AdaLN table backwards goes with it. Removing it is not only dead-code removal. `if bool(is_pad.any())` is a data-dependent branch, which has no fullgraph representation, and it was the reason whole-model `fullgraph=True` compilation and `torch.export` were skipped: 4 passed, 1 skipped tests -k "compile or export or aot" Both now run. The model tests go from 40 passed / 4 skipped to 41 passed / 2 skipped, against the same 8 pre-existing memory-offload failures. `attention_mask` stays on the processor, attention and block signatures — it is the signature every other processor in diffusers has, and a custom one may want it — but the model itself never passes anything but `None`, so attention runs unmasked over the one document and every backend stays available. Two of the four documentation examples have been re-run against this and come back bit-identical, frames and audio: `t2va` at the full 768x1344 canvas and `fl2va`. The two `ref2va` ones are still to run. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Apply suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> |
||
|
Sayak Paul
|
7156122b1d |
[ci] ignore the token in vq_diffusion converter. (#14275)
* ignore the token in vq_diffusion converter. * change to regular official download link |
||
Atharva Joshi
|
db44fe6638 |
Cosmos3 edge support (#14181)
* Add Cosmos3 Edge modular pipeline support * Remove Cosmos3 conversion script * Refine Cosmos3 Edge norm selection * Use internal Cosmos3 modular test fixture * Use real VAE in Cosmos3 pipeline tests * Align Cosmos3 Edge tests with conventions * Fix Cosmos3 callbacks on Python 3.10 * Address Cosmos3 modular test feedback * Address Cosmos3 pipeline test review comments |
||
tangyanf
|
408ce8db29 |
feat: add JoyImage edit plus (#14032)
* feat: add image edit plus
* refactor: remove debug code
* fix: address review issues for JoyImage Edit Plus
- Remove einops dependency: replace rearrange with reshape/permute
- Remove sglang-specific code from transformer forward
- Remove unused import inspect from transformer
- Fix hardcoded device_type="cuda" to use device.type
- Simplify scheduler sigma math: delegate to retrieve_timesteps
- Remove unused enable_denormalization parameter
- Fix callback latents variable binding
- Fix output_type="pt" to return stacked tensor
- Set return_dict default to True in transformer forward
- Add dummy objects for JoyImageEditPlus classes
- Add transformer and pipeline test files
* fix: add missing newline at end of pipeline_output.py
* fix: add missing newline at end of pipeline_output.py
* doc: add joyimage-edit-plus doc
* refactor: update code format
* refactor: update code formate
* refactor: merge edit-plus conversion script into convert_joyimage_edit_to_diffusers.py
JoyImage Edit and Edit Plus share identical VAE and transformer weight
layouts — only the target model class differs. Consolidate both into a
single script with a --model_type flag (edit | edit_plus) instead of
maintaining two nearly-duplicate files.
* revert: remove unnecessary attention_mask change from transformer_joyimage.py
* Apply style fixes
* fix: remove stale Copied-from annotations that diverged from source
* fix: address CI check failures for edit-plus PR
- Add missing Returns section to JoyImageEditPlusTransformer3DModel.forward docstring
- Fix alphabetical ordering of dummy classes in dummy_pt_objects.py
* fix: use VAE dtype instead of float32 cast in edit-plus pipeline
Replace `torch.autocast(dtype=torch.float32)` + `.float()` with
`.to(self.vae.dtype)` for both VAE encode and decode calls.
The previous approach caused dtype mismatch (float32 input vs bfloat16
bias) on CPU where autocast does not automatically cast conv weights,
breaking the CI `test_layerwise_casting_inference` test.
* refactor: address dg845 review comments on edit-plus
Transformer:
- Simplify RoPE by removing dead txt_freqs path
- Fix patch_size type hint to list[int]
- Use self.config.xxx instead of explicit attributes
Pipeline:
- Merge VaeImageProcessor into JoyImageEditImageProcessor
- Inline normalize/denormalize latents at call sites
- Use image_processor.preprocess() for reference images
- Remove duplicate resize_center_crop in prepare_latents
- Rename padded_latents to latents in __call__
- Support pre-computed latents parameter
- Remove redundant padding token zeroing (attention mask suffices)
- Use image_processor.postprocess() for output conversion
- Return (image,) tuple when return_dict=False
* fix: guard against division-by-zero in CFG norm rescaling
Add .clamp_min(1e-6) to noise_norm to prevent NaN when comb_pred
is all-zeros (e.g. padded patches), matching pipeline_joyimage_edit.py.
* style: run ruff format on edit-plus pipeline
* refactor: address dg845 second-round review comments
- Fix vae_image_processor -> image_processor in docs
- Remove empty __init__ from AttnProcessor
- Restore # Copied from on retrieve_timesteps and sync with source
- Add check_inputs method for basic input validation
- Remove unnecessary test skip for test_from_save_pretrained_dtype_inference
* Apply style fixes
---------
Co-authored-by: tangyanfei.8 <tangyanfei.8@jd.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
|
||
|
Sayak Paul
|
4f37263ab7 |
[chore] update to 2026 finally. (#14079)
update to 2026 finally. |
||
David Bertoin
|
a487e2f131 |
Add PRXPixelPipeline: pixel-space PRX text-to-image pipeline (#13928)
Adds a pixel-space variant of PRX that denoises raw RGB directly (no VAE), conditioned on a Qwen3-VL text encoder: - PRXTransformer2DModel: new optional config args `bottleneck_size` (two-layer img_in projection for large patch dims) and `resolution_embeds` (PRXResolutionEmbedder conditions the timestep modulation on the latent resolution) - PRXPipeline: support for subclass-tuned tokenizer max length, light text cleaning, x-prediction flow matching (x0 -> velocity conversion), and non-unit initial noise scale - PRXPixelPipeline: thin subclass wiring the above together (vae optional/None, vae_scale_factor=1, 1024px default) - conversion script support for the pixel checkpoint format - registration in __init__ files + dummy objects, docs autodoc entry, fast pipeline tests Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
MaciejBalaNV
|
212a98d4f2 |
Add Sound Encoder to Cosmos3 (#13911)
* Initial version of sound encoder Signed-off-by: Maciej Bala <mbala@nvidia.com> * Answered review comments Signed-off-by: Maciej Bala <mbala@nvidia.com> * Reshaped encoder Snake1d Signed-off-by: Maciej Bala <mbala@nvidia.com> * Updated test Signed-off-by: Maciej Bala <mbala@nvidia.com> * Fixed test Signed-off-by: Maciej Bala <mbala@nvidia.com> * Apply style fixes * Added docstring Signed-off-by: Maciej Bala <mbala@nvidia.com> --------- Signed-off-by: Maciej Bala <mbala@nvidia.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
rmatif
|
c943837899 |
Add Anima modular pipeline (#13732)
* Add Anima pipeline * Fix empty Anima negative prompts * Fix Anima registration * Clean up Anima conditioner * Refactor Anima to modular * Use modular loader in Anima docs * Move Anima text conditioner * Apply style fixes * Address Anima review nits * Fix Anima rotary autocast * Clean up Anima LoRA checks * Update Anima model id * Document Anima LoRA loader * Document Anima conditioner forward --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
Guian Fang
|
d94ad81b87 |
[AnyFlow] FAR: standalone causal-mask builder + torch.compile follow-up (#13792)
* [AnyFlow] FAR: standalone causal-mask builder + torch.compile follow-up
Follow-up to #13745. Extracts FAR mask construction to a module-level
helper and adds an `attention_mask` forward kwarg so
AnyFlowFARTransformer3DModel can be wrapped in
`torch.compile(fullgraph=True)`. The pipeline pre-builds the mask during
KV-cache prefill so users get end-to-end fullgraph compile.
* Public method `AnyFlowFARTransformer3DModel.build_attention_mask(...)`
(modes: "train", "cache") plus private module-level helper
`_build_anyflow_far_causal_block_mask(...)`.
* `_build_freqs` cache lookup/write bypassed under
`torch.compiler.is_compiling()` to avoid a Dynamo guard recompile on
the second compiled call (applied in bidi source; synced to FAR via
`# Copied from`).
* `TestAnyFlowFARTransformer3DCompile(TorchCompileTesterMixin)` —
recompilation_and_graph_break, repeated_blocks, and group_offloading
pass on H200; AOT is `@pytest.mark.skip`'d (torch.export rejects
BlockMask as a pytree input).
* Base `get_dummy_inputs` omits `attention_mask` so every non-compile
test class exercises the in-forward fallback; the compile class
overrides to inject a pre-built mask.
* Bit-exact: pre-built path vs internal-build fallback max|Δ|=0.0e+00.
* [AnyFlow] docs: full author list, repo demo examples, slimmer pipeline page
* Full author list and NVIDIA → NUS → MIT institution order; TL;DR +
abstract + Available Models bullets.
* Rewritten pipeline-selection tip describing both pipelines symmetrically.
* T2V / I2V / V2V examples now use the canonical 81-frame setup and the
demo prompts / conditioning assets shipped under
`NVlabs/AnyFlow/assets/evaluation/` (linked via raw.githubusercontent.com).
* Drop the inline "Optimizing Memory" and "torch.compile" sections — those
notes will live in the NVlabs/AnyFlow repo's own performance guide rather
than the diffusers pipeline reference.
* Sync zh user guide and the two model-API stubs.
* [AnyFlow] FAR: move chunk_partition default into transformer config
- AnyFlowFARTransformer3DModel.__init__ now accepts chunk_partition via
@register_to_config (default (1, 3, 3, 3, 3, 3, 3, 2) for the released
81-frame checkpoints, matching the field on Hub).
- AnyFlowFARPipeline.__call__ no longer requires chunk_partition; defaults
to self.transformer.config.chunk_partition. Per-call override still
supported for V2V / non-default num_frames.
- Drop the AnyFlowFARPipeline.default_chunk_partition class attribute.
- Update docs (en pipelines/models, zh using-diffusers) and the conversion
script to match.
* [AnyFlow] FAR pipeline: fix `timesteps` shadowing across chunks
Inside the per-chunk rollout loop, the local variable `timesteps` was
reassigned to `self.scheduler.timesteps` after `set_timesteps()`. On the
next chunk iteration the same name was passed back into
`set_timesteps(timesteps=...)`, where a non-None value enters the
*custom-schedule* branch — `apply_shift` re-runs on already-shifted
values, double-shifting the schedule for every chunk after the first.
Concretely, with `shift=5` and `num_inference_steps=4`:
- chunk 0 timesteps: [1000, 937.5, 833.3, 625] (correct)
- chunk 1+ timesteps: [1000, 986.8, 961.3, 892.9] (double-shifted)
The later steps drift toward `t=1000` instead of toward `t=0`, the
flow-map model is conditioned on the wrong source sigma, and the chunk
KV cache accumulates errors that show up as artifacts in later video
frames.
Fix: rebind the cached schedule to a fresh local name
(`scheduler_timesteps`) so the outer-scope `timesteps` kwarg (the
user-provided custom schedule, when any) stays untouched across chunks.
Layer-by-layer verification against the NVlabs reference implementation
on H200 (elephant prompt, seed 0, 4 NFE, 81 frames):
- chunk 0 inference: bit-exact (0.0 mean diff)
- chunk 1 step 0: 0.194 → 0.014 (-93%)
- chunk 7 last step: 0.564 → 0.274 (-51%)
* [AnyFlow] FAR: doc-builder line wrap for chunk_partition docstrings
Pure rewrap to satisfy `doc-builder style --max_len 119`. Two docstrings
introduced in
|
||
|
Atharva Joshi
|
a1c7df4801 |
Adding Cosmos 3 to Diffusers (#13818)
* Adding Cosmos 3 * removed dead code * Change customer TimeEmbedding Layer to DIffusers Time Embedding * removed dependency on hugging face transformers * refactor 1 * Fixed Attention Pattern * Removed from Pretrain overrides * Removing normalization from the audio Tokenizer * fixed diffusers checkpoint * fixed video save uint conversion * added forward hook for cpu offload case * removed dead params for sound tokenizer * renaming audio encoder for readability * ruff format * Fix checkpoint conversion script for sound tokenizer * Audio Decoder trim and removing some dead code * removing dead sequence packing code * refactor pipeline to diffusers style formatting * removing use of cosmos3 audio encoder * Revert "removing use of cosmos3 audio encoder" This reverts commit 1b8b99a95e5acc749b5c100637ed382de138c606. * refactor audio encoder * inline remaining sequence packing functions and lint * Removed GenerationDataClean class and Action logic * inlined default args * removed dead code and refactoring * drop pipeline-helper @no_grad, inline derive helper, move guidance check to check_inputs Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop unused list bookkeeping from PackedSequence attn_modes was never read; sample_lens collapses into the existing sequence_length int (we only pack a single sample at a time); split_lens collapses into a single und_len int (only split_lens[0] was ever read). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * extract pack-time build state from PackedSequence dataclass curr, _use_mrope, _mrope_temporal_offset, _mrope_reset_spatial were transient counters used while building the joint sequence, not part of the finalized output. Thread them through _pack_*_tokens as positional args/returns so the dataclass only carries fields the pipeline actually reads back. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop private-API isinstance/shape asserts These were build-mode guards on PackedSequence internals and shape checks on tensors the pipeline itself constructs, both flagged as noise in private code per reviewer. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * build PackedSequence tensors on target device, drop to_cuda Thread device through pack_input_sequence and _pack_*_tokens helpers so all torch.tensor/zeros/arange calls land on the target device directly. Move CPU-side mRoPE tensors over with .to(device) at the append site. Pass device to finalize so list-to-tensor conversion lands on device too. Delete PackedSequence.to_cuda() and the helper _modality_to_cuda; drop the corresponding call sites in __call__. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop SequencePlan, skip_text_tokens, and bos_token_id branch SequencePlan's has_text and has_vision were True at every construction site and has_sound was derivable from x0_tokens_sound is not None. condition_frame_indexes_vision is now passed directly as a List[int] arg to pack_input_sequence. Removed the skip_text_tokens flag (never True) and the dead bos_token_id shift branch in _pack_text_tokens (special_tokens never carries one in this pipeline). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * use retrieve_latents helper in _encode_video Copied from stable_diffusion_img2img matching cosmos2_5 convention so make fix-copies keeps it synced. Functionally identical to the prior .latent_dist.mode() call but handles latent_dist/latents attribute variants uniformly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop get_data_and_condition + data_batch dict scaffolding normalize_video_databatch_inplace, augment_image_dim_inplace and remove_padding_from_latent were no-ops once is_preprocessed=True (always set by this pipeline) and the pipeline never pads. get_data_and_condition just orchestrated those plus a never-taken multi-vision branch. Replaced the whole chain with a few lines inline in prepare_latents: build vision_tensor on device, call _encode_video, set fps_vision. prepare_latents no longer needs input_caption_key, input_video_key, input_image_key, or the prompt kwarg. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop _load_image_as_tensor; use VideoProcessor for conditioning frame Image loading is the caller's responsibility (load_image from diffusers.utils), matching the cosmos2_5 example. The pipeline registers a VideoProcessor in __init__ and calls preprocess() to resize + normalize caller-supplied PIL / np / tensor inputs to [1, 3, H, W] in [-1, 1]. prepare_latents fills the temporal dim in two lines (single frame at t=0, repeat-pad the rest) preserving the prior i2v behavior. Inference script updated to call load_image() before passing to the pipeline. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * move encode/decode helpers + transformer forward into Cosmos3OmniTransformer Pull encode_text / encode_vision / encode_sound_tokens / decode_vision / decode_sound_tokens (and their pure-tensor helpers patchify_and_pack_latents, unpatchify_and_unpack_latents, apply_timestep_embeds_to_noisy_tokens, _pack_sound_latents, _unpack_sound_latents) from the pipeline into Cosmos3OmniTransformer as methods. The transformer's forward(packed_seq) now runs the full per-step pass: encode text/vision/sound, rotary + layer loop, decode vision/sound — returns (preds_vision, preds_sound). The pipeline's CFG loop drops the encode_*/decode_* method calls and the manual und/gen split/concat; each pass is now a single self.transformer(packed_seq) call. No self.transformer.{embed_tokens, vae2llm, llm2vae, sound2llm, llm2sound, time_embedder, time_proj, sound_modality_embed} access remains in the pipeline. Cosmos3VLTextModel is kept as a structural wrapper for now — flattening it would break the published checkpoint layout. Tracked separately. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * remove Cosmos3VLTextModel; flatten transformer layout embed_tokens / layers / norm / norm_moe_gen / rotary_emb are now direct attributes of Cosmos3OmniTransformer. The converter strips the leading `model.` prefix from the source language_model state-dict so new conversions land at the flat layout natively. Published Hub artifact (nvidia/Cosmos3-Nano) needs its transformer safetensors + index.json re-keyed with the same prefix strip before this code can load it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * move save_img_or_video/save_wav to cosmos/export_utils.py Per reviewer guidance, custom video/audio export helpers belong in a pipeline-local export_utils.py (mirroring pipelines/ltx2/export_utils.py) rather than living inside the pipeline file. Pipeline imports trim the now-unused pathlib/numpy/export_to_video; inference example updated to import from the new location. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop @torch.no_grad on pipeline __call__ Diffusers pipeline convention: __call__ does not wear a torch.no_grad decorator; the responsibility for grad context sits with the caller. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * follow standard transformer conventions in Cosmos3OmniTransformer - Declare _no_split_modules, _repeated_blocks, _skip_layerwise_casting_patterns, _keep_in_fp32_modules, _supports_gradient_checkpointing on the transformer. - Wire self.gradient_checkpointing + the _gradient_checkpointing_func branch in forward so the flag is honest (models.md gotcha #3). - Add PeftAdapterMixin and AttentionMixin to the mixin set so LoRA loading and the attention-backend setters work. - CosmosAttnProcessor3_0 now declares _attention_backend / _parallel_config and forwards them to dispatch_attention_fn, matching the pattern in models.md and transformer_wan.py. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * restore @torch.no_grad() on Cosmos3OmniDiffusersPipeline.__call__ Reverting bea4eecf9 — removing the decorator causes GPU OOM during inference because the autograd graph accumulates across the full denoising loop (35 steps × dual cond/uncond passes × full transformer). pipelines.md gotcha #2 documents this exact failure mode and the convention is upheld by every sibling pipeline (pipeline_flux.py:652, pipeline_qwenimage.py:462, pipeline_wan.py:381, pipeline_ltx.py:535). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * drop training-only dummy projection paths in transformer decode helpers _decode_vision and _decode_sound both guarded a "no noisy tokens" branch that ran zeroed projections to keep the autograd graph intact. Those branches only fire when a pure-conditioning step has no MSE-loss tokens, which never happens in the inference pipeline — every workflow has at least one noisy vision frame, and _decode_sound is gated on has_sound which itself requires noisy sound tokens. Deleting per CLAUDE.md: "delete training-time code paths… only keep the inference path." Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * inline modality encode/decode helpers into transformer forward models.md "Coding style": all layer calls should be visible directly in forward — avoid helper functions that hide nn.Module calls. Inlines _encode_text / _encode_vision / _encode_sound / _decode_vision / _decode_sound into forward so embed_tokens, vae2llm, sound2llm, llm2vae, llm2sound, and time_embedder are all visible at the call site. Pure- tensor helpers (_patchify_and_pack_latents, _unpatchify_and_unpack_latents, _pack_sound_latents, _unpack_sound_latents, _apply_timestep_embeds_to_noisy_tokens) stay as methods since they don't hide layer state. Also drops the inference-unreachable guards while collapsing the helpers: the "vision is None" / "sound is None" / "mse_loss_indexes.numel() > 0" branches never fire because the pipeline always packs vision, only routes to the sound branch when sound is present, and condition_frame_indexes never covers the entire stream. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * trim dead flags and idioms in Cosmos3 pipeline __call__ - Drop assert config.use_moe — use_moe is never read by the model and asserts vanish under python -O; if !use_moe is unsupported the place to surface it is check_inputs, not a stripped assert. - Delete the joint_attn_implementation == "flex" path entirely (the include_end_of_generation_token branch in pack_input_sequence, the include_eog hoist, and both call-site kwargs). The published config is "two_way"; the flex branch and the end_of_generation special token were dead under every shipped checkpoint. - Drop torch._inductor.cudagraph_mark_step_begin() from the step loop — cudagraph stepping belongs in the caller's torch.compile wrapper, not the inference pipeline. - Replace four int(torch.prod(torch.tensor(shape))) idioms with math.prod(shape) — no tensor allocation, no .item() sync, and math is already imported. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * introduce Cosmos3Condition for prepare_latents return shape Replace the 7-tuple return from prepare_latents with a Cosmos3Condition dataclass that carries the encoded conditioning latents (vision + optional sound), their fps tensors, the conditioning frame indices, and the num_vision_items count. The denoising loop and _postprocess_latents now read these as named attributes instead of positional tuple unpacking. Addresses reviewer thread huggingface/diffusers-new-model-addition-cosmos#1 comment 3278569263 ("create something like Cosmos3Condition class for condition input"). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * restructure pack helpers as data-returning pipeline methods Replace the four module-level pack functions (_pack_text_tokens, _pack_vision_tokens, _pack_sound_tokens, pack_input_sequence) with methods on Cosmos3OmniDiffusersPipeline. The three per-segment methods now build and return their own data (text_ids/mrope_ids tuple for text, a populated ModalityData for vision/sound) instead of mutating a shared PackedSequence builder; pack_input_sequence orchestrates them. Other cleanups along the way: - Drop dead branches that the published config never exercises: use_mrope=False (model is always unified_3d_mrope), has_generation=False (always True), multi-vision-items (num_vision_items always 1), and the curr_rope_id non-mrope path. - Move the bf16 cast of per-step noisy tokens to before pack_input_sequence so the build-then-mutate pattern on packed_seq.vision.tokens disappears. - Drop the latent_patch_size / config hoists in __call__ that are now read directly inside the pack methods. Addresses reviewer threads huggingface/diffusers-new-model-addition-cosmos#1 comments 3278871807, 3278908766, 3278918514. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * collapse builder-pattern PackedSequence/ModalityData into flat dataclasses ModalityData and PackedSequence carried a list-or-tensor union for every field so they could double as builders during packing, with finalize() converting the lists to tensors at the end. Now that the pack methods each build their segment in one shot, finalize() is just a list->tensor conversion the pack methods can do themselves. - Rename ModalityData to _ModalityData (internal) with all-tensor fields (lists only for per-item entries like tokens / condition_mask). - Rename PackedSequence to Cosmos3PackedSequence and drop its finalize() method; fields are direct tensors at construction time. - _pack_vision_tokens / _pack_sound_tokens now build finalized _ModalityData directly via torch.arange / torch.tensor / torch.full. - pack_input_sequence builds the final Cosmos3PackedSequence in one return statement; no more two-stage build-then-finalize. - Update the transformer's forward docstring to reference the new name. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * exclude dtype from saved transformer config via ignore_for_config ModelMixin.from_pretrained injects dtype into init_dict (configuration_utils.py:289) whenever it appears in the loader's unused_kwargs, so Cosmos3OmniTransformer.__init__ has to accept dtype. But the default @register_to_config decorator was also serializing it into config.json on every save — leaving a stray "dtype": "bfloat16" key that doesn't describe the architecture, just the load-time runtime preference. Adding ignore_for_config = ["dtype"] keeps the decorator from registering dtype while still accepting it on the init signature. New saves omit dtype; existing checkpoints that have it log a warning at load time but the value is re-injected and __init__ ignores it, so loading is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * add Cosmos3 pipeline and transformer docs pages; retire JSON example inputs - New docs/source/en/api/pipelines/cosmos3.md with copy-pasteable text-to-image, text-to-video, image-to-video, and text-to-video-with-sound example snippets. The docs page is now the canonical reference for application code instead of the JSON-driven inputs/ directory. - New docs/source/en/api/models/cosmos3_omni_transformer.md describing the MoT dual-pathway architecture and showing a from_pretrained snippet. - Wire both pages into docs/source/en/_toctree.yml. - Export Cosmos3OmniDiffusersPipeline from the top-level diffusers package (matching every sibling pipeline) and add the corresponding dummy class for torch/transformers-unavailable environments. - Delete examples/cosmos3/inputs/omni/{t2i,t2v,i2v}.json and rewrite inference_cosmos3.py to take --prompt / --vision-path / --num-frames directly as CLI args. The script stays as a development smoke-test runner; canonical usage now lives in the docs. - Refresh examples/cosmos3/README.md to point at the docs page and reflect the new CLI surface. Addresses reviewer thread huggingface/diffusers-new-model-addition-cosmos#1 comment 3278258567. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * restore full prompts in cosmos3 docs from original example JSONs The text-to-video and image-to-video example prompts were shortened when porting from examples/cosmos3/inputs/omni/*.json into the docs page. Restore them verbatim from the JSONs so the docs reflect the prompts the model was actually demonstrated against, and so users copying from the docs get the same conditioning the example was tuned for. Also align the text-to-video-with-sound example: it now reuses the exact same prompt as the text-to-video block with only enable_sound=True added, instead of a hand-written waterfall prompt. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Renamed Cosmos3 module attributes Signed-off-by: Maciej Bala <mbala@nvidia.com> * bugfix Signed-off-by: Maciej Bala <mbala@nvidia.com> * Removed unnecessary helper function; added extra comments Signed-off-by: Maciej Bala <mbala@nvidia.com> * Moved from encode_prompt to tokenize_prompt Signed-off-by: Maciej Bala <mbala@nvidia.com> * Bring back video system prompt Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove multi frame conditioning for now Signed-off-by: Maciej Bala <mbala@nvidia.com> * Removed default negative prompts from code Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove Cosmos3condition; simplify sequence pack Signed-off-by: Maciej Bala <mbala@nvidia.com> * Clean up multiple parameters Signed-off-by: Maciej Bala <mbala@nvidia.com> * Simplify decode video; remove remainings of batching Signed-off-by: Maciej Bala <mbala@nvidia.com> * simple renames Signed-off-by: Maciej Bala <mbala@nvidia.com> * Refactored schedulers for sound Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove unnecessary autocast Signed-off-by: Maciej Bala <mbala@nvidia.com> * Update sound example Signed-off-by: Maciej Bala <mbala@nvidia.com> * Simplify loops in transformer_cosmos3.py Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove unused config attributes Signed-off-by: Maciej Bala <mbala@nvidia.com> * Cleanup audio decoder Signed-off-by: Maciej Bala <mbala@nvidia.com> * Reuse encoder_video from LTX2 for Cosmos3 Signed-off-by: Maciej Bala <mbala@nvidia.com> * Fixed a few nits Signed-off-by: Maciej Bala <mbala@nvidia.com> * Moved to RMSNorm for Cosmos3 Signed-off-by: Maciej Bala <mbala@nvidia.com> * Remove meta_tensor usage Signed-off-by: Maciej Bala <mbala@nvidia.com> * Improved rope handling Signed-off-by: Maciej Bala <mbala@nvidia.com> * Improved prompt templates Signed-off-by: Maciej Bala <mbala@nvidia.com> * Added extra docs for templatete Signed-off-by: Maciej Bala <mbala@nvidia.com> * remove dataclasses * Cleanup after merging Signed-off-by: Maciej Bala <mbala@nvidia.com> * Added guardrails Signed-off-by: Maciej Bala <mbala@nvidia.com> * bugfixed guardrails Signed-off-by: Maciej Bala <mbala@nvidia.com> * Bugfix guardrails v2 Signed-off-by: Maciej Bala <mbala@nvidia.com> * Simplified input_timestep Signed-off-by: Maciej Bala <mbala@nvidia.com> * Add TODO Co-authored-by: YiYi Xu <yixu310@gmail.com> * simplify conditional mask generation * Inlined _postprocess_latents Signed-off-by: Maciej Bala <mbala@nvidia.com> * removed pack_input_sequence helper * restore export utils with deprecation warning * moved sampling rate to pipeline attribute * inlined sound and image condition mask * seperating static and timestep based sound and vision token packing * unpack transformer args * enabled selection of cosmos3 super * Update src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Apply suggestions from code review Co-authored-by: YiYi Xu <yixu310@gmail.com> * fixed sound and vision conditioning use from prepare latents * ruff format and doc builder * ran fix copies * move typing to python3.10 * move special token application from pack_text_tokens to tokenize_prompt * rename packing methods to process methods * remove guidance_scale check * fix nits * respect vae dtype in the pipeline * use vae dtype for vae normalization stats * skip CFG if guidance_scale is 1 * Remove unnecessary parameter Signed-off-by: Maciej Bala <mbala@nvidia.com> * fix CFG for sound * bugfix for sound CFG * ruff format * Fix apply_chat_template return dict arg to return BatchEncoding * added option to select attention processor * docs: refresh Cosmos 3 pipeline intro Replace the terse architectural lede with the launch-style positioning (unified WFM for Physical AI, consolidating Predict/Reason/Transfer into one omni-model) and split out "What's new" and "Available checkpoints" sections so the page leads with capability rather than repo IDs. * docs: document Cosmos3OmniPipeline.__call__ arguments Add the missing Args block to Cosmos3OmniPipeline.__call__ so utils/check_forward_call_docstrings.py passes — covers all 21 parameters from prompt through enable_safety_check, plus a Returns section describing the output dataclass. * style: doc-builder reflow on Cosmos3OmniPipeline.__call__ docstring --------- Signed-off-by: Maciej Bala <mbala@nvidia.com> Co-authored-by: Yuliya Zhautouskaya <yzhautouskay@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Maciej Bala <mbala@nvidia.com> Co-authored-by: Dima Zhylko <dzhylko@nvidia.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
|
Guian Fang
|
e39aecff57 |
Add AnyFlow Any-Step Video Diffusion Pipelines (Bidirectional + FAR Causal) (#13745)
* [Pipelines] AnyFlow: scaffold pipelines/anyflow + register all top-level imports
This is the lazy-loader scaffolding only. Body files (pipeline_anyflow.py,
pipeline_anyflow_causal.py, transformer_anyflow.py,
scheduling_flow_map_euler_discrete.py) come in subsequent commits.
* [Schedulers] AnyFlow: add FlowMapEulerDiscreteScheduler
The flow-map scheduler advances samples from timestep t to caller-provided
target r in a single Euler step, supporting any-step sampling on flow-map-
distilled checkpoints. It is a general-purpose scheduler — not specific to the
AnyFlow checkpoints.
Tests: 12 standalone tests covering instantiation, set_timesteps endpoints,
shift identity/monotonicity, step shape preservation, zero-interval identity,
one-shot sampling, train weight schemes, scale_noise endpoints.
Docs: api/schedulers/flow_map_euler_discrete.md
* [Models] AnyFlow: add AnyFlowTransformer3DModel
A 3D DiT extending the v0.35.1 Wan2.1 backbone with two config-toggled modules:
* FAR causal blocks (init_far_model=True): block-sparse causal attention via
flex_attention + compressed-frame patch embedding for frame-level
autoregressive generation (Gu et al., 2025, arXiv:2503.19325).
* Dual-timestep flow-map embedding (init_flowmap_model=True): adds a delta
timestep embedder enabling flow-map sampling z_t -> z_r over arbitrary
intervals (AnyFlow).
With both flags off, the model reduces to stock Wan2.1.
The class is intentionally self-contained rather than annotated with
'# Copied from diffusers.models.transformers.transformer_wan' because upstream
Wan has been refactored extensively since v0.35.1 (new WanAttention class,
different processor architecture).
Tests: 9 unit tests covering construction in 3 modes, bidi forward shape and
determinism, return_dict variants, save/load round-trip with and without
init_far_model, gradient checkpointing toggle.
Docs: api/models/anyflow_transformer3d.md
* [Pipelines] AnyFlow: add AnyFlowPipeline and AnyFlowCausalPipeline
* AnyFlowPipeline (pipeline_anyflow.py, ~590 LOC): bidirectional T2V using
flow-map sampling. Loads checkpoints from nvidia/AnyFlow-Wan2.1-T2V-{1.3B,14B}.
* AnyFlowCausalPipeline (pipeline_anyflow_causal.py, ~700 LOC): FAR-based
causal pipeline supporting T2V/I2V/TV2V via task_type kwarg. Loads checkpoints
from nvidia/AnyFlow-FAR-Wan2.1-{1.3B,14B}-Diffusers.
Both pipelines reuse stock WanLoraLoaderMixin, AutoencoderKLWan, UMT5EncoderModel,
and AutoTokenizer from upstream. The transformer is the AnyFlowTransformer3DModel
introduced in the previous commit. The scheduler is FlowMapEulerDiscreteScheduler.
Tests:
* tests/pipelines/anyflow/test_anyflow.py: PipelineTesterMixin fast tests +
slow integration test against nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers.
* tests/pipelines/anyflow/test_anyflow_causal.py: same structure for FAR variant.
Reference slices for slow integration tests are deferred to Phase 7
(Final quality pass) where the user runs them on a real GPU.
* [Docs] AnyFlow: add main pipeline documentation page
Modeled on the Helios pipeline doc (PR #13208). Sections: paper link + abstract,
supported checkpoints table, memory/speed optimization tabs, T2V/I2V/TV2V
examples for both bidirectional and causal variants, autodoc trailers.
* [Auto/Scripts] AnyFlow: register AutoPipelineForText2Video + add conversion script
* Register AnyFlowPipeline in AUTO_TEXT2VIDEO_PIPELINES_MAPPING.
* AnyFlowCausalPipeline is intentionally NOT registered for AutoPipeline because
its task switch (t2v / i2v / tv2v) is too rich for a single auto-resolve key.
* scripts/convert_anyflow_to_diffusers.py: convert .pt training checkpoints
(with 'ema' state dict) into a diffusers save_pretrained layout. Supports all
4 released NVIDIA AnyFlow variants. Replaces the omegaconf-based config in the
upstream repo with argparse to match other diffusers conversion scripts.
* [Quality] AnyFlow: ruff-format + regenerated dummy stubs
* ruff format pass on all 5 source files (long lines + trailing comma fixes)
* check_dummies.py --fix_and_overwrite regenerated:
- dummy_pt_objects.py: AnyFlowTransformer3DModel + FlowMapEulerDiscreteScheduler
- dummy_torch_and_transformers_objects.py: AnyFlowPipeline + AnyFlowCausalPipeline
Local fast tests: 21/21 passed
- 12 scheduler tests (FlowMapEulerDiscreteScheduler)
- 9 transformer tests (AnyFlowTransformer3DModel construction + bidi forward + save/load)
The pipeline fast tests in tests/pipelines/anyflow/ require a local dev install
that matches the diffusers main branch's transformers >= compatibility floor.
The reference slices for slow integration tests (real GPU + 1.3B/14B
checkpoints) are intentionally left as TODO stubs to be captured by the user
on a real GPU machine before opening the PR.
* [AnyFlow] address review feedback: bug fixes + DMD wording + EN/ZH tutorials
Critical bug fixes (verified against precision-validation review):
* pipeline_anyflow.py / pipeline_anyflow_causal.py: replace hardcoded
transformer_dtype = torch.bfloat16 with self.transformer.dtype, so
pipe.to("cpu") and PipelineTesterMixin save/load tests do not crash on a
dtype mismatch in the patch_embedding conv3d.
* transformer_anyflow.py: drop the duplicate `base = base = ...` assignment in
_build_causal_mask (was a copy-paste typo carried over from FAR-Dev).
* transformer_anyflow.py: drop unused `q_is_context` / `k_is_context` locals
and the `# noqa: F841` markers that were silencing the dead-store warning.
* transformer_anyflow.py: remove `CacheMixin` from the inheritance list — the
pipeline manages KV cache directly, the mixin's interface is unused.
* transformer_anyflow.py: guard the module-level `torch.compile(flex_attention)`
with try/except so the file imports cleanly on CPU CI / no-Triton machines.
* convert_anyflow_to_diffusers.py: replace ad-hoc print warnings with the
stdlib logger (warning_once-style) and a module-level basicConfig.
Documentation accuracy:
* AnyFlowCausalPipeline class docstring + main pipeline doc + EN/ZH tutorial:
drop the fictitious `task_type` / `image` / `video` arguments and document
the real API: pass `context_sequence={"raw": tensor}` (or `{"latent": ...}`)
to switch between T2V (None) / I2V (1-frame) / TV2V (4n+1-frame) modes.
* Pipeline class docstrings + main doc: explicitly describe AnyFlow's
two-stage LoRA distillation including DMD reverse-divergence supervision
with Flow-Map backward simulation in stage 2 (was previously implicit).
* training_rollout: add detailed docstring explaining its role as the
3-segment Flow-Map backward simulation entry point used during DMD training.
* Long-form tutorial doc `using-diffusers/anyflow.md` (EN, 239 LOC) and
Chinese mirror `docs/source/zh/using-diffusers/anyflow.md` (224 LOC) added
and registered in both `_toctree.yml` files.
Tests:
* Skip `test_attention_slicing_forward_pass` in both pipeline test classes
with a clear rationale (custom attention processor does not support slicing).
* All 21 standalone tests still pass (12 scheduler + 9 transformer).
Quality gates:
* `ruff check` clean across all AnyFlow files.
* `ruff format --check` reports 6 files already formatted.
* `python utils/check_copies.py` reports no diff.
Out of scope for this commit (deferred until reviewer feedback):
* Splitting AnyFlowTransformer3DModel into bidi + causal subclasses
* Unifying _forward_inference / _forward_cache return types
* Migrating model tests from plain unittest to BaseModelTesterConfig + mixins
* HF model card / config.json metadata updates on the nvidia/* repos
(push to Hub manually before opening the PR)
* [AnyFlow] rename Causal->FAR + explicit forward signature + dataclass output
Round 2 of review feedback. Three groups of changes; transformer state-dict
keys, module hierarchy, and tensor flow are unchanged so the H200 bit-exact
validation remains valid.
A. Pipeline rename (mechanical, no behavior change):
* Class: AnyFlowCausalPipeline -> AnyFlowFARPipeline (Causal in diffusers
usually means an attention mask; AnyFlow's variant is FAR autoregressive,
so the FAR name is more specific and matches the paper).
* File: pipeline_anyflow_causal.py -> pipeline_anyflow_far.py (git mv).
* Test file: test_anyflow_causal.py -> test_anyflow_far.py (git mv).
* All references updated in src/, tests/, docs/, scripts/, plus stale
anyflowcausalpipeline anchor links in tutorial markdown.
B. Pipeline test bug fixes (closes 19 fast-test failures reported by
precision-validation reviewer):
* pipeline_anyflow.py / pipeline_anyflow_far.py: __call__ now sets
self._num_timesteps = num_inference_steps before the rollout, so the
PipelineTesterMixin callback tests can read pipe.num_timesteps.
* tests/pipelines/anyflow/test_anyflow_far.py: drop the fictitious
task_type="t2v" kwarg that crashed every causal fast test (the FAR
pipeline selects mode via context_sequence, not a task_type arg).
C. Transformer architecture cleanups (review-driven, no tensor changes):
* Replace forward(*args, **kwargs) dispatcher with an explicit signature
listing every supported kwarg (hidden_states, timestep, r_timestep,
encoder_hidden_states, encoder_hidden_states_image, chunk_partition,
clean_hidden_states, clean_timestep, kv_cache, kv_cache_flag, is_causal,
attention_kwargs, return_dict). Helps IDE / type-checker / torch.compile
tracing.
* Drop SimpleNamespace returns. Add AnyFlowFARTransformerOutput
(BaseOutput dataclass with sample + kv_cache fields) for the two causal
paths that need to also propagate kv_cache (_forward_inference and the
newly return_dict-aware _forward_cache). _forward_train and
_forward_bidirection now consistently return Transformer2DModelOutput.
Pipeline call sites already use return_dict=False with positional
unpacking, so the fix is transparent there.
Out of scope (deferred until canonical-org HF metadata sync):
* Splitting AnyFlowTransformer3DModel into a bidi class plus an
AnyFlowFARTransformer3DModel subclass — touches register_to_config keys
and would require updating model_index.json on every released checkpoint.
* Promoting chunk_partition from register_to_config to a forward-time
argument (same reason).
* Renaming training_rollout to _denoise — would break callers in the
FAR-Dev on-policy trainer that produced the released checkpoints.
Local fast tests: 21/21 still pass (12 scheduler + 9 transformer).
ruff check, ruff format, and check_copies.py are all clean.
* [AnyFlow] wire callback_on_step_end through inference_range + add chunk_partition to FAR fast-test fixture
Two root causes for the 19 remaining PipelineTesterMixin failures, identified
by the H200 reviewer:
1. callback_on_step_end was accepted by __call__ but never invoked. Both
pipelines pass it through to training_rollout (and FAR additionally through
inference()), and inference_range now fires it after scheduler.step in
the standard inference branch:
if callback_on_step_end is not None:
callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs}
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = ...
negative_prompt_embeds = ...
`nonlocal prompt_embeds, negative_prompt_embeds` lets the callback rewrite
the closure-captured embeddings, matching upstream WanPipeline semantics.
The 3-segment grad_timestep training rollout does not invoke the callback;
it is intentionally training-only.
2. tests/pipelines/anyflow/test_anyflow_far.py::get_dummy_components built
the dummy transformer without a `chunk_partition`, leaving it None on the
model config and crashing the pipeline at `sum(self.transformer.config.chunk_partition)`.
Set `chunk_partition=[1, 1, 1]` in the fixture (3 chunks of 1 latent frame
each, matching the test's num_frames=9 -> 3 latent frames).
Local fast tests: 21/21 still pass.
ruff check, ruff format, and check_copies.py are all clean.
* [AnyFlow] Phase 2: split transformer + drop chunk_partition from config + rename helpers
Major architectural refactor that aligns the integration with diffusers conventions
ahead of the canonical-org Hub upload. State-dict keys, module hierarchy, and
tensor flow are unchanged so the H200 bit-exact validation remains valid; only
the on-disk transformer/config.json fields move.
Changes:
1. **Sibling transformer classes** replace the flag-driven single class:
* AnyFlowTransformer3DModel — bidirectional only. Drops compressed_patch_size /
full_chunk_limit / init_far_model / init_flowmap_model / chunk_partition
kwargs (always-on for AnyFlow distilled checkpoints).
* AnyFlowFARTransformer3DModel — adds far_patch_embedding + the 3 FAR forward
paths (train / cache-prefill / autoregressive inference).
* AnyFlowTimeTextImageEmbedding (the legacy single-time embedder used only by
the old setup_flowmap_model bootstrap) is removed; both classes now build
AnyFlowDualTimestepTextImageEmbedding directly in __init__.
* setup_flowmap_model / setup_far_model methods are removed; weight warm-start
for far_patch_embedding (trilinear interpolation from patch_embedding) moves
into AnyFlowFARTransformer3DModel.__init__.
2. **chunk_partition** is no longer a model config field. The FAR pipeline owns
the schedule:
* AnyFlowFARPipeline.default_chunk_partition = [1, 3, 3, 3, 3, 3, 3, 2]
matches the released 81-frame NVIDIA checkpoints.
* AnyFlowFARPipeline.__call__ / _denoise_rollout accept a chunk_partition
argument that overrides the default for non-default num_frames.
3. **training_rollout -> _denoise_rollout** rename across both pipelines and all
English / Chinese docs that referenced it. Signals the method is internal to
the pipeline driver, not a public training API.
4. **Conversion script + tests + docs + registries**:
* scripts/convert_anyflow_to_diffusers.py: VARIANTS dict picks the right
transformer class per variant; init_far_model / init_flowmap_model /
chunk_partition kwargs are removed from the from_pretrained call.
* Transformer test file split into AnyFlowTransformer3DModelTest and
AnyFlowFARTransformer3DModelTest classes.
* Pipeline test fixtures use the right class and pass chunk_partition via
get_dummy_inputs (3-frame schedule [1, 1, 1] for the 9-frame test).
* New docs page docs/source/en/api/models/anyflow_far_transformer3d.md;
anyflow_transformer3d.md rewritten for the bidi-only class.
* AnyFlowFARTransformer3DModel registered in src/diffusers/__init__.py,
src/diffusers/models/__init__.py, models/transformers/__init__.py and the
dummy_pt_objects.py stubs.
* docs/source/en/_toctree.yml: new entry for the FAR transformer page.
5. **Cleanups**:
* Pipeline __call__ no longer passes is_causal=False to the bidi forward (the
bidi class doesn't accept it).
* Pipeline class docstrings drop stale references to init_*_model flags.
Local tests: 22/22 pass (12 scheduler + 10 transformer covering both classes).
ruff check / format / check_copies clean.
Hub artifacts (model_index.json, transformer/config.json, scheduler config) need
to be regenerated for the released checkpoints; the HF update guide will be
delivered separately.
* [AnyFlow] Phase 3: convention compliance against .ai/AGENTS.md + .ai/models.md
Hard violations (per official diffusers guidelines):
* drop einops dependency — replace 25+ rearrange() calls with native
permute/reshape/unflatten in transformer + both pipelines
* device-gate torch.float64 — apply_rotary_emb and AnyFlowRotaryPosEmbed now
fall back to float32 / complex64 on MPS / NPU; freqs are lazily rebuilt
per-device via _build_freqs (matches transformer_wan / transformer_flux
pattern)
* migrate attention to dispatch_attention_fn — replace direct
F.scaled_dot_product_attention calls with dispatch_attention_fn (works
with sage / flash / native backends); introduce AnyFlowAttention(
AttentionModuleMixin) with _default_processor_cls / _available_processors;
rename processors to AnyFlowAttnProcessor / AnyFlowCrossAttnProcessor and
declare _attention_backend / _parallel_config class attrs
* drop dead config fields — qk_norm and added_kv_proj_dim are pruned from
both transformer __init__ signatures and AnyFlowTransformerBlock;
AnyFlowAttention is hardcoded to rms-norm-across-heads (the only scheme
the released checkpoints use) and has no add_k_proj path (T2V only)
* add _repeated_blocks = ["AnyFlowTransformerBlock"] to both transformer
classes for compile_repeated_blocks() support (matches Wan)
* annotate prepare_latents with `# Copied from diffusers.pipelines.wan.
pipeline_wan.WanPipeline.prepare_latents`; the pipeline-side rearrange
to (B, T, C, H, W) layout is moved to the call site
State-dict keys are preserved (legacy Attention had identical to_q / to_k /
to_v / to_out / norm_q / norm_k naming), so existing AnyFlow checkpoints load
bit-exactly into the new AnyFlowAttention class.
The HF Hub config-update guide is updated correspondingly: transformer/
config.json now drops qk_norm and added_kv_proj_dim alongside the previous
init_far_model / init_flowmap_model / chunk_partition removals.
22 fast CPU tests still pass; ruff format / ruff check / check_copies all
clean.
* [AnyFlow] FAR fast-test compat: rope 0-dim guard + flex_attention CPU/head-dim fallbacks + KV-cache dtype + num_timesteps
Phase 3 migrated bidi + cross-attention to dispatch_attention_fn but the FAR
causal path still calls flex_attention directly, which has hard requirements
(CPU compile, head_dim >= 16) that fail on PipelineTesterMixin's tiny dummy
components. Real ckpts (head_dim=128, CUDA) never hit these branches; bit-exact
numerical equivalence with FAR-Dev preserved on all 4 released ckpts (forward
0.00e+00, backward kernel-nondet only, ratio 1.000).
Code fixes:
1. AnyFlowRotaryPosEmbed._forward_compressed_frame / _forward_full_frame now
short-circuit to an empty tensor when num_frames / height / width is 0.
PipelineTesterMixin's dummy VAE has scale_factor_spatial=8, so a 16x16 raw
spatial input becomes a 2x2 latent which then floors to 0 against
compressed_patch_size=(1, 4, 4); the original
`freqs[:0].view(0, k, 1, -1)` reshape was ambiguous in that regime.
2. flex_attention dispatch: split the module-load
`torch.compile(flex_attention, dynamic=True)` into `_flex_attention_eager`
(always available) plus `_flex_attention_compiled`, with a tiny wrapper
that picks compiled for CUDA tensors and eager for CPU. Avoids
torch._inductor C++ codegen failures that broke fast tests after
`pipe.to("cpu")`. CUDA performance unchanged (L10 benchmark: 0.0% delta on
bidi 1.3B fwd, 0.0% delta on FAR causal 1.3B fwd).
3. AnyFlowAttnProcessor (FAR causal branch): when head_dim < 16
(flex_attention's hard minimum) zero-pad q/k/v's last dim to 16 and pass
`scale=1/sqrt(original_head_dim)` to flex_attention. Padded value rows
contribute 0, so trimming the output back is mathematically equivalent.
Released ckpts use head_dim=128 so the branch is never taken in production.
4. pipeline_anyflow_far.encode_kv_cache: replace the hardcoded
`latents.to(torch.bfloat16)` with `self.transformer.dtype`. The hardcoded
bf16 crashed conv3d on dummy fp32 components ("Input type (BFloat16) and
bias type (float) should be the same"); real bf16 ckpts are unaffected.
5. pipeline_anyflow_far._denoise_rollout sets
`self._num_timesteps = (len(chunk_partition) - num_context_chunks) * num_inference_steps`
before the chunk loop, so PipelineTesterMixin.test_callback_cfg's
`pipe.num_timesteps`-based assertion matches the actual number of callback
fires (chunks * NFE) instead of the previous hardcoded num_inference_steps.
Tests:
* test_callback_inputs cannot pass without changing FAR's chunk-wise output
semantics — it zeroes latents on the final step and asserts the *entire*
output buffer is zero, but only the active chunk's slice is overwritten in
a chunk-wise rollout. Marked `@unittest.skip` with a detailed rationale;
callback functionality itself is still covered by test_callback_cfg.
* Full pytest run on tests/pipelines/anyflow/ +
tests/models/transformers/test_models_transformer_anyflow.py +
tests/schedulers/test_scheduler_flow_map_euler_discrete.py: 81 passed,
0 failed, 11 skipped.
Quality gates:
* `ruff check` and `ruff format --check` clean across all AnyFlow files.
* `python utils/check_copies.py` clean.
* `python utils/check_dummies.py` clean.
* [AnyFlow] docs/code: paper-release tidy-up
User-facing alignment with the official HF Hub model card and the day-of-announcement
materials at https://huggingface.co/collections/nvidia/anyflow.
* Fill in the arXiv identifier 2605.13724 (5 paper links + 2 BibTeX entries).
* Rename TV2V → V2V across docs + pipeline_anyflow{,_far}.py so the diffusers
copy uses the same Video-to-Video terminology as the official model card.
* Add the [nvidia/anyflow](https://huggingface.co/collections/nvidia/anyflow)
HF collection link to the three tutorial intros.
* Drop the temporary "guyuchao/* staging" tip from the EN tutorial / API page
/ ZH tutorial — the nvidia/AnyFlow-*-Diffusers repos are now live.
* Wire up NVlabs/AnyFlow (training code) and nvlabs.github.io/AnyFlow (project
page) in place of the prior <github-org> / <project-page-url> placeholders.
* Cite the authors (Yuchao Gu, Guian Fang et al.) and NUS ShowLab × NVIDIA
affiliation in the main tutorial, API pipeline page, and both transformer
model pages; BibTeX uses the standard `and others` to elide the full list
until the next pass.
Working tree, CI gates, and tests after the change:
ruff format --check ✓
ruff check ✓
python utils/check_copies.py ✓
python utils/check_dummies.py ✓
pytest tests/models + tests/schedulers (22 fast) ✓
No production code logic changes — only docstring wording inside pipeline
files (TV2V → V2V).
* [AnyFlow] docs: drop in official BibTeX (full author list)
Replace the placeholder ``@article{gu2026anyflow, author = {Gu, Yuchao and
Fang, Guian and others}, ...}`` block in both the English and Chinese
tutorials with the canonical ``@misc{gu2026anyflowanystepvideodiffusion,
...}`` form from arxiv.org/abs/2605.13724, which lists all seven authors:
Yuchao Gu, Guian Fang, Yuxin Jiang, Weijia Mao, Song Han, Han Cai,
Mike Zheng Shou.
Docs-only.
* [AnyFlow] align with diffusers conventions + drop training-only code
Scheduler
- FlowMapEulerDiscreteScheduler.step now returns a
FlowMapEulerDiscreteSchedulerOutput dataclass (or tuple with return_dict=False)
and uses the conventional positional order (model_output, timestep, sample,
r_timestep).
- Drop training-only helpers: adaptive_weighting, set_train_weight,
get_train_weight, linear_timesteps_weights, and the weight_type config field.
- Add scale_model_input no-op for API parity; raise ValueError on missing
r_timestep.
Transformer
- Remove gate_track debug write inside
AnyFlowDualTimestepTextImageEmbedding.forward_timestep.
- Compile flex_attention lazily on first CUDA call instead of at import time.
- Replace assert with ValueError in build_block_mask.
- Resolve <arxiv-id> placeholders to 2605.13724.
Pipelines (AnyFlowPipeline + AnyFlowFARPipeline)
- Add EXAMPLE_DOC_STRING + @replace_example_docstring and full __call__
docstrings covering every argument.
- Move use_mean_velocity from __init__ to __call__ so save/load round-trips.
- Drop _denoise_rollout's grad_timestep branch (DMD on-policy training rollout),
the inner inference_range closure, and the redundant negative-prompt concat.
- Replace asserts with ValueError; wire show_progress to tqdm; rename inference
-> _inference; remove dead current_timestep property.
- Update scheduler.step call sites to the new signature.
- Trim class docstrings to inference-only language.
Pipeline output
- Add Apache 2.0 license header; switch to relative import.
Auto pipeline / conversion script
- Register AnyFlowFARPipeline in AUTO_IMAGE2VIDEO_PIPELINES_MAPPING and
AUTO_VIDEO2VIDEO_PIPELINES_MAPPING.
- Document the weights_only=False requirement in the conversion script.
Tests
- Scheduler tests use the new step signature and verify the Output dataclass
contract.
- Drop the four obsolete training-weight tests; drop weight_type kwarg from
pipeline test fixtures; remove internal milestone names from TODO comments.
Docs
- Resolve <arxiv-id> in the scheduler docs page.
- Trim DMD / on-policy distillation language in EN/ZH tutorials and the
pipelines page; the paper abstract quote is preserved verbatim.
* [AnyFlow] split FAR causal transformer into transformer_anyflow_far.py
Per @dg845's review on #13745: extract FAR causal modules into a dedicated
sibling file so each transformer variant reads in isolation. Shared submodules
are duplicated via `# Copied from` so `make fix-copies` keeps both in sync.
- `transformer_anyflow.py`: bidi-only. `AnyFlowAttnProcessor` no longer carries
the flex/KV-cache branch (was: dispatch in one branch, bare flex_attention in
the other); `AnyFlowRotaryPosEmbed` drops the compressed-frame helpers and
the `is_causal` arg; `AnyFlowDualTimestepTextImageEmbedding` drops its causal
branch. `AnyFlowTransformerBlock` keeps a single class with a new
`is_causal: bool = False` ctor flag that selects the self-attn processor —
the forward path is identical in both modes, only the processor differs.
- `transformer_anyflow_far.py`: new. Contains `AnyFlowFARTransformerOutput`,
`AnyFlowCausalAttnProcessor` (routed through `dispatch_attention_fn(backend=
"flex")` with a clear ValueError when a non-flex backend is configured; the
BlockMask is consumed only by the flex backend in `_native_flex_attention`),
`AnyFlowDualTimestepTextImageEmbeddingCausal`, `AnyFlowCausalRotaryPosEmbed`,
`AnyFlowFARTransformer3DModel`, and `# Copied from` clones of the shared
shared `AnyFlowAttention`/`AnyFlowCrossAttnProcessor`/`AnyFlowImageEmbedding`/
`AnyFlowTransformerBlock`/`AnyFlowAttnProcessor` modules.
Verified bit-exact against the pre-refactor branch on H200 (float32):
- bidi: L2 = 0.000e+00, max|Δ| = 0.000e+00
- FAR : L2 = 4.772e-06, max|Δ| = 3.576e-07
The FAR delta is fp32 accumulation noise from the dispatch path permuting
(B,L,H,D) ↔ (B,H,L,D) around the same `flex_attention` kernel.
Addresses review comments at transformer_anyflow.py:215, :261, :450, :622,
:671, :958.
* [AnyFlow] pipeline cleanup: video_processor, encode_video, inline rollout, kwarg rename
Per @dg845's review on #13745, applied to both bidi `AnyFlowPipeline` and
causal `AnyFlowFARPipeline`:
- Use `self.video_processor.preprocess_video(...)` instead of the manual
`* 2 - 1` normalize.
- Merge `vae_encode` + `encode_latents` + `_normalize_latents` into a single
`encode_video` method, mirroring `WanImageToVideoPipeline.encode_image`'s
flat structure.
- Inline `_denoise_rollout` into `AnyFlowPipeline.__call__`. For the FAR
pipeline, inline both `_denoise_rollout` and `_inference` as a nested loop
(outer over chunks, inner over denoising steps), mirroring
`WanAnimatePipeline.__call__`. `encode_kv_cache` is intentionally kept as a
method — it is one transformer call with a different `kv_cache_flag` mode
(cache-write), and inlining it would interleave two distinct forward
semantics in the same loop body and lose readability.
- Rename `context_sequence` → `video` (pixel-space) + `video_latents`
(pre-encoded), matching `WanVideoToVideoPipeline`. For the FAR pipeline,
the old `{"raw"/"latent"}` dict form is replaced by the two kwargs.
Mutually-exclusive validation raises `ValueError`.
Addresses review comments at pipeline_anyflow.py:358, :372, :393, :473 and
pipeline_anyflow_far.py:395, :489, :675.
* [AnyFlow] scheduler: N-length timesteps + step defaults r_timestep
Per @dg845's review on #13745:
- `set_timesteps(N)` now produces `N` timesteps backed by an internal
`sigmas[N+1]` linspace, matching `FlowMatchEulerDiscreteScheduler.set_
timesteps`. The final sigma (== 0) is the implicit r-endpoint of the last
step; the pipeline rollouts iterate `for i, t in enumerate(timesteps)`
without the old `[:-1]` slicing.
- `step(r_timestep=None)` now defaults to the next timestep on the schedule
(resolved via fp-tolerant `argmin` over `sigmas[:-1]`), instead of raising.
Any-step sampling is preserved when `r_timestep` is explicit. The raise
stays only for the case where the caller passes a `timestep` value that
isn't on the schedule and provides no `r_timestep` — there's no sensible
default in that case.
- Build sigmas in float64 on CPU then move to the target device, with a
float32 downcast for MPS / NPU (float64 isn't supported on those backends).
Pipeline rollout loops updated to compute `r = sigmas[i + 1] * num_train_
timesteps` for the model's `r_timestep` input and pass `r_timestep=None` to
`scheduler.step` (which resolves it from the schedule internally).
Addresses review comments at scheduling_flow_map_euler_discrete.py:107 and
:148.
* [AnyFlow] tests: regenerate via generate_model_tests.py; split bidi/FAR files
Per @dg845's review on #13745: replaced the hand-rolled transformer tests
with the standard mixin-based suite produced by `utils/generate_model_tests
.py`, and split the FAR causal model tests into their own file to mirror the
transformer file split.
- `tests/models/transformers/test_models_transformer_anyflow.py`: regenerated
bidi suite. Pulls in `ModelTesterMixin`, `MemoryTesterMixin`,
`TrainingTesterMixin`, `AttentionTesterMixin`, `TorchCompileTesterMixin` via
`BaseModelTesterConfig`, with `get_init_dict()` / `get_dummy_inputs()`
filled in for the small bidi config used in CI.
- `tests/models/transformers/test_models_transformer_anyflow_far.py`: new.
Same mixin set (TorchCompile is intentionally skipped — FAR's
`_build_causal_mask` uses `flex_attention.create_block_mask(_compile=False)`
which conflicts with the standard compile tester's assumptions; the bidi
file covers compile, FAR is bit-exact-validated end-to-end on H200 via the
pipeline replay). Also carries an `AnyFlowCausalAttnProcessor` smoke test
that exercises the backend gate (non-flex backends must raise) and asserts
the `AnyFlowFARTransformerOutput` dataclass exposes the expected fields.
Addresses review comments at test_models_transformer_anyflow.py:71 and :128.
* [AnyFlow] docs: update for video / video_latents kwarg rename
Following the pipeline kwarg refactor in e9d50b2, sweep the user-facing docs
to reflect the new API:
- `docs/source/en/api/pipelines/anyflow.md`: T2V / I2V / V2V code examples now
use `video=` instead of `context_sequence={"raw": ...}`. The "Generation
with AnyFlow (FAR Causal)" intro describes the new mutually-exclusive
`video` / `video_latents` selector.
- `docs/source/en/using-diffusers/anyflow.md`: the scenario selector table,
the "Image-to-video and video-to-video" walkthrough, and the closing note
about pre-encoded latents are all updated. `vae_encode` references are
replaced with `encode_video`.
* [AnyFlow] tests: skip FAR training tests on CPU (flex backward); align scheduler tests with N-length timesteps
- TestAnyFlowFARTransformer3DTraining: skip test_training / test_training_with_ema /
test_gradient_checkpointing_equivalence on CPU. FAR causal self-attn uses
torch.nn.attention.flex_attention whose backward kernel is GPU-only.
- test_scheduler_flow_map_euler_discrete: assert timesteps is N-length (not N+1) and
the sigma=0 r-endpoint lives in self.sigmas[-1]; test_step_one_shot_sampling now
exercises r_timestep=None (resolved from sigmas) since N=1 has no timesteps[1].
* [AnyFlow] docs: complete forward() Args: sections for check_forward_call_docstrings
main #13758 added utils/check_forward_call_docstrings.py which requires every signature
arg to appear as its own `name (...):` entry under Args:. Expand the bidi and FAR
transformer forward docstrings to list each parameter individually.
* [AnyFlow] apply 5/21 review suggestions (A: 1-click)
FAR transformer:
- AnyFlowCausalAttnProcessor: default _attention_backend = 'flex' (was None);
remove None from _SUPPORTED_BACKENDS. None previously fell through to SDPA
which silently ignored the BlockMask; failing loudly is the right default.
- dispatch_attention_fn call: read self._attention_backend instead of hardcoded
'flex', so '_native_flex' selection works.
- _build_freqs / _forward_full_frame: add '# Copied from' to bidi RoPE.
Pipelines:
- bidi + FAR docstrings: video shape (B, C, T, H, W) -> (B, T, C, H, W) to
match VideoProcessor.preprocess_video.
- FAR EXAMPLE_DOC_STRING: single-frame I2V tensor wrap uses unsqueeze(1) for the
T axis instead of unsqueeze(2).
- FAR encode_video: drop duplicated @torch.no_grad() decorator.
Tests:
- test_anyflow / test_anyflow_far: lift the test_save_load_optional_components
skip (the test actually passes).
- FAR processor smoke test: assert default backend is 'flex' (was 'None').
* [AnyFlow] apply 5/21 review suggestions (B: refactors)
Pipelines:
- check_inputs accepts video / video_latents and raises early on:
(a) mutual exclusion (was checked late in __call__);
(b) FAR's (num_frames - 1) % 4 == 0 constraint.
__call__ no longer carries duplicate validation.
- FAR pipeline: drop the show_progress kwarg and replace the single tqdm with
nested progress bars in the LLaDA-2 pattern — outer 'Chunks' (position=0)
and per-chunk inner 'Inference Steps' (position=1, leave=False) — both
picking up DiffusionPipeline._progress_bar_config (so set_progress_bar_config
controls them, including disable=None).
Scheduler:
- step() resolves source and target sigmas by indexing self.sigmas via the new
index_for_timestep(), instead of dividing the input timesteps by
num_train_timesteps. This keeps the math correct for any future schedule
whose timestep/sigma relationship is non-linear. For an off-schedule
r_timestep the code falls back to r / num_train_timesteps, so explicit
any-step sampling outside the schedule still works (and t off-schedule with
r=None still raises a clear ValueError, as before).
Numerical equivalence: for the shipped linspace+shift schedule the two
formulations are bit-identical (verified: max abs diff = 0.0 over an N=8,
shift=5 schedule).
* [AnyFlow] apply Claude bot review (5/21): 8 findings beyond dg845's list
Finding #1 — attention_kwargs plumbing:
Both transformers now decorate forward() with @apply_lora_scale('attention_kwargs')
(matches Wan); pipelines forward attention_kwargs to the transformer + encode_kv_cache,
and the unused parameter is dropped from the inner _forward_train / _forward_cache /
_forward_inference signatures. Pipeline docstrings updated to the standard wording.
Finding #2 — naming:
Rename far_cfg -> layout_cfg in the bidi transformer (the bidi path is not FAR; the
FAR transformer keeps far_cfg, which is accurate there).
Finding #3 — scheduler state machine:
Add _step_index, _begin_index, step_index property, begin_index property,
set_begin_index(), _init_step_index(). step() lazily initializes and advances the
counter so downstream callbacks / composable schedulers can observe rollout progress.
Sigma resolution remains a pure function of (timestep, r_timestep) — calling step()
twice with identical args still returns identical prev_sample (idempotent).
Finding #4 — redundant @torch.no_grad():
Drop the redundant decorators on bidi pipeline's encode_video and FAR pipeline's
encode_kv_cache (callers are already in __call__'s no-grad scope).
Finding #5 — dead code:
Remove the unreachable temb.ndim == 2 else branch from the bidi transformer's
output-norm path (condition_embedder.forward always returns a 3D temb).
Finding #6 — private rename:
forward_far_patchify[_inference] -> _forward_far_patchify[_inference] (only called
internally by _forward_train / _forward_cache / _forward_inference).
Finding #7 — pipeline comment numbering:
Bidi + FAR pipelines renumber steps so the # 4. slot is no longer skipped.
Finding #8 — mask-mod comment numbering:
_build_causal_mask numbered comments now run 1) 2) 3) ... (was 1) 3) 4) ...).
Tests:
- New test_step_index_advances + test_set_begin_index_anchors_step_index in the
scheduler test file exercise the new state machine.
- All existing pipeline / transformer / scheduler tests still pass (85 passed,
85 skipped on CPU).
Bit-exact: 8-step rollout vs the previous formulation, max abs diff = 0.0 (the new
sigma-lookup is byte-identical to t/num_train_timesteps on this schedule).
* [AnyFlow] scheduler: honour off-schedule any-step in _init_step_index; drop dead _resolve_next_timestep
Audit caught two issues in the previous scheduler commit:
1. The new state machine raised in _init_step_index whenever the first timestep
wasn't on the active schedule, contradicting the documented contract that
step() falls back to t/num_train_timesteps for off-schedule any-step
sampling. The fall-back numerics were intact but they were unreachable —
the init check fired first.
Fix: _init_step_index now initializes _step_index to 0 when the timestep is
off-schedule (still a valid observable counter for callbacks). step()'s
sigma resolution is untouched, so on-schedule rollouts stay bit-exact and
off-schedule any-step sampling actually runs again. Regression test:
test_step_off_schedule_anystep_supported.
2. _resolve_next_timestep had no remaining callers after the step() rewrite
inlined the same lookup. Removed (private helper, no external API).
* [AnyFlow] docs: align user guides with video shape + kwarg fixes
- en api/pipelines/anyflow.md: video shape (B, C, T, H, W) -> (B, T, C, H, W);
example tensor wrap uses unsqueeze(0).unsqueeze(1) and permute(0, 3, 1, 2)
to match VideoProcessor.preprocess_video's 5D contract.
- zh using-diffusers/anyflow.md: same shape fixes; also flip the I2V / V2V
examples from the obsolete context_sequence={...} dict to the current
video= / video_latents= kwargs; helper to_video_tensor returns (1, T, C, H, W);
add a note about mutual exclusion.
* [AnyFlow] tests: drop @slow integration test scaffolds for initial PR
.ai/skills/model-integration/SKILL.md is explicit: 'No integration / slow
tests in the initial PR — don't add anything gated on @slow / RUN_SLOW=1
yet.' Our two integration test classes were shape-only assertions with TODOs
for a future numeric reference, so dropping them loses no actual coverage —
the relevant rollouts are covered by H200 bit-exact replay outside the
pytest suite. Can land a follow-up PR after merge with proper numeric
reference slices once the maintainer is comfortable enabling slow tests.
* Apply style fixes
* [AnyFlow] apply 5/22 dg845 review: comment cleanups + custom sigmas/timesteps schedule
dg845 third pass — 7 of 9 comments applied; the 8th (custom sigmas/timesteps support)
matches FlowMatchEulerDiscreteScheduler conventions; the 9th (_build_causal_mask
refactor) is explicitly marked non-blocking and deferred to a follow-up that also
re-enables TorchCompileTesterMixin.
Comment cleanups:
- transformer_anyflow.py:704 temb output-norm comment: drop redundant 'no ndim==2 branch'.
- pipeline_anyflow.py:550 denoise loop comment: '# 6. Denoising loop'.
- pipeline_anyflow_far.py:684 denoise loop comment: '# 8. Denoising loop (outer over
chunks, inner over timesteps).'.
- pipeline_anyflow_far.py:702 drop trailing inline comment on `timesteps = scheduler.timesteps`.
- scheduling_flow_map_euler_discrete.py: clearer wording on the off-schedule `r_timestep`
error.
Custom schedule support:
- FlowMapEulerDiscreteScheduler.set_timesteps gains `sigmas` and `timesteps` kwargs
mirroring FlowMatchEulerDiscreteScheduler. Default behaviour is unchanged
(linspace + shift); the validation + length-N → length-N+1 terminal-0 append are
shared with the default path so on-schedule rollouts stay bit-exact.
- AnyFlowPipeline.__call__ and AnyFlowFARPipeline.__call__ accept `sigmas` and
`timesteps` kwargs, override num_inference_steps from their length, and forward
to set_timesteps (matches LTX2Pipeline pattern).
- New scheduler tests: test_set_timesteps_custom_sigmas and
test_set_timesteps_custom_timesteps cover both override paths.
Dtype skip on save/load:
- TestAnyFlowTransformer3D and TestAnyFlowFARTransformer3D now skip
test_from_save_pretrained_dtype_inference (parametrized over fp16/bf16), mirroring
WanTransformer3DModel's skip — the test's tolerance requirements are too high for
meaningful signal under AnyFlow's flow-map mixed-precision sampling.
* [AnyFlow] docs: apply hf-doc-builder line wrap (max_len 119)
CI doc-builder style check flagged 3 files with docstring lines >119 chars.
Ran 'doc-builder style src/diffusers docs/source --max_len 119' to autoformat;
content unchanged, line wrapping only.
* [AnyFlow] apply 5/22 follow-up review: new_zeros terminal sigma + cleanup
dg845 blocking suggestion (r3287274209):
- scheduling_flow_map_euler_discrete.py:185 — use `working_sigmas.new_zeros(1)`
instead of `torch.zeros(1, dtype=...)` so the appended terminal sigma inherits
both device and dtype from working_sigmas. The current working_sigmas always
starts on CPU so the device mismatch is latent, but new_zeros is the correct
defensive pattern and matches how the published FAR test fixtures run on CUDA.
Claude bot final-review follow-ups:
- transformer_anyflow_far.py: drop three stale `# step 3: generate attention mask`
comments left over from the original numbered-step structure (bot #6).
- pipeline_anyflow_far.py: annotate `encode_video` with
`# Copied from diffusers.pipelines.anyflow.pipeline_anyflow.AnyFlowPipeline.encode_video`
and align docstring + inline comment so `make fix-copies` keeps them in sync (bot #3).
Skipped (not real / judgment-call):
- bot #2 (private rename of `_forward_far_patchify*`) — already done in 84605d5;
bot was looking at a stale snapshot.
- bot #4 (check_inputs `# Copied from`) — FAR's check_inputs has an extra
`(num_frames - 1) % 4 == 0` constraint that doesn't map onto the bidi version,
so a clean `# Copied from` link would require restructuring. Bot called it a
consistency nit; leaving as-is.
- bot #5 (`encode_kv_cache` → `_encode_kv_cache`) — bot itself flagged this as
judgment-call territory; the helper is a coherent operation that advanced
inference callers may want to invoke directly.
---------
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
|
||
MQ
|
10302496a6 |
[feat] JoyAI-JoyImage-Edit support (#13444)
* [feat] JoyAI-JoyImage-Edit support * [fix] remove rearrange * [refactor] two pass when do cfg * [refactor] remove repa, use wantimetextembeding, refactor modulate code * [refactor] Joyimage Attention refactor * remove vae tiling and autocast * [fix] remove einops from setup.py * [refactor] Refactor JoyImageEditPipeline to use explicit arguments instead of namespace and remove _build_arg * [fix] remove deprecated method decode_latents * [refactor] refactor the image pre-processing logic into a separate VaeImageProcessor subclass * [refactor] add JoyImageAttention to align with Attention + AttnProcessor design and update conversion script for new weight key mapping (e.g. img_attn_qkv -> attn.img_attn_qkv) * [refactor] simplify bucket logic in JoyImageEditImageProcessor by replacing runtime generation with precomputed lookup tables * [fix] remove leftover training-only parameters * [fix] add layerwise casting and fp32 module patterns to JoyImageTransformer3DModel. Reference WanTransformer3DModel to fix layer casting errors during inference. * [test] add JoyImageEditPipeline fast tests and JoyImageEditTransformer3DModel model tests * [fix] fix some pipeline args to support batch inference * [fix] duplicate images to match batch size when fewer images than prompts in JoyImageEditPipeline * [fix] remove no longer used config parameters * Apply style fixes * [fix] remove unused dataclass and rewrite helpers as inline functions * [fix] make dummy objects for JoyImageEdit * [fix] allow test_torch_compile_repeated_blocks to pass * [fix] add examples on JoyImageEditPipeline * fix code style issues with ruff and black * Apply style fixes * [fix] change default num_inference_steps to 40 * [fix] use forward hook to extract pre-norm hidden states for transformers 5.x compatibility * [fix] change the assert to ValueError in pipeline * [fix] rename JoyImageTransformer3DModel to JoyImageEditTransformer3DModel, clean up anything about the alias * [fix] support gradient checkpointing * [refactor] simplify RoPE utilities, inline helpers, copy WanTimeTextImageEmbedding locally and remove unused parameters * [fix] remove _get_text_encoder_ckpt and qwen_processor * [fix] change nn.RMSNorm to FP32LayerNorm * [fix] small fixes for suggestions given by Claude * [refactor] build model using from _pretained instead of config * [refactor] auto-wrap prompt and support text-to-image in JoyImage Edit pipeline * make style, make quality and make fix-copies * [test] small fix to use vocab_size=1024 * [refactor] separate encode_prompt_multiple_images from encode_prompt, support prompt_embeds/prompt_embesd_mask/num_images_per_prompt in edit mode * [test] fix CI: use strict=False for xfail and add @require_torch_accelerator to group offloading test * [refactor] separate image_latents from latents in prepare_latents to align with flux2 * make style --------- Co-authored-by: zhangmaoquan.1 <zhangmaoquan.1@jd.com> Co-authored-by: huangfeice <huangfeice@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> Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
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>
|
||
Lancer
|
c507097764 |
[Bugfix] Fix shape mismatch in LongCatAudioDiTTransformer conversion (#13494)
Signed-off-by: Lancer <maruixiang6688@gmail.com> |
||
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> |
||
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> |
||
|
Junsong Chen
|
c6f72ad2f6 |
add ltx2 vae in sana-video; (#13229)
* add ltx2 vae in sana-video; * add ltx vae in conversion script; * Update src/diffusers/pipelines/sana_video/pipeline_sana_video.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * condition `vae_scale_factor_xxx` related settings on VAE types; * make the mean/std depends on vae class; --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
Ando
|
8ec0a5ccad |
feat: implement rae autoencoder. (#13046)
* feat: implement three RAE encoders(dinov2, siglip2, mae) * feat: finish first version of autoencoder_rae * fix formatting * make fix-copies * initial doc * fix latent_mean / latent_var init types to accept config-friendly inputs * use mean and std convention * cleanup * add rae to diffusers script * use imports * use attention * remove unneeded class * example traiing script * input and ground truth sizes have to be the same * fix argument * move loss to training script * cleanup * simplify mixins * fix training script * fix entrypoint for instantiating the AutoencoderRAE * added encoder_image_size config * undo last change * fixes from pretrained weights * cleanups * address reviews * fix train script to use pretrained * fix conversion script review * latebt normalization buffers are now always registered with no-op defaults * Update examples/research_projects/autoencoder_rae/README.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update src/diffusers/models/autoencoders/autoencoder_rae.py Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * use image url * Encoder is frozen * fix slow test * remove config * use ModelTesterMixin and AutoencoderTesterMixin * make quality * strip final layernorm when converting * _strip_final_layernorm_affine for training script * fix test * add dispatch forward and update conversion script * update training script * error out as soon as possible and add comments * Update src/diffusers/models/autoencoders/autoencoder_rae.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * use buffer * inline * Update src/diffusers/models/autoencoders/autoencoder_rae.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * remove optional * _noising takes a generator * Update src/diffusers/models/autoencoders/autoencoder_rae.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * fix api * rename * remove unittest * use randn_tensor * fix device map on multigpu * check if the key is missing in the original state dict and only then add to the allow_missing set * remove initialize_weights --------- Co-authored-by: wangyuqi <wangyuqi@MBP-FJDQNJTWYN-0208.local> Co-authored-by: Kashif Rasul <kashif.rasul@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> |
||
Dhruv Nair
|
47e8faf3b9 |
Clean up accidental files (#13202)
update |
||
|
Dhruv Nair
|
84ff061b1d |
[Modular] Save Modular Pipeline weights to Hub (#13168)
* update * update * update * update * update * update |
||
Miguel Martin
|
212db7b999 |
Cosmos Transfer2.5 Auto-Regressive Inference Pipeline (#13114)
* AR * address comments * address comments 2 |
||
|
Sayak Paul
|
2843b3d37a |
Sunset Python 3.8 & get rid of explicit typing exports where possible (#12524)
* drop python 3.8 * remove list, tuple, dict from typing * fold Unions into | * up * fix a bunch and please me. * up * up * up * up * up * up * enforce 3.10.0. * up * up * up * up * up * up * up * up * Update setup.py Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> * up. * python 3.10. * ifx * up * up * up * up * final * up * fix typing utils. * up * up * up * up * up * up * fix * up * up * up * up * up * up * handle modern types. * up * up * fix ip adapter type checking. * up * up * up * up * up * up * up * revert docstring changes. * keep deleted files deleted. * keep deleted files deleted. --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
|
Miguel Martin
|
a1816166a5 |
Cosmos Transfer2.5 inference pipeline: general/{seg, depth, blur, edge} (#13066)
* initial conversion script * cosmos control net block * CosmosAttention * base model conversion * wip * pipeline updates * convert controlnet * pipeline: working without controls * wip * debugging * Almost working * temp * control working * cleanup + detail on neg_encoder_hidden_states * convert edge * pos emb for control latents * convert all chkpts * resolve TODOs * remove prints * Docs * add siglip image reference encoder * Add unit tests * controlnet: add duplicate layers * Additional tests * skip less * skip less * remove image_ref * minor * docs * remove skipped test in transfer * Don't crash process * formatting * revert some changes * remove skipped test * make style * Address comment + fix example * CosmosAttnProcessor2_0 revert + CosmosAttnProcessor2_5 changes * make style * make fix-copies |
||
Vinh H. Pham
|
a2ea45a5da |
LTX2 distilled checkpoint support (#12934)
* add constants for distill sigmas values and allow ltx pipeline to pass in sigmas * add time conditioning conversion and token packing for latents * make style & quality * remove prenorm * add sigma param to ltx2 i2v * fix copies and add pack latents to i2v * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Infer latent dims if latents/audio_latents is supplied * add note for predefined sigmas * run make style and quality * revert distill timesteps & set original_state_dict_repo_idd to default None * add latent normalize * add create noised state, delete last sigmas * remove normalize step in latent upsample pipeline and move it to ltx2 pipeline * add create noise latent to i2v pipeline * fix copies * parse none value in weight conversion script * explicit shape handling * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * make style * add two stage inference tests * add ltx2 documentation * update i2v expected_audio_slice * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Apply suggestion from @dg845 Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update ltx2.md to remove one-stage example Removed one-stage generation example code and added comments for noise scale in two-stage generation. --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: Daniel Gu <dgu8957@gmail.com> |
||
YiYi Xu
|
61f175660a |
Flux2 klein (#12982)
* flux2-klein * Apply suggestions from code review Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Klein tests (#2) * tests * up * tests * up * support step-distilled * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * doc string etc * style * more * copies * klein lora training scripts (#3) * initial commit * initial commit * remove remote text encoder * initial commit * initial commit * initial commit * revert * img2img fix * text encoder + tokenizer * text encoder + tokenizer * update readme * guidance * guidance * guidance * test * test * revert changes not needed for the non klein model * Update examples/dreambooth/train_dreambooth_lora_flux2_klein.py Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * fix guidance * fix validation * fix validation * fix validation * fix path * space --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * style * Update src/diffusers/pipelines/flux2/pipeline_flux2_klein.py * Apply style fixes * auto pipeline --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: Linoy Tsaban <57615435+linoytsaban@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
dg845
|
c10bdd9b73 |
Add LTX 2.0 Video Pipelines (#12915)
* Initial LTX 2.0 transformer implementation * Add tests for LTX 2 transformer model * Get LTX 2 transformer tests working * Rename LTX 2 compile test class to have LTX2 * Remove RoPE debug print statements * Get LTX 2 transformer compile tests passing * Fix LTX 2 transformer shape errors * Initial script to convert LTX 2 transformer to diffusers * Add more LTX 2 transformer audio arguments * Allow LTX 2 transformer to be loaded from local path for conversion * Improve dummy inputs and add test for LTX 2 transformer consistency * Fix LTX 2 transformer bugs so consistency test passes * Initial implementation of LTX 2.0 video VAE * Explicitly specify temporal and spatial VAE scale factors when converting * Add initial LTX 2.0 video VAE tests * Add initial LTX 2.0 video VAE tests (part 2) * Get diffusers implementation on par with official LTX 2.0 video VAE implementation * Initial LTX 2.0 vocoder implementation * Use RMSNorm implementation closer to original for LTX 2.0 video VAE * start audio decoder. * init registration. * up * simplify and clean up * up * Initial LTX 2.0 text encoder implementation * Rough initial LTX 2.0 pipeline implementation * up * up * up * up * Add imports for LTX 2.0 Audio VAE * Conversion script for LTX 2.0 Audio VAE Decoder * Add Audio VAE logic to T2V pipeline * Duplicate scheduler for audio latents * Support num_videos_per_prompt for prompt embeddings * LTX 2.0 scheduler and full pipeline conversion * Add script to test full LTX2Pipeline T2V inference * Fix pipeline return bugs * Add LTX 2 text encoder and vocoder to ltx2 subdirectory __init__ * Fix more bugs in LTX2Pipeline.__call__ * Improve CPU offload support * Fix pipeline audio VAE decoding dtype bug * Fix video shape error in full pipeline test script * Get LTX 2 T2V pipeline to produce reasonable outputs * Make LTX 2.0 scheduler more consistent with original code * Fix typo when applying scheduler fix in T2V inference script * Refactor Audio VAE to be simpler and remove helpers (#7) * remove resolve causality axes stuff. * remove a bunch of helpers. * remove adjust output shape helper. * remove the use of audiolatentshape. * move normalization and patchify out of pipeline. * fix * up * up * Remove unpatchify and patchify ops before audio latents denormalization (#9) --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Add support for I2V (#8) * start i2v. * up * up * up * up * up * remove uniform strategy code. * remove unneeded code. * Denormalize audio latents in I2V pipeline (analogous to T2V change) (#11) * test i2v. * Move Video and Audio Text Encoder Connectors to Transformer (#12) * Denormalize audio latents in I2V pipeline (analogous to T2V change) * Initial refactor to put video and audio text encoder connectors in transformer * Get LTX 2 transformer tests working after connector refactor * precompute run_connectors,. * fixes * Address review comments * Calculate RoPE double precisions freqs using torch instead of np * Further simplify LTX 2 RoPE freq calc * Make connectors a separate module (#18) * remove text_encoder.py * address yiyi's comments. * up * up * up * up --------- Co-authored-by: sayakpaul <spsayakpaul@gmail.com> * up (#19) * address initial feedback from lightricks team (#16) * cross_attn_timestep_scale_multiplier to 1000 * implement split rope type. * up * propagate rope_type to rope embed classes as well. * up * When using split RoPE, make sure that the output dtype is same as input dtype * Fix apply split RoPE shape error when reshaping x to 4D * Add export_utils file for exporting LTX 2.0 videos with audio * Tests for T2V and I2V (#6) * add ltx2 pipeline tests. * up * up * up * up * remove content * style * Denormalize audio latents in I2V pipeline (analogous to T2V change) * Initial refactor to put video and audio text encoder connectors in transformer * Get LTX 2 transformer tests working after connector refactor * up * up * i2v tests. * up * Address review comments * Calculate RoPE double precisions freqs using torch instead of np * Further simplify LTX 2 RoPE freq calc * revert unneded changes. * up * up * update to split style rope. * up --------- Co-authored-by: Daniel Gu <dgu8957@gmail.com> * up * use export util funcs. * Point original checkpoint to LTX 2.0 official checkpoint * Allow the I2V pipeline to accept image URLs * make style and make quality * remove function map. * remove args. * update docs. * update doc entries. * disable ltx2_consistency test * Simplify LTX 2 RoPE forward by removing coords is None logic * make style and make quality * Support LTX 2.0 audio VAE encoder * Apply suggestions from code review Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Remove print statement in audio VAE * up * Fix bug when calculating audio RoPE coords * Ltx 2 latent upsample pipeline (#12922) * Initial implementation of LTX 2.0 latent upsampling pipeline * Add new LTX 2.0 spatial latent upsampler logic * Add test script for LTX 2.0 latent upsampling * Add option to enable VAE tiling in upsampling test script * Get latent upsampler working with video latents * Fix typo in BlurDownsample * Add latent upsample pipeline docstring and example * Remove deprecated pipeline VAE slicing/tiling methods * make style and make quality * When returning latents, return unpacked and denormalized latents for T2V and I2V * Add model_cpu_offload_seq for latent upsampling pipeline --------- Co-authored-by: Daniel Gu <dgu8957@gmail.com> * Fix latent upsampler filename in LTX 2 conversion script * Add latent upsample pipeline to LTX 2 docs * Add dummy objects for LTX 2 latent upsample pipeline * Set default FPS to official LTX 2 ckpt default of 24.0 * Set default CFG scale to official LTX 2 ckpt default of 4.0 * Update LTX 2 pipeline example docstrings * make style and make quality * Remove LTX 2 test scripts * Fix LTX 2 upsample pipeline example docstring * Add logic to convert and save a LTX 2 upsampling pipeline * Document LTX2VideoTransformer3DModel forward pass --------- Co-authored-by: sayakpaul <spsayakpaul@gmail.com> |
||
|
Miguel Martin
|
973a077c6a |
Cosmos Predict2.5 14b Conversion (#12863)
14b conversion |
||
Miguel Martin
|
b5309683cb |
Cosmos Predict2.5 Base: inference pipeline, scheduler & chkpt conversion (#12852)
* cosmos predict2.5 base: convert chkpt & pipeline - New scheduler: scheduling_flow_unipc_multistep.py - Changes to TransformerCosmos for text embeddings via crossattn_proj * scheduler cleanup * simplify inference pipeline * cleanup scheduler + tests * Basic tests for flow unipc * working b2b inference * Rename everything * Tests for pipeline present, but not working (predict2 also not working) * docstring update * wrapper pipelines + make style * remove unnecessary files * UniPCMultistep: support use_karras_sigmas=True and use_flow_sigmas=True * use UniPCMultistepScheduler + fix tests for pipeline * Remove FlowUniPCMultistepScheduler * UniPCMultistepScheduler for use_flow_sigmas=True & use_karras_sigmas=True * num_inference_steps=36 due to bug in scheduler used by predict2.5 * Address comments * make style + make fix-copies * fix tests + remove references to old pipelines * address comments * add revision in from_pretrained call * fix tests |
||
|
YiYi Xu
|
671149e036 |
[HunyuanVideo1.5] support step-distilled (#12802)
* support step-distilled * style |
||
Guo-Hua Wang
|
4f136f842c |
Add support for Ovis-Image (#12740)
* add ovis_image * fix code quality * optimize pipeline_ovis_image.py according to the feedbacks * optimize imports * add docs * make style * make style * add ovis to toctree * oops --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
|
YiYi Xu
|
6156cf8f22 |
Hunyuanvideo15 (#12696)
* add --------- Co-authored-by: yiyi@huggingface.co <yiyi@ip-26-0-161-123.ec2.internal> Co-authored-by: yiyi@huggingface.co <yiyi@ip-26-0-160-103.ec2.internal> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
Sayak Paul
|
5ffb73d4ae |
let's go Flux2 🚀 (#12711)
* add vae * Initial commit for Flux 2 Transformer implementation * add pipeline part * small edits to the pipeline and conversion * update conversion script * fix * up up * finish pipeline * Remove Flux IP Adapter logic for now * Remove deprecated 3D id logic * Remove ControlNet logic for now * Add link to ViT-22B paper as reference for parallel transformer blocks such as the Flux 2 single stream block * update pipeline * Don't use biases for input projs and output AdaNorm * up * Remove bias for double stream block text QKV projections * Add script to convert Flux 2 transformer to diffusers * make style and make quality * fix a few things. * allow sft files to go. * fix image processor * fix batch * style a bit * Fix some bugs in Flux 2 transformer implementation * Fix dummy input preparation and fix some test bugs * fix dtype casting in timestep guidance module. * resolve conflicts., * remove ip adapter stuff. * Fix Flux 2 transformer consistency test * Fix bug in Flux2TransformerBlock (double stream block) * Get remaining Flux 2 transformer tests passing * make style; make quality; make fix-copies * remove stuff. * fix type annotaton. * remove unneeded stuff from tests * tests * up * up * add sf support * Remove unused IP Adapter and ControlNet logic from transformer (#9) * copied from * Apply suggestions from code review Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: apolinário <joaopaulo.passos@gmail.com> * up * up * up * up * up * Refactor Flux2Attention into separate classes for double stream and single stream attention * Add _supports_qkv_fusion to AttentionModuleMixin to allow subclasses to disable QKV fusion * Have Flux2ParallelSelfAttention inherit from AttentionModuleMixin with _supports_qkv_fusion=False * Log debug message when calling fuse_projections on a AttentionModuleMixin subclass that does not support QKV fusion * Address review comments * Update src/diffusers/pipelines/flux2/pipeline_flux2.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * up * Remove maybe_allow_in_graph decorators for Flux 2 transformer blocks (#12) * up * support ostris loras. (#13) * up * update schdule * up * up (#17) * add training scripts (#16) * add training scripts Co-authored-by: Linoy Tsaban <linoytsaban@gmail.com> * model cpu offload in validation. * add flux.2 readme * add img2img and tests * cpu offload in log validation * Apply suggestions from code review * fix * up * fixes * remove i2i training tests for now. --------- Co-authored-by: Linoy Tsaban <linoytsaban@gmail.com> Co-authored-by: linoytsaban <linoy@huggingface.co> * up --------- Co-authored-by: yiyixuxu <yixu310@gmail.com> Co-authored-by: Daniel Gu <dgu8957@gmail.com> Co-authored-by: yiyi@huggingface.co <yiyi@ip-10-53-87-203.ec2.internal> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> Co-authored-by: apolinário <joaopaulo.passos@gmail.com> Co-authored-by: yiyi@huggingface.co <yiyi@ip-26-0-160-103.ec2.internal> Co-authored-by: Linoy Tsaban <linoytsaban@gmail.com> Co-authored-by: linoytsaban <linoy@huggingface.co> |
||
|
Junsong Chen
|
1afc21855e |
SANA-Video Image to Video pipeline SanaImageToVideoPipeline support (#12634)
* move sana-video to a new dir and add `SanaImageToVideoPipeline` with no modify; * fix bug and run text/image-to-vidoe success; * make style; quality; fix-copies; * add sana image-to-video pipeline in markdown; * add test case for sana image-to-video; * make style; * add a init file in sana-video test dir; * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana_video/test_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana_video/test_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * minor update; * fix bug and skip fp16 save test; Co-authored-by: Yuyang Zhao <43061147+HeliosZhao@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * add copied from for `encode_prompt` * Apply style fixes --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: Yuyang Zhao <43061147+HeliosZhao@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
dg845
|
d8e4805816 |
[WIP]Add Wan2.2 Animate Pipeline (Continuation of #12442 by tolgacangoz) (#12526)
--------- Co-authored-by: Tolga Cangöz <mtcangoz@gmail.com> Co-authored-by: Tolga Cangöz <46008593+tolgacangoz@users.noreply.github.com> |
||
Yashwant Bezawada
|
0fd58c7706 |
fix: correct import path for load_model_dict_into_meta in conversion scripts (#12616)
The function load_model_dict_into_meta was moved from modeling_utils.py to model_loading_utils.py but the imports in the conversion scripts were not updated, causing ImportError when running these scripts. This fixes the import in 6 conversion scripts: - scripts/convert_sd3_to_diffusers.py - scripts/convert_stable_cascade_lite.py - scripts/convert_stable_cascade.py - scripts/convert_stable_audio.py - scripts/convert_sana_to_diffusers.py - scripts/convert_sana_controlnet_to_diffusers.py Fixes #12606 |
||
|
Junsong Chen
|
b3e9dfced7 |
[SANA-Video] Adding 5s pre-trained 480p SANA-Video inference (#12584)
* 1. add `SanaVideoTransformer3DModel` in transformer_sana_video.py 2. add `SanaVideoPipeline` in pipeline_sana_video.py 3. add all code we need for import `SanaVideoPipeline` * add a sample about how to use sana-video; * code update; * update hf model path; * update code; * sana-video can run now; * 1. add aspect ratio in sana-video-pipeline; 2. add reshape function in sana-video-processor; 3. fix convert pth to safetensor bugs; * default to use `use_resolution_binning`; * make style; * remove unused code; * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * support `dispatch_attention_fn` * 1. add sana-video markdown; 2. fix typos; * add two test case for sana-video (need check) * fix text-encoder in test-sana-video; * Update tests/pipelines/sana/test_sana_video.py * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/sana/test_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/pipelines/sana/pipeline_sana_video.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/video_processor.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * make style make quality make fix-copies * toctree yaml update; * add sana-video-transformer3d markdown; * Apply style fixes --------- 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> |
||
|
YiYi Xu
|
a138d71ec1 |
HunyuanImage21 (#12333)
* add hunyuanimage2.1 --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
David Bertoin
|
dd07b19e27 |
Prx (#12525)
* rename photon to prx * rename photon into prx * Revert .gitignore to state before commit |
||
|
David Bertoin
|
cefc2cf82d |
Add Photon model and pipeline support (#12456)
* Add Photon model and pipeline support This commit adds support for the Photon image generation model: - PhotonTransformer2DModel: Core transformer architecture - PhotonPipeline: Text-to-image generation pipeline - Attention processor updates for Photon-specific attention mechanism - Conversion script for loading Photon checkpoints - Documentation and tests * just store the T5Gemma encoder * enhance_vae_properties if vae is provided only * remove autocast for text encoder forwad * BF16 example * conditioned CFG * remove enhance vae and use vae.config directly when possible * move PhotonAttnProcessor2_0 in transformer_photon * remove einops dependency and now inherits from AttentionMixin * unify the structure of the forward block * update doc * update doc * fix T5Gemma loading from hub * fix timestep shift * remove lora support from doc * Rename EmbedND for PhotoEmbedND * remove modulation dataclass * put _attn_forward and _ffn_forward logic in PhotonBlock's forward * renam LastLayer for FinalLayer * remove lora related code * rename vae_spatial_compression_ratio for vae_scale_factor * support prompt_embeds in call * move xattention conditionning out computation out of the denoising loop * add negative prompts * Use _import_structure for lazy loading * make quality + style * add pipeline test + corresponding fixes * utility function that determines the default resolution given the VAE * Refactor PhotonAttention to match Flux pattern * built-in RMSNorm * Revert accidental .gitignore change * parameter names match the standard diffusers conventions * renaming and remove unecessary attributes setting * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * quantization example * added doc to toctree * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/api/pipelines/photon.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * use dispatch_attention_fn for multiple attention backend support * naming changes * make fix copy * Update docs/source/en/api/pipelines/photon.md Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Add PhotonTransformer2DModel to TYPE_CHECKING imports * make fix-copies * Use Tuple instead of tuple Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * restrict the version of transformers Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/photon/test_pipeline_photon.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update tests/pipelines/photon/test_pipeline_photon.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * change | for Optional * fix nits. * use typing Dict --------- Co-authored-by: davidb <davidb@worker-10.soperator-worker-svc.soperator.svc.cluster.local> Co-authored-by: David Briand <david@photoroom.com> Co-authored-by: davidb <davidb@worker-8.soperator-worker-svc.soperator.svc.cluster.local> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: sayakpaul <spsayakpaul@gmail.com> |
||
Ali Imran
|
1b456bd5d5 |
docs: cleanup of runway model (#12503)
* cleanup of runway model * quality fixes |
||
|
YiYi Xu
|
f072c64bf2 |
ltx0.9.8 (without IC lora, autoregressive sampling) (#12493)
update Co-authored-by: Aryan <aryan@huggingface.co> |
||
Linoy Tsaban
|
b50014067d |
Add Wan2.2 VACE - Fun (#12324)
* support Wan2.2-VACE-Fun-A14B * support Wan2.2-VACE-Fun-A14B * support Wan2.2-VACE-Fun-A14B * Apply style fixes * test --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> |
||
|
YiYi Xu
|
a6d9f6a1a9 |
[WIP] Wan2.2 (#12004)
* support wan 2.2 i2v * add t2v + vae2.2 * add conversion script for vae 2.2 * add * add 5b t2v * conversion script * refactor out reearrange * remove a copied from in skyreels * Apply suggestions from code review Co-authored-by: bagheera <59658056+bghira@users.noreply.github.com> * Update src/diffusers/models/transformers/transformer_wan.py * fix fast tests * style --------- Co-authored-by: bagheera <59658056+bghira@users.noreply.github.com> |