update for code quality check - #6
Merged
Merged
Conversation
lawrence-cj
added a commit
that referenced
this pull request
Dec 23, 2024
2. add multi-scale dataset and dataloader; 3. add multi-scale training config files; 4. code update;
lawrence-cj
pushed a commit
that referenced
this pull request
Jan 12, 2026
* Initial LTX 2.0 transformer implementation * Add tests for LTX 2 transformer model * Get LTX 2 transformer tests working * Rename LTX 2 compile test class to have LTX2 * Remove RoPE debug print statements * Get LTX 2 transformer compile tests passing * Fix LTX 2 transformer shape errors * Initial script to convert LTX 2 transformer to diffusers * Add more LTX 2 transformer audio arguments * Allow LTX 2 transformer to be loaded from local path for conversion * Improve dummy inputs and add test for LTX 2 transformer consistency * Fix LTX 2 transformer bugs so consistency test passes * Initial implementation of LTX 2.0 video VAE * Explicitly specify temporal and spatial VAE scale factors when converting * Add initial LTX 2.0 video VAE tests * Add initial LTX 2.0 video VAE tests (part 2) * Get diffusers implementation on par with official LTX 2.0 video VAE implementation * Initial LTX 2.0 vocoder implementation * Use RMSNorm implementation closer to original for LTX 2.0 video VAE * start audio decoder. * init registration. * up * simplify and clean up * up * Initial LTX 2.0 text encoder implementation * Rough initial LTX 2.0 pipeline implementation * up * up * up * up * Add imports for LTX 2.0 Audio VAE * Conversion script for LTX 2.0 Audio VAE Decoder * Add Audio VAE logic to T2V pipeline * Duplicate scheduler for audio latents * Support num_videos_per_prompt for prompt embeddings * LTX 2.0 scheduler and full pipeline conversion * Add script to test full LTX2Pipeline T2V inference * Fix pipeline return bugs * Add LTX 2 text encoder and vocoder to ltx2 subdirectory __init__ * Fix more bugs in LTX2Pipeline.__call__ * Improve CPU offload support * Fix pipeline audio VAE decoding dtype bug * Fix video shape error in full pipeline test script * Get LTX 2 T2V pipeline to produce reasonable outputs * Make LTX 2.0 scheduler more consistent with original code * Fix typo when applying scheduler fix in T2V inference script * Refactor Audio VAE to be simpler and remove helpers (#7) * remove resolve causality axes stuff. * remove a bunch of helpers. * remove adjust output shape helper. * remove the use of audiolatentshape. * move normalization and patchify out of pipeline. * fix * up * up * Remove unpatchify and patchify ops before audio latents denormalization (#9) --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Add support for I2V (#8) * start i2v. * up * up * up * up * up * remove uniform strategy code. * remove unneeded code. * Denormalize audio latents in I2V pipeline (analogous to T2V change) (#11) * test i2v. * Move Video and Audio Text Encoder Connectors to Transformer (#12) * Denormalize audio latents in I2V pipeline (analogous to T2V change) * Initial refactor to put video and audio text encoder connectors in transformer * Get LTX 2 transformer tests working after connector refactor * precompute run_connectors,. * fixes * Address review comments * Calculate RoPE double precisions freqs using torch instead of np * Further simplify LTX 2 RoPE freq calc * Make connectors a separate module (huggingface#18) * remove text_encoder.py * address yiyi's comments. * up * up * up * up --------- Co-authored-by: sayakpaul <spsayakpaul@gmail.com> * up (huggingface#19) * address initial feedback from lightricks team (huggingface#16) * cross_attn_timestep_scale_multiplier to 1000 * implement split rope type. * up * propagate rope_type to rope embed classes as well. * up * When using split RoPE, make sure that the output dtype is same as input dtype * Fix apply split RoPE shape error when reshaping x to 4D * Add export_utils file for exporting LTX 2.0 videos with audio * Tests for T2V and I2V (#6) * add ltx2 pipeline tests. * up * up * up * up * remove content * style * Denormalize audio latents in I2V pipeline (analogous to T2V change) * Initial refactor to put video and audio text encoder connectors in transformer * Get LTX 2 transformer tests working after connector refactor * up * up * i2v tests. * up * Address review comments * Calculate RoPE double precisions freqs using torch instead of np * Further simplify LTX 2 RoPE freq calc * revert unneded changes. * up * up * update to split style rope. * up --------- Co-authored-by: Daniel Gu <dgu8957@gmail.com> * up * use export util funcs. * Point original checkpoint to LTX 2.0 official checkpoint * Allow the I2V pipeline to accept image URLs * make style and make quality * remove function map. * remove args. * update docs. * update doc entries. * disable ltx2_consistency test * Simplify LTX 2 RoPE forward by removing coords is None logic * make style and make quality * Support LTX 2.0 audio VAE encoder * Apply suggestions from code review Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Remove print statement in audio VAE * up * Fix bug when calculating audio RoPE coords * Ltx 2 latent upsample pipeline (huggingface#12922) * Initial implementation of LTX 2.0 latent upsampling pipeline * Add new LTX 2.0 spatial latent upsampler logic * Add test script for LTX 2.0 latent upsampling * Add option to enable VAE tiling in upsampling test script * Get latent upsampler working with video latents * Fix typo in BlurDownsample * Add latent upsample pipeline docstring and example * Remove deprecated pipeline VAE slicing/tiling methods * make style and make quality * When returning latents, return unpacked and denormalized latents for T2V and I2V * Add model_cpu_offload_seq for latent upsampling pipeline --------- Co-authored-by: Daniel Gu <dgu8957@gmail.com> * Fix latent upsampler filename in LTX 2 conversion script * Add latent upsample pipeline to LTX 2 docs * Add dummy objects for LTX 2 latent upsample pipeline * Set default FPS to official LTX 2 ckpt default of 24.0 * Set default CFG scale to official LTX 2 ckpt default of 4.0 * Update LTX 2 pipeline example docstrings * make style and make quality * Remove LTX 2 test scripts * Fix LTX 2 upsample pipeline example docstring * Add logic to convert and save a LTX 2 upsampling pipeline * Document LTX2VideoTransformer3DModel forward pass --------- Co-authored-by: sayakpaul <spsayakpaul@gmail.com>
lawrence-cj
pushed a commit
that referenced
this pull request
Jun 15, 2026
…gingface#13815) * feat(pipelines): add DreamLite text-to-image and image-edit pipelines Add ByteDance's DreamLite model family to diffusers. DreamLite is a UNet-based diffusion model that supports both text-to-image generation and reference-image editing through a shared 3-branch dual-CFG design. Two pipelines are shipped: * DreamLitePipeline - full 3-branch dual CFG (negative, reference, prompt); supports T2I and I2I editing at 1024x1024. * DreamLiteMobilePipeline - distilled single-branch variant for on-device inference; no CFG. New model code (all isolated under *_dreamlite.py / unet_dreamlite.py to avoid touching shared upstream files): * models/transformers/transformer_2d_dreamlite.py - DreamLite 2D transformer block. * models/unets/unet_dreamlite.py - DreamLiteUNetModel. * models/unets/unet_2d_blocks_dreamlite.py - DreamLite-specific down/up/mid blocks. * models/resnet_dreamlite.py - DreamLite ResNet variants. * models/attention_processor.py - add DreamLiteAttnProcessor2_0 (pure addition, no existing processor modified). Pipeline + tests + docs: * pipelines/dreamlite/{__init__.py, pipeline_dreamlite.py, pipeline_dreamlite_mobile.py, pipeline_output.py}. * tests/pipelines/dreamlite/{test_pipeline_dreamlite.py, test_pipeline_dreamlite_mobile.py} with the standard PipelineTesterMixin suite; setUp/tearDown auto-patches encode_prompt with a fake so MagicMock text encoders work without per-test boilerplate. * Skip 8 mixin tests that don't apply to DreamLite (MagicMock serialisation, custom attention processor, encode_prompt return shape, batch_size > 1 sweep), mirroring SD3 / Flux conventions. * docs/source/en/api/pipelines/dreamlite.md + _toctree.yml entry (alphabetically between DiT and EasyAnimate). * Register exports in 6 __init__.py files. Two real bugs surfaced by the mixin test suite are fixed in this commit: * num_images_per_prompt > 1: prompt_embeds and text_attention_mask are now repeated along the batch dimension in both pipelines' T2I and I2I branches before being passed to the UNet. * vae=None: __init__ now guards the encoder_block_out_channels lookup so encode_prompt can be tested in isolation per PipelineTesterMixin convention. SlowTests real-checkpoint resolution is set to 1024x1024 (the only size DreamLite is trained for). Test result: 27 passed, 50 skipped, 0 failed on CPU fast suite. make style && make quality: clean. * docs+tests(pipelines/dreamlite): pin Hub repos to `diffusers` branch The `carlofkl/DreamLite-{base,mobile}` Hub repos host two flavours of the same checkpoint: * `main` branch - keeps `model_index.json` pointing at ByteDance's internal package path so the original (non-diffusers) reference code can still load these weights. * `diffusers` branch - rewrites the `unet` entry of `model_index.json` to `["diffusers", "DreamLiteUNetModel"]` so this integration loads correctly from `diffusers`. This commit pins every `from_pretrained(...)` call shipped with the diffusers integration (docs examples, pipeline docstrings, SlowTests) to `revision="diffusers"`. Local-override env vars (DREAMLITE_BASE_PATH / DREAMLITE_MOBILE_PATH) still bypass the revision pin. * chore(pipelines/dreamlite): sync `# Copied from` blocks + dummy objects after rebase Mechanical changes after rebasing onto current `main`: * `pipeline_dreamlite.py::retrieve_timesteps` — re-synced from `diffusers.pipelines.flux.pipeline_flux.retrieve_timesteps` (PEP 604 type hints, expanded docstring, plus the new `accepts_timesteps` / `accept_sigmas` introspection guards). DreamLite's default code path uses `num_inference_steps` (uniform schedule) and never passes custom `timesteps` / `sigmas`, so the added guards are dead-code for this pipeline — behaviour is unchanged. * `dummy_pt_objects.py` / `dummy_torch_and_transformers_objects.py` — registered the dummy classes auto-generated by `make fix-copies` for `DreamLiteTransformer2DModel`, `DreamLiteUNetModel`, `DreamLitePipeline`, `DreamLiteMobilePipeline`, `DreamLitePipelineOutput`. Generated by `make fix-copies`. No hand edits. * docs(dreamlite): register attention processor + split combined docstring entries - Register DreamLiteAttnProcessor2_0 in docs/source/en/api/attnprocessor.md (fixes check_support_list.py). - Split combined 'height / width' and 'guidance_scale / image_guidance_scale' entries in the two pipeline docstrings; add a complete Args block to DreamLiteTransformer2DModel.forward (fixes check_forward_call_docstrings.py). No behavioral change. * refactor(dreamlite): address review feedback from huggingface#13815 - Inline the down/up block factories and define DreamLiteCrossAttn{,NoSelfAttn}{Down,Up}Block2D directly (review #1, #2) - Rename DownBlock2DDreamLite/UpBlock2DDreamLite to DreamLiteDownBlock2D/DreamLiteUpBlock2D to match diffusers naming conventions (review #3, #4) - Merge unet_2d_blocks_dreamlite.py into unet_dreamlite.py to mirror recent transformer model files (review #5) - Wire max_sequence_length into the tokenizer call for generate mode (review #6) - Replace hard-coded drop_idx values (64/34) with self.prompt_template_encode_*_start_idx attributes plus a comment explaining how the offsets are derived (review #7, #8) - Drop the manual Image.resize call and rely on VaeImageProcessor's LANCZOS default in preprocess(image, height, width) (review #9) - Use self.guidance_scale / self.image_guidance_scale properties in the CFG combine instead of the underscore-prefixed attributes (review #10, #11) - Inline retrieve_latents / retrieve_timesteps / calculate_shift in the mobile pipeline with `# Copied from` markers, removing the cross-pipeline imports (review #12) - Add `# Copied from` marker to _extract_masked_hidden in the mobile pipeline (review huggingface#13) * refactor(dreamlite): address dg845 follow-up review - Merge resnet_dreamlite.py (DepthwiseSeparableConv + ResnetBlock2DDreamLite) into unet_dreamlite.py and delete the standalone module (review #1) - Move DreamLiteAttnProcessor2_0 from attention_processor.py into unet_dreamlite.py to keep all DreamLite-specific code in one place; update docs autodoc reference accordingly (review #2) - Drop the PyTorch 2.0 hasattr/ImportError check in DreamLiteAttnProcessor2_0.__init__ (diffusers already requires torch>=2.0; matches Wan deprecation) (review #3) - Drop the deprecated `scale` argument handling from DreamLiteAttnProcessor2_0.__call__ (new model, no legacy callers) (review #4) - Switch SDPA call to dispatch_attention_fn so all diffusers attention backends (FlashAttention, FlashAttention-3, sageattention, etc.) are selectable (review #5) - Rename block dispatch keys in _get_{down,mid,up}_block_dreamlite to match the Python class names (DreamLiteCrossAttn{Down,Up}Block2D / DreamLiteCrossAttnNoSelfAttn{Down,Up}Block2D / DreamLiteUNetMidBlock2DCrossAttn / DreamLite{Down,Up}Block2D); default down/up/mid block_types in DreamLiteUNetModel and the test fixtures are updated to the new keys (review #6, #7); the carlofkl/DreamLite-{base,mobile} (diffusers branch) Hub configs are being updated in lock-step - Localize retrieve_latents inside pipeline_dreamlite.py with a `# Copied from` marker, removing the cross-pipeline import; mirrors the mobile pipeline (review #8) - Add a check_inputs() method to both DreamLitePipeline and DreamLiteMobilePipeline (mobile uses `# Copied from`); call it from __call__; pulls the image-type validation out of prepare_image_latents and adds prompt-type and h/w-divisibility checks (review #9) * fix(dreamlite): correct Q/K/V layout for dispatch_attention_fn dispatch_attention_fn expects (batch, seq, heads, head_dim) and handles the transpose internally; the previous code passed (batch, heads, seq, head_dim), which collided with the dispatch's internal transpose and broke inference (RuntimeError: tensor size mismatch at non-singleton dimension 1). * test(dreamlite): swap MagicMock for tiny real Qwen3-VL fixture Address dg845's review: rebuild the DreamLite fast-test fixture around a real (tiny) Qwen3VLForConditionalGeneration + Qwen3VLProcessor so the standard PipelineTesterMixin save/load, dtype, and offload tests run end-to-end against the actual encode_prompt code path. Override DreamLiteUNetModel.set_default_attn_processor to reinstall the GQA processor so mixin utilities that round-trip through it keep working. * Apply style fixes * fix(dreamlite): address blocking review issues from huggingface#13815 - Override _no_split_modules / _repeated_blocks on DreamLiteUNetModel with the actual DreamLite class names (BasicTransformerBlockDreamLite, ResnetBlock2DDreamLite, DreamLiteCrossAttnUpBlock2D, DreamLiteUpBlock2D) so device_map="auto" and compile_repeated_blocks() match correctly. - Keep attention masks as bool tensors in DreamLiteTransformer2DModel instead of converting them to dense additive float biases. The dense format hard-raises on flash / _flash_3 / _sage backends in dispatch_attention_fn (which requires dtype == torch.bool). - Add explicit parentheses around each clause in check_inputs's mixed and/or condition (both pipelines) for readability. - Replace nn.Module.__init__(self) with ModelMixin.__init__(self) in DreamLiteUNetModel.__init__ so mixin state (e.g. _gradient_checkpointing_func) is properly initialised. ConfigMixin / PushToHubMixin don't define their own __init__, so this covers the full chain without re-running UNet2DConditionModel.__init__. * fix(dreamlite): forward all processor outputs to Qwen3VL text encoder Recent versions of Qwen3VLProcessor add an mm_token_type_ids output, and Qwen3VLModel.compute_3d_position_ids raises ValueError whenever multimodal inputs are present (image_grid_thw is not None) but mm_token_type_ids is None. encode_prompt previously forwarded only input_ids / attention_mask / pixel_values / image_grid_thw, dropping the new field and breaking the fast pipeline tests against transformers main. Switch to ``self.text_encoder(**tk_out, output_hidden_states=True)`` (matching NucleusMoEImagePipeline) so all processor outputs are forwarded automatically and future additions don't regress this path. * Apply style fixes * docs(dreamlite): address final review nits from huggingface#13815 - Replace broken cat.png URL in editing examples (both base and mobile) with the standard `huggingface/documentation-images` source used elsewhere in the diffusers docs. - Promote the recommended guidance_scale=3.5 / image_guidance_scale=1.5 to the default values of DreamLitePipeline.__call__, and drop the now-redundant explicit args from the docs examples. - Switch the EXAMPLE_DOC_STRING examples in both pipelines from torch.float16 to torch.bfloat16 for consistency with the rest of the docs. --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
lawrence-cj
pushed a commit
that referenced
this pull request
Jun 15, 2026
* [.ai] add self-review skill, retire parity-testing skill, and tighten the agent guides - New `self-review` skill mirroring the `@claude` CI review (rubric from review-rules.md, call-path dead-code analysis), report-only, with the report flagging what to fix before submitting (blocking + dead code) vs what to leave for the actual review. - Remove the WIP `parity-testing` skill; preserve its pitfalls as `model-integration/pitfalls.md` (numerical-discrepancy reference). - model-integration: restructure around a grouped checklist, default-to-modular, an overall file-structure sketch (details deferred to the guides), a fresh-conversion `Model parity test` example (internal, not shipped), and a filled-in weight/checkpoint-conversion section. - Centralize the loading rule (from_pretrained / from_single_file, no custom loaders) in models.md; add per-folder File structure sections to models.md / pipelines.md; default-to-modular note in pipelines.md. - AGENTS.md: dedicated 'Self-review before a PR' and 'Reference guides' sections. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [.ai] simplify pitfalls #6 and drop the model-storage / injection-test entries Trim pitfall #6 to the essential point (small dtype diffs compound into a large final difference), remove the `/tmp` model-storage and incomplete-injection-test pitfalls, and renumber 1-16 with cross-references updated. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [.ai] drop parity-harness-specific pitfalls With the parity-testing skill gone, remove the stale-test-fixtures pitfall (saved tensors / cross-pipeline fixtures no longer apply) and de-jargon the noise-dtype detection note. Keeps the pitfalls list generic to numerical discrepancy. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [.ai] trim pitfalls to a concise possible-causes reference Drop the variable-shadowing and decoder-config pitfalls and the noise-dtype 'Detection' aside, tighten the remaining entries, renumber 1-12 (cross-refs updated), and reframe the intro as a non-checklist reference list of possible causes to consult only when outputs don't match. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Apply suggestion from @yiyixuxu * Apply suggestion from @yiyixuxu * [docs] update contributing guide for the self-review skill Replace the retired parity-testing skill with self-review in the skills list, and add a 'Self-review before opening' step to the AI-assisted contributions section: run the self-review skill / review-rules, fix blocking issues + dead code, and treat the @claude CI review as a non-authoritative helper (note any intentional skips in the PR for the reviewer). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Apply suggestions from code review Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * [.ai] fix dangling pitfalls ref and broaden self-review scope - Drop the broken 'pitfalls.md #10' reference in the conversion step (the /tmp model-storage pitfall was removed); save to a local path instead. - Self-review now reviews the whole diff, not just src/diffusers/ and .ai/ — a contributor should review their own tests/docs/scripts too (the CI's scoping is a safety measure for untrusted PRs). Reword to 'same rubric as the CI'. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
DN6
pushed a commit
that referenced
this pull request
Sep 18, 2026
* 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>
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
2 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.
update for code quality check