From 94f30ef317d4607a2d91ac222b2bdf4addb1ea66 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Thu, 15 Aug 2024 16:16:53 +0900
Subject: [PATCH] fix: resize_mask - robust processing

---
 modules/impact/config.py |  2 +-
 modules/impact/utils.py  | 13 ++++++++++++-
 pyproject.toml           |  2 +-
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/modules/impact/config.py b/modules/impact/config.py
index 68f6722..c1f95f5 100644
--- a/modules/impact/config.py
+++ b/modules/impact/config.py
@@ -1,7 +1,7 @@
 import configparser
 import os
 
-version_code = [6, 2]
+version_code = [6, 2, 1]
 version = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '')
 
 dependency_version = 22
diff --git a/modules/impact/utils.py b/modules/impact/utils.py
index 128eaa5..bf4c40b 100644
--- a/modules/impact/utils.py
+++ b/modules/impact/utils.py
@@ -537,6 +537,16 @@ def make_3d_mask(mask):
     return mask
 
 
+def make_4d_mask(mask):
+    if len(mask.shape) == 3:
+        return mask.unsqueeze(0)
+
+    elif len(mask.shape) == 2:
+        return mask.unsqueeze(0).unsqueeze(0)
+
+    return mask
+
+
 def is_same_device(a, b):
     a_device = torch.device(a) if isinstance(a, str) else a
     b_device = torch.device(b) if isinstance(b, str) else b
@@ -556,7 +566,8 @@ from torchvision.transforms.functional import to_pil_image
 
 
 def resize_mask(mask, size):
-    resized_mask = torch.nn.functional.interpolate(mask.unsqueeze(0), size=size, mode='bilinear', align_corners=False)
+    mask = make_4d_mask(mask)
+    resized_mask = torch.nn.functional.interpolate(mask, size=size, mode='bilinear', align_corners=False)
     return resized_mask.squeeze(0)
 
 
diff --git a/pyproject.toml b/pyproject.toml
index c93d8ab..9e183d7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-impact-pack"
 description = "This extension offers various detector nodes and detailer nodes that allow you to configure a workflow that automatically enhances facial details. And provide iterative upscaler."
-version = "6.2"
+version = "6.2.1"
 license = { file = "LICENSE.txt" }
 dependencies = ["segment-anything", "scikit-image", "piexif", "transformers", "opencv-python-headless", "GitPython", "scipy>=1.11.4"]