Update app.py

app.py

@@ -413,11 +413,12 @@ class MatAnyoneHandler:
     """
     MatAnyone loader + inference adapter.
 
-    Key points:
-    - Uses first-frame *soft probability* seed (1xHxW float in [0,1]), not an index mask.
-    - Calls InferenceCore.step with the prob map as a **positional** arg (some builds reject `prob=`).
-    - Tries `matting=True` when supported; falls back if the kwarg is not available.
-    - Always feeds CHW tensors for images (3,H,W) and 1xHxW for probs — no extra batch dims.
+    Strategy:
+    - Image → CHW (3,H,W), no batch/time dims.
+    - Seed mask → try **2D prob** (H,W) first to avoid 5 vs 6-D concat, then fall back to 1xHxW.
+    - Call InferenceCore.step with prob as a **positional** argument.
+    - Try with/without `matting=True` (some builds don't accept it).
+    - Subsequent frames call step(image) with no seed.
     """
     def __init__(self):
         self.core = None
@@ -427,21 +428,29 @@ def __init__(self):
     def _to_chw_float(self, img01: np.ndarray) -> "torch.Tensor":
         """img01: HxWx3 in [0,1] -> torch float 3xHxW on DEVICE"""
         assert img01.ndim == 3 and img01.shape[2] == 3, f"Expected HxWx3, got {img01.shape}"
-        t = torch.from_numpy(img01.transpose(2, 0, 1)).contiguous().float()  # 3xHxW
+        t = torch.from_numpy(img01.transpose(2, 0, 1)).contiguous().float()  # (3,H,W)
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
 
-    def _prob_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
-        """mask_u8: HxW uint8 -> torch float 1xHxW on DEVICE, resized to (w,h) if needed"""
+    def _prob_hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
+        """mask_u8: HxW uint8 -> torch float (H,W) on DEVICE, resized if needed"""
         if mask_u8.shape[0] != h or mask_u8.shape[1] != w:
             mask_u8 = cv2.resize(mask_u8, (w, h), interpolation=cv2.INTER_NEAREST)
-        prob = (mask_u8.astype(np.float32) / 255.0)[None, ...]  # 1xHxW
+        prob = (mask_u8.astype(np.float32) / 255.0)  # (H,W)
+        t = torch.from_numpy(prob).contiguous().float()
+        return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
+
+    def _prob_1hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
+        """mask_u8: HxW uint8 -> torch float (1,H,W) on DEVICE, resized if needed"""
+        if mask_u8.shape[0] != h or mask_u8.shape[1] != w:
+            mask_u8 = cv2.resize(mask_u8, (w, h), interpolation=cv2.INTER_NEAREST)
+        prob = (mask_u8.astype(np.float32) / 255.0)[None, ...]  # (1,H,W)
         t = torch.from_numpy(prob).contiguous().float()
         return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
 
     def _alpha_to_u8_hw(self, alpha_like) -> np.ndarray:
         """
         Accepts torch Tensor or numpy-like. Returns uint8 HxW (0..255).
-        Handles shapes (H,W), (1,H,W), or (K,H,W) -> picks first channel.
+        Handles (H,W), (1,H,W), or (K,H,W) by taking the first channel if needed.
         Also handles MatAnyone tuples/lists like (indices, probs) by taking the 2nd item.
         """
         if isinstance(alpha_like, (list, tuple)) and len(alpha_like) > 1:
@@ -503,19 +512,45 @@ def initialize(self) -> bool:
                 state.matanyone_error = f"MatAnyone init error: {e}"
                 return False
 
