Fix VAEEncodeForInpaint mask growing on non-CPU or non-float32 masks - #15736
Fix VAEEncodeForInpaint mask growing on non-CPU or non-float32 masks#15736BlakeB254 wants to merge 1 commit into
Conversation
`grow_mask_by > 0` builds its erosion kernel with `torch.ones((1, 1, grow_mask_by, grow_mask_by))`, which lands on CPU in float32 no matter what `mask` is. The very next line convolves it with `mask`, so whenever the incoming mask is not a CPU float32 tensor the node raises instead of encoding. Two ways users hit it: - Mask on CUDA - "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!" This is Comfy-Org#2556, open since January 2024. `--gpu-only` is a reliable way in, since it keeps intermediates on the GPU. - Half-precision mask - "Input type (torch.HalfTensor) and weight type (torch.FloatTensor) should be the same". Building the kernel with the mask's own device and dtype fixes both. The kernel is all ones, so nothing else about the result changes, and a CPU float32 mask takes exactly the path it did before. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
🎉 Thank you for your contribution, we really appreciate it! 🎉 Like many open source projects, we require contributors to sign our Contributor License Agreement (CLA). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:
CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility. ✍ To sign, please post a new comment on this PR with exactly the following text: ✍ I have read and agree to the Contributor License Agreement You can retrigger this bot by commenting recheck in this Pull Request. Posted by the CLA Assistant Lite bot. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📜 Recent review details
|
| Check name | Status | Explanation |
|---|---|---|
| Title check | ✅ Passed | The title clearly identifies the fix for mask growth with non-CPU and non-float32 masks. |
| Description check | ✅ Passed | The description explains the device and dtype failures, the fix, affected cases, safety, and verification. |
| Docstring Coverage | ✅ Passed | No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. |
| Linked Issues check | ✅ Passed | Check skipped because no linked issues were found for this pull request. |
| Out of Scope Changes check | ✅ Passed | Check skipped because no linked issues were found for this pull request. |
✨ Finishing Touches 💡 1
🛠️ Fix failing CI checks 💡
- Create stacked PR
- Commit on current branch
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
What
VAEEncodeForInpaintbuilds its mask-growing kernel aswhich is always CPU / float32, and convolves it with
maskon the next line. Whenever the incoming mask isn't CPU float32, the node raises instead of encoding. This builds the kernel with the mask's own device and dtype.Two ways users hit it
Mask on CUDA —
Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!That's #2556, open since January 2024, with a traceback pointing straight at this node; the only answer on the thread is a workaround (rebuild the node's internals by hand).
--gpu-onlyis a reliable way in, since it keeps intermediates on the GPU — the same flag @rattus128 identified as the minimum reproducer for the sibling report in #11412.Half-precision mask —
Input type (torch.HalfTensor) and weight type (torch.FloatTensor) should be the sameReproducing without a GPU
The dtype half needs no CUDA, so it's easy to check both before and after:
Why it's safe
The kernel is all ones, so moving where it's allocated changes nothing about the result —
mask_erosionis the same tensor it was, just computed next to its input. A CPU float32 mask takes exactly the path it took before; this only adds the cases that currently raise.Verified on Linux aarch64 (NVIDIA GB10 / DGX Spark, torch 2.11.0+cu130) for the dtype case and CPU-unchanged case. I did not re-run the 2024 workflow from #2556, so I'd describe this as fixing the reported failure mode rather than certifying that whole thread.