* [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>
* [LoRA] convert Z-Image LoRAs that only carry the fused `qkv` and bare `out` attention keys
A LoRA trained on Z-Image's original module names (musubi-tuner, LyCORIS)
has `attention.qkv` and `attention.out` and nothing else. The non-diffusers
converter dropped both on the assumption that split `to.q/k/v` and `to_out.0`
keys are also present (the Anime-Z layout), which emptied such a LoRA and then
raised on its leftover `.alpha` keys. Split the fused key into to_q/to_k/to_v
like the single-file converter does and map bare `out` to `to_out.0` when no
split keys exist; the Anime-Z case keeps skipping the redundant keys.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
* [LoRA] review follow-ups: shorter comment with the layout's origin, drop the private-method test
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---------
Co-authored-by: christopher5106 <christopher@scenario3d.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Both layout steps build position_ids in float64 on CPU and then move it to
the execution device with a plain .to(device), which raises on MPS. Pass the
dtype through maybe_adjust_dtype_for_device so the transfer lands as float32
on mps/npu/neuron and is unchanged everywhere else. The fp64 construction is
untouched, and MiniMaxH3RotaryPosEmbed.forward already casts position_ids to
float32, so no device sees different rotary angles.
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* 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>
* [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>
A tag is mutable: whoever controls the action's repository can move
`v4` to different code, and the next run picks it up with no diff to
review. Pinning to the commit SHA freezes the code that runs today.
Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com>
A job with no `permissions:` block inherits whatever the repository
hands out. Each block added here grants what that job's own steps
need and nothing more.
Co-authored-by: hf-security-analysis[bot] <265538906+hf-security-analysis[bot]@users.noreply.github.com>
Use hf auth login in the LPL example error message
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
- TransformerSpatioTemporalModel.forward: no num_frames parameter
- get_2d_rotary_pos_embed_lumina: parameters are len_h / len_w, not grid
- AutoencoderKLHunyuanImageRefiner / AutoencoderKLHunyuanVideo15.tiled_decode:
no return_dict parameter; they return a plain tensor
Co-authored-by: revanth-045 <revanthreddych159@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
* fix: component device placement for split-device pipelines
DiffusionPipeline.device/`_execution_device` (#14383) already prefer any
non-CPU, non-meta component so a pipeline with components split across
devices (e.g. text encoder on CPU, denoising backbone on an accelerator)
reports the right execution device. But several component call sites
still moved a tensor to that execution `device` before handing it to a
specific submodule, instead of the submodule's own device -- so encoding
a prompt, encoding an IP-Adapter image, or decoding latents would crash
with a device-mismatch error as soon as that submodule wasn't already
sitting on the same device as everything else.
Fixes three component types across the major pipeline families
(StableDiffusionPipeline, StableDiffusionXLPipeline,
StableDiffusion3Pipeline, FluxPipeline, Flux2(Klein)Pipeline, WanPipeline,
WanImageToVideoPipeline):
- Text encoder(s): `encode_prompt` / `_get_t5_prompt_embeds` /
`_get_clip_prompt_embeds` now read the encoder's own `.device` instead
of trusting the passed-in `device`.
- VAE: the `vae.encode`/`vae.decode` calls in each `__call__`, plus
`Flux2Pipeline.prepare_image_latents` (whose output also wasn't being
moved back to the execution device).
- image_encoder (IP-Adapter): `encode_image`, plus
`run_safety_checker`'s feature-extractor input.
Also fixes a cross-device bug in `WanImageToVideoPipeline.prepare_latents`
and `Flux2KleinInpaintPipeline.prepare_latents`, where a vae-encoded
tensor and the freshly-sampled noise tensor could end up on different
devices before being combined by the scheduler.
Fixed the canonical/root pipelines by hand, then ran `make fix-copies` to
propagate the text-encoder/image-encoder fixes to every `# Copied from`
derivative.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
* review:reduce scope
* fix: fix not failing test while offloading
---------
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
`_offload_to_disk` releases the onloaded tensors at the end, returning their device memory to the allocator while the compute stream may still be reading them. The memory is then reused by the next onload, so the in-flight kernel reads foreign data and the output becomes NaN. `_offload_to_memory` already synchronizes for this reason, do the same on the disk path.
Reproducer: tests/models/autoencoders/test_models_autoencoder_vidtok.py::TestAutoencoderVidTokMemory::test_group_offloading_with_disk[leaf_level-False]
* 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>
