Skip to content

Commit

Permalink
Fix potential tensor device issue with ImageCompositeMasked.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Dec 21, 2023
1 parent a1e1c69 commit 6781b18
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion comfy_extras/nodes_mask.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from nodes import MAX_RESOLUTION

def composite(destination, source, x, y, mask = None, multiplier = 8, resize_source = False):
source = source.to(destination.device)
if resize_source:
source = torch.nn.functional.interpolate(source, size=(destination.shape[2], destination.shape[3]), mode="bilinear")

Expand All @@ -20,7 +21,7 @@ def composite(destination, source, x, y, mask = None, multiplier = 8, resize_sou
if mask is None:
mask = torch.ones_like(source)
else:
mask = mask.clone()
mask = mask.to(destination.device, copy=True)
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear")
mask = comfy.utils.repeat_to_batch_size(mask, source.shape[0])

Expand Down

0 comments on commit 6781b18

Please sign in to comment.