298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-11 02:15:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix issue with 3d masks.

Browse Source
This commit is contained in:
comfyanonymous 2024-11-13 07:18:30 -05:00
parent 3748e7ef7a
commit 3b9a6cf2b1
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
comfy/utils.py
View File

@ -853,19 +853,18 @@ def reshape_mask(input_mask, output_shape):
dims = len(output_shape) - 2 dims = len(output_shape) - 2
if dims == 1: if dims == 1:
mask = input_mask
scale_mode = "linear" scale_mode = "linear"
if dims == 2: if dims == 2:
mask = input_mask.reshape((-1, 1, input_mask.shape[-2], input_mask.shape[-1])) input_mask = input_mask.reshape((-1, 1, input_mask.shape[-2], input_mask.shape[-1]))
scale_mode = "bilinear" scale_mode = "bilinear"
if dims == 3: if dims == 3:
if len(input_mask.shape) < 5: if len(input_mask.shape) < 5:
mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1])) input_mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1]))
scale_mode = "trilinear" scale_mode = "trilinear"
mask = torch.nn.functional.interpolate(mask, size=output_shape[2:], mode=scale_mode) mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)
if mask.shape[1] < output_shape[1]: if mask.shape[1] < output_shape[1]:
mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]] mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]]
mask = comfy.utils.repeat_to_batch_size(mask, output_shape[0]) mask = comfy.utils.repeat_to_batch_size(mask, output_shape[0])