QwenImage21Pipeline: editing an image the model generated degrades at output_resolution=1024 regardless of seed (not noise replay, MPS)

[ck71]

Follow-up to #14824, which was closed as "noise replay" (same seed/stream/shape for generation and edit). I see a degradation at output_resolution=1024 that does not depend on the seed, so replay does not explain it.

Setup: diffusers 0121a91 and current main 0377f0c (same result), torch 2.14.0, transformers 5.17.0, bf16, MPS (M5 Max, macOS 26.6), 40 steps. Script: one edit per process, below.

input	seed	result	RGB std (input 0.259)	PSNR vs input
image generated by Qwen-Image-2.1 (via mflux, MLX RNG, seed 42), red circle drawn in	42	instruction followed (bicycle appears) but whole image grey and embossed	0.090	12.1 dB
same	12345	same degradation	0.099	11.9 dB
real photo (Kodak kodim01, 768×512)	42	clean, correct edit	0.176 (input 0.162)	17.9 dB
image generated by the diffusers pipeline (seed 700467539, 1376×768)	12345	edit applied, image darker, edge energy ×2.0	0.129 (input 0.167)	16.5 dB

The generation seed/stream differs from the edit's in rows 1, 2 and 4 (row 1 is even a different framework's RNG), so the edit's initial noise is not the draw that produced the image. At 512 and 768 all of these edits are clean. This may be what @peterc described in #14824 as a separate issue.

Script (one edit per process)

"""One edit at output_resolution=1024, run in its own process."""
import sys, time, json, numpy as np, torch
from PIL import Image
from diffusers import QwenImage21Pipeline
src, prompt, seed, out = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4]
img = Image.open(src).convert("RGB")
pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", dtype=torch.bfloat16).to("mps")
pipe.set_progress_bar_config(disable=True)
t = time.time()
arr = pipe(image=img, prompt=prompt, num_inference_steps=40, generator=torch.Generator("mps").manual_seed(seed),
           output_resolution=1024, output_type="np").images[0]
rgb = (np.clip(arr[..., :3], 0, 1) * 255).round().astype("uint8")
res = Image.fromarray(rgb); res.save(out)
# RGB std as a contrast measure (washed-out ~0.09, clean ~0.26 on this image) and edge energy relative to the input
def grad(a): a = a.mean(-1); return np.abs(np.diff(a, axis=0)).mean() + np.abs(np.diff(a, axis=1)).mean()
inp = np.asarray(img.resize(res.size)).astype(float) / 255; o = rgb.astype(float) / 255
print(json.dumps({"out": out, "s": round(time.time() - t), "std": round(float(o.std()), 3), "std_in": round(float(inp.std()), 3),
                  "edge_ratio": round(float(grad(o) / grad(inp)), 2), "psnr": round(float(10 * np.log10(1 / ((o - inp) ** 2).mean())), 1)}))

Usage: python edit1024.py <image> <prompt> <seed> <out.png>

[vinitsonawane45]

I'd like to investigate this issue.

The fact that the degradation is reproducible across different input images and edit seeds, while 512/768 remain clean, makes the 1024-resolution path particularly interesting. I also agree that the reported results make noise replay unlikely to be the primary explanation.

I'll try to reproduce this on MPS and trace the 1024 path, focusing on the image resize/VAE encode → latent preparation → denoising → VAE decode stages and checking for numerical differences or invalid values specific to bfloat16/MPS.

If I can reproduce it reliably, I'll narrow it down to the smallest failing operation and prepare a minimal regression test/fix rather than changing the sampling behavior blindly.

[peteklapka]

Independent repro of this, hit it in production yesterday. M2 Ultra (192GB unified), macOS, torch 2.14.0, diffusers @ 0377f0c — same commit you tested. We're editing a Qwen-Image-2.1-generated 1024² RGBA image via a long-lived server process; every edit at output_resolution=1024 comes back with the global grey/embossed/desaturated texture, instruction correctly applied.

To rule out our stack we ran the pipeline directly (no server wrapper), same prompt/seed/input, three configurations: fp16 + input.convert("RGB"), bf16 + convert("RGB"), and bf16 with RGBA passed through untouched — all three embossed at 1024. Same setup at output_resolution=768 and 512: both clean. So it's not dtype, not alpha handling, not MPS-specific in our case.

One extra data point for #14859: we made four consecutive 1024 edits through one resident pipeline (our server model) and got emboss on all four — no NaNs — so the NaN-on-second-edit behavior may be torch-version- or input-specific rather than universal.

[vinitsonawane45]

Thanks for the additional reproduction — that's very useful.

The fp16/bf16, RGB/RGBA, direct-pipeline, and repeated-edit results make the dtype, alpha handling, and long-lived-process hypotheses much less likely. The fact that 512/768 remain clean while 1024 consistently produces the same embossed/desaturated output points me more strongly toward a resolution-dependent preprocessing/latent path.

I'll focus the investigation on comparing the 768 vs 1024 execution paths, particularly the resize/target-dimension calculation, VAE encode/decode, latent packing/preparation, and spatial/positional inputs to the transformer. I'll try to identify the first stage where the 1024 tensors diverge abnormally before proposing a fix.

I'll report back once I have a reproducible stage-level difference.

[ck71]

Thanks @vinitsonawane45. I think I found the variable that matters, and it corrects my original description: it is not model-generated vs. real images, it is the aspect ratio at 1024. My first report compared a 3:2 real photo against square generated images, which made it look like an input-origin issue.

Same setup as before (diffusers 0121a91, torch 2.14.0, bf16, MPS, M5 Max 128 GB, macOS 26.6), 40 steps, one call per fresh process, 43 runs in total (34 at 1024, 9 at 768). Contrast = RGB std of the output (input in brackets).

output at output_resolution=1024	runs	result
1248×832 / 832×1248 (3:2, 2:3), all condition images 3:2/2:3 too	14 (real photos, a generated image, 2–3 condition images, 2 seeds)	all clean, e.g. 0.293 (0.316), 0.247 (0.259)
1184×896 / 896×1184 (4:3, 3:4)	12	washed out / embossed (e.g. 0.097 vs 0.315, edge energy ×1.7–×5) or all NaN
1024×1024 (1:1)	2	embossed (edge energy ×4.0) or all NaN
1376×768 / 768×1376 (16:9, 9:16)	2	9:16 strongly embossed, 16:9 milder but lighting/colors drift