-    # ----- video matting using first-frame PROB mask (PATCHED) -----
+    # ----- robust call helpers ------------------------------------------------
+    def _call_step_seed(self, img_chw: "torch.Tensor",
+                        prob_hw: "torch.Tensor",
+                        prob_1hw: "torch.Tensor"):
+        """
+        Try a few safe permutations to satisfy different MatAnyone builds.
+        Order chosen to avoid the 5-vs-6D concat error seen in group_modules.py.
+        """
+        # 1) image (3,H,W), prob (H,W)
+        try:
+            return self.core.step(img_chw, prob_hw)
+        except (TypeError, RuntimeError):
+            pass
+
+        # 2) image (3,H,W), prob (1,H,W)
+        try:
+            return self.core.step(img_chw, prob_1hw)
+        except (TypeError, RuntimeError):
+            pass
+
+        # 3) image (3,H,W), prob (H,W), matting kw
+        try:
+            return self.core.step(img_chw, prob_hw, matting=True)
+        except (TypeError, RuntimeError):
+            pass
+
+        # 4) image (3,H,W), prob (1,H,W), matting kw
+        return self.core.step(img_chw, prob_1hw, matting=True)
+
+    # ----- video matting using first-frame PROB mask (PATCHED) ----------------
     def process_video(self, input_path: str, mask_path: str, output_path: str) -> str:
         """
         Produce a single-channel alpha mp4 matching input fps & size.
 
         First frame:
-          - Generate soft prob (1,H,W) from SAM2 mask and pass as positional arg to step().
-          - Try step(image, prob, matting=True); if TypeError, call step(image, prob).
-
+          - Build both (H,W) and (1,H,W) soft prob tensors from SAM2 mask.
+          - Call step(...) via _call_step_seed (tries 2D first to dodge 6D concat).
         Remaining frames:
-          - Try step(image, matting=True); fallback to step(image).
-
-        Returns: path to alpha.mp4
+          - Call step(image) with no seed.
         """
         if not self.initialized or self.core is None:
             raise RuntimeError("MatAnyone not initialized")
@@ -532,25 +567,19 @@ def process_video(self, input_path: str, mask_path: str, output_path: str) -> st
         w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
-        # soft seed prob (1,H,W) in [0,1]
         seed_mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
         if seed_mask is None:
             cap.release()
             raise RuntimeError("Seed mask read failed")
-        prob_1hw = self._prob_from_mask_u8(seed_mask, w, h)  # (1,H,W) float
+
+        prob_hw  = self._prob_hw_from_mask_u8(seed_mask, w, h)    # (H,W)
+        prob_1hw = self._prob_1hw_from_mask_u8(seed_mask, w, h)   # (1,H,W)
 
         # temp frames
         tmp_dir = TEMP_DIR / f"ma_{int(time.time())}_{random.randint(1000,9999)}"
         tmp_dir.mkdir(parents=True, exist_ok=True)
         memory_manager.register_temp_file(str(tmp_dir))
 
-        def _step_with_prob(image_chw: "torch.Tensor", prob_1hw_t: "torch.Tensor"):
-            """Call step with positional prob; fall back if 'matting' kwarg unsupported."""
-            try:
-                return self.core.step(image_chw, prob_1hw_t, matting=True)
-            except TypeError:
-                return self.core.step(image_chw, prob_1hw_t)
-
         frame_idx = 0
 
         # --- first frame (with soft prob) ---
@@ -558,11 +587,12 @@ def _step_with_prob(image_chw: "torch.Tensor", prob_1hw_t: "torch.Tensor"):
         if not ok or frame_bgr is None:
             cap.release()
             raise RuntimeError("Empty first frame")
+
         frame_rgb01 = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
         img_chw = self._to_chw_float(frame_rgb01)  # (3,H,W)
 
         with torch.no_grad():
-            out_prob = _step_with_prob(img_chw, prob_1hw)
+            out_prob = self._call_step_seed(img_chw, prob_hw, prob_1hw)
 
         alpha_u8 = self._alpha_to_u8_hw(out_prob)
         cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
@@ -573,14 +603,16 @@ def _step_with_prob(image_chw: "torch.Tensor", prob_1hw_t: "torch.Tensor"):
             ok, frame_bgr = cap.read()
             if not ok or frame_bgr is None:
                 break
+
             frame_rgb01 = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
             img_chw = self._to_chw_float(frame_rgb01)
 
             with torch.no_grad():
                 try:
-                    out_prob = self.core.step(img_chw, matting=True)
+                    out_prob = self.core.step(img_chw)           # simplest path
                 except TypeError:
-                    out_prob = self.core.step(img_chw)
+                    # Extremely old/new variants: try permissive kw
+                    out_prob = self.core.step(img_chw, matting=True)
 
             alpha_u8 = self._alpha_to_u8_hw(out_prob)
             cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)