Update app.py

app.py

@@ -409,6 +409,244 @@ def create_mask(self, image_rgb: np.ndarray) -> Optional[np.ndarray]:
 # CHAPTER 6: MATANYONE HANDLER (First-frame PROB mask)
 # ==============================================================================
 
+class MatAnyoneHandler:
+    """
+    Fixed MatAnyone loader + inference adapter.
+    
+    Key fix: Only pass tensor inputs to MatAnyone.core.step() since the
+    internal pad_divide_by function expects tensors, not numpy arrays.
+    """
+    def __init__(self):
+        self.core = None
+        self.initialized = False
+
+    # ----- tensor helpers -----
+    def _to_chw_float(self, img01: np.ndarray) -> "torch.Tensor":
+        """img01: HxWx3 in [0,1] -> torch float (3,H,W) on DEVICE (no batch)."""
+        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()  # (3,H,W)
+        return t.to(DEVICE, non_blocking=CUDA_AVAILABLE)
+
+    def _prob_hw_from_mask_u8(self, mask_u8: np.ndarray, w: int, h: int) -> "torch.Tensor":
+        """mask_u8: HxW -> torch float (H,W) in [0,1] on DEVICE (no batch, no channel)."""
+        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)  # (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":
+        """Optional: 1xHxW (channel-first, still unbatched)."""
+        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 / numpy / tuple(list) outputs.
+        Returns uint8 HxW (0..255). Squeezes common shapes down to HxW.
+        """
+        if isinstance(alpha_like, (list, tuple)) and len(alpha_like) > 1:
+            alpha_like = alpha_like[1]  # (indices, probs) -> take probs
+
+        if isinstance(alpha_like, torch.Tensor):
+            t = alpha_like.detach()
+            if t.is_cuda:
+                t = t.cpu()
+            a = t.float().clamp(0, 1).numpy()
+        else:
+            a = np.asarray(alpha_like, dtype=np.float32)
+            a = np.clip(a, 0, 1)
+
+        a = np.squeeze(a)
+        if a.ndim == 3 and a.shape[0] >= 1:  # (1,H,W) -> (H,W)
+            a = a[0]
+        if a.ndim != 2:
+            raise ValueError(f"Alpha must be HxW; got {a.shape}")
+
+        return np.clip(a * 255.0, 0, 255).astype(np.uint8)
+
+    def initialize(self) -> bool:
+        if not TORCH_AVAILABLE:
+            state.matanyone_error = "PyTorch required"
+            return False
+        with memory_manager.mem_context("MatAnyone init"):
+            try:
+                _reset_hydra()
+                repo_path = ensure_repo("matanyone", "https://github.com/pq-yang/MatAnyone.git")
+                if not repo_path:
+                    state.matanyone_error = "Clone failed"
+                    return False
+                try:
+                    from matanyone.inference.inference_core import InferenceCore
+                    from matanyone.utils.get_default_model import get_matanyone_model
+                except Exception as e:
+                    state.matanyone_error = f"Import error: {e}"
+                    return False
+
+                ckpt = CHECKPOINTS / "matanyone.pth"
+                net = None
+                if ckpt.exists():
+                    net = get_matanyone_model(str(ckpt), device=DEVICE)
+                else:
+                    url = "https://github.com/pq-yang/MatAnyone/releases/download/v1.0.0/matanyone.pth"
+                    if download_file(url, ckpt, "MatAnyone"):
+                        net = get_matanyone_model(str(ckpt), device=DEVICE)
+
+                if net is None:
+                    state.matanyone_error = "Model load failed"
+                    return False
+
+                self.core = InferenceCore(net)
+                self.initialized = True
+                state.matanyone_ready = True
+                return True
+            except Exception as e:
+                state.matanyone_error = f"MatAnyone init error: {e}"
+                return False
+
+    # ----- FIXED: tensor-only call helpers --------------------------------
+    def _try_step_variants_seed(self,
+                                img_chw_t: "torch.Tensor",
+                                prob_hw_t: "torch.Tensor",
+                                prob_1hw_t: "torch.Tensor"):
+        """
+        Try multiple MatAnyone.step() signatures with TENSOR INPUTS ONLY.
+        MatAnyone's internal functions expect tensors, not numpy arrays.
+        
+        Order (most to least preferred):
+          1) step(CHW_tensor, HW_tensor)
+          2) step(CHW_tensor, HW_tensor, matting=True)
+          3) step(CHW_tensor, 1HW_tensor)
+          4) step(CHW_tensor, 1HW_tensor, matting=True)
+        """
+        trials = [
+            ( (img_chw_t, prob_hw_t), {} ),
+            ( (img_chw_t, prob_hw_t), {"matting": True} ),
+            ( (img_chw_t, prob_1hw_t), {} ),
+            ( (img_chw_t, prob_1hw_t), {"matting": True} ),
+        ]
+        last_err = None
+        for (args, kwargs) in trials:
+            try:
+                return self.core.step(*args, **kwargs)
+            except Exception as e:
+                last_err = e
+                # Keep trying next variant
+        raise last_err  # bubble up the most informative final error
+
+    def _try_step_variants_noseed(self, img_chw_t: "torch.Tensor"):
+        """
+        Variants when no seed is provided on subsequent frames.
+        TENSOR INPUT ONLY.
+        """
+        trials = [
+            ( (img_chw_t,), {} ),
+            ( (img_chw_t,), {"matting": True} ),
+        ]
+        last_err = None
+        for (args, kwargs) in trials:
+            try:
+                return self.core.step(*args, **kwargs)
+            except Exception as e:
+                last_err = e
+        raise last_err
+
+    # ----- video matting using first-frame PROB mask --------------------------
+    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: pass a soft seed prob (~HW) alongside the image.
+        Remaining frames: call step(image) only.
+        """
+        if not self.initialized or self.core is None:
+            raise RuntimeError("MatAnyone not initialized")
+
+        out_dir = Path(output_path)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        alpha_path = out_dir / "alpha.mp4"
+
+        cap = cv2.VideoCapture(input_path)
+        if not cap.isOpened():
+            raise RuntimeError("Could not open input video")
+
+        fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
+        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+        # soft seed prob - prepare tensor versions only
+        seed_mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+        if seed_mask is None:
+            cap.release()
+            raise RuntimeError("Seed mask read failed")
+
+        prob_hw_t  = self._prob_hw_from_mask_u8(seed_mask, w, h)    # (H,W) torch
+        prob_1hw_t = self._prob_1hw_from_mask_u8(seed_mask, w, h)   # (1,H,W) torch
+
+        # 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))
+
+        frame_idx = 0
+
+        # --- first frame (with soft prob) ---
+        ok, frame_bgr = cap.read()
+        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_t = self._to_chw_float(frame_rgb01)     # (3,H,W) torch
+
+        with torch.no_grad():
+            out_prob = self._try_step_variants_seed(
+                img_chw_t, prob_hw_t, prob_1hw_t
+            )
+
+        alpha_u8 = self._alpha_to_u8_hw(out_prob)
+        cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
+        frame_idx += 1
+
+        # --- remaining frames (no seed) ---
+        while True:
+            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_t = self._to_chw_float(frame_rgb01)
+
+            with torch.no_grad():
+                out_prob = self._try_step_variants_noseed(img_chw_t)
+
+            alpha_u8 = self._alpha_to_u8_hw(out_prob)
+            cv2.imwrite(str(tmp_dir / f"{frame_idx:06d}.png"), alpha_u8)
+            frame_idx += 1
+
+        cap.release()
+
+        # --- encode PNGs → alpha mp4 ---
+        list_file = tmp_dir / "list.txt"
+        with open(list_file, "w") as f:
+            for i in range(frame_idx):
+                f.write(f"file '{(tmp_dir / f'{i:06d}.png').as_posix()}'\n")
+
+        cmd = [
+            "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
+            "-f", "concat", "-safe", "0",
+            "-r", f"{fps:.6f}",
+            "-i", str(list_file),
+            "-vf", f"format=gray,scale={w}:{h}:flags=area",
+            "-pix_fmt", "yuv420p",
+            "-c:v", "libx264", "-preset", "medium", "-crf", "18",
+            str(alpha_path)
+        ]
+        subprocess.run(cmd, check=True)
+        return str(alpha_path)
 # =============================================================================
 # CHAPTER 6: MATANYONE HANDLER (Robust unbatched calls + fallbacks)
 # =============================================================================