• Same image, only cropped differently: kodim04 cropped to 2:3 is clean, the same photo cropped to 1:1 or 3:4 is embossed; kodim01 3:2 clean, 4:3 embossed; a 3:4 phone photo cropped to 2:3 becomes clean. I wrote down six predictions before running them, all six held.
• Input size does not matter: the same 3:4 photo pre-resized to 512, 768, 1024 or 1600 px is broken every time.
• 768 is clean in every aspect ratio I tried (1:1, 4:3, 3:4, 3:2, 2:3; 9 runs).
• NaN also in a fresh process, on the first call (not only from the second call as in QwenImage21Pipeline on MPS: every edit after the first in the same process returns all-NaN at output_resolution=1024 #14859): multi-image calls with a 3:4 or 1:1 image among the conditions.
• Every condition image counts, not only the output shape: with width=832, height=1248 (2:3) and two condition images, a 2:3 first image plus a 3:4 second image gives all NaN (2 seeds); both cropped to 2:3 is clean (2 seeds). Two other mixes with the off-ratio image first happened to be clean, so I would not rely on order.

Since each condition image is resized to its own aspect ratio at output_resolution² and its shape goes into img_shapes, my guess is something shape-dependent in the positional/RoPE path or latent packing at 1024. For reference, the latent patch grids are 78×52 (3:2, clean), 74×56 (4:3), 64×64 (1:1), 86×48 (16:9). I have not traced it further.

Workaround on my side: at 1024 I center-crop every image to 3:2 or 2:3 and pass the matching width/height.

Minimal repro with a public Kodak image

"""Same photo, two crops, edited at output_resolution=1024. Expected: both clean.
Observed on MPS: 2:3 clean, 1:1 embossed. Run once per crop in a fresh process."""
import io, sys, urllib.request
import numpy as np, torch
from PIL import Image
from diffusers import QwenImage21Pipeline

crop = sys.argv[1]  # "2x3" or "1x1"
img = Image.open(io.BytesIO(urllib.request.urlopen("https://r0k.us/graphics/kodak/kodak/kodim04.png").read())).convert("RGB")
if crop == "1x1":
    img = img.crop((0, 128, 512, 640))  # 512x768 -> 512x512
pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", dtype=torch.bfloat16).to("mps")
arr = pipe(image=img, prompt="Change the red hat to a navy blue knitted beanie. Keep the face and everything else identical.",
           num_inference_steps=40, generator=torch.Generator("mps").manual_seed(1),
           output_resolution=1024, output_type="np").images[0]
out = np.clip(arr[..., :3], 0, 1)
inp = np.asarray(img.resize(out.shape[1::-1])).astype(float) / 255
print(crop, out.shape[:2], "NaN:", int(np.isnan(arr).sum()), "std out/in:", round(float(out.std()), 3), round(float(inp.std()), 3))
Image.fromarray((out * 255).round().astype("uint8")).save(f"kodim04-{crop}.png")

Numbers from my equivalent runs (same image, crop, prompt and seed; the script above is a condensed version of them): 2:3 → 832×1248, std 0.176 (input 0.188), clean; 1:1 → 1024×1024, std 0.131 (input 0.176), embossed.

[vinitsonawane45]

Thanks, this gives a much clearer direction.

I’d like to investigate this further and prepare a focused fix with a regression test for the aspect-ratio-dependent 1024 behavior.

Before I open a PR, is it okay for me to proceed with the implementation based on these findings?

[ck71]

Yes, please go ahead, and thanks for picking this up. I'm just the reporter, so the maintainers will have the final say on the PR, but I'm happy to test any branch on MPS (M5 Max) against the same set of cases: the 3:2 vs. 1:1/4:3 crops, the mixed-ratio condition images that produced NaN, and the 768 baseline. One note on the regression test: since CI probably doesn't run on MPS, it may be worth checking whether the problem reproduces on CPU or CUDA at 1024 with a 1:1 or 4:3 input. I've only run it on MPS so far.

[vinitsonawane45]

@ck71 I added a minimal diagnostic script for issue #14858 that instruments the live QwenImage21Pipeline path.

It compares:

• 1024 3:2 — known clean
• 1024 4:3 — known broken
• 1024 1:1 — known broken
• 768 4:3 — known clean

It records tensor shapes/stats through VAE encode → latent packing → img_shapes → RoPE → Q/K/V → attention → transformer → scheduler → decode, and compares the default/cache/no-cache attention paths.

Could you run diagnose_qwenimage21_issue_14858.py on your M5 Max using the same images/prompts from the reproduction?

The default run uses 2 steps for a faster diagnostic; --steps 40 runs the full reproduction. The trace is written to qwenimage21_issue_14858.jsonl.

The main thing I'm looking for is the first operation where the 1024 4:3 or 1:1 case diverges or becomes non-finite, compared with the clean 3:2/768 cases.

I don't want to make a speculative source change. Once we identify that boundary, I'll implement the smallest fix and add a focused regression test.

[ck71]

Happy to run it, but I can't find the script: there's no link or attachment in your comment, and I don't see a fork, branch or gist. Could you share it (a gist or a branch is fine)? I'll run it with the same Kodak crops (3:2, 4:3, 1:1 at 1024, plus 4:3 at 768) and post the trace.

[vinitsonawane45]

Thanks @ck71! I’ve pushed the diagnostic script to my fork here:

https://github.com/vinitsonawane45/diffusers/tree/fix/qwenimage21-issue-14858-diagnostic

Commit: 6cb32387103d5b505ad6e294c9a0435810b743fb

The script is diagnose_qwenimage21_issue_14858.py.

Please run it with the same Kodak crops you used:

• 1024 3:2
• 1024 4:3
• 1024 1:1
• 768 4:3

The default run uses 2 steps for a quicker diagnostic; --steps 40 runs the full reproduction.

It writes the trace to qwenimage21_issue_14858.jsonl. The main thing I'm looking for is the first operation where the 1024 4:3/1:1 cases diverge or become non-finite compared with the clean 3:2/768 cases.

Once you post the trace, I’ll use that to narrow down the exact failing operation before making any source changes.

[ck71]

@vinitsonawane45 thanks for sharing the script. I read through it before running anything; in the meantime I found the boundary with a layer-by-layer comparison against CPU, details below. Two notes on the script, in case you want to keep using it: it runs all 12 case/mode combinations in one process, which on MPS triggers the separate state problem from #14859 (NaN from the second 1024 call on), and the broken cases here never become non-finite (they are washed out, not NaN), so "first non-finite" won't find them. Happy to run it anyway if you still want the trace.

Root cause: F.pad on 5-D tensors is wrong on MPS once the trailing dims reach 65,536 elements, and the Qwen-Image 2.1 VAE encoder hits it in QwenImage21AvgDown3D. It is an MPS bug, not a model or pipeline bug: the same 1:1 edit at 1024 on CPU (fp32, 2 steps) gives output std 0.232, on MPS 0.141; the clean 3:2 case is 0.241 on CPU and 0.233 on MPS.

How I narrowed it down (torch 2.14.0, diffusers 0121a91, M5 Max, macOS 26.6):

1. Hooks on the live pipeline, MPS vs. CPU: the transformer's prefill looks the same in clean and broken cases; the first stage that already differs is the VAE encode of the condition image (latent std 1.02 on MPS vs. 1.17 on CPU for 1:1; 1.148 vs. 1.158 for 3:2).
2. Layer-by-layer comparison of the VAE encoder, MPS bf16 vs. CPU fp32 (1:1 at 1024): every layer is within ~1 % up to down_blocks.1.avg_shortcut, whose output is all zeros on MPS (std 0.0 vs. 1.216). Same for down_blocks.2.avg_shortcut. Everything after that is off by 20–85 %.
3. Inside QwenImage21AvgDown3D.forward the 8-D view/permute/contiguous/mean is fine; the temporal F.pad(x, (0, 0, 0, 0, pad_t, 0)) is what breaks.

Pure PyTorch repro:

import torch, torch.nn.functional as F
for w in (65535, 65536):
    x = torch.arange(1, w + 1, dtype=torch.float32).reshape(1, 1, 1, 1, w)
    print(w, torch.equal(F.pad(x, (0, 0, 0, 0, 1, 0)), F.pad(x.to("mps"), (0, 0, 0, 0, 1, 0)).cpu()))
# 65535 True
# 65536 False

It fails for any padded axis of a 5-D tensor at that size (T, H or W), in fp32 and bf16; the same sizes as 3-D or 4-D tensors are fine.

Why it looked like an aspect-ratio rule: in down_blocks.2 the spatial size is 16× the number of image tokens. 3:2 at 1024 is 312×208 = 64,896 (< 65,536, correct), 1:1 is 256×256 = 65,536, 4:3 is 296×224 = 66,304, 16:9 is 344×192 = 66,048 (all broken). At 768 every ratio stays below the limit there. down_blocks.1 is above the limit in every case (e.g. 624×416 for 3:2, 448×336 at 768), so the encoder output on MPS is always somewhat off (relative error vs. CPU ≈ 0.15 at 768, ≈ 0.55 for 3:2 at 1024); the model apparently tolerates one broken shortcut but not two.

Fix candidate (avoids F.pad, same result on CPU):

pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
if pad_t:
    x = torch.cat([x.new_zeros(x.shape[0], x.shape[1], pad_t, *x.shape[3:]), x], dim=2)

With only this change, the two broken cases from before are clean at 1024 and 40 steps: kodim04 1:1 std 0.182 (input 0.176, before 0.131), a 3:4 phone photo 0.311 (input 0.314, before 0.097), no embossing.

For a regression test that runs in CI without MPS, checking QwenImage21AvgDown3D against a reference that does not use F.pad probably won't catch it on CPU/CUDA, since the bug is in the MPS kernel. The encoder has other F.pad calls (e.g. the causal conv padding); conv_in was within 0.3 % in my comparison, so I did not dig further, but a general audit of 5-D F.pad on MPS may be worthwhile. I have not checked whether the NaN in #14859 has the same origin. I'd also be glad to report the F.pad issue to PyTorch unless someone already knows of an existing issue.

[vinitsonawane45]

Thanks @ck71 for the detailed investigation and for confirming the root cause.

I’ve incorporated the torch.cat-based temporal padding workaround into PR #14902 along with a regression test for the expected padding behavior.

Please have a look at the PR when you get a chance. If you think any changes, additional tests, or adjustments are needed, I’m happy to make them.