Skip to content

Fix VAEEncodeForInpaint mask growing on non-CPU or non-float32 masks - #15736

Open
BlakeB254 wants to merge 1 commit into
Comfy-Org:masterfrom
BlakeB254:fix/vae-encode-inpaint-mask-device
Open

Fix VAEEncodeForInpaint mask growing on non-CPU or non-float32 masks#15736
BlakeB254 wants to merge 1 commit into
Comfy-Org:masterfrom
BlakeB254:fix/vae-encode-inpaint-mask-device

Conversation

@BlakeB254

Copy link
Copy Markdown

What

VAEEncodeForInpaint builds its mask-growing kernel as

kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))

which is always CPU / float32, and convolves it with mask on 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 CUDAExpected 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-only is 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 maskInput type (torch.HalfTensor) and weight type (torch.FloatTensor) should be the same

Reproducing without a GPU

The dtype half needs no CUDA, so it's easy to check both before and after:

import math, torch

def grow(mask, n=6):
    k = torch.ones((1, 1, n, n))                      # current
    # k = torch.ones((1, 1, n, n), device=mask.device, dtype=mask.dtype)   # this PR
    return torch.clamp(
        torch.nn.functional.conv2d(mask.round(), k, padding=math.ceil((n - 1) / 2)), 0, 1)

grow(torch.rand((1, 1, 64, 64)))                      # ok
grow(torch.rand((1, 1, 64, 64), dtype=torch.float16)) # RuntimeError on master, ok with this PR

Why it's safe

The kernel is all ones, so moving where it's allocated changes nothing about the result — mask_erosion is 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.

