Skip to content

Commit 767ee30

Browse files
authored
ZImageFunControlNet: Fix mask concatenation in --gpu-only (#11421)
This operation trades in latents which in --gpu-only may be out of the GPU The two VAE results will follow the --gpu-only defined behaviour so follow the inpaint image device when calculating the mask in this path.
1 parent 3ab9748 commit 767ee30

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

comfy_extras/nodes_model_patch.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -348,7 +348,7 @@ def encode_latent_cond(self, control_image=None, inpaint_image=None):
348348
if self.mask is None:
349349
mask_ = torch.zeros_like(inpaint_image_latent)[:, :1]
350350
else:
351-
mask_ = comfy.utils.common_upscale(self.mask.view(self.mask.shape[0], -1, self.mask.shape[-2], self.mask.shape[-1]).mean(dim=1, keepdim=True), inpaint_image_latent.shape[-1], inpaint_image_latent.shape[-2], "nearest", "center")
351+
mask_ = comfy.utils.common_upscale(self.mask.view(self.mask.shape[0], -1, self.mask.shape[-2], self.mask.shape[-1]).mean(dim=1, keepdim=True).to(device=inpaint_image_latent.device), inpaint_image_latent.shape[-1], inpaint_image_latent.shape[-2], "nearest", "center")
352352

353353
if latent_image is None:
354354
latent_image = comfy.latent_formats.Flux().process_in(self.vae.encode(torch.ones_like(inpaint_image) * 0.5))

0 commit comments

Comments
 (0)