QwenImage21Pipeline on MPS: every edit after the first in the same process returns all-NaN at output_resolution=1024

[ck71]

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>

[vinitsonawane45]

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]

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:

pos.	job	NaN	current_allocated_memory / driver_allocated_memory
1	text-to-image 1024×1024	0	32.5 / 46.3 GB
2	edit, output_resolution=512	0	32.5 / 46.2 GB
3	text-to-image 1024×1024 (RGBA prompt)	0	32.5 / 47.6 GB
4	edit, output_resolution=768	all NaN	32.5 / 47.4 GB
5	text-to-image 1024×1024, same seed as 1	0, bit-identical to 1	32.5 / 48.2 GB

Job 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

import argparse, os
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
import numpy as np, torch
from PIL import Image
from diffusers import QwenImage21Pipeline

p = argparse.ArgumentParser(); p.add_argument("image"); a = p.parse_args()
base = Image.open(a.image).convert("RGB")  # any 1024x1024 photo
rgba = ("This is an RGBA image with transparency. A hand-painted wooden shop sign reading \"MOOL\". "
        "The image has alpha channel and the background is transparent.")
jobs = [
    ("t2i-1024", dict(prompt="A lighthouse on a rocky coast at dusk, photorealistic", width=1024, height=1024, output_resolution=1024), 1),
    ("edit-512", dict(prompt="Make the awning dark red. Keep everything else identical.", image=base, width=512, height=512, output_resolution=512), 2),
    ("t2i-1024-rgba", dict(prompt=rgba, width=1024, height=1024, output_resolution=1024), 3),
    ("edit-768", dict(prompt="Add a black cat on the bench. Keep everything else identical.", image=base, width=768, height=768, output_resolution=768), 4),
    ("t2i-1024-again", dict(prompt="A lighthouse on a rocky coast at dusk, photorealistic", width=1024, height=1024, output_resolution=1024), 1),
]
pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", dtype=torch.bfloat16).to("mps")
pipe.set_progress_bar_config(disable=True)
for name, kw, seed in jobs:
    arr = pipe(num_inference_steps=20, generator=torch.Generator("mps").manual_seed(seed), output_type="np", **kw).images[0]
    print(name, "NaN:", int(np.isnan(arr).sum()), "of", arr.size)

Usage: python reuse_repro.py <1024x1024 image>