From 75fbd9c9335896530d0f07196af486dd8569f439 Mon Sep 17 00:00:00 2001 From: Alexis Rolland Date: Fri, 8 May 2026 18:25:36 +0800 Subject: [PATCH] Move mask to CPU --- comfy/bg_removal_model.py | 2 +- comfy_extras/nodes_compositing.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/bg_removal_model.py b/comfy/bg_removal_model.py index cb7c2ee53..b23f06a5f 100644 --- a/comfy/bg_removal_model.py +++ b/comfy/bg_removal_model.py @@ -47,7 +47,7 @@ class BackgroundRemovalModel(): out = self.model(pixel_values=pixel_values) out = torch.nn.functional.interpolate(out, size=(H, W), mode="bicubic", antialias=False) - mask = out.sigmoid() + mask = out.sigmoid().float().cpu() if mask.ndim == 3: mask = mask.unsqueeze(0) if mask.shape[1] != 1: diff --git a/comfy_extras/nodes_compositing.py b/comfy_extras/nodes_compositing.py index ade6f6408..720efc629 100644 --- a/comfy_extras/nodes_compositing.py +++ b/comfy_extras/nodes_compositing.py @@ -203,7 +203,7 @@ class JoinImageWithAlpha(io.ComfyNode): @classmethod def execute(cls, image: torch.Tensor, alpha: torch.Tensor) -> io.NodeOutput: batch_size = max(len(image), len(alpha)) - alpha = 1.0 - resize_mask(alpha.to(image.device), image.shape[1:]) + alpha = 1.0 - resize_mask(alpha.to(image), image.shape[1:]) alpha = comfy.utils.repeat_to_batch_size(alpha, batch_size) image = comfy.utils.repeat_to_batch_size(image, batch_size) return io.NodeOutput(torch.cat((image[..., :3], alpha.unsqueeze(-1)), dim=-1))