From 3b9a6cf2b11094f92228b121c6a0d466ba5d5246 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 13 Nov 2024 07:18:30 -0500 Subject: [PATCH] Fix issue with 3d masks. --- comfy/utils.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/comfy/utils.py b/comfy/utils.py index 04926c1e..985cd9a1 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -853,19 +853,18 @@ def reshape_mask(input_mask, output_shape): dims = len(output_shape) - 2 if dims == 1: - mask = input_mask scale_mode = "linear" 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" if dims == 3: 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" - 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]: mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]] mask = comfy.utils.repeat_to_batch_size(mask, output_shape[0])