mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-24 20:59:27 +03:00
diffusers-serve
165 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
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>
|
||
Jingya HUANG
|
1d5d056ecf |
[TPU] TorchTPU backend integration - eager / torch.compile / tp (#14039)
* feat: add torchtpu * feat:draft TorchTPU support * fix: wan overflow issue + compile mode error on sdxl * doc: enhance with TorchTPU doc * doc: enhance with TorchTPU doc * style: remove unused imports flagged by ruff Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs: remove Debug Eager and Fused Eager sections from tpu.md Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat: add TP for TPU * style: fix import sorting in TPU test scripts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: ruff format TPU test scripts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: style * fix: test for native 4 devices * fix: propagate TPU device fixes to Flux/Flux2/Wan-family copies; register tpu.md; drop redundant execution_device check - Propagate the text_encoder.device-based fix (introduced for TPU CPU-offload support) from FluxPipeline/Flux2KleinPipeline/WanPipeline into their `# Copied from` copies (flux/*, flux2_klein_inpaint, visualcloze, anyflow, chronoedit, lucy_edit, skyreels_v2/*). SDXL-family copies of StableDiffusionXLPipeline.encode_prompt are intentionally left untouched; they'll be handled in a follow-up PR that fixes device placement for every pipeline component (not just text encoders). - Register docs/source/en/optimization/tpu.md in _toctree.yml (was breaking the docs build: "not present in the table of contents"). - Remove the redundant "prefer non-CPU, non-meta component" loop from DiffusionPipeline._execution_device: PR #14383 already fixed this in DiffusionPipeline.device, which _execution_device falls back to. Verified on TPU hardware that _execution_device still resolves correctly for a split-placement pipeline after the removal. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * tests: cleanup * fix: fix compile mode * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * test: remove flux2 e2e test * doc: apply suggestions * doc: apply suggestions * review: remove monkey patch * review: revert neuron-specific changes in the tests * review: apply suggestions * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: add tpu to tp doc * review: remove unnecessary for tpu * test: assert every _tp_plan parameter is sharded after a TP load Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * removal: delete redundant tp shard helpers * removal: drop TPU workarounds no longer needed after recent main changes * doc: keep the FLUX.2-dev text encoder off a single TPU chip in the TP example * doc: encode the prompt on CPU in the TPU TP example so it runs end to end * test: tighten TPU TP tolerance and simplify the TPU TP worker * doc: use enable_model_cpu_offload() in the TPU eager example * doc: shard both FLUX.2-dev text encoder and transformer in the TPU TP example * doc: run the text encoders on TPU in the compiled example * test: drop the redundant TPU sync in the TP worker * test: trim comments in the TPU TP tests * test: run the TPU TP test with mp.spawn like the CUDA one * test: pin the TPU TP test to 4 chips * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: use DistributedConfig and dtype, and drop the torch_tpu import in the TPU examples * test: skip the TPU TP test for models without a _tp_plan * doc: correct the warning of recompilation * fix: clone the SD3 timestep so torch.compile doesn't recompile every step, thus we can use it for the tpu.md compile mode snippet * doc: update advice for compile mode when one chip doesn't fit * test: pass the TPU TP tolerances as test arguments * review: apply suggestions to the doc * review: view compilation patch for tpu-only --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: JingyaHuang <huang_jingya@outlok.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Christopher
|
78bf2c6266 |
[LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) (#14163)
* [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) Adds loading of LoKr (LyCORIS Kronecker product) adapters: - `load_lora_adapter` detects `lokr_` keys and injects a peft `LoKrConfig`, inferred from the tensor shapes via `_create_lokr_config` (decompose factor, per-module rank/alpha patterns). - State dict conversions for the formats in the wild: ai-toolkit Z-Image (dotted diffusers paths under `diffusion_model.`), ai-toolkit BFL Flux2 (fused qkv), LyCORIS underscore format, and bare dotted diffusers paths. - BFL fused-QKV LoKr cannot be split exactly into separate Q/K/V Kronecker factors, so `Flux2LoraLoaderMixin.load_lora_weights` fuses the model's QKV projections and maps the adapter 1:1 (exact). - Alpha follows the LyCORIS convention: scaling applies only to rank-decomposed factors and is baked into the weights at conversion. Fixes #13221 * [LoRA] LoKr review follow-ups: adapter-agnostic error message, refuse fused-QKV load over an unfused adapter - "Invalid adapter checkpoint. We currently support LoRA and LoKr." replaces the message that still said "LoRA checkpoint" and described the substring check. - Loading a fused-QKV LoKr checkpoint now refuses when an adapter is already injected on to_q/to_k/to_v or add_{q,k,v}_proj: fuse_qkv_projections() would replace those modules and orphan it. Covered by a test. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * [LoRA] rename the LoKr tests to test_lokr_conversion.py The file exercises checkpoint conversion and loading for LoKr adapters, not LoRA. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [LoRA] LoKr review follow-ups: scope to Z-Image/Flux2, model-level fused QKV, LoKrTesterMixin - Only the Z-Image and Flux2 loader mixins accept LoKr checkpoints, the two that convert them; the other mixins are back to their LoRA-only check. - The fused-QKV handling moves to `_maybe_fuse_qkv_projections_for_lokr` in peft_utils, called from `load_lora_adapter`, so the model fuses its projections whichever entry point loads the adapter. The quantized-model error explains why. - `_bake_lokr_alpha` becomes `_bake_lokr_alpha_` as it edits the dict in place. - The BFL Flux2 converter no longer accepts expanded diffusers block names under `diffusion_model.`: no LoKr checkpoint uses that layout; diffusers-named checkpoints have no prefix and take the generic converter. - The LyCORIS converter names the keys it does not recognize. - Tests move from tests/lora to a `LoKrTesterMixin` in tests/models/testing_utils, with the checkpoint formats tested on the Flux2 and Z-Image model test classes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [LoRA] mark ZImageLoraLoaderMixin.load_lora_weights as copied from Flux2 --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Sayak Paul
|
bc5e3bd9c5 |
Kandinsky6 (#14949)
* implement Kandinsky 6 - add TI2VA and SR pipelines - additonally implement PiflowScheduler and MMAudioVAE * update docs for Kandinsky 6 * remove redundant bigvgan code and manual checkpoint loading inside model code * add missing _no_split_modules * use attention processor in SR * fix unguarded dependencies - removed: einops, pydantic - lazy: torchvision, av, librosa * remove dead manual weight loading code; flattent MMAudioVAE structure * move attention processors to transformer files * remove duplicate code in K6 latent_upscaler * address 1st batch of comments * patch transformer k6 * patch SR * patch pipeline style * replace TimestepEmbedding and fix batching * refactor ti2va dit and pipeline * refactor piflow scheduler * refactor SR * patch sr transformer autocast * major refactor * untrack notebooks directory Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * patch mmaudio VAE * patch mmaudio VAE * fix inits * fix flex attn * refactoring * remove __all__ * separate MMAudio VAE and Vocoder * fix naming * refactor checkpointing, magcache, transformer, sr vae for K6 * refactor * address apply_scale_shift_norm problem * update docs * apply style chanes * address review feedback on transformer, SR VAE, and flex attention Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * address comments * fix critical bugs * Revert checkpoint-breaking module renames from |
||
|
Sayak Paul
|
771779849b |
[tests] remove test_pipelines_common.py (#14669)
* migrate imports * migrate stable diffusion adapter tests * up * move last remaining bits. * simplify test policy wording * move around * move flux ip adapter testing stuff to flux module of test * fix(ci): harden GitHub Actions workflows (#14669) (#14787) fix(ci): harden workflow files flagged on #14669 Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com> * clarify comment. * remove name from batch_input --------- Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com> |
||
|
Sayak Paul
|
c2798cc785 |
[tests] refactor the pipeline-level attention processor tests (#14701)
* refactor the pipeline-level attention processor tests * simplify stuff. * remove utils.py |
||
Dhruv Nair
|
578c9b2c66 |
Add single file support for Krea 2 (#14914)
update |
||
|
Dhruv Nair
|
4de185d6ec |
Add single file support for Minimax H3 (#14839)
update Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
yzhautouskay
|
78405ec9a1 |
[Cosmos3] Fix Transfer SeaCache artifacts with control CFG (#14897)
* Fix Cosmos 3 Transfer SeaCache indicators * Clarify SeaCache indicator guidance; tests refactor --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
fef717ffb0 |
[core] propagate sage attention updates. (#14584)
* propagate sage attention updates. * add sage blackwell. * up * style. * address feedback * style. |
||
Hidemitsu Ikeda
|
bdc2bea37a |
Fix TaylorSeer cache crash when the feature shape changes between steps (#14831)
* Fix TaylorSeer cache crash when the feature shape changes between steps * Add QwenImage 2.1 x TaylorSeer tests covering the KV-cache prefill shape change --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Dhruv Nair
|
f2edf6ea30 |
Add single file support for Qwen 2.1 (#14836)
update |
||
|
yzhautouskay
|
344d6e7300 |
Add SeaCache support for Cosmos3 pipelines (#14663)
* Add optional FP32 sampling state for Cosmos3 * Add cache execution contexts to Cosmos3 pipelines * Add SeaCache support for Cosmos3 * Cap consecutive SeaCache reuses * Align SeaCache defaults with Cosmos3 inference * Enable SeaCache by default for Cosmos3 * Document SeaCache for Cosmos3 * Use FP32 sampling state by default for Cosmos3 * Fix SeaCache reconfiguration, transfer state isolation, and FP32 sampling * Remove SeaCache stats and ablation instrumentation * Move SeaCache to transformer-level and disable by default * Align SeaCache and tests with model-level cache patterns * Make Cosmos3 sampling state always FP32 * Add pipe definition to the docs snippet * Remove redundant hook registry cache resets * Use raw vision latents for SeaCache indicators * Make Cosmos3 SeaCache safe for regional compilation * Fix SeaCache with Cosmos3 model parallelism * Support single-stream models in SeaCache * Add Wan T2V support to SeaCache * pass denoising step info to cache hooks via cache_context Replaces SeaCache's pipeline-state callbacks: the pipeline attaches the loop coordinates (step_index, sigma, num_inference_steps) to the denoiser call with `cache_context(name, ...)`, carried by a frozen CacheContext dataclass and stored on each stateful hook's StateManager (the registry no longer holds context). Co-Authored-By: Dhruv Nair <dhruv.nair@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Fix Wan cache metadata and Cosmos modular cleanup Pass denoising step metadata to Wan T2V conditional and unconditional cache contexts. Restore stateful cache reset after Cosmos modular pipeline calls, restore the required Literal import, and update regression tests and documentation. * Refactor Cosmos3 SeaCache decoder hook * Apply CI formatting updates * Fix distilled system prompt fallback * Preserve Cosmos3 fp32 sampling state --------- Co-authored-by: yiyixuxu <yixu310@gmail.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
naykun
|
6256aa7666 |
Add Qwen-Image 2.1 (#14804)
* feat: add Qwen-Image 2.1 pipeline with block-causal attention and KV cache New classes: - QwenImage21Transformer2DModel: single-stream transformer with block-causal attention and t=0 modulation for the text and condition-image prefix - AutoencoderKLQwenImage21: 64-channel VAE (z_dim=64, decoder_base_dim=144) - QwenImage21Pipeline: text-to-image and image-conditioned generation Attention: - Block-causal: the joint text/image sequence is causal while each image block (condition and target) stays internally bidirectional. Built as a compiled flex_attention BlockMask, which keeps the score matrix block-sparse and makes 2048x2048 feasible. - flex_attention is optional. Without it the mask is approximated by a two-pass prefill (prefix causally, then the target image over the cached prefix). Exact for text and for the target image, approximate for condition images. KV cache: - The text and condition-image prefix is modulated from t=0, so its activations do not change across denoising steps and its keys and values are cached after the first step. Later steps only recompute the target image's tokens. Also: separate text-to-image and image-conditioned prompt templates with image-pad token downsampling, and plain classifier-free guidance. Includes model tests, docs, and full registration. * fix style * refactor: address PR review feedback for Qwen-Image 2.1 - Remove `causal_block` config flag (always on for released checkpoint) - Split attention into QwenImage21FlexAttnProcessor and QwenImage21SDPAAttnProcessor - Replace two-pass approximate SDPA prefill with exact multi-pass prefill (each image block gets bidirectional attention, text gets causal mask) - Refactor KV cache to QwenImage21KVCache/QwenImage21KVLayerCache classes with explicit kv_cache_mode="extract"/"cached"/"extend" - Lazy-compile flex_attention on first use (fixes OOM on uncompiled path) - Fix edit pipeline: remove broken _downsample_image_pad_tokens, add mm_token_type_ids for transformers 5.x, auto-convert RGB to RGBA - Use @apply_lora_scale decorator, extract _IMG_TOKENS_PER_SLOT constant - Add # Copied from markers for retrieve_latents and _encode_vae_image - Fix mutable default feat_idx=[0] in all 8 VAE forward methods - Update docs: add usage snippet, remove stale causal_block references - Delete examples/qwenimage21/ (snippet moved to docs) * fix: KV cache pinned the whole prefill sequence at batch size 1 The "extract" branch stored the prefix as `key[:, cache_write_slice].contiguous()`. At batch size 1 that slice already counts as contiguous, because PyTorch ignores size-1 dimensions in the check, so `contiguous()` returned the same view and the cache pinned the full prefill K/V for every step of the denoising loop: 8.0 GiB of resident memory at 2048x2048 across 32 layers. At batch size 2 and above the slice is not contiguous, the copy happens, and the leak disappears, so no test caught it. Store a `clone()` instead, and assert in the tests that the cached prefix owns its storage. Measured at 2048x2048 on one H100, bf16, batch 1, 20 steps: peak memory for a full pipeline call drops from 64.5 GiB to 56.5 GiB. * move the block-causal segmentation into the SDPA processor Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WpU4nuugzCw8T1c4tL6N2A * refactor: let the processor pick the block-causal path Builds on the previous commit by @yiyixuxu, which moved the segmentation into the processor. - Drop the internal `torch.compile(flex_attention)`. The convention is that the caller compiles, so `QwenImage21FlexAttnProcessor` goes through `dispatch_attention_fn(..., backend="flex")` and warns once when it finds an uncompiled `flex_attention` — that falls back to a dense fp32 score matrix and runs out of memory at high resolution. - `_attention_backend` is `None` on the flex processor. It is what `set_attention_backend()` sets and only applies to the cached decode steps; the prefill needs the flex kernel for its `BlockMask` and is not configurable. The processor raises from `__init__` when flex_attention is unavailable. - Pad the sequence axis directly instead of transposing around `F.pad`, so the padded tensors stay contiguous, which the compiled flex kernel requires. - Rename `QwenImage21SDPAAttnProcessor` to `QwenImage21AttnProcessor` and make it the default. Once compiled, flex is 1.5% faster end to end at 2048x2048 (31.8s vs 32.3s over 20 steps) and 3.8% faster with two condition images, in exchange for 37s of compilation; uncompiled it cannot render 2048x2048 at all. A default that only works when the caller compiles is the wrong trade, so flex is documented as the opt-in path instead. - Derive the prefix segment boundaries once per forward rather than once per layer. They only depend on `image_ids` and `prefix_len`, so the per-layer version repeated the same `tolist()` device sync 32 times. `forward` passes down whichever representation the installed processors read, and builds neither for a processor that does not need it. - Replace `test_non_flex_backend_rejected_when_causal`, which no longer describes the intended behaviour, with a check that `set_attention_backend` only affects decode. Measured on one H100, bf16, batch 1, 20 steps at 2048x2048: the default path runs out of the box in 32.7s at 56.5 GiB peak, and `set_attn_processor(QwenImage21AttnProcessor())` now works at all — it used to hand the flex `BlockMask` to a non-flex kernel and raise. * fix: VAE class defaults did not describe the released model `scale_factor_spatial` was 8 while the encoder applies four spatial downsamples: `encode` takes a 1024x1024 image to a (1, 64, 1, 64, 64) latent, so the ratio is 16. Every tile-to-latent conversion divides by it, so tiling silently produced a wrong-shaped latent — a 2048x2048 encode came out as 168x168 instead of 128x128. The in/out channel defaults were 3 while this VAE takes four channels, so the class could not be instantiated from its own defaults. Add the model test file that was missing, covering the ratio against the architecture and the shape of a tiled encode. Tile values are not compared: each tile starts the causal convolution feature cache fresh, which is a property of the tiling implementation rather than of these defaults. * address review feedback: copies, pipeline, docs VAE: - Mark the classes that are byte-identical to their Wan originals with `# Copied from` (`DupUp3D`, `WanUpsample`, `WanRMS_norm`, `WanAttentionBlock`, `patchify`, `unpatchify`). Adopting the upstream `RMS_norm.forward` in the process also picks up a fix we had missed: it normalizes in fp32 for fp16/bf16/fp8 inputs. - Drop the `non_linearity` argument, which was always "silu", from the blocks that take it, and validate `AvgDown3D`'s channel divisibility with a `ValueError` before the fields are assigned rather than with an `assert` after. Pipeline: - Take `calculate_dimensions` verbatim from the edit pipeline so it can carry a `# Copied from`, which is what fixes `check_repository_consistency`: the marker on `_encode_vae_image` was one blank line out of sync. - Move the `QwenImage21KVCache` import to the top of the module. Docs: - Apply @stevhliu's suggestions. `models.autoencoders.vae.AutoencoderKLOutput` does not exist, so that autodoc reference was broken; the module is `autoencoder_kl`. - Describe the two attention processors instead of the old "with and without the flex backend" split, and add the snippet for opting into flex, which has to be compiled. `make style` also reflowed a few docstrings from the previous commit. * address review feedback: validate before doing work The `kv_cache` checks in the transformer's `forward` ran after the input projections, the joint sequence build, the rotary embeddings and the modulation, so a bad `kv_cache_mode` was only reported once that work had been done. They now run first. Following the same point through the pipeline turned up a check that could never fire: `check_inputs` warns when `height` and `width` are not divisible by `vae_scale_factor * 2`, but the rounding happened before the call, so the values it saw were always divisible. It now runs before the rounding, and `output_resolution=1000` warns and yields 992x992. * fix: num_images_per_prompt > 1 raised in the pipeline `encode_prompt` expands `prompt_embeds` and its mask to `batch_size * num_images_per_prompt`, but returns `image_pad_mask` unexpanded, and the target slots appended to that mask were sized from `latents`, which is expanded. Any call with `num_images_per_prompt > 1` therefore died in the concatenation: RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 1 but got size 2 for tensor number 1 in the list. Size the slots from each mask's own batch instead. The transformer reads the layout from row 0 because samples share it, so the mask does not need expanding. Verified end to end: `num_images_per_prompt=2`, a list of two prompts, both together, and each of those with classifier-free guidance and with a condition image. * Use the recommended sampling defaults: 40 steps, no guidance Qwen-Image 2.1 is meant to be sampled in 40 steps without classifier-free guidance, so `num_inference_steps` defaults to 40 and `true_cfg_scale` to 1.0. The other QwenImage pipelines default to 50 and 4.0, which is why this differs from its siblings. It also removes a warning from every default call: `true_cfg_scale=4.0` with no negative prompt took the "guidance is not enabled" branch. Passing a `negative_prompt` without raising `true_cfg_scale` still warns, which is the case worth warning about. The docs and the example docstring rely on the defaults now instead of passing a step count, and the docs state the recommendation. * docs: complete the forward and __call__ docstrings `utils/check_forward_call_docstrings.py` on main checks that every argument in a forward/__call__ signature has a docstring entry and that a non-None return type has a Returns section. It landed after this branch's base, so it only started running here once the copy check stopped failing ahead of it. Add the missing entries: `sample_posterior` and `generator` plus Returns on the VAE's forward, `attention_kwargs` and `return_dict` plus Returns on the transformer's, and the four embedding arguments and `callback_on_step_end_tensor_inputs` on the pipeline's __call__. * add pipeline tests for qwenimage 2.1 (#6) * add more copied froms (#5) * docs: describe multiple condition images, and inline the flex warning The pipeline page now has a section on passing several condition images, which is where @sayakpaul asked for it, and the flex section emphasises that the processor wants a compiled model. `_warn_if_flex_attention_is_uncompiled()` is inlined at its only call site, as requested. A class-level flag keeps it to one warning per process: the default processor is constructed per attention module, so 32 instances would otherwise each warn, and the logger has no `warning_once`. * fix-copies: drop the AvgDown3D marker `AvgDown3D` validates its channel divisibility with a `ValueError` before assigning its fields, where Wan still asserts after, so the copy is not consistent and `check_copies` fails on it. The other ten markers from #5 are fine and stay. * fix: feed the transformer the pre-norm text hidden state The transformer was trained on the last decoder layer's output of the text encoder, before the encoder's final RMSNorm. Up to transformers 4.x that is what `hidden_states[-1]` holds. From transformers 5.0 the output capturing ties that entry to `last_hidden_state`, so it comes back normalized instead, and the transformer reads something a third of the way off — visible first as garbled text in the rendered image. Neutralize the final norm for the encoder call with a forward hook that returns the module's input, so `hidden_states[-1]` is the layer output on either version. Nothing else is touched: the weights stay untouched, which matters because they are on `meta` under offloading, and no version check is needed. On the released checkpoint the prompt embeddings now match the pre-norm value exactly and the rendered image is pixel-identical to it, where before the whole image shifted by 5.35/255 on average. * fix: match the checkpoint's text conditioning, and repair the unexercised paths Aligned with the text encoder the checkpoint was trained with: - The image marker is `<image1>`, `<image2>`, … as in training, not `Picture 1: `. The two tokenize to different lengths (4 tokens against 5), so the conditioning the transformer read was a sequence that never occurred in training, and the rendered image moves 4.02/255 on average. With the marker corrected the output is pixel-identical to the training template. This also retires the `random.choice` over four spellings of that word, which ran on the global RNG and left image-conditioned generation irreproducible from `generator`. - Condition images reach the vision encoder with their alpha composited over white, as in training. The VAE still reads all four channels. - The processor pads on the left, as in training. The joint sequence is re-padded on the right either way, so this only changes the positions the encoder itself sees for a batch of prompts of different lengths. - An empty prompt becomes a space. Qwen has no bos token, so the encoder would otherwise have nothing to read. Prompt embeddings: - The 2D prompt mask was repeated with `repeat(1, n, 1)`, which prepends an axis and tiles the rows where the 3D embeddings interleave theirs. With more than one prompt and more than one image per prompt, each sample was denoised against another prompt's padding. - Supplying `prompt_embeds` raised: `image_pad_mask` only comes from `encode_prompt`'s own encoding, and the target slots appended to it were sized from `latents`. It is synthesized for text embeddings now, and required when the embeddings cover condition images. Supplying embeddings without a mask raised too. - `has_neg_prompt` no longer requires `negative_prompt_embeds_mask`, so a caller who passes `encode_prompt`'s own output back in keeps guidance. An unpadded prompt returns `None` for the mask, the pattern `pipeline_qwenimage.py` uses, since a mask that carries no information costs the backends that reject one. Condition images: - A tensor or ndarray `image` raised on `image.size`. They are normalized to PIL up front; a latents tensor is rejected with a message, because the text encoder has to see the image, and a per-prompt nested list with another, because one flat set applies to the whole batch. That also removes the half-wired path where a latents tensor was silently dropped from both the prompt and the latents. - A list of prompts with a condition image raised a bare `StopIteration`: every prompt's template repeats the placeholders, but the processor was handed one set of images. Denoising and validation: - Interrupting on the first step used to `continue`, skipping the step that prefills the KV cache and leaving the next one to decode from an empty one. It breaks out now. - `kv_cache_mode` without a `kv_cache` is rejected instead of failing later on a shape. - `prepare_latents` checks the generator list before spending a VAE encode per image. - The VAE's image convolution names its limitation instead of asserting: it folds the single frame away and has no temporal context, so it cannot take a feature cache. Removed: `_downsample_image_pad_tokens`, `_max_length`, the `_drop_idx_ti2i` alias, and the KV cache's `is_populated` and `clear()`, none of which anything reached. Training collapses each run of `<|image_pad|>` to one token; skipping that is equivalent here, since those positions are overwritten by the VAE latents either way. The text encoder hook now points at huggingface/transformers#48087, which lets the config untie `hidden_states[-1]` from 5.18 and will make the hook unnecessary. * fix callback test (#2) * Add a TODO at the text encoder hook It can be replaced with `tie_last_hidden_states=False` in the text encoder's config once huggingface/transformers#48087 ships in a stable release. --------- Co-authored-by: yiyixuxu <yixu310@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
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> |
||
|
Sayak Paul
|
daad99a04a |
[tests] remove quanto tests. (#14254)
* remove quanto tests. * deprecate quanto and remove tests |
||
|
Sayak Paul
|
0aa743c08d |
[tests] migrate lora tests for pipelines to use new mixins (#14268)
* start migrating lora tests. * massive refactor. * up * remove planning doc * address review: tighten lora bias tests * move run_pipe() * collapse. * up * done * update * only use diffusers api contract in tests. * fix: pass tmp_path through AuraFlow test_lora_B_bias override Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
|
Jingya HUANG
|
7d2e86a85a |
[core] Support tensor parallelism for model inference (CUDA, Neuron) (#13718)
* draft:add neuron as a legit backend
* feat: neuron-specific changes in the pipeline
* tests: eager tests
* draft: start with tp for flux2
* fix: style
* fix:apr_02 beta
* feat:add wan
* fix:pixart
* fix: rewrite flux swiglu activation to avoid gather op in neuron IR
* test: pixart compile mode on neuron
* cleanup & fix style
* merge: another change
* review: cleanup+suggestions
* fix: CIs style
* tests: add test units for tp
* fix: in case of text-encoder(s) on CPU
* review:cleanup+add test
* fix: style
* doc: remove it for now
* Add from_single_file support for SkyReelsV2 and ChronoEdit transformers (#13946)
SkyReels-V2 and ChronoEdit are both built on Wan, and their transformers have
the same keys as WanTransformer3DModel, so they reuse
convert_wan_transformer_to_diffusers (like WanVACE / WanAnimate). This lets the
community GGUF builds load directly.
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
* multi-GPU VAE Fix for Cosmos 3 (#13924)
fix(cosmos3): pin VAE latent norm buffers to encode output device
Under sharded placement (device_map="balanced"), vae.encode() runs on the
VAE's own device while the mean/inv_std buffers were pinned to x.device,
causing a cross-device RuntimeError. Compute raw_mu first, then pin the
normalization buffers to its device so all tensors share one device.
Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* docs: fix repeated word typo in set_timesteps docstring (#13876)
* docs: fix repeated word typo in set_timesteps docstring
Removed the duplicate word "schedule" from the docstring for the sigmas argument in EulerDiscreteScheduler.set_timesteps.
* Update scheduling_euler_discrete.py
* Apply style fixes
---------
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
* clean some stuff to simplify code.
* clean more to remove permutation related shenanigans.
* revert: put torch.chunk back
* Update docs/source/en/training/distributed_inference.md
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Address Sayak's TP review nits (docstrings + stale comments)
- _modeling_parallel: fix stale comment claiming attention processors read
tp_degree at runtime (no longer true after the PackedColwise/Rowwise
refactor; processors are TP-agnostic via head_dim).
- tensor_parallel: single backticks instead of double in docstrings, add
input/output examples to _blocks_to_block_sizes and _resolve_tp_plan, and
drop the stale "applies the Flux2 fused-weight permutations" line from
apply_tensor_parallel (permuters were removed).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Fix stale permuter reference in Flux2ParallelSelfAttnProcessor comment
The proportional QKV/MLP slice is now guaranteed by PackedColwiseParallel
block-by-block sharding, not the removed fused-weight permuter.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Derive TP backend inside apply_tensor_parallel from the mesh device type
Drops the `backend` parameter: apply_tensor_parallel already holds the TP
mesh, so it derives neuron-vs-default from `config._mesh.device_type` itself.
The caller no longer computes it. Documents why torch_device can't be used
(Neuron reports as "cpu" via torch's accelerator API). Removes the now-unused
is_torch_neuronx_available import from modeling_utils.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Simplify TP backend derivation to the mesh device type alone
Drop the redundant is_torch_neuronx_available() guard: a DeviceMesh("neuron", ...)
only exists when running on Neuron, so tp_mesh.device_type is the single source
of truth.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Remove incorrect "Neuron reports as cpu" rationale from TP backend docstring
Verified on the current native torch_neuronx drop that Neuron does surface as
the torch accelerator (torch._C._get_accelerator().type == "neuron",
get_device() == "neuron"); the old "reports as cpu" claim was stale. The
mesh-device-type derivation is unchanged — it's the device being sharded onto.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Drop stale "XLA" labels from Neuron device/runtime comments
Neuron is a native torch backend now (not XLA-routed): remove the "(XLA)"
parenthetical from the randn-on-CPU note in torch_utils and the "XLA runtime"
mention from the Neuron TP test docstring. Behavior unchanged.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* make style: reflow tensor_parallel docstrings
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Revert "revert: put torch.chunk back"
This reverts commit
|
||
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> |
||
|
Dhruv Nair
|
aeda394db7 |
[CI] Allow single file tests to run as fast tests (#14324)
* update * update * update * update * [FA2] Use `kernels-community/aiter-flash-attn-ck` Hub kernel for ROCm FlashAttention (drop `aiter` dependency) (#14436) Use Hub kernel Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * add HF_TOKEN to pr workflow * fix token * clean up * Apply suggestions from code review Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: BADAOUI Abdennacer <106801897+Abdennacer-Badaoui@users.noreply.github.com> |
||
|
Sayak Paul
|
425b113c21 |
[tests] refactor pipeline-level quantization tests (#14435)
* fix(torchao): route root-level tensors around safetensors reconstruction Models with parameters at the root of the module tree (e.g. Wan's `scale_shift_table`) crashed torchao's `unflatten_tensor_state_dict` when loading serialized checkpoints, since flattened tensor names are assumed to carry a module prefix. Filter such tensors and their metadata entries out of the reconstruction and merge them back unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(bnb): defer 8-bit weights split from their SCB stats across shards Sharded serialization can place an 8-bit weight and its `SCB` statistics in different shard files, in which case the shard-by-shard loader failed with "Missing quantization component `SCB`". Hold the incomplete half of the pair back until its counterpart arrives with a later shard, mirroring the torchao pending mechanism, and disable parallel shard loading for prequantized 8-bit checkpoints. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * [tests] extend model-level quantization tester mixins Migrate remaining model-level coverage from tests/quantization into the tester mixins so it runs for every wired model: - base: buffer-placement assertions in the device-map test, and an opt-in sharded-serialization test enabled by setting `sharded_serialization_config` - bnb: serialization across all configs (sharded included), dtype assignment and adapter training for 8-bit, device moves preserving the memory footprint, corrupted-state-dict loading error, and a fixed modules-to-not-convert test (BitsAndBytesConfig only exposes llm_int8_skip_modules; the old test passed an unsupported kwarg and only survived by being skipped) - torchao: custom device maps with cpu/disk offload, generalized from the Flux-specific test - gguf: the diffusers-format single-file loading path, wired into the Flux model tests Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * [tests] fix model-level quantization test wiring - SD3.5: the quantized testers reused the random-init dummy inputs (4 latent channels, fp32) while the tiny Hub checkpoint has in_channels=8 and the quantizers load the model in half precision; give them matching inputs and relax the 4-bit memory expectation for the tiny checkpoint. - QwenImage / Flux2: the quantized testers had no Hub checkpoint wired at all, so every test errored; point them at hf-internal-testing/tiny-qwenimage-pipe and tiny-flux2 with matching inputs. - NucleusMoE: no tiny checkpoint exists on the Hub yet; comment the testers out like the LTX ones. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * [tests] split tests/quantization into pipeline-level and backend-level tiers Pipeline-level quantization tests (pipeline quality slices, cpu offload, LoRA loading, compile, PipelineQuantizationConfig) move to tests/pipelines/testing_utils/quantization.py, marked per backend so the nightly CI can select them with `pytest -m`. tests/quantization keeps only backend-level tests that fit neither tier: config validation, utility warnings, and GGUF CUDA kernel correctness. Tests already covered by the model-level mixins are dropped. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * [ci] select nightly quantization tests by marker across test tiers Each backend job now runs `pytest -m <marker>` over tests/models, tests/quantization, and tests/pipelines/testing_utils/quantization.py, giving the model-level mixin tests a nightly home with the backend dependencies installed. The torchao job additionally installs mslk. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: sayak@huggingface.co <sayak@ip-10-53-90-45.ec2.internal> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
|
Sayak Paul
|
356c2fc7e5 |
[tests] skip parallelism tests for some models. (#14505)
* skip parallelism tests for some models. * clarify comment * skip others too. |
||
|
Sayak Paul
|
d5baa4fb54 |
tests: fix cuda model tests. (#14464)
* tests: fix cuda model tests. * up * up * up * add to docstrings. * restore previous beehaviour. --------- 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>
|
||
YiYi Xu
|
3a2f35d4ef |
Wan-Animate-2 (authored by @kelseyee) (#14413)
* support wan-animate-2 --------- Co-authored-by: kelseyee <971704395@qq.com> Co-authored-by: yiyi@huggingface.co <yiyi@huggingface.co> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: sayakpaul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
50e7158093 |
[tests] fix torchao tests (#14258)
fix torchao tests Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> |
||
|
Sayak Paul
|
3b68f8d087 |
[core] support cp in h3. (#14407)
* support cp in h3. * empty |
||
|
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> |
||
Álvaro Somoza
|
8b33b47332 |
[quantization] SDNQ core loading (#14277)
* initial * fix code quality * dependency test fix * apply review suggestions * add extras and version check * fix gating and prefer local cache for modular * simplify transformers registration and make the version warning silenceable * remove pipeline level test * only model level testing * layers to quantizze, change em dashes * env flag for transformers --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
jiqing-feng
|
52110fbbf6 |
Fix LoRA hot-swapping recompilation with different_shapes_for_compilation (#14297)
[tests] fix hot-swapping recompilation with different_shapes_for_compilation
`test_hotswapping_compiled_model_linear` and
`test_hotswapping_compiled_model_both_linear_and_other` failed for every model
setting `different_shapes_for_compilation`, e.g.:
RecompileError: tensor 'hidden_states' size mismatch at index 1.
expected 16, actual 32
The dummy inputs are `(batch, height * width, channels)` with `channels = 16`
and the first traced shape is `(4, 4)`, so the image sequence length equals the
channel count. Duck shaping assigns both dims the same symbol, which `img_in`
(an `nn.Linear` with constant `in_features`) then specializes to 16, forcing a
recompile on the next shape.
`use_duck_shape = False` was already set in `test_compile_on_different_shapes`
and `test_hotswapping_compile_on_different_shapes` (#11327) but the other two
multi-shape hot-swapping tests were missed. Move it into
`_check_model_hotswap` so all of them are covered, and drop the now-redundant
duplicate.
Also drop the two `xfail(strict=True)` markers on
`TestQwenImageTransformerLoRAHotSwap` that were masking this bug.
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
|
||
Atharva Joshi
|
86e6dac536 | Fix Cosmos3 Edge generator K normalization (#14246) | ||
|
Sayak Paul
|
994a92dedd |
[tests] fix cosmos3 tests (#14229)
fix cosmos3 tests |
||
|
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 |
||
|
Sayak Paul
|
5f33428194 |
update licensing for others to 2026. (#14128)
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> |
||
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>
|
||
Cheung Ka Wai
|
5be68ced49 |
fix _flash_3_varlen_hub mask handling (#14115)
* fix fa3 mask * add test case --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
72eb60c2da |
[tests] fix test_from_save_pretrained_dtype_inference (#13872)
* fix test_from_save_pretrained_dtype_inference * fix * add warning to to() * address feedback * fix * fix * fix * logger.debug -> logger.warning * Update src/diffusers/models/modeling_utils.py Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> * fix --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
|
Sayak Paul
|
9d52996c4d |
[tests] implement base model output caching in model-level tests (#14059)
* implement base model output caching in model-level tests * single quotes * memory --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> |
||
|
Sayak Paul
|
6a71b6e332 |
Fix model cuda tests (#13975)
* port final set of model tests and others * fix extracter. * fix cuda tests for models. * Revert "fix extracter." This reverts commit |
||
|
YiYi Xu
|
7104cb43c3 |
Add Krea 2 (K2) text-to-image pipeline and transformer (#14045)
Adds Krea2Transformer2DModel and Krea2Pipeline: a flow-matching text-to-image MMDiT with grouped-query attention, Qwen3-VL text conditioning, and the Qwen-Image VAE. Co-authored-by: Elea Zhong <elea@krea.ai> Co-authored-by: Naga Sai Abhinay Devarinti <devarintinagasaiabhinay@gmail.com> |
||
Hz_Zhang
|
6e8c8c4cac |
Fix Ideogram4MRoPE collapsing under torch.autocast (compute rotary in float32) (#13922)
* Fix `Ideogram4MRoPE` collapsing under `torch.autocast` (compute rotary in float32)
Ideogram4 builds image-token positions as IMAGE_POSITION_OFFSET (65536) + (t, h, w).
`Ideogram4MRoPE.forward` casts its operands to float32, but the rotary matmul (and
cos/sin) is on autocast's downcast list, so under torch.autocast("cuda", bfloat16) —
common in training and pipeline code — it runs in bfloat16 anyway. bfloat16's step at
65536 is 512, so every image position in a <=512 grid rounds to the same value: all
image tokens get identical rotary embeddings, spatial information is lost, and the
decoded image degenerates to a flat color.
Wrap the frequency computation in torch.autocast(enabled=False) so the rotary
embeddings are always computed in float32, matching how transformers guards its RoPE
modules. Added a regression test that fails on main and passes with the fix.
Fixes #13920
* Compute the rotary frequencies in float64 instead of disabling autocast
Per review: replace the torch.autocast(enabled=False) guard with a float64 computation,
which autocast does not downcast — matching the float64 rope path used elsewhere (Flux).
The autocast and float32 paths stay bit-identical (max|delta|=0).
* Disable autocast for Ideogram4 rope matmul instead of using float64
Per review, use torch.autocast(enabled=False) around the rotary matmul (as
the original implementation did) rather than computing in float64, and adopt
the clearer comment describing the bfloat16 collapse at the 65536 offset.
* Disable autocast for ernie_image and helios rope einsum
Extend the Ideogram4 fix: ernie_image's `rope` and helios's
`get_frequency_batched` build rotary freqs with an unguarded float32
einsum over raw position ids. Under an ambient autocast the einsum runs in
bfloat16 on CUDA, which cannot represent consecutive integers past 256, so
positions degrade — the same bug, matching the guards mochi/omnigen already
have. Wrap each in torch.autocast(enabled=False).
* Disable autocast for Cosmos3 VL-text rope matmul
Cosmos3VLTextRotaryEmbedding builds its interleaved-mrope freqs with an
unguarded position-id matmul (same shape as Ideogram4), so an ambient autocast
downcasts it to bfloat16 and collapses positions past 256. Wrap in
torch.autocast(enabled=False).
* Tighten autocast(enabled=False) to just the rope matmul
Per review, scope the guard to the precision-sensitive position-id matmul
in Ideogram4 and Cosmos3 rather than the whole freqs block (ernie_image and
helios already wrap only the einsum).
---------
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
|
||
|
Sayak Paul
|
5e7540fc5c |
[tests] port final set of model tests and others (#13974)
* port final set of model tests and others * fix extracter. |
||
|
Dhruv Nair
|
3d3eedf720 |
[CI] Refactor Z Image Transformer Tests (#13253)
* update * update * update * update * [CI] Refactor SD3 Transformer Test (#13340) * update * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * refactor unet tests (3d_condition, motion, controlnetxs) (#13897) * refactor unet_3d_condition tests * refactor unet_motion tests * refactor unet_controlnetxs tests * refactor unet_1d tests (#13898) * refactor unet_1d tests * use per-sample output_shape for unet_1d tests --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * refactor unet_2d tests (#13901) Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * [chore] log quant config to the user_agent (#13850) log quant config to the user_agent * Integrate AutoRound into Diffusers (#13552) * support auto_round Signed-off-by: Xin He <xin3.he@intel.com> * add document and unit tests Signed-off-by: Xin He <xin3.he@intel.com> * fix CI Signed-off-by: Xin He <xin3.he@intel.com> * Apply suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * update document and overwrite the default quantization_config with specified backend. Signed-off-by: Xin He <xin3.he@intel.com> * add UT and fix bug Signed-off-by: Xin He <xin3.he@intel.com> * update per comments Signed-off-by: Xin He <xin3.he@intel.com> * update per comments Signed-off-by: Xin He <xin3.he@intel.com> * fix compile error in doc Signed-off-by: Xin He <xin3.he@intel.com> * Apply style fixes * small nits * Add auto_round dependency to the versions table Signed-off-by: Xin He <xin3.he@intel.com> * fix make deps_table_check_updated Signed-off-by: Xin He <xin3.he@intel.com> * fix CI Signed-off-by: Xin He <xin3.he@intel.com> --------- Signed-off-by: Xin He <xin3.he@intel.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * [tests] refactor UNet model tests to align with the new pattern (#13153) * refactor unet2d condition model tests. * fix tests * up * fix * Revert "fix" This reverts commit |
||
|
jiqing-feng
|
784fa62652 |
Use device_map="auto" in single file tests to support large models on limited GPU memory (#13816)
* fix flux tests OOM on 24G GPU Signed-off-by: jiqing-feng <jiqing.feng@intel.com> * revert wrong change Signed-off-by: jiqing-feng <jiqing.feng@intel.com> --------- Signed-off-by: jiqing-feng <jiqing.feng@intel.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Xin He
|
c31bb1c148 |
Integrate AutoRound into Diffusers (#13552)
* support auto_round Signed-off-by: Xin He <xin3.he@intel.com> * add document and unit tests Signed-off-by: Xin He <xin3.he@intel.com> * fix CI Signed-off-by: Xin He <xin3.he@intel.com> * Apply suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * update document and overwrite the default quantization_config with specified backend. Signed-off-by: Xin He <xin3.he@intel.com> * add UT and fix bug Signed-off-by: Xin He <xin3.he@intel.com> * update per comments Signed-off-by: Xin He <xin3.he@intel.com> * update per comments Signed-off-by: Xin He <xin3.he@intel.com> * fix compile error in doc Signed-off-by: Xin He <xin3.he@intel.com> * Apply style fixes * small nits * Add auto_round dependency to the versions table Signed-off-by: Xin He <xin3.he@intel.com> * fix make deps_table_check_updated Signed-off-by: Xin He <xin3.he@intel.com> * fix CI Signed-off-by: Xin He <xin3.he@intel.com> --------- Signed-off-by: Xin He <xin3.he@intel.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Dhruv Nair
|
2c7efb9534 |
[CI] Refactor SD3 Transformer Test (#13340)
* update * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Dhruv Nair
|
c07f09cf9f |
[CI] Refactor Skyreels, Lumina, Ominigen, Mochi transformer tests (#13348)
* update * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Dhruv Nair
|
47538fc21d |
[CI] Refactor Chroma , LongCat and HiDream Transformer Tests (#13345)
* update * update * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Cheung Ka Wai
|
86dab152ea |
Fix the QwenImage Attention mask under Ulysses SP (#13756)
* fix mask * handle mask locally * update according to comment * fix ulysses_anything as well * Update src/diffusers/models/transformers/transformer_qwenimage.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * add accuracy test * fix controlnet * refactor worker * refactor correctness test * fix style --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
f3d42be118 |
fix chronoedit tests on PRs (#13870)
up |