Fix Qwen VAE spatial chunking failure on tall/high-resolution images#2398
Fix Qwen VAE spatial chunking failure on tall/high-resolution images#2398Fengxiaoxiao-001 wants to merge 4 commits into
Conversation
PR Description This PR fixes a deterministic VAE precompute failure in sd-scripts/library/qwen_image_autoencoder_kl.py when processing high-resolution images with spatial chunking enabled. Problem During VAE latent precomputation, some valid inputs trigger a Conv2d runtime error: Calculated padded input size per channel: (2 x 1857) Kernel size (3 x 3) can't be greater than actual input size This is not a GPU OOM issue. The failure is caused by the chunking logic producing an intermediate tensor whose height is too small for a subsequent 3x3 convolution. Root Cause The current chunking strategy can create boundary cases where the last spatial chunk leads to an invalid feature map shape during VAE encoding, especially with certain high-resolution images and small chunk sizes. Fix Adjust the chunking logic to avoid producing invalid intermediate tensor shapes. Ensure every chunk remains valid for the downstream 3x3 convolution. Preserve spatial chunking for memory efficiency without cropping or resizing images. Keep the fix localized to the VAE implementation so existing training and inference behavior remains unchanged except for improved robustness. Why this matters Prevents crashes during dataset precomputation on valid image sizes. Avoids unnecessary fallback to crop/resize paths. Preserves full-resolution training data. Improves reliability for large images and low-memory chunking configurations. Validation Reproduced the failure with a valid 1280x1856 image. Confirmed the issue was not caused by GPU OOM. Verified that the updated logic avoids invalid Conv2d input sizes.
|
Thank you for this PR. There are so many changes that it is difficult to review. Could you please revert the changes that were solely for formatting? Also, it would be helpful if you could write your comments in English. |
|
Thank you for the review. I have reverted the formatting-only changes and kept the PR focused on the functional 2D image-only VAE implementation. I also updated the comments to concise English comments. |
|
Thank you for update the PR, and the detailed report — I was able to reproduce the bug and confirm it. After digging into the chunking logic, I found that the crash actually comes from an incorrect output-size calculation, not the chunk boundary handling. The following two lines are wrong for y_height = org_shape[2] // self.stride[0]
y_width = org_shape[3] // self.stride[1]They should use PyTorch's actual Conv2d output-size formula (which also accounts for padding): y_height = (org_shape[2] + 2 * self.original_padding[0] - self.kernel_size[0]) // self.stride[0] + 1
y_width = (org_shape[3] + 2 * self.original_padding[1] - self.kernel_size[1]) // self.stride[1] + 1With just this change, the existing chunking algorithm produces output that is bit-identical to the non-chunked path, with no crashes. I verified this across a wide range of image heights and chunk sizes for both the Since I'd like to keep the change minimal — the fuller rewrite also alters the meaning of That said, I'd be very happy to merge this if you'd like to update the PR to use these two lines instead (after confirming it works on your end). Thanks again for catching and reporting the issue! |
|
Thank you for thoroughly reviewing my PR and figuring out the real root cause. |
PR Description
This PR fixes a deterministic VAE precompute failure in sd-scripts/library/qwen_image_autoencoder_kl.py when processing high-resolution images with spatial chunking enabled.
Problem
During VAE latent precomputation, some valid inputs trigger a Conv2d runtime error:
Calculated padded input size per channel: (2 x 1857) Kernel size (3 x 3) can't be greater than actual input size This is not a GPU OOM issue. The failure is caused by the chunking logic producing an intermediate tensor whose height is too small for a subsequent 3x3 convolution.
Root Cause
The current chunking strategy can create boundary cases where the last spatial chunk leads to an invalid feature map shape during VAE encoding, especially with certain high-resolution images and small chunk sizes.
Fix
Adjust the chunking logic to avoid producing invalid intermediate tensor shapes. Ensure every chunk remains valid for the downstream 3x3 convolution. Preserve spatial chunking for memory efficiency without cropping or resizing images. Keep the fix localized to the VAE implementation so existing training and inference behavior remains unchanged except for improved robustness. Why this matters
Prevents crashes during dataset precomputation on valid image sizes. Avoids unnecessary fallback to crop/resize paths.
Preserves full-resolution training data.
Improves reliability for large images and low-memory chunking configurations. Validation
Reproduced the failure with a valid 1280x1856 image. Confirmed the issue was not caused by GPU OOM.
Verified that the updated logic avoids invalid Conv2d input sizes.