Fix QwenImage21 MPS temporal padding corruption

[vinitsonawane45]

Summary

Avoid F.pad for temporal padding in QwenImage21AvgDown3D.

On MPS, torch.nn.functional.pad can produce incorrect results for sufficiently large 5-D tensors. Qwen-Image 2.1 can hit this condition during VAE encoding at certain image aspect ratios and resolutions, causing severe output degradation.

This change replaces the temporal F.pad operation with equivalent zero-padding using torch.cat.

Root cause

The issue was reproduced on MPS when the affected 5-D tensor dimensions reached the relevant size threshold. The corruption occurs inside the temporal padding operation in QwenImage21AvgDown3D.

The problem is aspect-ratio dependent. At output_resolution=1024, several aspect ratios can cross the affected tensor-size boundary, while others remain below it.

Changes

• Replace only the temporal F.pad operation with torch.cat and a zero-initialized tensor.
• Keep the existing spatial padding unchanged.
• Add a CPU equivalence test for the padding behavior.
• Add an MPS-specific regression test at the 256 × 256 / 65,536 boundary where the original implementation produces incorrect results.

Validation

• Ruff lint: passed
• Ruff format: passed
• Python compilation: passed
• git diff --check: passed
• CPU equivalence test: passed
• MPS regression test: passed on an Apple M5 Max
• The MPS regression test fails with the original implementation and passes with this fix.
• The issue reporter validated the actual PR branch (7ef0f38) on an M5 Max.
• The previously failing 1024×1024 edit produces a clean output with the fix and matches the reporter's patched validation.

The full pytest collection was blocked locally by an existing environment incompatibility involving huggingface_hub (resolve_revision import error); no dependency changes were made as part of this fix.

Related issues

Fixes #14858
Related to #14859

[ck71]

Thanks for turning this around so quickly. The code change matches what I validated (same logic, applied as a monkeypatch on 0121a91), so that part looks good to me. A few notes:

1. The regression test passes with the old code too. It uses a 1×1×3×1×1 tensor on CPU, far below the size where MPS breaks, so it only checks that the new implementation equals the old one. A test that actually fails without the fix needs MPS and a large enough tensor. This one fails on the old code and passes with your change on my M5 Max (torch 2.14.0):

@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="the F.pad bug is specific to MPS")
def test_temporal_downsample_matches_cpu_on_mps(self):
    torch.manual_seed(0)
    downsample = QwenImage21AvgDown3D(in_channels=192, out_channels=384, factor_t=2, factor_s=2)
    # 256 x 256 = 65,536 trailing elements: the size where temporal F.pad breaks on MPS
    sample = torch.randn(1, 192, 1, 256, 256)
    torch.testing.assert_close(downsample(sample.to("mps")).cpu(), downsample(sample))

It won't run in CI without MPS, but it will on any Apple Silicon dev machine. Keeping your CPU equivalence test next to it makes sense.

2. Correction to my earlier comment. I wrote that it fails "for any padded axis of a 5-D tensor"; I had only checked that on a single-channel tensor (1, 1, 1, 1, 65536). With 96 channels at 512×512, padding H or W alone is correct and only the temporal pad is wrong. The other F.pad in this file (the causal conv, line 180) runs on a 4-D tensor after squeeze(2) and is not affected (checked for conv_in and both down blocks). So your change covers the right spot.

3. Validation numbers (same logic as the PR, MPS bf16 vs. CPU fp32):

• VAE encode, relative error vs. CPU: 1.2–1.5 % at 512, 768 and 1024 in 1:1, 3:4 and 2:3 (without the fix: ~15 % at 768, 15–63 % at 1024). The decoder is unaffected (0.3 %).
• Edits at 1024, 40 steps, no cropping: 1:1, 4:3, 3:4 and 16:9 all clean, plus a two-image group edit with two 3:4 photos (which returned all-NaN before); 1:1, 3:4 and the group edit with 2 seeds, the others with 1.

To be precise about the PR description: I validated the same change as a monkeypatch, not your branch itself. Happy to run the branch directly if that helps the review.

[vinitsonawane45]

Thanks @ck71 for the detailed feedback. I’ve added the requested MPS-specific regression test at the 256×256 / 65,536 boundary while keeping the existing CPU equivalence test.

The update is pushed as commit 7ef0f381. The production fix remains unchanged.

Please take another look when you have a chance, and I’m happy to make any further changes you recommend.

[ck71]

Thanks! I ran the branch itself (7ef0f38) on the M5 Max this time: both TestQwenImage21AvgDown3D tests pass, and with the fix reverted the MPS test fails while the CPU one passes, as intended. End-to-end with the branch's source and no monkeypatch, the 1:1 edit at 1024 is clean and identical to my patched run (output std 0.182, input 0.176). Looks good from my side; the maintainers have the final word. This change also makes #14859 go away (details there), so the PR could reference that issue too.

[vinitsonawane45]

Thanks @ck71 for testing the actual branch and confirming both the regression test and end-to-end behavior on M5 Max. I’ve updated the PR description to keep #14859 as a related issue while keeping #14858 as the primary fix.