diff --git a/nodes/mask_nodes.py b/nodes/mask_nodes.py
index ddd81b4..7951ffe 100644
--- a/nodes/mask_nodes.py
+++ b/nodes/mask_nodes.py
@@ -7,10 +7,10 @@ import numpy as np
 from contextlib import nullcontext
 import os
 import json
-import model_management
-from nodes import MAX_RESOLUTION
+# import model_management
+ # from nodes import MAX_RESOLUTION
 
-from ..utility.utility import tensor2pil, pil2tensor
+from utility.utility import tensor2pil, pil2tensor
 
 script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
@@ -30,17 +30,30 @@ class BatchImageToMask:
     DESCRIPTION = "Converts RGB images to binary masks using grayscale conversion and thresholding, with optional morphological dilation"
 
     def batch_convert_to_mask(self, images, threshold, dilation_amount):
-        B, H, W, C = images.shape
-        masks_tensor = torch.zeros((B, H, W), device=images.device)
-        
-        # Convert each RGB image to a binary mask
-        for i in range(B):
-            img = images[i]
-            gray = img.mean(dim=-1)  # Simple RGB to grayscale conversion
-            mask = (gray > threshold).float()  # Binarize using threshold
-            masks_tensor[i] = mask
+        # Get shape - handle both [B, H, W] and [B, C, H, W] formats
+        if len(images.shape) == 3:
+            B, H, W = images.shape  # Already in [B, H, W] format
+            # Create a tensor to hold the masks
+            masks_tensor = torch.zeros((B, H, W), device=images.device)
             
-                    # Apply morphological dilation if requested
+            # Convert each image to a binary mask
+            for i in range(B):
+                mask = (images[i] > threshold).float()  # Binarize using threshold
+                masks_tensor[i] = mask
+        else:
+            # Handle [B, C, H, W] format (standard ComfyUI IMAGE)
+            B, C, H, W = images.shape
+            masks_tensor = torch.zeros((B, H, W), device=images.device)
+            
+            # Convert each RGB image to a binary mask
+            for i in range(B):
+                img = images[i]  # Shape: [C, H, W]
+                # Convert to grayscale by taking mean across channels
+                gray = img.mean(dim=0)  # Shape: [H, W]
+                mask = (gray > threshold).float()  # Binarize using threshold
+                masks_tensor[i] = mask
+            
+        # Apply morphological dilation if requested
         if dilation_amount > 0:
             c = 1  # non-tapered corners
             kernel = np.array([[c, 1, c],
@@ -48,13 +61,13 @@ class BatchImageToMask:
                              [c, 1, c]])
             
             dilated_masks = []
-            for m in mask_tensor:
+            for m in masks_tensor:
                 output = m.cpu().numpy().astype(np.float32)
                 for _ in range(dilation_amount):
                     output = scipy.ndimage.grey_dilation(output, footprint=kernel)
                 dilated_masks.append(torch.from_numpy(output))
             
-            mask_tensor = torch.stack(dilated_masks)
+            masks_tensor = torch.stack(dilated_masks).to(images.device)
 
         return (masks_tensor,)
 
@@ -76,6 +89,10 @@ class MapTrajectoriesToSegmentedMasks:
         # Parse trajectories JSON
         paths_list = json.loads(trajectories)
         
+        # Handle case where masks is a tuple (common in ComfyUI node system)
+        if isinstance(masks, tuple):
+            masks = masks[0]  # Extract the mask tensor from the tuple
+        
         if len(paths_list) != masks.shape[0]:
             raise ValueError(f"Number of trajectories ({len(paths_list)}) must match number of masks ({masks.shape[0]})")
         
@@ -120,3 +137,4 @@ class MapTrajectoriesToSegmentedMasks:
         
         return (masks, translated_trajectories_json)
 
+