[BugFix] Fix modulate_index shape error in Qwen-Image-Edit Task

[mxuax]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
This PR fixes issue #1094.

Purpose

Fix modulate_index shape error in Qwen-Image-Edit Task
This bug is due to zero_cond_t being true when editing, and the current sp_plan does not include the sharding for the variable modulate_index created. I added a new submodule to handle this

Test Plan

Command:
CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/offline_inference/image_to_image/image_edit.py --model "Qwen/Qwen-Image-Edit-2511" --image input.png --prompt "Add a sunset sky background" --output output_u2r2.png --ulysses_degree 2 --ring_degree 2 --enforce_eager --num_inference_steps 20

input image:

Test Result

Success with the right image.

---

Essential Elements of an Effective PR Description Checklist

• [ - ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• [ - ]The test plan, such as providing test commands.
• [ - ] The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

💡 Codex Review

vllm-omni/vllm_omni/diffusion/attention/backends/flash_attn.py

Lines 26 to 30 in 05a1d17

	if not HAS_FLASH_ATTN:
	raise ImportError(
	"FlashAttentionBackend requires Flash Attention. "
	"Please install one of: fa3-fwd, flash-attention, or flash-attn. "
	"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"

Gate FlashAttention import error on CUDA/ROCm only

Raising ImportError at module import time makes the FLASH_ATTN backend unusable on NPU, even when mindiesd is present. The NPU platform selects FLASH_ATTN when mindiesd is available, but HAS_FLASH_ATTN is computed from FA2/FA3 CUDA/ROCm packages, which are typically absent on NPU, so importing this module now fails before the forward_npu path can be used. This breaks attention on NPU deployments that previously worked; consider deferring the check to forward_cuda or guarding it with a platform check so NPU can still load the backend.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[mxuax]

This BugFix is ready. @ZJY0516 @hsliuustc0106

[hsliuustc0106]

please add the benchmark results from benchmark/diffusion folder

[Copilot]

Pull request overview

This PR fixes a shape mismatch error in the Qwen-Image-Edit task when using sequence parallelism (USP=2 or higher). The bug occurred because the modulate_index tensor was not being sharded correctly when zero_cond_t is enabled in image editing models.

Changes:

• Introduced ModulateIndexPrepare module to encapsulate modulate_index creation logic and enable proper sequence parallel sharding
• Updated _sp_plan configuration to shard the modulate_index output
• Refactored the forward method to use the new module instead of inline tensor creation

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ZJY0516]

I think this is what we need to catch in CI

[gcanlin]

LGTM. It works on NPU as well.

[mxuax]

please add the benchmark results from benchmark/diffusion folder

This PR fixes the bug for edit, but the benchmark is for t2i. Thus, I think it won't influence the result.