Add Qwen-Image 2.1 - #14804
Open
naykun wants to merge 18 commits into
Open
Add Qwen-Image 2.1#14804naykun wants to merge 18 commits into
naykun wants to merge 18 commits into
Conversation
…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.
- 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)
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.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WpU4nuugzCw8T1c4tL6N2A
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.
`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.
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.
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.
`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.
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.
`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__.
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`.
`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 huggingface#5 are fine and stay.
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.
…ised 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.
github-actions
Bot
added
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds support for Qwen-Image 2.1, a unified text-to-image and image-to-image model.
We're happy to open-source Qwen-Image 2.1, the most balanced and best value-for-compute model in the Qwen-Image
family so far.
Thanks @yiyixuxu, @sayakpaul and @stevhliu for the reviews and the help getting this in shape.