diff --git a/nodes.py b/nodes.py
index 82fa62c..dd98ce8 100644
--- a/nodes.py
+++ b/nodes.py
@@ -231,7 +231,7 @@ class GenerateNAID:
                 action = "infill"
                 image, mask, add_original_image = option["infill"]
                 params["image"] = image_to_base64(resize_image(image, (width, height)))
-                params["mask"] = naimask_to_base64(resize_to_naimask(mask, (width, height)))
+                params["mask"] = naimask_to_base64(resize_to_naimask(mask, (width, height), "4" in model))
                 params["add_original_image"] = add_original_image
 
             if "vibe" in option:
diff --git a/pyproject.toml b/pyproject.toml
index 1e27e89..26c1d0e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_naidgenerator"
 description = "This extension helps generate images through NAI."
-version = "1.0.6"
+version = "1.0.7"
 license = { file = "LICENSE" }
 dependencies = ["python-dotenv", "argon2-cffi"]
 
diff --git a/utils.py b/utils.py
index 3482f82..4c92820 100644
--- a/utils.py
+++ b/utils.py
@@ -121,12 +121,14 @@ def resize_image(image, size_to):
     s = s.movedim(1,-1)
     return s
 
-def resize_to_naimask(mask, image_size=None):
+def resize_to_naimask(mask, image_size=None, is_v4=False):
     samples = mask.movedim(-1,1)
     w, h = (samples.shape[3], samples.shape[2]) if not image_size else image_size
     width = int(np.ceil(w / 64) * 8)
     height = int(np.ceil(h / 64) * 8)
     s = comfy.utils.common_upscale(samples, width, height, "nearest-exact", "disabled")
+    if is_v4:
+        s = comfy.utils.common_upscale(s, width*8, height*8, "nearest-exact", "disabled")
     s = s.movedim(1,-1)
     return s