mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-24 20:59:27 +03:00
refactor-pab-fastercache
2039 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Sayak Paul
|
6a8b5b7f78 | Merge branch 'main' into refactor-pab-fastercache | ||
Dhruv Nair
|
ea981b8b7d |
[CI] Clean up test_modeling_common.py (#14988)
clean up |
||
Vinh H. Pham
|
7c6eb78306 |
[hooks] Support Flux2 in MagCache (#14756)
* [hooks] Support Flux2 in MagCache MagCache could not be enabled on `Flux2Transformer2DModel` (Flux2 and Flux2 Klein): 1. `Flux2TransformerBlock` and `Flux2SingleTransformerBlock` were not registered in `TransformerBlockRegistry`, so `apply_mag_cache` raised `ValueError: Model class ... not registered`. 2. Once registered, the tail hook raised a shape error. The head hook records the image-only input of `transformer_blocks[0]`, but the Flux2 model concatenates the text and image streams before the single-stream blocks and calls them with `encoder_hidden_states=None`, so the last single block returns the fused `[text; image]` sequence. The residual `tail_output - head_input` could not be formed. Register the two Flux2 blocks (the single block with `return_encoder_hidden_states_index=None`, matching how the model calls it) and add `_align_to_hidden_states`, which keeps only the trailing image tokens of the fused tensor when computing the residual. The cached residual is therefore image-only, which is what the reference MagCache4FLUX implementation does, and it matches the head input on skipped steps. This also replaces a dead "fallback" branch in the tail hook that subtracted mismatched shapes anyway. Tests: `MagCacheTesterMixin` for `Flux2Transformer2DModel`, `Flux2Pipeline` and `Flux2KleinPipeline`. FirstBlockCache has the same head/tail structure and still fails on Flux2; left for a follow-up. * [hooks] Slice the fused tail output in the existing mismatch branch Implement the Flux2 residual fix inside the tail hook's sequence-length mismatch branch instead of a separate `_align_to_hidden_states` helper. The surrounding structure stays as it was: the `diff == 0` case still covers a batch-size mismatch via broadcasting, and only the mismatched-length case changes to subtract the trailing image tokens. --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
|
Sayak Paul
|
017c81e7b7 | Merge branch 'main' into refactor-pab-fastercache | ||
Wang, Yi
|
839af3a043 |
Update FlashAttention Hub integration and guard CUDA-only tests (#15024)
* Update FlashAttention Hub integration and guard CUDA-only tests Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * fmt Signed-off-by: Wang, Yi <yi.a.wang@intel.com> --------- Signed-off-by: Wang, Yi <yi.a.wang@intel.com> |
||
Javier De Jesus
|
08d69774f0 |
[utils] preserve method metadata in apply_forward_hook (#13803)
* [utils] preserve method metadata in apply_forward_hook Add @functools.wraps to the inner wrapper so it inherits __name__, __doc__, __module__ and the signature of the wrapped method. Without this, methods decorated with @apply_forward_hook (e.g. AutoencoderKL's encode and decode) lose their docstring and signature, so doc-builder cannot render them through the [[autodoc]] directive and they show up as "wrapper" with a generic (*args, **kwargs) signature on the live API docs. The same affects VQModel, ConsistencyDecoderVAE and the other autoencoders that use the decorator. Add regression tests asserting that the wrapped method exposes the right metadata (name, docstring, signature including default values, __wrapped__) and that the underlying offload hook still fires through the wrapper. Fixes #6575 * Trim apply_forward_hook test to a single assertion |
||
Sergio Graciá
|
916a5791ed |
[Z-Image] Apply the LoRA scale passed via joint_attention_kwargs (#15015)
Fixes #15007 Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Junsong Chen
|
1628cb1c33 |
Add SANA-WM camera-controlled image-to-video pipeline (#13881)
* feat(sana-wm): add diffusers-style SANA-WM camera-controlled I2V pipeline
Adds the public SANA-WM bidirectional camera-controlled image-to-video
model as a first-class diffusers pipeline + transformer. Layout mirrors
``sana_video``: the model lives under ``src/diffusers/models/transformers/``
as a near-single-file (kernels split off so the ``@triton.jit`` decorators
don't drown the model body); the pipeline lives under
``src/diffusers/pipelines/sana_wm/``.
Files added:
src/diffusers/models/transformers/
├── transformer_sana_wm.py # SanaWMTransformer3DModel + blocks + helpers
└── transformer_sana_wm_kernels.py # fused Triton kernels + camera math
src/diffusers/pipelines/sana_wm/
├── __init__.py
├── pipeline_sana_wm.py
├── pipeline_output.py
├── refiner.py
└── cam_utils.py
Pipeline architecture:
* Stage 1: 1600M ``SanaWMTransformer3DModel`` DiT with bidirectional
GDN-Triton linear attention + UCPE camera-control branch, LTX-style
flow-matching Euler scheduler with per-token timesteps.
* Stage 2: LTX-2 sink-bidirectional Euler refiner (3 distilled sigma
steps, reuses diffusers' ``LTX2VideoTransformer3DModel`` +
``LTX2TextConnectors`` + Gemma-3 text encoder).
* Decode through the LTX-2 VAE (``AutoencoderKLLTX2Video``).
One-line usage:
pipe = SanaWMPipeline.from_pretrained(
"Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
torch_dtype=torch.bfloat16,
).to("cuda")
out = pipe(image=img, prompt="...", action="w-80,jw-40,w-40",
intrinsics=[fx, fy, cx, cy])
End-to-end smoke test (stage-1 + refiner + VAE decode) passes on H100.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* feat(sana-wm): align pipeline with merged sana_video style; fix mp4 export
transformer_sana_wm.py:
* License header switched to the "HuggingFace Team and SANA-WM Authors"
style used by merged sana_video.
* Imports rewritten in stdlib -> third-party -> diffusers order; use
diffusers `from ...utils import logging` instead of stdlib `logging`.
* Fix 9 `Optional[X]` annotations written as `X or None` (Python's `or`
short-circuits and silently returns `X`).
* Fix two `assert (cond, msg)` tuple-asserts in PatchEmbedMS3D.forward
that always pass (SyntaxWarning at import time).
* Remove duplicate `__all__` declarations (the second silently overwrote
the first).
* Remove dead `reset_bn` (imports a nonexistent `packages.apps.utils`,
would crash on call).
* Remove the duplicate `logger = logging.getLogger(__name__)` further
down in the file.
transformer_sana_wm_kernels.py:
* License header normalized; collapse three duplicate triton/torch import
blocks into one.
pipeline_sana_wm.py:
* License header normalized.
* `_decode_latents` now returns `(T, H, W, 3)` float in [0, 1], matching
the diffusers convention used by `VideoProcessor`. Returning uint8
silently broke `export_to_video`: it does `frame * 255` assuming float
input, so uint8 overflows to `(-x) mod 256` and inverts colors.
* `__call__` converts to PIL/uint8 only when `output_type="pil"`.
* Intrinsics argument now accepts (4,), (F, 4), (3, 3), and (F, 3, 3)
forms (auto-extracts fx, fy, cx, cy from a 3x3 K) and auto-trims to
`num_frames` when a longer-than-needed trajectory is passed.
* Inline `retrieve_timesteps` with the standard `# Copied from
diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps`
marker, matching merged sana_video.
* Docstrings + EXAMPLE_DOC_STRING updated to reflect the new return type.
pipeline_output.py:
* Update `frames` field docstring to describe the new float [0, 1] return.
refiner.py, cam_utils.py, scripts/sana_wm/convert_sana_wm_to_diffusers.py:
* License headers normalized.
Docs:
* New `docs/source/en/api/pipelines/sana_wm.md` and
`docs/source/en/api/models/sana_wm_transformer3d.md`, modeled on
sana_video.md / sana_video_transformer3d.md, wired into
`docs/source/en/_toctree.yml` under Models and Pipelines.
5s end-to-end smoke test (81 frames @ 16fps, 30 stage-1 steps + 3-step
LTX-2 refiner) passes on 1x H100 80GB with `enable_model_cpu_offload`.
Round-trip diff vs raw float frames is 2.06/255 mean (h264 lossy noise),
confirming the export_to_video fix.
* feat(sana-wm): port chunk-causal AR refiner mode (RefinerChunkRunner + KV cache hooks)
The first cleanup pass only kept the legacy single-shot refiner path. That
path is what the model was *not* trained on — its docstring even says
"feeding the full sequence at once is out-of-distribution" — and its cost
is O(T^2) attention over the full latent volume, which made longer videos
unusable (~21 min per refiner step at 321 frames on an H100).
Port the chunk-causal AR mode from the upstream reference so the refiner
matches the training contract:
* `refine_latents` now defaults to `block_size=3, kv_max_frames=11`
(the canonical AR recipe). Pass `block_size=None` to fall back to the
legacy single-shot path.
* New `_refine_latents_ar` + `_RefinerChunkRunner` orchestrate the sliding
window: pre-capture pre-RoPE sink K/V on `z_sana[:source_sink_frames]`
at sigma=0, then for each `block_size`-frame chunk run a 3-step Euler
with prefix `{sink_k_pre, sink_v, sink_pe, history_k, history_v}` and
capture post-RoPE K/V to feed the next window. History is bounded to
`kv_max_frames - source_sink_frames` so per-block compute is constant.
* New `_predict_x0_active_block` runs the transformer on the active block
only (Q from active, K/V from prefix+active).
* New `_capture_block_kv` runs sigma=0 forward with a pre_rope/post_rope
capture flag set on each `attn1`.
* New `_forward_video_only_with_rope` takes a pre-built RoPE so each block
can use absolute frame positions in the source video.
* `_streaming_self_attention` extended with the `_kv_cache_capture`,
`_tf_capture_kv`, `_tf_kv_prefix` hook contract that AR mode uses to
inject and capture K/V on each block.
* New helpers: `_build_rotary_emb_for_absolute_positions`,
`_set_kv_prefix_on_blocks`, `_clear_kv_prefix_on_blocks`,
`_set_capture_flag_on_blocks`, `_collect_captured_kv_from_blocks`.
* `_encode_prompt` now also moves the Gemma-3 text encoder back to CPU
after producing the embeds — otherwise it stays resident through the
entire AR loop and gates how much GPU memory the refiner transformer
has left.
Module-level docstring updated to document both modes; existing
single-shot path preserved verbatim.
* feat(sana-wm): block-level checkpoint for AR refiner (resume after preemption)
The AR refiner is expensive (~3-5 min per block) and the refinement loop
ran end-to-end has no in-progress state to recover, so a SLURM preemption
mid-refinement loses all progress. With the canonical
``block_size=3, kv_max_frames=11`` setup, refining a 50s video is 34
blocks of work that has to make it through without preemption on a
backfill queue.
Add per-block atomic checkpointing:
* ``SanaWMLTX2Refiner.refine_latents(checkpoint_dir=Path)`` and
``_refine_latents_ar`` accept a directory. After each completed AR
block, the AR loop writes ``checkpoint_dir/state.pt`` atomically
(tmp + os.replace).
* The payload is ``{block_idx_done, n_blocks, sink_size, block_size,
output_shape, output, runner_state}``. ``runner_state`` is a CPU snapshot
of the runner's ``_sink_kv_pre``, ``_history_kv_post``,
``_history_frames`` and ``torch.Generator`` state.
* On entry, if ``state.pt`` exists with a compatible shape signature, the
AR loop loads the persisted output tensor + runner state and resumes
from ``block_idx_done + 1`` instead of recomputing from scratch.
* ``SanaWMPipeline.__call__(refiner_checkpoint_dir=...)`` plumbs the
directory through to the refiner.
Checkpoint size: ~output_volume + sink_KV (~360MB for 50 layers) +
rolling history KV (~3-4GB at full capacity) — saved once per block,
total per-block save overhead ~10s on lustre.
* test(sana-wm): add CPU unit tests + slow GPU integration stub
* CPU unit tests for cam_utils helpers (action DSL → c2w, intrinsics
rescale-for-crop, resize+center-crop, snap_num_frames 8k+1 rounding).
* Public-surface registration tests (top-level diffusers symbols,
SanaWMPipelineOutput dataclass shape, refiner signature has AR defaults
+ checkpoint_dir, pipeline __call__ accepts c2w/action/intrinsics/
refiner_checkpoint_dir).
* @slow @require_torch_accelerator integration stub for an end-to-end I2V
against the public checkpoint, currently @unittest.skip — wires up the
nightly GPU path without exploding regular CI.
SanaWMTransformer3DModel has hardcoded depth/hidden_size/num_heads inside
its inner SanaMSVideoCamCtrl (not exposed through register_to_config), so
the usual PipelineTesterMixin small-config fast tests aren't applicable
without a transformer refactor (followup PR).
* feat(sana-wm): make triton optional + auto-fallback to pure-PyTorch attention
`transformer_sana_wm_kernels.py` previously did a hard `import triton`
at the top of the file. That blocked importing the SANA-WM transformer
on any environment without Triton (CPU-only, ROCm without Triton,
older Triton, etc.), even though the model has pure-PyTorch attention
classes for every `*Triton` variant.
Make Triton optional and have the dispatcher transparently fall back:
* Wrap `import triton` / `import triton.language as tl` in try/except.
When unavailable, install a shim where `@triton.jit` is a no-op so
the kernel function definitions still load (they just aren't compiled
by Triton). Module-level `triton.X` / `tl.X` lookups return a
self-shimming sentinel so signature parsing doesn't blow up either.
* Add `is_triton_available()` + `_require_triton(entry_point)`. The four
Triton-backed entry points called by the model (`fused_qk_inv_rms`,
`fused_bigdn_func`, `cam_prep_func`, `cam_scan_bidi_chunkwise`) now
raise a clear RuntimeError on a Triton-less host with a hint to use
the pure-PyTorch attention variants — but the dispatcher does this
automatically (see below) so users shouldn't ever see it.
* Delete the leftover duplicate `import torch / triton / triton.language`
block at line 262 (left over from the upstream port).
* Register `BidirectionalGDNUCPESinglePathLiteLA` in `ATTENTION_BLOCKS`
so the fallback chain can find it.
* New `_resolve_attention_block(name, role)` walks the requested class's
MRO at dispatch time. If Triton isn't usable AND the requested class
name ends in `Triton`, route to the closest registered non-`Triton`
ancestor (BidirectionalGDNUCPESinglePathLiteLABothTriton ->
BidirectionalGDNUCPESinglePathLiteLA, etc.) and log a one-shot warning.
* Rewire both `SanaVideoMSCamCtrlBlock` dispatch sites to use
`_resolve_attention_block` for the GDN+UCPE camera branch and the main
attention branch (the `BidirectionalSoftmaxUCPESinglePathLiteLA` branch
doesn't use Triton at all so it stays hard-coded).
Tests:
* `test_kernels_module_imports_with_triton_hidden` — reloads the kernels
module with `sys.modules['triton'] = None` and verifies the module
imports, `is_triton_available()` is False, and the pure-PyTorch helpers
remain callable.
* `test_resolve_attention_block_cpu_fallback` — on a CPU-only host, the
three `*Triton` attn types resolve to the correct non-Triton ancestor.
* `test_triton_entry_point_raises_clean_error_without_triton` — verifies
the `_require_triton` guard yields a RuntimeError that mentions Triton.
* fix(sana-wm): make optional deps lazy + register transformer in __init__
Three CI checks were failing on the PR:
1. `check_code_quality` (43 ruff errors): mix of unused imports / import
sorting / E731 lambdas (auto-fixable) plus a handful of F821 dead-code
references inherited from the upstream research codebase (`xformers.*`
inside `if _xformers_available:` blocks, an undefined `BlockHook` type
annotation, two `x_sa`/`mlp_out` references in a block forward whose
live assignment was already overridden by subclasses). Ran `ruff check
--fix --unsafe-fixes` + `ruff format`, fixed the type annotation
manually, and added targeted `# noqa: F821` markers on the conditionally
unreachable lines.
2. `check_torch_dependencies`: `transformer_sana_wm.py` hard-imported
`einops`, `fla`, `timm`, `termcolor`. The minimum-deps CI environment
doesn't have them, and diffusers' lazy loader rewrites `ModuleNotFoundError`
as `RuntimeError` so `test_pipeline_imports` blew up. Wrapped each of
the four optional imports in a try/except shim — `rearrange`/
`ShortConvolution`/`DropPath`/`Attention_`/`Mlp` become placeholders
that raise a clear `ImportError` on construction, `colored` falls back
to plain text. Class bodies that subclass these still parse at module
load, so `import diffusers.models.transformers.transformer_sana_wm`
succeeds anywhere. Same treatment for the kernels file's
`from einops import rearrange, repeat`.
3. `build_pr_documentation`: doc-builder imported `SanaWMTransformer3DModel`
from `diffusers.models.transformers` (not the diffusers top level) and
that subpackage's `__init__.py` was missing the entry. Added the import.
* style(sana-wm): apply make style + fix-copies
* `doc-builder style src/diffusers docs/source --max_len 119` rewraps
docstrings in the six SANA-WM files (transformer, kernels, pipeline,
refiner, output, cam_utils) to the repo-wide 119-column limit. No
behaviour change — purely whitespace inside docstrings.
* `make fix-copies` regenerates `dummy_pt_objects.py` and
`dummy_torch_and_transformers_objects.py` to add `DummyObject` stubs
for the three new public classes (`SanaWMTransformer3DModel`,
`SanaWMPipeline`, `SanaWMLTX2Refiner`), so `from diffusers import …`
gives the standard "missing backend" message on installs without
torch / transformers.
Verified: `make quality` passes (ruff check, ruff format check,
doc-builder style check_only, check_doc_toc). Test suite still
15 passed / 1 skipped.
* refactor(sana-wm): drop einops dependency, inline with torch ops
Per @dg845's review on `transformer_sana_wm.py:44`. The 9 call sites all
match well-known patterns that have one-liner torch equivalents:
rearrange(s, "b d -> (b d)") -> s.reshape(b * d)
rearrange(x, "(b d) d2 -> b (d d2)", ...) -> x.reshape(b, d * d2)
rearrange(R, "b t h w i j -> b t h w j i") -> R.transpose(-1, -2)
repeat(x, "b h w c -> b t h w c", t=T) -> x.unsqueeze(1).expand(-1, T, -1, -1, -1)
repeat(x, "b t c -> b t h w c", h=H, w=W) -> x[:, :, None, None, :].expand(-1, -1, H, W, -1)
repeat(x, "... -> b ...", b=B) -> x.unsqueeze(0).expand(B, *x.shape)
repeat(x, "H W C -> B T H W C", B, T) -> x[None, None].expand(B, T, -1, -1, -1)
repeat(x, "B H W C -> B T H W C", T) -> x.unsqueeze(1).expand(-1, T, -1, -1, -1)
Each replacement is bit-identical to the einops original — verified
against a fresh `einops` install on random tensors before swapping. The
optional `from einops import ...` shim block is gone from both
`transformer_sana_wm.py` and `transformer_sana_wm_kernels.py`.
* refactor(sana-wm): remove dead code per @dg845's review
Drop 23 unused symbols (1451 lines) from `transformer_sana_wm.py` that
aren't reachable from the public SANA-WM checkpoint's
`SanaMSVideoCamCtrl` -> `SanaVideoMSCamCtrlBlock` -> GDN/UCPE attention
path. Each was verified to have zero call sites outside of its own
definition (or only within other-deleted items).
Classes:
* `SanaMS`, `SanaMSBlock` — alternative `Sana` subclass + block, not
used by the SANA-WM checkpoint (which goes through `SanaMSVideoCamCtrl`).
* `ChunkCausalAttention`, `CachedCausalAttention`,
`ChunkedLiteLAReLURope`, `LiteLAReLURope` — chunk-causal / cached
attention variants and their common base; SANA-WM uses the bidi GDN
path. `LiteLAReLURope` had only the three (now-deleted) subclasses
referencing it.
* `PAGCFGIdentitySelfAttnProcessorLiteLA`,
`PAGIdentitySelfAttnProcessorLiteLA`, `SelfAttnProcessorLiteLA`,
`SelfAttnProcessorLiteLAReLURope` — PAG processors; we don't expose
PAG in the SANA-WM pipeline.
* `ChunkGLUMBConvTemp`, `CachedGLUMBConvTemp`, `MBConvPreGLU` —
alternative FFN/conv blocks; the checkpoint uses `GLUMBConvTemp`.
* `MaskFinalLayer`, `DecoderLayer` — alternative final layers; the
checkpoint uses `T2IFinalLayer`.
* `LabelEmbedder`, `CaptionEmbedderDoubleBr` — alternative embedders;
the checkpoint uses `CaptionEmbedder`.
Helpers:
* `set_grad_checkpoint`, `prepare_prompt_ar`, `resize_and_crop_tensor`,
`generate_temporal_head_mask_mod`, `is_chunk_causal_request`,
`get_chunk_index_from_config` — training-only or
chunk-causal-mode helpers with no inference call sites.
Verified after the diff:
* `make quality` clean.
* CPU test suite 15 passed / 1 skipped (slow GPU integration).
* `import diffusers ; SanaWMPipeline / SanaWMTransformer3DModel /
SanaWMLTX2Refiner` resolve identically.
* fix(sana-wm): defer transformers import + document mask/return_dict
Two PR-CI failures, both ours:
* `check_torch_dependencies`: line 33 hard-imported `transformers`, which
isn't present in the minimum-deps environment. Move the lone
`AutoModelForCausalLM` use site inside `initialize_gemma_params` (a
training-only helper).
* `check_repository_consistency`: `SanaWMTransformer3DModel.forward`'s
docstring was missing entries for `mask` and `return_dict`. Added.
* refactor(sana-wm): address review feedback on SanaWMPipeline
Per @dg845's review comments on pipeline_sana_wm.py:
* Read the VAE spatial/temporal strides from the ``vae`` component
(``vae_spatial_compression_ratio`` / ``vae_temporal_compression_ratio``)
instead of hardcoding 32/8, mirroring LTX2Pipeline.
* Inline the stage-1 sampling loop into ``__call__`` and factor the noise
init into a standard ``prepare_latents`` method.
* Use ``self.scheduler`` for the flow-matching Euler steps instead of
constructing a new scheduler per call.
* Drive the sampling loop with ``self.progress_bar`` (respects
``set_progress_bar_config``) instead of a bare tqdm.
* Move input validation/normalization into a ``check_inputs`` method.
* Move first-frame resize+center-crop, [-1, 1] normalization and the
intrinsics rescale into a ``SanaWMImageProcessor(VaeImageProcessor)``
subclass (new ``image_processor.py``).
* Add ``generator`` as a ``__call__`` argument (``seed`` kept as a
convenience shortcut).
* Post-process decoded latents with ``VideoProcessor.postprocess_video``
rather than a hand-rolled conversion; drop the ``.cpu()`` casts on the
``output_type="latent"`` path.
* Move the pipeline README into the docs (docs/.../sana_wm.md); delete
the in-package README.
* refactor(sana-wm): make SanaWMLTX2Refiner a standalone DiffusionPipeline
Per @dg845's review: the refiner has components (transformer, connectors,
text encoder, tokenizer) and runs a denoising loop, so it fits better as a
pipeline than a ModelMixin.
* SanaWMLTX2Refiner now subclasses DiffusionPipeline and registers its
components via ``register_modules`` (dropping the bespoke
``from_pretrained`` / ``save_pretrained``); standard load/save now handle
the ``refiner/`` subfolder.
* Add a ``FlowMatchEulerDiscreteScheduler`` component (shift=1.0) and drive
the Euler steps through ``scheduler.step`` / ``scheduler.scale_noise``
(single-shot and per-AR-block) instead of hand-rolled updates. Numerically
equivalent to the previous flow-matching update.
* Rename the entry point ``refine_latents`` -> ``__call__``; add a ``device``
arg so the parent can hand it the execution device without a bulk move.
* SanaWMPipeline: keep the refiner as an optional nested component; free the
parent's GPU weights before running it (it manages its own sub-module
placement) and bring the VAE back for decode.
* Conversion script emits the new refiner layout (model_index.json +
scheduler/ + tokenizer/); test asserts the refiner is a DiffusionPipeline
with the canonical AR ``__call__`` defaults.
Validated end-to-end on 1xH100 (stage-1 + refiner on the official demo,
coherent video output).
* docs(sana-wm): document the refiner __call__ `device` arg
Fixes the check_repository_consistency failure — `SanaWMLTX2Refiner.__call__`
gained a `device` parameter but its docstring wasn't updated.
* refactor(sana-wm): clean up dead code + reuse shared utils in the DiT
Addresses several of @dg845's transformer review comments (the safe,
output-preserving subset — GPU smoke gives byte-identical output):
* Reuse diffusers' shared `FP32LayerNorm` (models/normalization.py) and
`get_1d_rotary_pos_embed` (models/embeddings.py); delete the local copies.
* Remove dead inference code paths:
- the `if self.diagonal_mask is not None:` flex-attention block
(`diagonal_mask` is always `None`) + the now-unused
`create_block_mask_cached` helper and `create_block_mask` import;
- the `SANA_FSDP2_BLOCK_TIMING` block-timing/profiling scaffolding;
- the `save_qkv` / `qkv_store_buffer` visualization hooks (never
enabled at inference), at both the attention and model level.
* Collapse `SanaVideoMSCamCtrlBlock.forward_frame_aware` into `forward`
(the pipeline/refiner always pass >=3D timesteps, so the non-frame-aware
branch was dead — it even referenced undefined locals).
* Drop the 191-line `SanaMSVideoCamCtrl.load_state_dict` shape-remapping
override — it's never reached by the shipped convert/inference flow
(`nn.Module.load_state_dict` on the wrapper doesn't call it), and the
release checkpoint already ships correctly-shaped weights.
* Delete the unused `PatchEmbedMS` module.
Net -503 lines; no numerics change (stage-1 + refiner smoke output mean
identical to the pre-refactor run).
* refactor(sana-wm): merge the 3 DiT classes into one SanaWMTransformer3DModel
Per @dg845's review: `SanaWMTransformer3DModel` was a thin wrapper over
`SanaMSVideoCamCtrl`, which in turn subclassed `Sana` and overwrote most of
it. Fold all three into a single `SanaWMTransformer3DModel(ModelMixin,
ConfigMixin)`:
* The `@register_to_config` __init__ signature is unchanged (so config.json
is identical); the body builds the modules directly on `self` instead of
a nested `self._inner`.
* Only the surviving `Sana.__init__` pieces are kept (t_embedder,
cfg_embedder, attention_y_norm, config attrs, initialize_weights); the
parts the subclass overwrote are gone.
* `forward` takes the diffusers signature (hidden_states / timestep /
encoder_hidden_states / encoder_attention_mask / return_dict) and returns
`Transformer2DModelOutput`, folding in the old wrapper's arg-renaming.
* Deleted now-unreachable code: `class Sana`, `class SanaMSVideoCamCtrl`,
`class SanaBlock`, `class PatchEmbed` (the video model uses
`PatchEmbedMS3D`), the `add_inner_prefix` helper, and the dead
`sincos`/`flux_rope` pos-embed branches in forward (release uses
`wan_rope`). Net -468 lines.
This drops the `_inner.` state-dict prefix, so the conversion script no
longer adds it. State-dict is otherwise identical: the merged model's
`state_dict()` has exactly the same 871 param keys as before (verified),
and a stage-1 + refiner GPU smoke on the public checkpoint gives
byte-identical output (frame mean 0.5560, matching the pre-merge run).
Addresses the class-merge review comment and removes the unused
`SanaBlock`/`PatchEmbed` modules.
* refactor(sana-wm): drop the no-op custom grad-checkpoint wrapper
`auto_grad_checkpoint` gated on a `grad_checkpointing` attribute that was
never set, so it always fell through to `module(*args, **kwargs)` — i.e. a
no-op. Call the transformer blocks directly instead and delete the unused
`auto_grad_checkpoint` / `checkpoint_sequential` helpers and the
`torch.utils.checkpoint` import.
(Inference is unchanged — the block never read the `use_reentrant` kwarg
the wrapper passed. Full training-time gradient checkpointing via the
standard `_gradient_checkpointing_func` would need the block's dynamic
kwargs — camera_conditions / prope_fns / chunk_index — threaded through,
so it's left as a follow-up for this inference-focused release.)
* refactor(sana-wm): drop the fla-core and termcolor dependencies
Per @dg845 / @sayakpaul: the SANA-WM transformer required `fla-core` to be
constructible at all (and fla's `ShortConvolution` cannot even run on CPU —
its dispatch does `torch.cpu.device(...)`, which doesn't exist).
`ShortConvolution` was only ever used as a depthwise *causal* conv1d with
`activation=None`, so it is replaced by a ~20-line self-contained PyTorch
module with the same parameter layout (`weight` of shape `(C, 1, K)`, no
bias) and the same `(output, cache)` return signature. No `fla-core` and no
`kernels` dependency is needed now, and the model can be built and run on
CPU. Verified on an H100 in bf16 at the real config (hidden 2240, kernel 4):
`max|Δ|` vs fla is exactly `0.0` for every shape tested, and the state dict
keeps the identical 871 keys.
Also:
* Use the existing `is_timm_available()` utility for the timm imports. The
`else` branch keeps a placeholder class rather than raising at module
scope, because several layers subclass these symbols and this module is
imported eagerly by `diffusers.models` — a module-level raise would break
plain `import diffusers` (and the `check_torch_dependencies` CI job) when
timm isn't installed. The error is instead raised on construction.
* Drop `termcolor` entirely: replace the `self.logger = print` +
`colored(...)` pattern with the module-level `logger`, so there is no
optional dependency left to gate.
* Only attempt the optional null-embedding load when `null_embed_path` is
actually set (it is unset for the public checkpoint, so this previously
logged a spurious failure on every construction) and load it with
`weights_only=True`.
* fix(sana-wm): address review nits in the pipeline / refiner / cam utils
* refiner: load the AR resume checkpoint with `weights_only=True`. The
payload is only tensors / ints / tuples / dicts plus the generator's uint8
state, so it round-trips safely under the restricted unpickler.
* refiner: fix `_tokens_per_frame` to divide by `patch_size_t` rather than
multiply — `_pack_latents` emits `(T // patch_size_t) * (H // p) * (W // p)`
tokens, so one latent frame contributes `(H // p) * (W // p) / patch_size_t`.
No-op for LTX-2 (`patch_size_t=1`) but the history trimming would have kept
too many tokens otherwise.
* pipeline: replace `torch.cuda.empty_cache()` with diffusers'
backend-agnostic `empty_device_cache(device.type)`.
* pipeline: extract the offload probe into `_model_cpu_offload_active()`,
matching the `hasattr(self, "_all_hooks") and len(self._all_hooks) > 0`
idiom `DiffusionPipeline` uses internally, instead of a bare truthiness
check on the attribute.
* pipeline_output: add the missing Apache-2.0 license header.
* cam_utils: drop the unnecessary `+ 1e-6` when normalizing the forward /
right vectors — the branch is only taken when the norm is already > 0, so
the epsilon just introduced a small systematic bias.
* refactor(sana-wm): strip research-repo code from the DiT, fix model conventions
From the `self-review` skill run against `.ai/models.md` / `.ai/AGENTS.md`.
`transformer_sana_wm.py`: 6555 -> 4592 lines.
Conventions:
* `_no_split_modules` was `["blocks"]` — an attribute name, but accelerate
matches on class name, so it matched nothing and `device_map="auto"` could
split a block across devices and crash. Now `["SanaVideoMSCamCtrlBlock"]`.
* Add `_repeated_blocks` (enables `compile_repeated_blocks()`) and
`_skip_layerwise_casting_patterns`. `_keep_in_fp32_modules` is deliberately
left unset with a note: the blocks apply the timestep modulation inline, so
keeping `scale_shift_table` / `t_embedder` in fp32 upcasts the hidden states
and feeds fp32 activations to bf16 weights (caught by a GPU smoke run).
* Expose `num_layers` / `hidden_size` / `num_attention_heads` / `patch_size`
through `register_to_config` instead of hardcoding the release architecture,
so a tiny variant can be built for tests. Defaults are the released values,
so `config.json` and the state dict are unchanged.
* `torch.float64` -> `torch.float32` on the live RoPE paths (gotcha 5),
`torch.empty` -> `torch.zeros` for parameter init (gotcha 6), and stop
reading `self.proj.weight.dtype` to cast activations (gotcha 4).
* `WanRotaryPosEmbed.freqs` is a non-persistent buffer instead of a plain
attribute reassigned inside `forward` (which broke `.to()` and compile).
Dead code (`AGENTS.md`: "delete training-time code paths, experimental flags,
and ablation branches entirely — only keep the inference path"):
* All weight-init / transfer-learning helpers — `from_pretrained` overwrites
them, and they also printed ~20 lines of noise on every construction.
* `CaptionEmbedder.initialize_gemma_params` (fetched `google/gemma-2b-it`
out-of-band at runtime, and referenced an attribute that never exists),
`token_drop`, and the training branch of its forward.
* The cam-debug statistics apparatus, the `save_block_output` hooks (whose
`get_block_output` was never defined), and `block_hook`.
* Both unreachable recurrence variants, `_maybe_drop_cam_branch`, the xformers
branches (`_xformers_available` was a literal `False`, defined twice, and
`xformers` was never imported), 3 env-var escape hatches, 9 unused classes,
12 unused module-level helpers, and the sincos family — which also removes
the NumPy import, satisfying the no-NumPy-in-forward rule.
* Inline the `fp32_attention` mechanism: it was set on every submodule via
`model.apply` and read at 18 sites, so it was folded to the shipped
always-on behaviour rather than deleted.
Ephemeral comments (commit SHAs from a private tree, references to files that
don't exist in diffusers, FSDP2 rationale, stale docstrings) removed.
State dict is unchanged: 871/871 keys match the released checkpoint.
* fix(sana-wm): honour scheduler shift, use randn_tensor, drop cluster scaffolding
From the `self-review` skill run against `.ai/pipelines.md` / `.ai/testing.md`.
`flow_shift` was a dead knob. `__call__` set `self.scheduler.config.shift`,
but `FlowMatchEulerDiscreteScheduler.set_timesteps` reads `self.shift` (i.e.
`self._shift`), so the documented argument had no effect and every run used
the checkpoint's `shift=9.8`. It also half-mutated a `FrozenDict` — the freeze
guard checks a name-mangled attribute, so the assignment silently desynced
`config.shift` from `config["shift"]` on a shared component. Removed the
argument so the scheduler owns the shift (`pipelines.md` gotcha 3), which also
drops a per-call mutation of a registered component (gotcha 7).
Other pipeline fixes:
* `torch.randn` -> `randn_tensor` in `prepare_latents` and the refiner
(gotcha 10). `__call__` advertises `generator: Generator | list[Generator]`,
but `torch.randn` raises on a generator list and on a CPU generator with a
CUDA device, so that path could not work.
* Delete the SLURM preemption/resume feature (`checkpoint_dir`,
`_atomic_save_state`, `_capture_state` / `_restore_state`) and the
single-shot refiner path the docstring itself called a debugging fallback —
roughly 270 lines of research-cluster scaffolding.
* `_empty_cuda_cache` (CUDA-only) -> `empty_device_cache`;
`@torch.inference_mode()` -> `@torch.no_grad()` and removed from inner
helpers the decorator already covers (gotcha 2).
* Remove dead `_callback_tensor_inputs` (no `callback_on_step_end` exists),
`_exclude_from_cpu_offload` (a no-op — the base class already skips
non-`nn.Module` components), `_kv_max_frames`, and `latents.detach()`.
Tests: drop the `@slow` integration stub — `testing.md` says integration and
slow tests don't belong in the initial PR.
`refiner.py`: 1286 -> 1017 lines.
* refactor(sana-wm): drop the timm dependency
Per @yiyixuxu's review. The three timm imports were all trivially replaceable:
* `Mlp` — both local subclasses fully overrode `forward` and only inherited
`fc1` / `act` / `drop1` / `fc2` / `drop2`, so it is now a self-contained
module and the redundant `class Mlp(Mlp)` wrapper is gone.
* `Attention` — `GDN` inherited only `qkv` and `proj` from it (it already
defines its own `q_norm` / `k_norm`), so those two layers are declared
inline and `GDN` subclasses `nn.Module`.
* `DropPath` — `drop_path` was hardcoded to `0.0`, so this was always
`nn.Identity()`. Removed along with the rest of the plumbing.
No `dummy_timm_objects.py` is needed since the dependency is gone rather than
optional, and `setup.py` is untouched (other models still use timm).
State dict is unchanged (871/871 keys match the released checkpoint) and the
end-to-end GPU smoke gives byte-identical output (frame mean 0.5554).
* refactor(sana-wm): drop the Triton kernels and apply review feedback
Per @yiyixuxu's inline review. Removes `transformer_sana_wm_kernels.py`
entirely (3234 lines) — the Triton kernels will be published to
`kernels-community` and wired back in a follow-up PR.
* Delete the three `*Triton` attention classes. Each only overrode `forward`
to call a fused kernel; their pure-PyTorch parents (`BidirectionalGDN`,
`BidirectionalGDNUCPESinglePathLiteLA`) compute the same result. The
Triton-availability probe and the MRO-walking fallback go with them.
The released `config.json` still names the `*Triton` variants, so
`ATTENTION_BLOCKS` maps those strings onto the pure-PyTorch classes.
* Move the nine pure-PyTorch camera-math helpers that were reachable from
the model into `transformer_sana_wm.py`; drop the rest with the kernels
file (they were only reachable from the deleted Triton path).
Other review items:
* Use `get_activation()` from `..activations`; delete the local activation
and normalization registries (`build_norm` was only ever called with
`None`).
* Delete `GLUMBConv` and make `GLUMBConvTemp` self-contained. `ConvLayer` is
slimmed and renamed `SanaWMConvLayer`, but has to stay a module — the
checkpoint keys are `mlp.inverted_conv.conv.weight`, so collapsing it to a
bare `nn.Conv2d` would drop the `.conv` level.
* Remove the closure-returning helpers the review flagged as hard to read
and compile-hostile: `prepare_prope_fns` and friends now return plain
tensors via `_prepare_ucpe_ray_transforms` + `_apply_ucpe_transform`, and
the `_register_block` decorator factory is an explicit dict update.
* Remove all seven `@torch.compile` decorators — users can reach for
`compile_repeated_blocks()` instead.
* Drop `partial` (pass `chunk_size` at the call site), `DWMlp`, the tuple
helpers, `get_same_padding`, `modulate`, an unused `apply_rotary_emb`,
and the validation of internal-only arguments.
* Tests: drop `SanaWMTritonFallbackTests` along with the mechanism it covered.
State dict is unchanged (871/871 keys) and the end-to-end GPU smoke passes
with correct video output.
* refactor(sana-wm): inline the modulation helper, drop the backward-only shim
Two more items from @yiyixuxu's inline review:
* Inline `t2i_modulate` at its four call sites.
* Delete `_IdentityForwardContiguousBackward` / `_contiguous_backward`. It is
the identity in forward and only exists to hand a contiguous gradient to
the backward pass, so it is dead weight in an inference-only port.
State dict unchanged (871/871 keys); GPU smoke output matches the previous
run exactly (frame mean 0.5560).
* Update src/diffusers/models/transformers/transformer_sana_wm.py
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* Update src/diffusers/models/transformers/transformer_sana_wm.py
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* style(sana-wm): drop the blank line left by removing the camera-branch init
* refactor(sana-wm): address @dg845's review round
Transformer:
* Stop stashing patch-grid shape on `self` during `forward` (`self.f/h/w`) and
thread it through as locals; `unpatchify` now takes it explicitly.
* Replace the 7 `assert`s with `ValueError`s.
* Drop `attn_drop`/`proj_drop` from `MultiHeadCrossAttention` (training-only,
and `attn_drop` was never applied) plus a stale training-era banner comment.
Pipeline / refiner:
* Don't mutate the components handed to the pipeline. The VAE tiling +
framewise settings move to the docs, `padding_side="right"` is passed per
tokenizer call, and the `.eval()` calls are gone (no `self.training`
branches remain, and `from_pretrained` already returns eval-mode modules).
* Gate `cam_utils`' optional imports on `is_torchvision_available()` and a new
`is_pi3_available()` helper.
* Annotate `SanaWMLTX2Refiner.__init__` and inline `_refine_latents_ar` into
its single caller.
Conversion script:
* Move to `scripts/` alongside the other Sana converters.
* Raise on missing/unexpected keys instead of printing, so a bad mapping can't
silently emit a broken transformer.
State dict unchanged (871/871 keys). GPU smoke on the public checkpoint is
byte-identical to the previous run (frame mean 0.5560), including with the VAE
settings applied by the caller rather than the pipeline.
* refactor(sana-wm): rename WanRotaryPosEmbed to SanaWMRotaryPosEmbed
The name implied this was Wan's rotary embedding, but it isn't: the per-axis
split is configurable through `fhw_dim`, and the frequencies stay complex in a
single `freqs` buffer instead of being split into real cos/sin buffers. So it
can't carry a `# Copied from`. Renamed, with a docstring recording why.
The buffer is `persistent=False`, so the state dict is unchanged (871/871).
* test(sana-wm): add transformer model tests, migrate pipeline tests to pytest
* `tests/models/transformers/test_models_transformer_sana_wm.py` — generated
with `utils/generate_model_tests.py` and filled in, following
`test_models_transformer_sana_video.py`. The tiny config sets
`softmax_every_n=2` so one block exercises the GDN camera branch and the
other the softmax variant. Dummy inputs supply the conditioning the forward
requires: `encoder_attention_mask`, `(B, F, 20)` camera conditions, and
`chunk_plucker`.
* `tests/pipelines/sana_wm/test_sana_wm.py` — rewritten in the pytest style of
`tests/pipelines/sana_video/test_sana_video.py`: no `unittest`, bare
asserts, module-level imports, and `parametrize` in place of loop-style
cases (15 test functions become 31 cases).
`AttentionTesterMixin` is skipped because the model calls
`F.scaled_dot_product_attention` directly rather than going through a
diffusers attention processor.
* fix(sana-wm): make the transformer layerwise-casting and offload safe
* `TimestepEmbedder.dtype` used `next(self.parameters()).dtype`, which reports
the storage dtype under layerwise casting rather than the compute dtype. Use
`get_parameter_dtype`, which is layerwise-casting aware. (The model-level
`self.dtype` already routes through it via `ModelMixin`.)
* Stop reading submodule `.weight`/`.bias` inside `forward`. Group-offload
hooks fire on a module's `forward`, so reading its parameters directly
leaves them offloaded and trips a device mismatch. The frame-gate and
output-gate helpers now call their submodules, and the fused camera QKV
projection becomes three `q/k/v` calls -- algebraically the same as one
GEMM over the concatenated weights.
* Don't mutate the tokenizer in `SanaWMLTX2Refiner._encode_prompt`; pass
`padding_side="left"` per call, matching `SanaWMPipeline`.
* Add a `torch.compile` test with `recompile_limit=2` -- the repeated block
compiles once per attention variant.
GPU smoke unchanged (frame mean 0.5560).
* refactor(sana-wm): explicit arguments instead of **kwargs, generator instead of seed
Replace the `**kwargs` bag threaded through model -> block -> attention with
explicit keyword arguments. Only three runtime keys were ever read at the
leaves (`frame_valid_mask`, `precomputed_gates`, `ucpe_ray_transforms`); the
rest were forwarded and silently swallowed. `camera_embedding`,
`chunk_index`, `chunk_index_global` and `chunk_split_strategy` turned out to
be pure dead plumbing -- written into the per-block kwargs dicts and never
read by any attention or MLP forward -- so they are gone.
The pipeline also grows the standard diffusers arguments:
* `prompt_embeds` / `prompt_attention_mask` / `negative_prompt_embeds` /
`negative_prompt_attention_mask` on `encode_prompt` and `__call__`.
* `seed` / `refiner_seed` are replaced by `generator` / `refiner_generator`.
`generator=torch.Generator(device).manual_seed(42)` reproduces exactly what
`seed=42` used to build, so results are unchanged.
State dict unchanged (871/871). CPU old-vs-new equality on a tiny config over
both attention variants is exact (`torch.equal`, max diff 0.0), and the GPU
smoke on the public checkpoint still gives frame mean 0.5560.
* refactor(sana-wm): drop the unreachable chunk-split strategies
@yiyixuxu was right that the model only ever runs uniform chunking; my earlier
reply defending the strategies was wrong. The two sites that actually chunk
both call `normalize_chunk_index(None, T, chunk_size)` with three positional
arguments, so `chunk_split_strategy` always took its `"uniform"` default. The
configured `first_chunk_plus_one` reached the per-block kwargs dict and was
then swallowed by the attention forwards' `**kwargs` without ever being read.
Flattening that bag into explicit arguments is what surfaced it.
Also note the `chunk_size` those call sites use is `chunk_gdn_chunk_size`, a
different attribute from the `chunk_size` that was being threaded through.
So `chunk_index_from_chunk_size`, `normalize_chunk_index`,
`is_uniform_chunking` and `compute_chunk_sizes` are gone, the uniform
boundaries are inlined at both call sites, and `chunk_split_strategy` is
dropped from the model and block constructors.
The released `config.json` still carries the key; loading is unaffected
(`extract_init_dict` ignores it) but it logs an "not expected and will be
ignored" warning, so it should come out of the checkpoint config on the next
export.
State dict unchanged (871/871); GPU smoke still frame mean 0.5560, which
confirms those branches were never taken.
* refactor(sana-wm): use the shared RMSNorm
@yiyixuxu asked whether the shared `RMSNorm` could be used here and my earlier
reply overstated the obstacles. Re-checking each one:
* `scale_factor` is not a blocker. `attention_y_norm.weight` is one of the 871
checkpoint keys, so `from_pretrained` overwrites whatever the constructor
initialised -- exactly the same reasoning that removed the other
hand-written inits in this PR. It only ever affected from-scratch models.
* `norm_dim` is not a blocker either; it is always the default `-1`.
* The numerics do differ: ours ran the normalisation and the weight multiply
in fp32, while the shared class computes only the variance in fp32 and does
the scaling and weight multiply in the input/weight dtype.
Measured rather than argued: the end-to-end GPU smoke on the public
checkpoint moves from frame mean 0.5560 to 0.5561, and the decoded video is
visually identical. That is well inside bf16 noise, so the local class is not
worth keeping.
The now-unused `y_norm_scale_factor` config argument goes too. As with
`chunk_split_strategy`, the released `config.json` still carries it and will
log an "not expected and will be ignored" warning until the checkpoint config
is re-exported.
* refactor(sana-wm): migrate MultiHeadCrossAttention to the diffusers attention API
* Split the attention maths into a `SanaWMCrossAttnProcessor` and make
`MultiHeadCrossAttention` an `AttentionModuleMixin`, so the processor can be
swapped and the standard tooling applies.
* Use `dispatch_attention_fn` instead of calling
`F.scaled_dot_product_attention` directly, which lets the model pick up
alternative attention backends. This also drops the two transposes, since
the dispatcher takes `(batch, seq, heads, dim)`.
* Build a boolean padding mask rather than an additive float one. The old
`(1 - mask) * -10000.0` form isn't supported by the varlen backends.
Projection names are unchanged (`q_linear`, `kv_linear`, `proj`, `q_norm`,
`k_norm`), so the state dict stays at 871/871 keys. The GPU smoke on the
public checkpoint is unchanged at frame mean 0.5561, so `-inf` masking makes
no difference against the old `-10000.0` at this precision.
* feat(sana-wm): give the refiner its own transformer model
Per @yiyixuxu and @dg845: the stage-2 refiner no longer reaches into
`LTX2VideoTransformer3DModel`. It used to drive AR refinement by setting
attributes on LTX-2's submodules (`attn._tf_kv_prefix`, `_kv_cache_capture`,
`_tf_capture_kv`), which meant a second `forward` implementation smuggled in
from the pipeline.
New `transformer_sana_wm_refiner.py`, following the `transformer_wan_vace.py`
pattern of importing the reusable pieces rather than duplicating them:
* `SanaWMRefinerKVLayerCache` / `SanaWMRefinerKVCache`, modelled on
`Flux2KVCache`, holding the per-layer sink and history K/V.
* `SanaWMLTX2RefinerTransformerBlock`, which subclasses
`LTX2VideoTransformerBlock` and overrides only `forward`, so the submodule
structure stays byte-identical.
* `SanaWMLTX2RefinerTransformer3DModel`, taking `kv_cache` / `kv_cache_mode`
explicitly through `forward`.
`refiner.py` drops from 982 to 638 lines: the four `*_kv_prefix_on_blocks` /
`*_captured_kv_*` helpers, `_forward_video_block`, `_streaming_self_attention`
and `_forward_video_only_with_rope` are all gone, along with a dead
`n_context_tokens > 0` branch that neither call site could reach.
The state dict is unchanged at 3510 keys, identical to
`LTX2VideoTransformer3DModel` at the same config, so only the class name in
the checkpoint metadata has to change. CPU equivalence against the old path
is exact across four AR blocks -- covering sink injection, history growth and
history trimming -- with a negative control confirming the cache path is
actually exercised. The GPU smoke on the public checkpoint is unchanged at
frame mean 0.5561.
* refactor(sana-wm): run the refiner as a separate pipeline
Per @dg845, the two stages are no longer nested. `SanaWMPipeline` produces
stage-1 latents and `SanaWMLTX2Refiner` consumes them:
latents = SanaWMPipeline(...)
video = SanaWMLTX2Refiner(...)
The refiner takes an optional `vae` -- pass the base pipeline's so the weights
are shared, the way SDXL shares components between base and refiner -- and
decodes to video when given one. The manual device juggling in
`SanaWMPipeline.__call__` is gone with the nesting that forced it; each
pipeline now manages its own placement through the standard offload hooks.
Un-nesting turned up a real bug: `torch_dtype` never reached the nested
sub-pipeline, so `SanaWMPipeline.from_pretrained(..., torch_dtype=bfloat16)`
silently ran the whole of stage 2 in float32 -- transformer, text encoder and
connectors alike. That is where the "~87 GB of refiner weights" in the old
comment came from. Loaded on its own the refiner honours the requested dtype,
which halves its footprint; the GPU smoke moves from frame mean 0.5561 to
0.5548 and the decoded video is visually identical.
The conversion script writes the refiner to its own output directory
(`--dst-refiner`, default `<dst>-refiner`) rather than a `refiner/`
subfolder. `DiffusionPipeline.from_pretrained` has no `subfolder` argument, so
a nested folder silently loads the *base* pipeline's same-named component
folders instead -- exactly the kind of wrong-weights failure that only
surfaces because `connectors/` happens not to exist at the top level.
* Apply style fixes
* docs(sana-wm): document the refiner's output_type argument
`check_forward_call_docstrings` flagged `SanaWMLTX2Refiner.__call__` as
missing an entry for `output_type`, which the un-nesting added.
* refactor(sana-wm): support only what the released checkpoint uses
@yiyixuxu asked whether we can support just the released configuration and
flatten what's left. We can — the checkpoint only ever instantiates two
attention classes, so the other five in the chain were dead weight.
* Collapse `GDN -> BidirectionalGDN`, `GDN -> _GDNUCPEBase`,
`BidirectionalGDNUCPELiteLA -> ...PostUCPERenorm -> ...SinglePathLiteLA`
into a self-contained `BidirectionalGDNUCPESinglePathLiteLA` plus
`_SoftmaxUCPESinglePathLiteLA`. Inheritance depth goes 5 -> 1.
* Blocks take the class, not a string: `SanaVideoMSCamCtrlBlock` gains an
`attn_cls` argument, and the model picks per layer with a plain loop. That
removes the `ATTENTION_BLOCKS` registry, `_resolve_attention_block`,
`_GDN_TO_SOFTMAX_CAMCTRL`, `_inject_softmax_layers`, the `attn_type` /
`camctrl_type` validation, and the unreachable no-camera branch.
* `attn_type` and `camctrl_type` leave the config signature. As with
`chunk_split_strategy` and `y_norm_scale_factor`, the released
`config.json` still carries them and loading logs a benign "not expected"
warning until the checkpoint is re-exported.
* Also drop the zero-init on `GLUMBConvTemp.t_conv` (overwritten by
`from_pretrained`) and inline the one-line `restore_shape` helper. The
model docs lose a paragraph describing an `_inner.` prefix and an
`add_inner_prefix` helper that no longer exist.
3908 -> 3153 lines. State dict unchanged at 871/871 keys with identical
ordering, and 15 GDN + 5 softmax blocks as before. Old-vs-new CPU equality is
exact (`torch.equal`, max diff 0.0) on both a tiny config and the released
geometry, across a `softmax_every_n` sweep and on the backward pass, with
negative controls confirming the comparison isn't vacuous. The GPU smoke is
byte-identical too: refined latents mean 0.035024 / std 0.887467, frame mean
0.5548.
* refactor(sana-wm): make the softmax attention class standalone
Drops the last bit of inheritance between the two attention classes, per
@yiyixuxu. `_SoftmaxUCPESinglePathLiteLA` now subclasses `nn.Module`
directly.
It turned out to be a simplification rather than duplication: of the parent's
thirteen methods it only ever called three (`_apply_output_gate`,
`_prepare_frame_valid_masks`, `_stabilize_cam_transforms`), and the other ten
were dead for it. It is also always built with `conv_kernel_size=0`, so the
whole short-convolution half of the inherited `__init__` -- `conv_q/k/v`,
the `*_cam` variants, `kernel_func` -- was permanently `None` or unused and
is now gone.
State dict unchanged: 871/871 keys in the same order, and the softmax blocks
still emit the same 24 keys each. Old-vs-new equality is exact
(`torch.equal`, max abs diff 0.0) on two input configurations, with negative
controls from 0.109 to 14.62 confirming the comparison isn't vacuous; the
inlined `__init__` even consumes the RNG stream identically, so seeded
from-scratch weights are bit-for-bit the same. GPU smoke unchanged at
refined-latent mean 0.035024 / std 0.887467.
* feat(sana-wm): export the refiner weights in bfloat16
The refiner's Gemma-3 text encoder was the one component still stored in
float32, at 45.4 GB. Since the pipeline is loaded with
`torch_dtype=torch.bfloat16`, those weights were cast on every load anyway,
so storing them cast changes nothing and halves the download:
text_encoder 45.4 GB -> 22.7 GB
refiner total 84 GB -> 61 GB
The VAE stays in float32, as the docs recommend. Verified rather than
assumed: the GPU smoke is bit-identical (refined latents mean 0.035024 / std
0.887467, frame mean 0.5548), and it now runs in 124s instead of 384s because
there is far less to load.
* refactor(sana-wm): pipeline clean-up from @yiyixuxu's review
`SanaWMLTX2Refiner`:
* Dissolve `_RefinerChunkRunner` (191 lines) into `__call__`, so the AR loop
reads top-to-bottom like `pipeline_wan_animate.py`, and swap the bare
`tqdm` for `self.progress_bar(...)`.
* Actually support `enable_model_cpu_offload()`. The docs already advertised
it, but three things blocked it: the `device` default resolved through
`next(self.transformer.parameters()).device`, which is `cpu` under offload;
`_encode_prompt` called the inner Gemma backbone directly, bypassing the
hook that `cpu_offload_with_hook` installs on the top-level module; and
`_decode_latents` read `self.vae.device`. All three now go through
`self._execution_device`, and the manual `.to(...)` juggling is gone.
* TODO on the `LTX2TextConnectors` import pointing at
huggingface/diffusers#14749.
Two argument removals fall out: `progress=` (use
`set_progress_bar_config(disable=True)`, which is what every other pipeline
does) and `device=`.
`SanaWMPipeline`:
* `check_inputs` only validates now; the camera-trajectory normalisation
moves to `prepare_camera_trajectory`.
* Add the standard `guidance_scale` / `do_classifier_free_guidance` /
`num_timesteps` / `current_timestep` / `interrupt` properties.
* `return_dict=False` returned only `(frames,)`, silently dropping `c2w` and
`latent`. It now returns the same three values either way -- the flag picks
the container, not the contents.
* Take a `PIL.Image`, not a path, and drop the `_model_cpu_offload_active`
helper that had no callers left after the un-nesting.
The Pi3X intrinsics estimator leaves `diffusers` entirely: it pulled in an
extra dependency and fetched a second checkpoint without the user ever
calling `from_pretrained`. It is packaged as a remote modular block hosted
alongside the checkpoint, and the docs show `ModularPipeline.from_pretrained`
instead. `is_pi3_available` goes with it.
Equivalence for the refiner rewrite is exact across six configurations --
varying sink size, block size, KV window and seed so history growth and
trimming both fire -- `torch.equal` with max abs diff 0.0, negative controls
from 1.9e-2 to 5.55. GPU smoke on the public checkpoint is unchanged at
refined-latent mean 0.035024 / std 0.887467.
* refactor(sana-wm): model-side clean-up from @yiyixuxu's review
* Split each attention class's main and camera branches into their own
modules, so `forward` is a flat list of layer calls. The branch modules
deliberately hold no parameters: if they owned `qkv` / `q_proj_cam` /
... PyTorch would prefix by attribute name and rename 390 of the 871
checkpoint keys. (`state_dict` duplicates a module aliased in two places
rather than de-duplicating it, so aliasing isn't a way out either.) The
compute is split; the ownership isn't.
* `transformer_sana_wm_refiner.py` is self-contained -- nine `# Copied from`
markers replace the `transformer_ltx2` import, at whole-function,
whole-class and method-level granularity as each piece required. That costs
+831 lines, since the copied block `__init__` drags in both attention
processors.
* The refiner's streaming self-attention is a real attention processor,
`SanaWMLTX2RefinerAttnProcessor`, and the block calls `self.attn1(...)`. The
processor swap happens in the model `__init__` after the blocks are built,
because `attn1` is constructed inside the `# Copied from` body.
* Inline the KV-cache mode strings, drop `plucker_proj`'s zero-init.
State dicts unchanged: 871/871 and 3510/3510, ordered lists identical.
Equivalence is exact -- max abs diff 0.0 -- against the real released
checkpoint for the attention split, and across five refiner configurations
driven through two AR blocks, with negative controls throughout.
`check_copies` passes. GPU smoke unchanged at refined-latent mean 0.035024 /
std 0.887467.
* test(sana-wm): cover the GDN branch, which was silently dead in these tests
`SanaWMTemporalShortConvolution` allocates its weight with `torch.zeros` and
nothing initialises it any more -- the hand-written init came out earlier in
this PR because `from_pretrained` overwrites it. That is true for real
checkpoints, but a from-scratch model gets a zero `conv_k`, so the GDN key is
identically zero and the whole branch returns zeros. Every test in this file
builds a random model, so they were all passing while exercising none of the
GDN maths.
Two tests: one pins the zero-init as deliberate rather than accidental, the
other fills the short convolutions with noise and asserts the output moves.
Checked that the second one bites -- forcing the short convolution to return
zeros makes it fail.
* [SANA-WM] Address review: reuse diffusers embeddings, grad checkpointing, callbacks, offload fixes
- Reuse Timesteps/TimestepEmbedding; drop unused cfg_embedder / y_embedding / model_max_length
- Rename Mlp/PatchEmbed/CaptionEmbedder with SanaWM prefix
- Support gradient checkpointing in SanaWMTransformer3DModel
- Keep complex RoPE freqs out of buffers so Module.to(dtype) doesn't drop the imaginary part
- Disable autocast around camera geometry (torch.cross)
- Pipeline: callback_on_step_end, Literal output_type, _execution_device, maybe_free_model_hooks
- Fix output docstring, bump copyright to 2026
- Port tests to new pytest mixins, add refiner transformer tests
* [SANA-WM] Drop cam_utils and registration tests
* [SANA-WM] Add return type annotations to forward methods
---------
Co-authored-by: HaoyiZhu <hyizhu1108@gmail.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
|
||
|
Wang, Yi
|
883d51ef02 |
Make Stable Diffusion device-map comparisons deterministic (#15010)
Description: The balanced device-map tests compare images from single-device and multi-device inference. On XPU, repeated fp16 inference can produce different results even with the same seed, causing the pixel comparison to fail. Enable PyTorch deterministic algorithms for the two image-comparison tests and restore the previous setting afterward. This keeps other tests unaffected. Validation: Both affected tests passed on a multi-XPU runner (RUN_NIGHTLY=1 RUN_SLOW=1). Signed-off-by: Wang, Yi A <yi.a.wang@intel.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Lucas Ruan
|
e1a6c67dc7 | Fix Anima text encoder exports and standalone core denoising (#14943) | ||
YiYi Xu
|
85c9fa4cae |
[cli] detect composite block classes in custom_blocks (#15001)
`diffusers-cli custom_blocks` only recognized classes whose base was spelled exactly `ModularPipelineBlocks`, so a `SequentialPipelineBlocks` or `AutoPipelineBlocks` subclass (the usual thing to publish) was reported as "could not be retrieved". Accept the composite subclasses as well, and match on the class name so dotted bases like `diffusers.modular_pipelines.SequentialPipelineBlocks` work too. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> |
||
Jingya HUANG
|
1d5d056ecf |
[TPU] TorchTPU backend integration - eager / torch.compile / tp (#14039)
* feat: add torchtpu * feat:draft TorchTPU support * fix: wan overflow issue + compile mode error on sdxl * doc: enhance with TorchTPU doc * doc: enhance with TorchTPU doc * style: remove unused imports flagged by ruff Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs: remove Debug Eager and Fused Eager sections from tpu.md Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * feat: add TP for TPU * style: fix import sorting in TPU test scripts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: ruff format TPU test scripts Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: style * fix: test for native 4 devices * fix: propagate TPU device fixes to Flux/Flux2/Wan-family copies; register tpu.md; drop redundant execution_device check - Propagate the text_encoder.device-based fix (introduced for TPU CPU-offload support) from FluxPipeline/Flux2KleinPipeline/WanPipeline into their `# Copied from` copies (flux/*, flux2_klein_inpaint, visualcloze, anyflow, chronoedit, lucy_edit, skyreels_v2/*). SDXL-family copies of StableDiffusionXLPipeline.encode_prompt are intentionally left untouched; they'll be handled in a follow-up PR that fixes device placement for every pipeline component (not just text encoders). - Register docs/source/en/optimization/tpu.md in _toctree.yml (was breaking the docs build: "not present in the table of contents"). - Remove the redundant "prefer non-CPU, non-meta component" loop from DiffusionPipeline._execution_device: PR #14383 already fixed this in DiffusionPipeline.device, which _execution_device falls back to. Verified on TPU hardware that _execution_device still resolves correctly for a split-placement pipeline after the removal. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * tests: cleanup * fix: fix compile mode * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * test: remove flux2 e2e test * doc: apply suggestions * doc: apply suggestions * review: remove monkey patch * review: revert neuron-specific changes in the tests * review: apply suggestions * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: add tpu to tp doc * review: remove unnecessary for tpu * test: assert every _tp_plan parameter is sharded after a TP load Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * removal: delete redundant tp shard helpers * removal: drop TPU workarounds no longer needed after recent main changes * doc: keep the FLUX.2-dev text encoder off a single TPU chip in the TP example * doc: encode the prompt on CPU in the TPU TP example so it runs end to end * test: tighten TPU TP tolerance and simplify the TPU TP worker * doc: use enable_model_cpu_offload() in the TPU eager example * doc: shard both FLUX.2-dev text encoder and transformer in the TPU TP example * doc: run the text encoders on TPU in the compiled example * test: drop the redundant TPU sync in the TP worker * test: trim comments in the TPU TP tests * test: run the TPU TP test with mp.spawn like the CUDA one * test: pin the TPU TP test to 4 chips * Update docs/source/en/optimization/tpu.md Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * doc: use DistributedConfig and dtype, and drop the torch_tpu import in the TPU examples * test: skip the TPU TP test for models without a _tp_plan * doc: correct the warning of recompilation * fix: clone the SD3 timestep so torch.compile doesn't recompile every step, thus we can use it for the tpu.md compile mode snippet * doc: update advice for compile mode when one chip doesn't fit * test: pass the TPU TP tolerances as test arguments * review: apply suggestions to the doc * review: view compilation patch for tpu-only --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: JingyaHuang <huang_jingya@outlok.com> Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Christopher
|
78bf2c6266 |
[LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) (#14163)
* [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) Adds loading of LoKr (LyCORIS Kronecker product) adapters: - `load_lora_adapter` detects `lokr_` keys and injects a peft `LoKrConfig`, inferred from the tensor shapes via `_create_lokr_config` (decompose factor, per-module rank/alpha patterns). - State dict conversions for the formats in the wild: ai-toolkit Z-Image (dotted diffusers paths under `diffusion_model.`), ai-toolkit BFL Flux2 (fused qkv), LyCORIS underscore format, and bare dotted diffusers paths. - BFL fused-QKV LoKr cannot be split exactly into separate Q/K/V Kronecker factors, so `Flux2LoraLoaderMixin.load_lora_weights` fuses the model's QKV projections and maps the adapter 1:1 (exact). - Alpha follows the LyCORIS convention: scaling applies only to rank-decomposed factors and is baked into the weights at conversion. Fixes #13221 * [LoRA] LoKr review follow-ups: adapter-agnostic error message, refuse fused-QKV load over an unfused adapter - "Invalid adapter checkpoint. We currently support LoRA and LoKr." replaces the message that still said "LoRA checkpoint" and described the substring check. - Loading a fused-QKV LoKr checkpoint now refuses when an adapter is already injected on to_q/to_k/to_v or add_{q,k,v}_proj: fuse_qkv_projections() would replace those modules and orphan it. Covered by a test. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * [LoRA] rename the LoKr tests to test_lokr_conversion.py The file exercises checkpoint conversion and loading for LoKr adapters, not LoRA. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [LoRA] LoKr review follow-ups: scope to Z-Image/Flux2, model-level fused QKV, LoKrTesterMixin - Only the Z-Image and Flux2 loader mixins accept LoKr checkpoints, the two that convert them; the other mixins are back to their LoRA-only check. - The fused-QKV handling moves to `_maybe_fuse_qkv_projections_for_lokr` in peft_utils, called from `load_lora_adapter`, so the model fuses its projections whichever entry point loads the adapter. The quantized-model error explains why. - `_bake_lokr_alpha` becomes `_bake_lokr_alpha_` as it edits the dict in place. - The BFL Flux2 converter no longer accepts expanded diffusers block names under `diffusion_model.`: no LoKr checkpoint uses that layout; diffusers-named checkpoints have no prefix and take the generic converter. - The LyCORIS converter names the keys it does not recognize. - Tests move from tests/lora to a `LoKrTesterMixin` in tests/models/testing_utils, with the checkpoint formats tested on the Flux2 and Z-Image model test classes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * [LoRA] mark ZImageLoraLoaderMixin.load_lora_weights as copied from Flux2 --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
12b843e5d3 | Merge branch 'main' into refactor-pab-fastercache | ||
Patrick Ribbsaeter
|
d961a388fd |
fix(image_processor): maintain 3-tuple return contract in InpaintProcessor.preprocess when mask is None (#14481)
fix(image_processor): preserve 3-tuple return contract when mask is None (#14470) Ensure InpaintProcessor.preprocess returns (image, None, postprocessing_kwargs) when mask is None, preserving the expected 3-tuple return contract and signature consistency across all preprocess branches. Consolidate unit tests into a parameterized test covering with-mask, without-mask, and crop configurations. |
||
Subodh
|
e0415c41f8 |
Fix active adapter detection for text encoders (#14961)
* Fix active adapter detection for text encoders Signed-off-by: Subodh Bansode <subodh171011@gmail.com> * test: move text encoder adapter regression test Signed-off-by: Subodh Bansode <subodh171011@gmail.com> --------- Signed-off-by: Subodh Bansode <subodh171011@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
2192286010 | preserve deprecated functionality. | ||
|
Sayak Paul
|
14b4c7155d | resolve conflictsa | ||
|
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> |
||
naykun
|
da1d3829cf |
Configure fixed sampling sigmas on QwenImage21Pipeline (#14950)
* feat(qwenimage21): configure sampling sigmas on the pipeline * docs(qwenimage21): clarify default and runtime sampling sigmas |
||
|
Sayak Paul
|
c2798cc785 |
[tests] refactor the pipeline-level attention processor tests (#14701)
* refactor the pipeline-level attention processor tests * simplify stuff. * remove utils.py |
||
|
Sayak Paul
|
cff9dafbbf |
[modular] Echo team joy future academy jd echo by @Gelercatty (#14941)
* Add Echo modular pipeline * Address Echo modular pipeline review feedback * Inline Echo input list normalization * Align Echo latent boundaries and batch expansion * Polish Echo documentation per review * Load official Echo checkpoint and separate block tests * consolidate testing. * formatting --------- Co-authored-by: Yanwen Ma <gelercat117@gmail.com> |
||
|
Sayak Paul
|
ad8168c5c0 | deprecate current_timestep_callback instead of removing. | ||
|
Sayak Paul
|
2e6bc80744 | Merge branch 'main' into refactor-pab-fastercache | ||
|
Dhruv Nair
|
578c9b2c66 |
Add single file support for Krea 2 (#14914)
update |
||
|
Dhruv Nair
|
acbabca385 |
Consolidate torch device backend dispatch (#14792)
* introduce TorchDeviceBackend * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Dhruv Nair
|
4de185d6ec |
Add single file support for Minimax H3 (#14839)
update Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Dhruv Nair
|
a91ff967f5 |
[Modular] Avoid downloading weights when loading from an existing local path (#14797)
* avoid redownload from local path * update * update * update --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Jingya HUANG
|
863092ffec |
[core] Shard tensor-parallel checkpoints on load and save (#14544)
* [core] Shard tensor-parallel checkpoints on load and save Stream each rank's slice of a tensor-parallel checkpoint straight off disk instead of materializing the full checkpoint on every rank and resharding it afterwards, and gather the shards back on save. - `from_pretrained(..., parallel_config=TensorParallelConfig(...))` resolves the shard specs on the still-meta model, then slices each safetensors tensor before the dtype cast, so host memory peaks at ~1/tp_degree of the checkpoint. - `save_pretrained` all-gathers the DTensors into an ordinary checkpoint, or writes a distributed checkpoint with `dcp=True` so no full tensor is ever formed. The writing `tp_degree` is recorded, since a packed weight's stored layout is interleaved by it. - Factor the plan interpretation out of the Neuron pre-shard path into shared `TPShardSpec` / `resolve_tp_shard_specs` / `_local_shard` / `_hooks_only_styles` helpers, so both backends and both the load and save paths shard identically. * Raise when tensor parallelism is combined with quantization, offloading or LoRA Addresses the remaining two items of the review on #13718: tensor parallelism was rejected alongside quantization and `device_map` only on the `from_pretrained` streaming path, while `enable_parallelism` — which the quantization error message itself recommended — accepted a quantized, offloaded or adapter-injected model and sharded it anyway. - Add `_check_tp_model_state`, called from `apply_tensor_parallel`, the one chokepoint every TP entry point funnels through. It rejects a model that is quantized, group-offloaded, placed by accelerate (`device_map` or CPU offload), or has PEFT layers injected. Placed before the device-type check so the reported reason is the useful one. - Guard the reverse order too: `enable_group_offload`, the two pipeline CPU-offload methods, and `load_lora_adapter` now refuse a tensor-parallel model. - `save_pretrained` refuses a quantized tensor-parallel model. Previously the `dcp=True` branch returned before the quantizer's serialization step, writing shards with no quantization metadata and no error. - The DCP load guard checked the `quantization_config` kwarg only, so a pre-quantized checkpoint directory loaded silently; check the config's own entry too, and add the missing `_tp_plan` check that otherwise surfaced as a raw `AttributeError`. - Correct the `from_pretrained` message and the doc sentence that pointed at `enable_parallelism` as a way to shard a quantized model. The new tests are the first tensor-parallel tests that need neither an accelerator nor more than one rank: every case asserts a raise before any collective, so they run single-process on gloo. * review: remove tp save/ dcp related * review: add a new helper . * review: moved the TP checks into enable_parallelism * review: move the warning under caution block * review: add loading time numbers * review: move tp size dividende check * review: delete test * review: consolidate tp checks * review: consolidate the tp checks, all in tensor_parallel.py * review: improve conditionals * review: debug info * review: check weights format before instantiating the model * review: restore duplicated _find_mismatched_keys since irrelevant * revert: delete duplicated helper * review: update doc for loading benchmark of flux2-dev * review: improve the raising check * review: add an args section * Test tensor-parallel loading from sharded checkpoints --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
yzhautouskay
|
78405ec9a1 |
[Cosmos3] Fix Transfer SeaCache artifacts with control CFG (#14897)
* Fix Cosmos 3 Transfer SeaCache indicators * Clarify SeaCache indicator guidance; tests refactor --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Sayak Paul
|
360959d1cf | Merge branch 'main' into refactor-pab-fastercache | ||
|
Sayak Paul
|
fef717ffb0 |
[core] propagate sage attention updates. (#14584)
* propagate sage attention updates. * add sage blackwell. * up * style. * address feedback * style. |
||
|
sayakpaul
|
0a1f36769c | feat: implement cache_context for pab and fastercache. | ||
BeyondBirthday07
|
5ff8e59ff9 |
Fix LoRA loading for flattened CLIP text encoders (#14872)
Co-authored-by: Aznix07 <sruhilmodi13@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
Álvaro Somoza
|
51a454be9a |
[MInimax H3] Fix VAE decode (#14754)
* intiial fix * fix code quality * fix failing pr test * fix gpu test * don't override the requested dtype * add test * apply review --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> |
||
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> |
||
Linoy Tsaban
|
e0118ade2f |
[LoRA] add LoRA training for Qwen Image 2.1 (#14808)
* [LoRA] add LoRA training for Qwen Image 2.1 Add DreamBooth LoRA training for Qwen-Image 2.1, text-to-image and image-to-image, with fast tests and a README section. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * default rank and alpha to 16 and document the ratio `--rank` and `--lora_alpha` are independent arguments, so raising the rank alone leaves the update scaled by `lora_alpha / rank`. Default both to 16, which keeps the scale at 1, and add a README section explaining the ratio. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * fix validation when embeddings are precomputed or only the final pass runs Two fixes from review: `QwenImage21ValidationPipeline` substituted the cached image-pad mask after calling the base `encode_prompt`, but that call raises for a missing mask when it is handed `prompt_embeds` together with a condition image, so the substitution never ran and image-to-image validation failed. Pass the mask in before the call instead. `--final_validation_prompt` was accepted by the final-inference guard, but the prompt embeddings were only built under `--validation_prompt`, so the final pass reached a pipeline whose text encoder is `None` with no prompt at all. Build the embeddings for whichever prompt is set. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * fix encoding under --offload without --cache_latents `vae.encode` sat outside the `offload_models` block, so the context manager had already moved the VAE back to the CPU by the time it ran, while the prepared dataloader hands the batch over on the accelerator: RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same Move the call inside the block and cover that flag combination, which is the only path that encodes pixels inside the training loop. The image-to-image trainer already had it right. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * add LoRA tests for the Qwen-Image 2.1 pipeline Follow the Flux layout and reuse `LoraTesterMixin` and `LoraMemoryTesterMixin`, so the adapters these trainers produce are covered by the pipeline's own loading, fusing and memory-offload tests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * fix caching a prompt mask that `encode_prompt` returns as None `encode_prompt` drops the prompt mask when nothing in the batch is padded, which is the common case for `--caption_column` datasets since captions in a bucket often tokenize to the same length. The cache is filled per sample, so slicing that `None` raised during latent caching: prompt_embeds_mask_cache[idx] = prompt_embeds_mask[i : i + 1] TypeError: 'NoneType' object is not subscriptable Materialize the mask before filling the cache, and cover the path with a test: none of the existing ones passed `--caption_column`, which is why this only showed up on a real dataset. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * pass rank and alpha explicitly in the README dog example The dog example was written for the old 4/4 defaults; now that both default to 16, spell out 4/4 so the documented command keeps training the size it was tuned for. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * drop copy-paste leftovers from other trainers `--output_dir` defaulted to `hidream-dreambooth-lora`, the weight-decay help mentioned UNet params, and the VAE cast comment talked about the Flux VAE. All three were inherited verbatim from the trainer this was derived from. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * lower the README dog example to lr 1e-4 At 2e-4 the example memorizes the five training images: prompted scenes collapse to the training backdrop rather than following the prompt. 1e-4 is what the image-to-image example already uses. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011qnMKe5MVc7B4XXHGPvXuZ * skip the LoRA scale test for the Qwen-Image 2.1 dummy Halving the LoRA scale moves this dummy transformer's output by less than the tolerance the assertion allows, so the two outputs compare equal and the test fails on CPU. Skip it with the measurement recorded rather than loosening a tolerance shared with every other pipeline's LoRA tests. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> |
||
|
Sayak Paul
|
0121a91f9d |
test and docs for xet dedup based file sharing. (#14835)
* test and docs for xet dedup based file sharing. * up[ |
||
|
Dhruv Nair
|
f2edf6ea30 |
Add single file support for Qwen 2.1 (#14836)
update |
||
Shivam Shrirao
|
80c7ed262a |
Fix flash/sage varlen prep under torch.compile with dynamic shapes (#14568)
* Fix flash/sage varlen prep under torch.compile with dynamic shapes * Handle zero-length sequences in varlen offsets * Add varlen dynamic compile test * Remove private helper test |
||
|
Wang, Yi
|
7a4e1e3c5f |
extend tests/others/test_cli_commands.py::TestRunCommand::test_group_… (#14697)
extend tests/others/test_cli_commands.py::TestRunCommand::test_group_offload_arg tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py::TestDreamLiteMobilePipelineIntegration tests/pipelines/dreamlite/test_pipeline_dreamlite.py::TestDreamLitePipelineIntegration Signed-off-by: Wang, Yi <yi.a.wang@intel.com> |
||
|
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> |
||
|
Sayak Paul
|
d77d530445 |
[torchao]: fix how FqnConfig is resolved and handled (#14686)
* fix: how fqnconfig passed to torchaoconfig is resolved. * testing with quantize_() * change torchao version in setup.py as well. |
||
Akshan Krithick
|
83107dcbb9 |
fix modular pipeline class resolution for families outside the auto-pipeline mappings (#14736)
* resolve modular pipeline classes for families outside the auto-pipeline mappings and fail clearly without blocks * test modular class resolution outside the auto-pipeline mappings * register ltx in the auto video mappings instead of resolving by pipeline folder * drop the unreachable workflow check and rename the resolution test * move the resolution asserts into the fallback test --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> |
||
|
Sayak Paul
|
04e724ce51 |
[tests] partially fix the sd2 upscale tests (#14689)
partially fix the sd2 upscale tests |
||
|
Sayak Paul
|
d30c748f5f |
tests: add lora tester mixin for modular and migrate minimax h3. (#14605)
* tests: add lora tester mixin for modular and migrate minimax h3. * enable lora testing modular ci. * remove seed. * up |
||
Suryansh Sijwali
|
c5469b7ceb |
Invalidate HookRegistry child-registries cache on enable/disable cache (#14093)
* Invalidate HookRegistry child-registries cache on enable/disable cache * Test the enable-cache-after-inference flow through a pipeline * Make the FirstBlockCache regression test drive two explicit cache_context forwards * Drop the pipeline-level enable-cache-after-inference test --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> |
||
|
Akshan Krithick
|
253be5e645 |
add modular blockset for wan 2.2 vace (#14432)
* add modular blockset for wan 2.2 vace * drop transformer from vace encoder components, adapt tests to new mixin structure * regenerate stale vace blocks docstring * generalize wan22 loop denoiser for vace, drop the vace copy * Delete image.png * fold trim_latents into the vace core denoise step --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> |