Repository navigation
QwenImage21Pipeline on MPS: every edit after the first in the same process returns all-NaN at output_resolution=1024 #14859
Description
Repeated image edits with one pipeline instance return all-NaN from the second call on, but only at output_resolution=1024. At 512 and 768 three consecutive edits are clean. gc.collect(), torch.mps.synchronize() and torch.mps.empty_cache() between calls do not help.
Setup: diffusers 0121a91 and current main 0377f0c (identical), torch 2.14.0, bf16, MPS, M5 Max 128 GB, macOS 26.6.
Minimal repro (script below, 20 steps, one pipeline, three edits of the same 1024² image):
| output_resolution | edit 1 | edit 2 | edit 3 |
|---|---|---|---|
| 512 | 0 NaN | 0 NaN | 0 NaN |
| 768 | 0 NaN | 0 NaN | 0 NaN |
| 1024 | 0 NaN | 4,194,304 of 4,194,304 NaN | all NaN |
Workaround: a fresh process per edit at 1024.
Minimal repro
"""Minimal repro: every edit after the first in the same process returns NaN on MPS at 1024.
Three edits of the same image with one pipeline instance, with gc and torch.mps.empty_cache()
in between. Expected: no NaN. Observed at output_resolution=1024: first clean, then all NaN.
"""
import gc, sys, time
import numpy as np, torch
from PIL import Image
from diffusers import QwenImage21Pipeline
import diffusers
image = Image.open(sys.argv[1]).convert("RGB")
res = int(sys.argv[2]) if len(sys.argv) > 2 else 512
pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", dtype=torch.bfloat16).to("mps")
pipe.set_progress_bar_config(disable=True)
print("diffusers", diffusers.__version__, "| torch", torch.__version__, "| output_resolution", res, flush=True)
for i, prompt in enumerate(["Change the scene to a rainy evening.", "Make the sign read HELLO.", "Add a black cat on the bench."], 1):
t = time.time()
arr = pipe(image=image, prompt=prompt, num_inference_steps=20, generator=torch.Generator("mps").manual_seed(i),
output_resolution=res, output_type="np").images[0]
print(f"edit {i}: {time.time() - t:.0f} s, NaN values: {int(np.isnan(arr).sum())} of {arr.size}", flush=True)
del arr; gc.collect(); torch.mps.synchronize(); torch.mps.empty_cache()Usage: python nan_repro.py <image> <output_resolution>
Activity
vinitsonawane45 commentedon Sep 23, 2026
I’d like to investigate this issue. The fact that the first 1024² edit succeeds while subsequent edits return all-NaN in the same process, while 512/768 remain stable, suggests there may be a stateful MPS/backend interaction rather than a simple input issue.
I’m going to reproduce it and narrow down where NaNs first appear (transformer, scheduler, or VAE), while also checking whether the behavior changes with pipeline recreation, dtype, and PyTorch versions.
If I can identify a minimal root cause and a regression test, I’d be happy to work on a fix.
ck71 commentedon Sep 24, 2026
Thanks @vinitsonawane45 for looking into this. One more data point that may help narrow down the state: it is not limited to repeated 1024 edits. A 768 edit also returns all-NaN when 1024 text-to-image jobs (no condition image) ran earlier in the same process.
Same setup (diffusers 0121a91, torch 2.14.0, bf16, MPS, M5 Max 128 GB, macOS 26.6), one pipeline instance, five calls in a row, 20 steps each:
current_allocated_memory/driver_allocated_memoryoutput_resolution=512output_resolution=768Job 4 alone in a fresh process is clean. The calls without NaN are bit-identical to the same call in a fresh process (job 5 matches job 1, job 3 matches a fresh run), so the leftover state only shows up in the image-conditioned path. Together with the original table (three 512 or three 768 edits in a row stay clean), my guess is that a preceding call at a larger sequence length leaves something behind that a later edit picks up, but I have not isolated which of jobs 1–3 is required.
Script
Usage:
python reuse_repro.py <1024x1024 image>ck71 commentedon Sep 29, 2026
Update: this appears to have the same root cause as #14858. The
F.padworkaround inQwenImage21AvgDown3D(PR #14902) also makes the NaNs here go away.Same setup as above (diffusers 0121a91, torch 2.14.0, bf16, MPS, M5 Max, macOS 26.6), one pipeline instance per sequence:
output_resolution=1024(repro script from the first post, 20 steps)The run without the fix was repeated today right after the one with the fix, on the same machine and environment, so the NaNs have not disappeared for some unrelated reason.
What I can't explain yet is why a wrong VAE encode in one call turns the next call into NaN; I only know that with the encode fixed, neither the first nor the later calls produce NaN in these sequences. So I'd treat this as "fixed by the same change, mechanism of the carry-over not understood". @vinitsonawane45, if you agree, #14902 could reference this issue as well, and I'm happy to rerun both sequences on the PR branch.
vinitsonawane45 commentedon Sep 29, 2026
Thanks @ck71 for checking this and for the additional validation.
Yes, I agree that #14859 should be referenced in #14902 since the same change fixes the NaNs in the tested sequences as well.
I’ll keep the PR focused on the
QwenImage21AvgDown3Dtemporal padding issue and won’t speculate about the exact mechanism behind the cross-call NaNs until that is better understood.Happy to make any further changes if the maintainers would like additional tests or documentation.