`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>
@github-actions

Copy link
Copy Markdown
Contributor

🎉 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:

  • Confirm that you own your contribution.
  • Keep the right to reuse your own code.
  • Grant us a copyright license to include and share it within our projects.

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.

@coderabbitai

coderabbitai Bot commented Aug 19, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 2f826beb-d899-4aee-9dae-b6ac1263e916

📥 Commits

Reviewing files that changed from the base of the PR and between c67885b and d935f48.

📒 Files selected for processing (1)
  • nodes.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

📜 Recent review details
⚠️ CI failures not shown inline (2)

GitHub Actions: CLA Assistant / cla-assistant: Fix VAEEncodeForInpaint mask growing on non-CPU or non-float32 masks

Conclusion: failure

View job details

##[group]Run contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08
 with:
   path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
   remote-organization-name: comfy-org
   remote-repository-name: comfy-cla
   path-to-signatures: signatures/cla.json
   branch: main
   allowlist: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot]
   custom-notsigned-prcomment: 🎉 Thank you for your contribution, we really appreciate it! 🎉
Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). 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:
- Confirm that you own your contribution.
- Keep the right to reuse your own code.
- Grant us a copyright license to include and share it within our projects.
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:** ✍
   custom-pr-sign-comment: I have read and agree to the Contributor License Agreement
   custom-allsigned-prcomment: ✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
   use-dco-flag: false
   lock-pullrequest-aftermerge: true
   suggest-recheck: true
 env:
   GITHUB_***REDACTED_SECRET_ASSIGNMENT***
   PERSONAL_ACCESS_***REDACTED_SECRET_ASSIGNMENT***
 ##[endgroup]
 CLA Assistant GitHub Action bot has started the process
 (node:2078) [DEP0040] DeprecationWarn...

GitHub Actions: CLA Assistant / 0_cla-assistant.txt: Fix VAEEncodeForInpaint mask growing on non-CPU or non-float32 masks

Conclusion: failure

View job details

##[group]Run contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08
 with:
   path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
   remote-organization-name: comfy-org
   remote-repository-name: comfy-cla
   path-to-signatures: signatures/cla.json
   branch: main
   allowlist: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot]
   custom-notsigned-prcomment: 🎉 Thank you for your contribution, we really appreciate it! 🎉
Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). 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:
- Confirm that you own your contribution.
- Keep the right to reuse your own code.
- Grant us a copyright license to include and share it within our projects.
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:** ✍
   custom-pr-sign-comment: I have read and agree to the Contributor License Agreement
   custom-allsigned-prcomment: ✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
   use-dco-flag: false
   lock-pullrequest-aftermerge: true
   suggest-recheck: true
 env:
   GITHUB_***REDACTED_SECRET_ASSIGNMENT***
   PERSONAL_ACCESS_***REDACTED_SECRET_ASSIGNMENT***
 ##[endgroup]
 CLA Assistant GitHub Action bot has started the process
 (node:2078) [DEP0040] DeprecationWarn...
🧰 Additional context used
📓 Path-based instructions (6)
**/*

📄 CodeRabbit inference engine (AGENTS.md)

**/*: Keep changes small, direct, and limited to the narrowest necessary code path and smallest number of files.
Prefer practical fixes, minimal dependencies, and existing repository patterns; remove obsolete, dead, unreachable, or unused code.
Preserve existing APIs, node names, model-loading behavior, file layout, and workflow compatibility unless replacement is explicitly intended.
Core ComfyUI must not add outbound internet requests, telemetry, tracking, reporting, remote configuration, or background network activity. User-authorized model downloads are limited to the requested artifact and must exclude telemetry and unrelated metadata.

Files:

  • nodes.py
**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Keep state and capability flags on the object that owns the behavior. Prefer explicit parent-owned attributes over probing child objects with getattr; use child checks only when the child owns the delegated behavior.
Preserve shared method signatures, argument order, return shapes, side effects, and error behavior unless every affected caller and interface is intentionally updated.
Do not add unused compatibility parameters, flags, attributes, constructor options, or model-specific options to shared helpers; keep one-off behavior at the integration boundary.
Normalize third-party return conventions at integration boundaries so core code receives the expected type and shape; avoid undocumented caller-side unwrapping.
Do not add torch.no_grad, torch.inference_mode, or inference-mode wrappers. Do not add model freeze/unfreeze toggles; only disable globally enabled inference mode when a training path requires gradients.
Remove inference-only training behavior such as dropout while preserving checkpoint and state-dict compatibility; use nn.Identity when deleting a module would alter keys or ordering.
Keep imports at module scope except established optional-backend probes or imports required to avoid cycles; avoid unnecessary try/except blocks and use specific exceptions with useful fallbacks.
Do not add workarounds for unsupported library versions, especially PyTorch exception-and-float-cast retries, unless a comment names the exact versions still requiring them.
Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently degrading output.
Match local style, keep comments sparse and useful, and remove comments that merely restate obvious code.
Treat dtype, device placement, VRAM use, and offloading as correctness concerns across CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low-VRAM environments.
Prefer existing ComfyUI and Comfy Kitchen operations, quantization helpers, cast/offload helpe...

Files:

  • nodes.py
**/*.{py,json}

📄 CodeRabbit inference engine (AGENTS.md)

Treat legacy combo, io.Combo, and io.DynamicCombo values affecting filesystem access as untrusted; revalidate them at load/save boundaries with folder_paths, containment checks, or fixed allowlists.

Files:

  • nodes.py
**/*.{py,md,txt,json}

📄 CodeRabbit inference engine (AGENTS.md)

Keep warning and info messages short and actionable, remove noisy or misleading logging, and make documentation edits concise, factual, and tied to changed behavior.

Files:

  • nodes.py
**

⚙️ CodeRabbit configuration file

**: IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing a with: block),
treat it as unchanged. Contributors should not feel obligated to address
pre-existing issues outside the scope of their contribution.

Files:

  • nodes.py
nodes.py

⚙️ CodeRabbit configuration file

nodes.py: Core node definitions (2500+ lines). Focus on:

  • Backward compatibility of NODE_CLASS_MAPPINGS
  • Consistency of INPUT_TYPES return format

Files:

  • nodes.py
🧠 Learnings (1)
📚 Learning: 2026-02-21T14:01:41.482Z
Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.

Applied to files:

  • nodes.py
🔇 Additional comments (1)
nodes.py (1)

438-438: LGTM!


📝 Walkthrough

Walkthrough

VAEEncodeForInpaint.encode now creates the mask-growth convolution kernel on mask.device. The kernel keeps its existing shape and dtype. The convolution no longer receives a kernel implicitly created on the default device.

Merge Risk: ⚪ Minimal · up to d935f

This localized change preserves existing CPU float32 behavior while allowing inpainting masks on other devices or dtypes to encode successfully; no actionable merge-blocking risk remains beyond normal checks and review.

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

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.

1 participant