feat: add reference image KV cache for QwenImageDiT

[bingchenlll]

Summary

• Add block-level static image KV cache support in QwenImageDiT to reuse K/V for reference (input) image tokens across denoising steps.
• Refactor Qwen image attention/block flow to support cached-image KV execution while keeping normal path unchanged.
• Add pipeline-level use_image_kv_cache override, with cache clear/restore guards to avoid stale cache leakage between runs.

Why

• Reduce redundant attention computation on static image-token regions in later denoising passes.
• Make cache behavior controllable per inference call for easier testing and rollout.
• Now only do forward pass at the first denoise step, then cache the K/V are later step attention operations.
• With step-distilled loras (4/8 steps), it saves 35% inference cost at 1024px resolution without quality loss.

Tested

• Run Qwen image generation with cache disabled (use_image_kv_cache=False) and verify output path is unchanged.
• Run with cache enabled (use_image_kv_cache=True) and verify generation succeeds.
• Compare enabled vs disabled outputs for quality sanity (no obvious regressions).
• Verify cache is cleared after run and restored correctly when override is used.

Notes

• This PR adds comments in the transformer loop clarifying why cache control is passed to each block.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant optimization for the QwenImageDiT model by implementing a static image Key-Value (KV) cache. The primary goal is to reduce computational overhead during image generation by reusing attention keys and values for static image tokens across multiple denoising steps. This enhancement not only improves inference efficiency, particularly for high-resolution images, but also provides granular control over cache behavior per inference call, facilitating easier testing and deployment without compromising output quality.

Highlights

• Image KV Cache: Added block-level static image KV cache support in QwenImageDiT to reuse K/V for reference (input) image tokens across denoising steps.
• Refactoring: Refactored Qwen image attention/block flow to support cached-image KV execution while keeping the normal path unchanged.
• Cache Control: Implemented pipeline-level use_image_kv_cache override, with cache clear/restore guards to prevent stale cache leakage between runs.
• Performance Improvement: Reduced redundant attention computation on static image-token regions, leading to a 35% inference cost saving at 1024px resolution without quality loss for step-distilled loras (4/8 steps).

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• diffsynth_engine/models/qwen_image/qwen_image_dit.py
  • Added dataclass import and defined ImageTokenCache to store static image KV data.
  • Refactored QwenDoubleStreamAttention by extracting several helper methods for projecting QKV, normalizing QK, and applying rotary embeddings, improving code modularity.
  • Introduced forward_with_cached_image_kv method in QwenDoubleStreamAttention to handle attention computations when static image KV is cached.
  • Added _image_token_cache attribute and a suite of private methods (_resolve_cache_state, _prepare_cached_image, _align_modulate_index_to_image, _apply_img_mlp_residual, _apply_txt_mlp_residual, _get_static_token_indices, _build_static_kv_cache) to QwenImageTransformerBlock for managing and utilizing the image token cache.
  • Modified QwenImageTransformerBlock.forward to conditionally use the new caching logic based on use_image_token_cache and the cache's readiness.
  • Added use_image_token_cache parameter to QwenImageDiT's constructor and as an instance attribute.
  • Updated QwenImageDiT.forward to pass the use_image_token_cache flag to individual transformer blocks.
  • Implemented clear_image_token_caches and set_image_token_cache_enabled methods in QwenImageDiT for external control over the cache state.
• diffsynth_engine/pipelines/qwen_image.py
  • Added use_image_kv_cache as an optional parameter to the __call__ method, allowing pipeline-level control over the image KV cache.
  • Integrated logic to enable, disable, and clear the image KV cache on the dit model within the __call__ method, ensuring proper state management and cleanup using a try...finally block.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a K/V cache for static image tokens in QwenImageDiT to optimize inference performance. The implementation includes refactoring the attention mechanism, adding caching logic within the transformer blocks, and providing pipeline-level controls for this new feature. My review identified a critical bug where static image tokens are inadvertently dropped when the cache is active, which would lead to incorrect model outputs. I have provided a detailed code suggestion to resolve this issue. Additionally, I've pointed out a minor redundancy in the pipeline's cache management logic within the finally block and suggested a cleaner approach. While the caching feature is a valuable performance enhancement, addressing the critical bug is essential before merging.

[bingchenlll]

@akaitsuki-ii Please check when available, thanks so much!

[akaitsuki-ii]

@bingchenlll Thank you and I think this is a cool idea! But I have some questions:

1. Without causal constraints in autoregressive model, the use of image KV cache in DiT is an approximate optimization. Does it introduce any loss of image quality? Is there any show cases?
2. We no longer maintain the old pipelines and here is a preview of the future release. Could you please port the changes to this branch?
3. We have some WIP fused kernels conflicting with the rotary kernel, need more consideration on this part