[Bug] Fix QwenImageEditPlus Series on NPU by zhangtao0408 · Pull Request #13017 · huggingface/diffusers · GitHub
Skip to content

[Bug] Fix QwenImageEditPlus Series on NPU - #13017

Merged
DN6 merged 13 commits into
huggingface:mainfrom
zhangtao0408:fix_npu_related_error
Feb 17, 2026
Merged

[Bug] Fix QwenImageEditPlus Series on NPU#13017
DN6 merged 13 commits into
huggingface:mainfrom
zhangtao0408:fix_npu_related_error

Conversation

@zhangtao0408

@zhangtao0408 zhangtao0408 commented Jan 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes: #13015
Fixes: #13016

Test Codes:

import torch
import torch_npu
import torch.distributed as dist

import os, time
from PIL import Image
from diffusers import QwenImageEditPlusPipeline, ContextParallelConfig
from diffusers.utils import load_image

# Initialize Env
rank = int(os.getenv("RANK", 0))
world_size = int(os.getenv("WORLD_SIZE", 1))

if world_size > 1 and not dist.is_initialized():
	dist.init_process_group(backend="hccl")
	rank = dist.get_rank()
	device = torch.device("npu", rank % torch.npu.device_count())
	torch.npu.set_device(device)
else:
    device='npu'

image1 = load_image("https://github.com/vipshop/cache-dit/raw/main/examples/data/edit2509_1.jpg")
image2 = load_image("https://github.com/vipshop/cache-dit/raw/main/examples/data/edit2509_2.jpg")
prompt = "The magician bear is on the left, the alchemist bear is on the right, facing each other in the central park square"

pipe = QwenImageEditPlusPipeline.from_pretrained(
    "/PATH/TO/Qwen-Image-Edit-2509",
    torch_dtype=torch.bfloat16
).to(device)
pipe.transformer.set_attention_backend("_native_npu")

pipe.set_progress_bar_config(disable=rank != 0)
pipe.enable_model_cpu_offload(device=device)

if world_size > 1:
    pipe.transformer.enable_parallelism(
        config=ContextParallelConfig(ulysses_degree=world_size)
    )

with torch.inference_mode():
    # Inference
    torch.npu.synchronize()
    start_time = time.time()
    output = pipe(
        image=[image1, image2],
        prompt=prompt,
        generator=torch.Generator(device="cpu").manual_seed(0),
        true_cfg_scale=4.0,
        negative_prompt=" ",
        num_inference_steps=20,
        num_images_per_prompt=1,
        height=1024,
        width=1024,
    )
    torch.npu.synchronize()
    end_time = time.time()
    
    inference_time = end_time - start_time
    if rank == 0:
        output_image = output.images[0]
        output_image.save(f"qwen-image-ulysses{world_size}-time{inference_time:.2f}s.png")
        print(f"image saved at qwen-image-ulysses{world_size}-time{inference_time:.2f}s.png")
  • Run the code
# 1 card
python3 qwen_image_edit_test.py

# 4 card
python3 qwen_image_edit_test.py

Results in comment

Before submitting

Who can review?

cc @yiyixuxu @sayakpaul @asomoza @DN6

@zhangtao0408

zhangtao0408 commented Jan 22, 2026

Copy link
Copy Markdown
Contributor Author

sayakpaul
sayakpaul reviewed Jan 22, 2026
View reviewed changes

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, left some comments.

Comment thread src/diffusers/models/attention_dispatch.py Outdated
Comment on lines +1132 to +1135
if (
attn_mask is not None
and torch.all(attn_mask != 0).item()
):

@sayakpaul sayakpaul Jan 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if (
attn_mask is not None
and torch.all(attn_mask != 0).item()
):
if attn_mask is not None and torch.all(attn_mask != 0):

Won't it work?

@zhangtao0408 zhangtao0408 Jan 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Won't it work?

diffusers/src/diffusers/models/attention_dispatch.py

Lines 1131 to 1136 in 5c92a77

# Skip Attention Mask if all values are 1, `None` mask can speedup the computation
if (
attn_mask is not None
and torch.all(attn_mask != 0).item()
):
attn_mask = None

Thanks for the reply!

Since NPU FA does not support the [B, Seq_len_kv] mask shape passed by QwenImageEditPlus, and the unsqueeze/expand operations slow down execution, I added logic to bypass these steps when the mask is all 1s. This optimization significantly improves speed in context parallel, as shown in the test results below:

Stage Card End to End Time(s)
Skip expand mask (set to None) 1 108.22
Skip expand mask (set to None) 4 49.83
Expand mask 1 108.62
Expand mask 4 57.74

@sayakpaul sayakpaul Jan 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's fine. I am asking if this condition would work (i.e., no item()):
if attn_mask is not None and torch.all(attn_mask != 0):

@zhangtao0408 zhangtao0408 Jan 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, that worked. I've removed item() and pushed the update.

Comment thread src/diffusers/models/attention_dispatch.py Outdated
# Skip Attention Mask if all values are 1, `None` mask can speedup the computation
if (
attn_mask is not None
and torch.all(attn_mask != 0).item()

@sayakpaul sayakpaul Jan 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above.

Comment thread src/diffusers/models/transformers/transformer_qwenimage.py
Comment on lines -167 to +171
per_sample_len = torch.where(has_active, active_positions.max(dim=1).values + 1, torch.as_tensor(text_seq_len))
per_sample_len = torch.where(
has_active,
active_positions.max(dim=1).values + 1,
torch.as_tensor(text_seq_len, device=encoder_hidden_states.device)
)

@sayakpaul sayakpaul Jan 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like an unrelated change? If so, could you undo it?

@zhangtao0408 zhangtao0408 Jan 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change is to fix #13015.

sayakpaul
sayakpaul reviewed Jan 22, 2026
View reviewed changes
Comment thread src/diffusers/models/attention_dispatch.py Outdated
Comment on lines +2450 to +2458
if (
attn_mask is not None
and attn_mask.ndim == 2
and attn_mask.shape[0] == query.shape[0]
and attn_mask.shape[1] == key.shape[1]
):
B, Sq, Skv = attn_mask.shape[0], query.shape[1], key.shape[1]
attn_mask = ~attn_mask.to(torch.bool)
attn_mask = attn_mask.unsqueeze(1).expand(B, Sq, Skv).unsqueeze(1).contiguous()

@sayakpaul sayakpaul Jan 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it make sense to have a small utlity named _maybe_modify_attn_mask_npu() so that it can be reused in the two places (here and above)?

@zhangtao0408 zhangtao0408 Jan 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your suggestion! I've updated the _maybe_modify_attn_mask_npu() method.

diffusers/src/diffusers/models/attention_dispatch.py

Lines 1114 to 1135 in 020a232

def _maybe_modify_attn_mask_npu(
query: torch.Tensor,
key: torch.Tensor,
attn_mask: Optional[torch.Tensor] = None
):
# Skip Attention Mask if all values are 1, `None` mask can speedup the computation
if (attn_mask is not None and torch.all(attn_mask != 0)):
attn_mask = None
# Reshape Attention Mask: [batch_size, seq_len_k] -> [batch_size, 1, sqe_len_q, seq_len_k]
# https://www.hiascend.com/document/detail/zh/Pytorch/730/apiref/torchnpuCustomsapi/docs/context/torch_npu-npu_fusion_attention.md
if (
attn_mask is not None
and attn_mask.ndim == 2
and attn_mask.shape[0] == query.shape[0]
and attn_mask.shape[1] == key.shape[1]
):
B, Sq, Skv = attn_mask.shape[0], query.shape[1], key.shape[1]
attn_mask = ~attn_mask.to(torch.bool)
attn_mask = attn_mask.unsqueeze(1).expand(B, Sq, Skv).unsqueeze(1).contiguous()
return attn_mask

@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Jan 23, 2026

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

This was referenced Jan 24, 2026
Closed
Merged
Draft
@zhangtao0408

zhangtao0408 commented Feb 4, 2026

Copy link
Copy Markdown
Contributor Author

Hi @sayakpaul , thanks for reviewing! Could you let me know what’s needed to get this PR merged? I’m happy to make any adjustments. Thanks!

@sayakpaul

sayakpaul commented Feb 4, 2026

Copy link
Copy Markdown
Member

@bot /style

@github-actions

github-actions Bot commented Feb 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

DN6
DN6 reviewed Feb 16, 2026
View reviewed changes
Comment thread src/diffusers/models/attention_dispatch.py Outdated
@DN6
DN6 merged commit bcbbded into huggingface:main Feb 17, 2026
10 of 11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@DN6 DN6 DN6 left review comments
@sayakpaul sayakpaul sayakpaul approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

5 participants

@zhangtao0408 @HuggingFaceDocBuilderDev @sayakpaul @DN6 @TaoZhang-Work