Skip to content

Fix Qwen VAE spatial chunking failure on tall/high-resolution images#2398

Open
Fengxiaoxiao-001 wants to merge 4 commits into
kohya-ss:mainfrom
Fengxiaoxiao-001:patch-1
Open

Fix Qwen VAE spatial chunking failure on tall/high-resolution images#2398
Fengxiaoxiao-001 wants to merge 4 commits into
kohya-ss:mainfrom
Fengxiaoxiao-001:patch-1

Conversation

@Fengxiaoxiao-001

Copy link
Copy Markdown

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.

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.
@kohya-ss

kohya-ss commented Jul 8, 2026

Copy link
Copy Markdown
Owner

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.

@Fengxiaoxiao-001

Copy link
Copy Markdown
Author

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.

@kohya-ss

Copy link
Copy Markdown
Owner

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 stride != 1:

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] + 1

With 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 stride=1, padding=1 and stride=2, padding=0 convolutions used in the VAE.

Since I'd like to keep the change minimal — the fuller rewrite also alters the meaning of spatial_chunk_size, which I'd prefer to avoid — I'm going to close this PR in favor of the two-line fix above.

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!

@Fengxiaoxiao-001

Fengxiaoxiao-001 commented Jul 13, 2026

Copy link
Copy Markdown
Author

Thank you for thoroughly reviewing my PR and figuring out the real root cause.
I will replace those two lines following your suggestion. Once validated, I will push the updated commit.
It’s a valuable lesson for me. Thanks again for giving this minimal‑fix solution instead of a large‑scale rewrite.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants