Update models/loaders/matanyone_loader.py

models/loaders/matanyone_loader.py

@@ -217,23 +217,37 @@ def guarded_method(*args, **kwargs):
                     return torch.ones((1, 512, 512), dtype=torch.float32) * 0.5
                 
                 try:
-                    # Coerce shapes
+                    # Coerce shapes - ensure we REALLY squeeze out extra dimensions
                     img_nchw = ensure_image_nchw(image, want_batched=True)
                     
+                    # CRITICAL FIX: Force squeeze all unnecessary dimensions
+                    while img_nchw.ndim > 4:
+                        if img_nchw.shape[0] == 1:
+                            img_nchw = img_nchw.squeeze(0)
+                        elif img_nchw.shape[1] == 1:
+                            img_nchw = img_nchw.squeeze(1)
+                        else:
+                            break
+                    
                     if idx_mask:
                         m_fixed = ensure_mask_for_matanyone(mask, idx_mask=True)
                     else:
                         m_fixed = ensure_mask_for_matanyone(mask, idx_mask=False, threshold=0.5)
                     
-                    # Log shapes for debugging
+                    # Log actual shapes being passed
                     logger.debug(f"MatAnyone input - image: {img_nchw.shape}, mask: {m_fixed.shape}, idx: {idx_mask}")
                     
+                    # For MatAnyone, we need CHW not NCHW for unbatched
+                    if img_nchw.ndim == 4 and img_nchw.shape[0] == 1:
+                        img_chw = img_nchw[0]  # Remove batch dimension
+                    else:
+                        img_chw = img_nchw
+                    
                     # Try unbatched first (most common)
                     try:
                         new_kwargs = dict(kwargs)
-                        # CRITICAL: Use unbatched (CHW) not batched for first attempt
-                        new_kwargs["image"] = img_nchw.squeeze(0) if img_nchw.shape[0] == 1 else img_nchw[0]  # CHW
-                        new_kwargs["mask"] = m_fixed.squeeze(0) if m_fixed.shape[0] == 1 else m_fixed  # HW or CHW
+                        new_kwargs["image"] = img_chw  # CHW
+                        new_kwargs["mask"] = m_fixed.squeeze(0) if m_fixed.ndim > 2 and m_fixed.shape[0] == 1 else m_fixed
                         new_kwargs["idx_mask"] = bool(idx_mask)
                         
                         result = original_method(**new_kwargs)