- Shard the transformer with `TensorParallelConfig` on CUDA, TPU and Neuron, with
remote-code transformers loaded via `--trust-remote-code`.
- Require `--seed`; reject `--device-map`, `--cpu-offload`, `--lora`,
`--context-parallel` and modular pipelines.
- Wrap `--remote` runs in torchrun.
- Load standard pipelines without an indexed `device_map` (e.g. `cuda:<LOCAL_RANK>`)
and place components after loading, fixing `--context-parallel` on ranks > 0.
* 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>
* 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 28c5326da
Commit 28c5326da (today, 14:03 +0300) renamed Kandinsky6FusedTransformerDecoderBlock's
videoT/audioT submodules to video_dec_block/audio_dec_block, and replaced both
Kandinsky6TimeEmbeddings and Kandinsky6SRTimeEmbeddings's hand-rolled in_layer/out_layer
with diffusers' Timesteps/TimestepEmbedding, purely for naming clarity. No checkpoint
conversion was re-run to match, so every published K6 checkpoint (unchanged since Sep 27,
confirmed via identical blob SHA256 across recent Hub commits) still ships the old names.
Verified against kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers and
kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers: this was silently dropping ~65% of the main
transformer's weights (every visual_transformer_blocks.*.{video,audio}T.* parameter) and
the SR transformer's time embeddings, then crashing with "Cannot copy out of meta tensor"
on enable_model_cpu_offload(). After this revert, all checkpoint keys match exactly
(4471/4471 for the main transformer, 459/459 for the SR transformer).
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
* Revert checkpoint-breaking latent-upscaler renames from 372f29088
Same class of regression as 9ae71f0d0, this time in Kandinsky6SRLatentUpscalerX2Branch
and Kandinsky6SRLatentUpscalerOutputHead: today's "address comments" commit (18:48 +0300)
dropped the private_ prefix from private_mid_blocks/private_upsample/private_blocks/
private_output_proj, and converted OutputHead from nn.Sequential to a plain nn.Module with
named norm/activation/conv submodules, again without a matching checkpoint re-conversion.
Verified against kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers's latent_upscaler: this left
every Kandinsky6SRLatentUpscalerX2Branch parameter (and the OutputHead-based output_proj
used standalone by the x2/x4 branches) unfilled on the meta device, crashing
enable_model_cpu_offload() the same way as the main transformer did before 9ae71f0d0.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
* move vocoder and LU
* add coppied from resolve i2v cond mode
* rename variables
* fix lu rename bug
* Update src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* Update src/diffusers/models/transformers/transformer_kandinsky6.py
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* address yiyixuxu's comments
* fix audio channels bug
* fix tests
* Fix CI failures on the Kandinsky6 PR
- tests/others/test_utils.py: assert the caller file path without assuming the checkout directory is named `diffusers`.
- .github/workflows/pr_tests.yml: raise PYTEST_TIMEOUT to 300 s for the example and pipeline CPU steps. Both run several workers with four threads each on an 8-vCPU runner, and the checkpointing example tests run several training subprocesses in sequence.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
* Avoid staging LFS dedup race in the org push test
The second push in TestModelPushToHub.test_push_to_hub_in_organization sent the same model bytes as the first push. The staging server can deduplicate those bytes against the first push's LFS object and then reject the commit with "LFS pointer pointed to a file that does not exist". Change the weights before the second push so its bytes differ.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
* Revert "Avoid staging LFS dedup race in the org push test"
This reverts commit 35e6cc4ef.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
* Revert "Fix CI failures on the Kandinsky6 PR"
This reverts commit 60cad4341.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
* Apply suggestion from @yiyixuxu
* Apply suggestion from @yiyixuxu
* Apply batched suggestions from code review
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* remove review-only files from the Kandinsky6 PR
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS
* Apply suggestion from @yiyixuxu
* restore .gitignore to match main
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS
* simplify the mono audio handling in _write_audio
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS
* remove the no-op set_attention_backend calls from the Kandinsky6 docs
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS
* give the Kandinsky6 transformer its own output class instead of reusing LTX-2's
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS
* fix the stale pipeline exports and add return annotations for the Kandinsky6 pipelines
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS
* address feedback from @leffff
---------
Co-authored-by: Denis Koposov <denis.koposov@phystech.edu>
Co-authored-by: leffff <levnovitskiy@gmail.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Lev Novitskiy <57654885+leffff@users.noreply.github.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
* [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>
* 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>
* 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>
* docs: fix latent -> latents in LTX latent upsampler docstrings
__call__ takes 'latents' (the docstring said 'latent', and its own
description says 'Input latents to normalize').
* docs: fix 'promt' typo in kandinsky5 default prompt templates
The default system template sent to the Qwen text encoder misspelled
'prompt engineer' as 'promt engineer' (4 pipelines).
Note: this changes the token sequence of the default template, so
default outputs will differ numerically from before / from the
upstream Kandinsky5 reference. If exact parity with the reference
repo is preferred, the reference itself would need the same fix.
* docs: fix docstring parameter names that do not match signatures
- flux controlnet encode_prompt documented a clip_skip parameter that
does not exist (copy-paste from SD-style docs; flux has no CLIP
layer-skipping)
- hunyuan_video1_5 pipelines documented num_images_per_prompt while
the parameter is num_videos_per_prompt
- kandinsky5 _encode_prompt_qwen/_encode_prompt_clip documented
num_videos_per_prompt, which neither method accepts
* docs(zh): fix English anchors pointing at translated Chinese headings
speed-memory-optims.md linked ./memory#model-offloading and
#group-offloading while the zh memory page translates those headings
(模型卸载 / 组卸载); using_peft_for_inference.md linked #hotswapping
while the zh heading is 热切换. Repointed at the Chinese slugs.
* docs: fix PAG anchor and the StableDiffusionPipelineSafe row in the SD overview
- pag.md linked #pagappliedlayers; the heading '### pag_applied_layers'
keeps its underscores (#pag_applied_layers)
- the SD overview table linked ./stable_diffusion_safe, a doc page that
no longer exists anywhere under docs/source; the row now names the
deprecated pipeline as plain text instead of a dead link
* docs: fix dead Masked Im2Im TOC anchors in community pipelines README
The TOC rows linked #stable-diffusion-masked-im2im /
#stable-diffusion-xl-masked-im2im; the actual headings are
'Masked Im2Im Stable Diffusion Pipeline' and 'Masked Im2Im Stable
Diffusion Pipeline XL' (lines 3184/3219).
* docs(examples): drop duplicated 'with' in T5/Gemma sequence-length help strings
16 example scripts' --max_sequence_length help read 'to use with with
the T5 text encoder' (or the Gemma variant); surfaced in --help output.
* fix(links): keep the pag_applied_layers anchor in doc-builder form
doc-builder slugifies headings by dropping underscores, so the heading
'### pag_applied_layers' anchors as #pagappliedlayers — the form the
check-links CI expects. GitHub-style underscores broke it.
Signed-off-by: simpleqt <89645338+simpleqt@users.noreply.github.com>
* fix(kandinsky5): recompute prompt_template_encode_start_idx after the promt fix
Correcting 'promt' -> 'prompt' in the system templates shifts each
template's token count by one (verified with the Qwen2.5-VL tokenizer:
41->40, 55->54, 129->128, 129->128). Without this, the hardcoded start
indices slice the encoded prompt one token off.
---------
Signed-off-by: simpleqt <89645338+simpleqt@users.noreply.github.com>
* Add opt-in mixed W8A8/W8A16 denoising for Cosmos3 ModelOpt FP8 checkpoints.
Keep native ModelOpt GEMM on middle steps and dequant-linear W8A16 on the first/last steps so CFG cond/uncond share one precision per scheduler step.
* Read Cosmos3 mixed-precision schedule from the checkpoint policy.
Enable first/last W8A16 only when transformer/config.json declares diffusion_step_policy, so distilled FP8 stays native W8A8 instead of inheriting a hardcoded 3+3 window.
* Harden Cosmos3 mixed-precision loading and document official Hub fp8 schedules.
Read the checkpoint runtime policy from on-disk transformer/config.json when the live ModelOpt config omits it, fail closed on incomplete policies, and allow FP32 activations on the W8A16 path.
* Address review: simplify FP8 mixed docs, no-op non-ModelOpt backends, drop unit tests.
W8A8/W8A16 is explained without listing call-site overrides; TorchAO and other
quantizers keep their native forwards; focused tests are removed until Hub fp8
usage is clearer.
* Document FP8 mixed default vs none trade-offs without ModelOpt restore boilerplate.
Keep serialized restore in the ModelOpt guide and spell out that all-W8A8 is faster but can flicker on multi-step video.
* Apply style fixes
* Regenerate Cosmos modular auto docstrings for mixed-precision inputs.
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
* Add LTX-2.5 DFR pipeline with generated keyframe slots
Ports DFRPipeline from the Lightricks reference. Stage 1 generates video plus
single-pixel-frame keyframe slots at a fraction of the requested resolution on
a VAE-aligned segment grid; both are spatially latent-upsampled and stage 2
re-denoises at twice that resolution with the slots re-attached and an optional
spatial detailing IC-LoRA active for that stage only. Optional temporal x2/x4
refine rounds tile the canvas at shared keyframes and densify with ancestral
Euler. With spatial_upscalings=2 a full-resolution detailing epilogue follows
the rounds.
The transformer already stored keyframes_abs_pos_embedding for load/save; this
wires it into the forward through a new video_keyframes_mask argument, which
only a keyframes-aware pipeline passes, so other pipelines are unaffected.
The epilogue denoises the whole canvas in one loop and tiles the transformer
call inside it, so every Euler step steps a canvas whose tiles have already
agreed on their overlaps. Spatial tiles blend under a trapezoidal mask, since
neither side of a height or width border holds a known answer. Temporal tiles
are cut on the keyframe seams the last refine round stitched on: both windows
reproduce a shared keyframe there, so the later one drops its run-up under a
rectangular mask rather than averaging it. Conditionings are attached once on
the whole canvas and filtered per tile at the token level, and a keyframe two
windows share is one token they both read.
The epilogue is handed its keyframes rather than asked to generate them. Each
carry plane is decoded on its own -- the VAE is causal, so a stacked decode
would bleed neighbours -- then Lanczos-stretched x2 in RGB and encoded again at
the output resolution, and pinned fully clean. Only the video latent is
spatially upsampled.
Conditioning fps is snapped to 60 above 30 rather than merely capped there, at
every stage. RoPE time is pixel_frame / fps, and the transformer is trained
around 24/25/30 and 60; a temporal round taking 24 fps to 48 lands between
those, and it shows as stutter at the latent borders. Playback fps is
unchanged, so 24 fps with one round still ships 48 fps.
A condition's index is read on the canvas num_frames asks for, and the moment
it names is carried onto each refine round's longer canvas by scaling its pixel
position by 2**round. The scaled position does not generally land on a latent
boundary, so it travels as a pixel index rather than through the public latent
index; a keyframe conditioning is appended as extra tokens instead of being
spliced into the base grid, so it does not need to.
height and width must be divisible by 2**spatial_upscalings times the VAE's
spatial compression ratio, which makes 4K 3840x2176 rather than 3840x2160. That
rule is checked ahead of the looser one every LTX-2 pipeline applies, so the
error names the divisor a DFR caller actually has to satisfy.
Four details are easy to get wrong, and each is covered by a test after showing
up as a visible seam at a tile handover:
- The ancestral step injects noise into every token, so the conditioning blend
has to be re-applied afterwards. Skipping it lets the strength-0.95 anchor
keyframes erode over the schedule, and those anchors are the only thing
pinning adjacent tiles onto the same content.
- Velocity is converted to x0 with each token's own noise level, not the scalar
schedule sigma: a token held at strength s sits at (1 - s) * sigma.
- Each tile draws its ancestral noise from a generator seeded
seed + 1000 * round + tile, kept separate from the main generator so the
draws do not consume state the next tile's initial noising reads.
- Two tiles invent the slot that falls in the later one's dropped lead-in. The
stitch keeps the earlier tile's frames there, so the earlier tile's copy is
the one the canvas holds, and the one the next round must anchor on.
* Split the DFR pipeline into composable stages
Addresses the review on #14567: use the same compose pattern as the other LTX
two-stage pipelines, get ancestral Euler from an existing scheduler, and put the
temporal rounds in their own pipeline so the schedule is not switched mid-call.
`LTX2DFRPipeline.__call__` is now one denoise pass at `height` x `width`. Callers
compose stage 1, `LTX2LatentUpsamplePipeline`, stage 2 and each temporal round,
and a documented recipe is the copy-paste 1080p path. The recipe knobs
(`spatial_upscalings`, `temporal_upscalings`, `detailing_lora_adapter_name`) and
the required upsampler components are gone; `height`/`width` are this pass, not
the final output.
`ancestral_euler_step` is replaced by `LTXEulerAncestralRFScheduler.step` plus a
re-application of the conditioning blend, which ancestral noise would otherwise
erode on the strength-0.95 seam anchors. The new `LTX2DFRTemporalRefinePipeline`
owns that scheduler and one round; stage 1, stage 2 and the epilogue stay on
`FlowMatchEulerDiscreteScheduler`. It refuses any other scheduler rather than
silently taking a deterministic step and returning a softer canvas.
Pack/unpack, `prepare_latents`, `denoise` and `encode_conditions` move to
`LTX2DFRCoreMixin` so neither pipeline subclasses the other. Public latents are
raw on both sides of every boundary, `output_type="latent"` returns the untrimmed
canvas so a slot on the pad is not dropped, and `trim_canvas` does the trim
before decode. `LTX2DFRPipelineOutput` carries `keyframes` and
`keyframe_positions`, which cannot be re-derived after a round.
Verified against the pre-split implementation: bit-exact on dummy components in
fp32 and bf16, with and without the IC-LoRA reference, over one and two rounds;
and within one bf16 ulp on the real checkpoint, where the only difference is that
the split normalizes upsampled latents in fp32 rather than bf16.
Also in this pass:
- `__call__` takes `video_tiles` (the `epilogue_tiles` layout) instead of a
resolved token plan. Resolving one needs the RoPE coordinates that only exist
once `prepare_latents` has run, so a caller could not build the plan at all.
- `rebuild_epilogue_keyframes` is public and takes and returns raw latents. The
composed epilogue needs it, so it was public API in practice while named
private, and its normalized return forced callers into `_denormalize_latents`.
- Drop prompt enhancement and `num_videos_per_prompt` from the temporal pipeline.
Enhancement belongs to stage 1 -- re-running it would denoise the canvas under
a different prompt than the one that generated it -- and the batch is set by
the incoming latent canvas, so `num_videos_per_prompt > 1` only ever raised.
- Fix the docs recipe: `requested_frames` counted latent frames where
`trim_canvas` wants pixel frames, truncating a 241-frame render to 25; the
three pipelines share components, so place them together instead of offloading
one and leaving `temporal_latent_upsampler` off the device; and the detailing
IC-LoRA is applied at 0.5, the strength the reference hardcodes.
- Pass `crf=0` on the conditions in the refine-round test. The default CRF sends
the image through H.264 re-compression, which needs PyAV, so the test failed on
any environment without it while testing nothing about re-compression.
* Drop LTX2DFRCoreMixin in favor of # Copied from
Each DFR pipeline now owns pack/unpack, prepare_latents, and denoise so the
classes stay self-contained. Shared helpers copy from LTX2Pipeline /
LTX2ConditionPipeline, and DFR-specific methods copy from LTX2DFRPipeline.
dfr_core.py keeps only the constants and canvas helpers.
* inline _finalize_output
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* remove dfr_core.py
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* make return_dict=False return a tuple of 4, consistent with diffusers convention; update docs and pipeline output accordingly
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* remove the dfr_layout test suite
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---------
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
* Move LTX2Guidance from /modular_pipelines/ltx2/ to guiders/ to allow LTX-2 guiders to be correctly loaded from checkpoint
* Add regression test for LTX-2 guiders declared as from_pretrained components
Covers the round trip this move fixes: a checkpoint that declares `guider` and
`audio_guider` in `modular_model_index.json` must record a resolvable
(library, class_name) pair and reload its non-default guidance scales. The
existing `test_modular_index_consistency` skips components without a
`pretrained_model_name_or_path`, so `from_config` guiders were uncovered.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Migrate LTX2Guidance docs to guider docs after move
* Improve LTX2Guidance docs and add arg docstrings
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Fibo Edit: multi-reference conditioning, batching, timesteps
* docs: simplify Fibo Edit page, fix EXAMPLE_DOC_STRING to the VLM prompt flow
* docs: point Fibo Edit at the Gemini prompt-to-JSON block and the 1.5 base/turbo ids
briaai/FIBO-edit-prompt-to-JSON is retired in favor of
briaai/FIBO-edit-gemini-prompt-to-JSON, which handles multiple reference
images and masks. Also corrects the checkpoint ids to their canonical
briaai/Fibo-Edit-1.5-base and briaai/Fibo-Edit-1.5-turbo spellings.
* apply review feedback on reference handling and the attention mask
- collapse the mask branch to `if not attention_mask.all():` in both Fibo pipelines
- encode references via `self.image_processor.preprocess(...)`, dropping the manual
numpy normalization and the `_vae_safe_size` helper
- normalize `image` into a reference list once in `__call__`, pass it to `check_inputs`
* make style
* make style
* LTX-2.5 modular pipeline init docs
* Update LTX-2.5 modular docs now that the modular pipeline automatic downloads the Gemma 4 prompt enhancer and processor
* fix some modular docs typos
* fix more typos
* Update docs/source/en/api/pipelines/ltx2.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Remove unused transformers imports from LTX-2.5 modular examples
---------
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
docs: fix typos in docs, comments and docstrings
Prose, comments and docstrings only. No executable line changed.
docs/:
- optimization/memory.md: maxmium -> maximum; heigh -> height (in the
channels_last layout tuple, which read "(batch size, heigh, width,
channels)")
- api/pipelines/bria_fibo.md: proffesional -> professional
- api/pipelines/wan.md: involed -> involved
src/ comments and docstrings:
- guiders/magnitude_aware_guidance.py: supression -> suppression
- models/attention.py: spliting -> splitting
- models/attention_dispatch.py: abritrary -> arbitrary (x2)
- models/transformers/transformer_ltx2.py: timstamps -> timestamps (x2)
- pipelines/ltx2/pipeline_ltx2{,_image2video}.py: corrct -> correct
- pipelines/deprecated/unidiffuser/modeling_uvit.py: tbe -> the
- pipelines/visualcloze/pipeline_visualcloze_{combined,generation}.py:
"whe they are lists" -> "when they are lists" in a user-facing
ValueError message
Deliberately not touched:
- pipeline_kandinsky*.py "promt"/"scren" live inside prompt_template,
which is fed to the text encoder verbatim. Changing them would change
generation output.
- ANE (Apple Neural Engine), MoT (Mixture-of-Transformers): initialisms.
- "racoon": appears in prompt examples whose reference images were
generated with that exact prompt.
- "Shuting Wang": a paper author's name.
* restructure guides as skill references
* reorganize AGENTS.md and references
* clean up
* clean up
* clean up
* fix AGENTS.md references and add check_ai to Makefile
* clean up
* update paths
* clean up
* draft:add neuron as a legit backend
* feat: neuron-specific changes in the pipeline
* tests: eager tests
* draft: start with tp for flux2
* fix: style
* fix:apr_02 beta
* feat:add wan
* fix:pixart
* fix: rewrite flux swiglu activation to avoid gather op in neuron IR
* test: pixart compile mode on neuron
* cleanup & fix style
* merge: another change
* review: cleanup+suggestions
* fix: CIs style
* tests: add test units for tp
* fix: in case of text-encoder(s) on CPU
* review:cleanup+add test
* fix: style
* doc: remove it for now
* Add from_single_file support for SkyReelsV2 and ChronoEdit transformers (#13946)
SkyReels-V2 and ChronoEdit are both built on Wan, and their transformers have
the same keys as WanTransformer3DModel, so they reuse
convert_wan_transformer_to_diffusers (like WanVACE / WanAnimate). This lets the
community GGUF builds load directly.
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
* multi-GPU VAE Fix for Cosmos 3 (#13924)
fix(cosmos3): pin VAE latent norm buffers to encode output device
Under sharded placement (device_map="balanced"), vae.encode() runs on the
VAE's own device while the mean/inv_std buffers were pinned to x.device,
causing a cross-device RuntimeError. Compute raw_mu first, then pin the
normalization buffers to its device so all tensors share one device.
Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* docs: fix repeated word typo in set_timesteps docstring (#13876)
* docs: fix repeated word typo in set_timesteps docstring
Removed the duplicate word "schedule" from the docstring for the sigmas argument in EulerDiscreteScheduler.set_timesteps.
* Update scheduling_euler_discrete.py
* Apply style fixes
---------
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
* clean some stuff to simplify code.
* clean more to remove permutation related shenanigans.
* revert: put torch.chunk back
* Update docs/source/en/training/distributed_inference.md
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Address Sayak's TP review nits (docstrings + stale comments)
- _modeling_parallel: fix stale comment claiming attention processors read
tp_degree at runtime (no longer true after the PackedColwise/Rowwise
refactor; processors are TP-agnostic via head_dim).
- tensor_parallel: single backticks instead of double in docstrings, add
input/output examples to _blocks_to_block_sizes and _resolve_tp_plan, and
drop the stale "applies the Flux2 fused-weight permutations" line from
apply_tensor_parallel (permuters were removed).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Fix stale permuter reference in Flux2ParallelSelfAttnProcessor comment
The proportional QKV/MLP slice is now guaranteed by PackedColwiseParallel
block-by-block sharding, not the removed fused-weight permuter.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Derive TP backend inside apply_tensor_parallel from the mesh device type
Drops the `backend` parameter: apply_tensor_parallel already holds the TP
mesh, so it derives neuron-vs-default from `config._mesh.device_type` itself.
The caller no longer computes it. Documents why torch_device can't be used
(Neuron reports as "cpu" via torch's accelerator API). Removes the now-unused
is_torch_neuronx_available import from modeling_utils.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Simplify TP backend derivation to the mesh device type alone
Drop the redundant is_torch_neuronx_available() guard: a DeviceMesh("neuron", ...)
only exists when running on Neuron, so tp_mesh.device_type is the single source
of truth.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Remove incorrect "Neuron reports as cpu" rationale from TP backend docstring
Verified on the current native torch_neuronx drop that Neuron does surface as
the torch accelerator (torch._C._get_accelerator().type == "neuron",
get_device() == "neuron"); the old "reports as cpu" claim was stale. The
mesh-device-type derivation is unchanged — it's the device being sharded onto.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Drop stale "XLA" labels from Neuron device/runtime comments
Neuron is a native torch backend now (not XLA-routed): remove the "(XLA)"
parenthetical from the randn-on-CPU note in torch_utils and the "XLA runtime"
mention from the Neuron TP test docstring. Behavior unchanged.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* make style: reflow tensor_parallel docstrings
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Revert "revert: put torch.chunk back"
This reverts commit 7ea75f771a.
* fix: sharding for neuron + validate flux 1
* fix: change for check_repository_consistency
* feat: support qwen image as well
* Update docs/source/en/training/distributed_inference.md
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* review: keep bria and nucleusmoe out
* review: keep bria and nucleusmoe out
* review: keep bria and nucleusmoe out
* review: address agent review
* review: apply suggestions
* Update docs/source/en/training/distributed_inference.md
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* review: change example tp size to 4
* review: add flux tp test
* review: address comments
* review: improve the doc
* review:apply qwen rotary_emb per device
* review:remove redundant device def after PR #14383
* review: doc title suggestion
* review: restore copy comment
* review: remove cp divisibility explanation from doc on parallelism guide
* fix style
---------
Co-authored-by: Hz_Zhang <47402297+HaozheZhang6@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Atharva Joshi <atjoshi@nvidia.com>
Co-authored-by: Atharva Joshi <atjoshi@smc521ge-0036.ipp2a2.colossus.nvidia.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Ramkumar R <ramkumarashvanth09@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
* Stable Audio 3 First Version Commit, require model dict loading on GPUs
* Fixing bugs for sa3 VAE
* Make quality fix
* Fix the unit test failure
* Fix SA3 pipeline noise: learned text-padding embedding + no cross-attn mask
The diffusers Stable Audio 3 pipeline produced noise instead of music because
the cross-attention conditioning was built differently from the reference:
- The reference text conditioner uses padding_mode="learned": padded text
positions (~245 of 256 for a short prompt) are filled with a trained
`padding_embedding`, and the DiT attends to *all* positions (its
cross-attention mask is intentionally disabled). Our pipeline instead zeroed
padded positions and masked them out, wiping ~95% of the conditioning signal.
Changes:
- transformer_stable_audio3: add learned `prompt_padding_embedding`; in forward,
replace masked text positions with it (in cond_token_dim space, before
to_cond_embed) then attend to the full context, matching the reference.
- pipeline_stable_audio_3: stop zeroing padded positions (the DiT now handles
them); default `silence_padding_duration` to 6.0 (reference headroom default).
- convert_..._to_diffusers: convert
`conditioner.conditioners.prompt.padding_embedding` into the DiT.
- scheduling_ping_pong: pin schedule endpoints (sigmas[0]=1.0, sigmas[-1]=0.0)
to match the reference LogSNRShift endpoint preservation.
- run_..._inference: coerce float16->float32 on CPU.
- tests: update transformer state-dict expectations (523 tensors, add
prompt_padding_embedding) and ping-pong sigma endpoint assertions.
Verified against the reference with identical noise + identical conditioning:
full 8-step ping-pong trajectories agree to 4e-6 (final latent). Removed the
dev-only parity scripts (verify_*_dit_parity, verify_*_vae_parity) from the tree.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
* Style fix
* Autopipeline available
* Adding document description for testing examples.
* Solving fix for serge review bot suggestion
* Design review comments fix
* CI/CD Fix for ping pong scheduler md file
* Fix doc build: remove autodoc for nonexistent StableAudio3EulerSchedulerOutput
* Fix doc build: qualify non-exported autodoc references with module paths
* Fix details in scheduler and some nits
* suggestion fix
* fix-copies: propagate prepare_latents dtype fix to audio2audio/inpaint pipelines
---------
Co-authored-by: Buffett Liu <jengliu@linkedin.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* [CLI] Clean up and improvements to `diffusers-cli`
Modular pipelines:
- `custom_blocks` also writes `modular_model_index.json`, so a packaged repo is
loadable by `ModularPipeline.from_pretrained` (and runnable by `run`), not just
importable as blocks.
- Add `--cpu-offload auto`, which offloads through the pipeline's
`ComponentsManager`, plus `--offload-margin` to set its `memory_reserve_margin`.
- Add `--workflow` passthrough; warn and ignore it on standard pipelines.
- Route `model_index.json` repos whose `_class_name` is modular to the modular path.
Output saving:
- Resolve every media field on the output object (`images`, `frames`, `audios`,
`audio`) instead of the first match, keeping the batch dimension. Fixes video
output being misrouted through the image branch and LTX2's audio being dropped.
- Fix the `output_type="pt"` video transpose: `postprocess_video` returns
`(B, F, C, H, W)`, not `(B, C, F, H, W)`.
- Save `(B, C, samples)` waveforms as multi-channel wavs (StableAudio).
- Write each video's frames to `<stem>-frames/<NNNN>.png`. Frames need no video
backend, so they replace the `.pt` fallback when `export_to_video` fails.
- Report frame directories as a single path so hundreds of frames don't flood the
terminal; `--push-to` and the sandbox download preserve subfolder structure.
- Raise on unrecognized payloads instead of dumping JSON/pickle.
Remote:
- Add `imageio`, `imageio-ffmpeg`, and `peft` to the default sandbox deps.
- Raise sandbox `start_timeout` to 600s; the 120s default expires while a cold node
pulls the pytorch base image.
* Apply suggestion from @sayakpaul
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* update
* update
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
- `qualty` -> `quality`
- `infered` -> `inferred` (x3)
- `bakend` -> `backend` (x3, in the trailing `# <--- Set attention
bakend to Flex` comments; the actual `set_attention_backend` calls
were already spelled correctly)
- `wothout` -> `without`
Prose and comments only, no code or behavior changes.
Modular Diffusers is no longer experimental: remove the import-time warning
and the docstring warnings from ModularPipeline, the block base classes, and
every model's modular pipeline, and drop it from the docs overview.
ComponentsManager and the Mellon integration are still moving, so they keep
a warning — added one to the ComponentsManager docs page and to
mellon_node_utils, which had none.
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
* Add LoRA support for MiniMax-H3
* Address review: Copied-from loader, LoRA badge, docstring caveats, drop generic scale test
* Move the MiniMax-H3 LoRA tests to tests/lora
* Support the remaining published MiniMax-H3 LoRA layouts
Three more published H3 LoRAs did not load. Two were key-layout gaps and the
third was the reason neither was noticed: a layout that reaches no module at
all used to return without an exception.
- musubi-tuner writes one flattened `lora_unet_` name per module. It is
un-flattened against the H3 module vocabulary, since `qkv_proj`,
`token_refiner`, `final_layer` and the output heads carry underscores that
are not path separators, and then goes through the existing kohya path.
- one producer publishes its own converter's output: diffusers module names
with peft's `.default.` infix left in and no component prefix. The infix is
dropped and the prefix added.
- a state dict that filters to nothing in both partitions now warns instead of
loading as a silent no-op, in the wording `load_lora_adapter` uses for the same
situation.
A fourth loaded, but at the wrong strength. One producer ships no `.alpha` scalars
and records the alpha it trained with in the file's own `__metadata__` instead,
under `alpha`. Its 8-step turbo LoRA pairs that entry's 8 with rank 128, an
effective scale of 0.0625, so synthesizing `alpha == rank` applied the adapter 16x
too strong. That entry is now read as the uniform network alpha. Per-module scalars
still win when a file carries both, since the converter has already folded them into
the weights, and a non-numeric value is warned about and ignored. `__metadata__` only
exists on a file, so `_fetch_state_dict` hands it back on request.
* Simplify the MiniMax-H3 LoRA comments
* Trim the MiniMax-H3 LoRA tests to loading and effectiveness
* Reference the alpha convention in the MiniMax-H3 LoRA docstring
* Drop a redundant format comment
* Test only diffusers-written state dicts
* Rely on get_peft_kwargs for alpha-less files
* Drop the peft version guards after the main cleanup
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Add MiniMax Music 3 transformer, pipeline sub-models, and conversion script
* Add MiniMax Music 3 pipeline, tests, and docs
* Wire example docstring, point diffusers-format references at MiniMaxAI/MiniMax-Music-3
* Use MiniMaxAI/MiniMax-Music3 as the official diffusers weights repo id
* Apply doc-builder style
* Use dtype in examples, demonstrate the structured prompt format
* Use soundfile in examples like the other audio pipelines
* Convert to a modular pipeline and move models under src/diffusers/models
Per review: the conditioner moves to models/condition_embedders, the RVQ depth
decoder and vocoder to models/transformers and models/autoencoders, and the
standard pipeline is replaced by MiniMaxMusic3ModularPipeline (anima-style
blocks, helios-style chunk loop, guider-abstracted CFG with a zeros
unconditional branch). Index configs now reference all components under the
diffusers library. Also applies the transformer review suggestions
(Transformer2DModelOutput, inlined rotary forward, explicit processor default).
* Point modular index loading specs at the Hub repo in the converter
* Apply doc-builder style to the modular blocks docstrings
* Address review-bot comments: overlap comment accuracy, single pipeline-level progress bar, zero-frame duration guard, import ordering, docstring cross-reference
* Support CPU offloading in the AR step: trigger offload hooks before submodule calls (minimax_h3 workaround), clear error when the AR models cannot co-reside
* Docs: PR-install tip and memory-usage section with measured offloading numbers
* Address review comments on the block assembly
Take the individual models rather than `components` in `_generate_depth_codes` and
`_embed_audio_frame`; `num_codebooks` / `audio_vocab_size` come off the depth decoder's
config, so the two pipeline properties that existed only for them are gone.
Assemble the blockset as three top-level children, each owning models and runnable on its
own — generate the semantics once and denoise it with different settings:
semantic_generator MiniMaxMusic3SemanticGenerationStep [tokenize, generate]
denoise MiniMaxMusic3CoreDenoiseStep [prepare_chunks, denoise]
decode MiniMaxMusic3VocoderDecodeStep
`MiniMaxMusic3TextEncoderStep` is renamed `MiniMaxMusic3TokenizeStep` (it only tokenizes)
and the autoregressive leaf `MiniMaxMusic3AutoregressiveStep`, freeing its old name for
the sequential.
Point the tests at `hf-internal-testing/tiny-minimax-music3-modular-pipe`; the previous
fixture repo's component specs pointed at a local path, so nothing loaded off that machine.
Drop the PR-branch install tip from the docs.
20s and 60s generations are bit-identical to before these changes.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---------
Co-authored-by: yiyixuxu <yixu310@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
