Update models/loaders/matanyone_loader.py

models/loaders/matanyone_loader.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
 """
 MatAnyone Loader - Stable Callable Wrapper for InferenceCore (extra-dim stripping)
 =================================================================================
@@ -11,16 +9,11 @@
     e.g. [B,T,C,H,W] -> [C,H,W]  (use first slice when B/T > 1 with a warning)
     e.g. [B,C,H,W]   -> [C,H,W]
     e.g. [H,W,C,1]   -> [H,W,C]
-- Optional CUDA mixed precision (fp16/bf16)
 - Robust alpha extraction -> (H,W) float32 [0,1]
 """
+from __future__ import annotations
 
-import os
-import time
 import logging
-import tempfile
-import traceback
-from pathlib import Path
 from typing import Optional, Dict, Any, Tuple, Union
 
 import numpy as np
@@ -28,6 +21,12 @@
 
 logger = logging.getLogger(__name__)
 
+try:
+    # Official import path
+    from matanyone.inference.inference_core import InferenceCore
+except Exception:  # keep import error defered until load()
+    InferenceCore = None  # type: ignore
+
 
 # ------------------------------ Helpers ------------------------------
 
@@ -75,28 +74,23 @@ def _ensure_chw_float01(image: Union[np.ndarray, torch.Tensor], *, name: str = "
 
     if torch.is_tensor(image):
         t = image
-        # Convert 4D (rare if caller passes) once more
         if t.ndim == 4:
             t = _strip_leading_extras_to_ndim(t, 3)
 
         if t.ndim == 3:
             c0, c1, c2 = t.shape
             if c0 in (1, 3, 4):
-                # CHW
-                pass
+                pass  # CHW
             elif c2 in (1, 3, 4):
-                # HWC -> CHW
-                t = t.permute(2, 0, 1)
+                t = t.permute(2, 0, 1)  # HWC -> CHW
             else:
-                # Ambiguous, assume HWC-like and take first channel after moving to CHW
                 logger.warning(f"{name}: ambiguous 3D shape {tuple(t.shape)}; attempting HWC->CHW then selecting first channel.")
                 t = t.permute(2, 0, 1)
                 if t.shape[0] > 1:
                     t = t[0]
-                    t = t.unsqueeze(0)  # back to 1HW
+                    t = t.unsqueeze(0)
         elif t.ndim == 2:
-            # HW -> 1HW
-            t = t.unsqueeze(0)
+            t = t.unsqueeze(0)  # 1HW
         else:
             raise ValueError(f"{name}: unsupported tensor dims {tuple(t.shape)} after stripping.")
 
@@ -107,21 +101,20 @@ def _ensure_chw_float01(image: Union[np.ndarray, torch.Tensor], *, name: str = "
         logger.debug(f"{name}: {orig_shape} -> {tuple(t.shape)} (CHW)")
         return t
 
-    # numpy path
     arr = np.asarray(image)
     if arr.ndim == 4:
         arr = _strip_leading_extras_to_ndim(arr, 3)
 
     if arr.ndim == 3:
-        if arr.shape[0] in (1, 3, 4):   # CHW
-            pass
-        elif arr.shape[-1] in (1, 3, 4):  # HWC -> CHW
-            arr = arr.transpose(2, 0, 1)
+        if arr.shape[0] in (1, 3, 4):
+            pass  # CHW
+        elif arr.shape[-1] in (1, 3, 4):
+            arr = arr.transpose(2, 0, 1)  # HWC -> CHW
         else:
             logger.warning(f"{name}: ambiguous 3D shape {arr.shape}; trying HWC->CHW and selecting first channel.")
-            arr = arr.transpose(2, 0, 1)  # HWC->CHW
+            arr = arr.transpose(2, 0, 1)
             if arr.shape[0] > 1:
-                arr = arr[0:1, ...]      # 1HW
+                arr = arr[0:1, ...]
     elif arr.ndim == 2:
         arr = arr[None, ...]  # 1HW
     else:
@@ -144,24 +137,20 @@ def _ensure_1hw_float01(mask: Union[np.ndarray, torch.Tensor], *, name: str = "m
     if torch.is_tensor(mask):
         m = mask
         if m.ndim == 3:
-            # 1HW or CHW or HWC-like
             if m.shape[0] == 1:
                 pass  # 1HW
             elif m.shape[-1] == 1:
                 m = m.permute(2, 0, 1)  # HW1 -> 1HW
             else:
-                # If multi-channel, take first
                 logger.warning(f"{name}: multi-channel {tuple(m.shape)}; using first channel.")
-                # Assume CHW or HWC-like already normalized earlier; prefer leading as channel
                 if m.shape[0] in (3, 4):
                     m = m[0:1, ...]
                 elif m.shape[-1] in (3, 4):
                     m = m.permute(2, 0, 1)[0:1, ...]
                 else:
-                    # Ambiguous -> take first along first axis and ensure 1HW
                     m = m[0:1, ...]
         elif m.ndim == 2:
-            m = m.unsqueeze(0)  # 1HW
+            m = m.unsqueeze(0)
         else:
             raise ValueError(f"{name}: unsupported tensor dims {tuple(m.shape)} after stripping.")
 
@@ -172,7 +161,6 @@ def _ensure_1hw_float01(mask: Union[np.ndarray, torch.Tensor], *, name: str = "m
         logger.debug(f"{name}: {orig_shape} -> {tuple(m.shape)} (1HW)")
         return m
 
-    # numpy path
     arr = np.asarray(mask)
     if arr.ndim == 3:
         if arr.shape[0] == 1:
@@ -182,13 +170,13 @@ def _ensure_1hw_float01(mask: Union[np.ndarray, torch.Tensor], *, name: str = "m
         else:
             logger.warning(f"{name}: multi-channel {arr.shape}; using first channel.")
             if arr.shape[0] in (3, 4):
-                arr = arr[0:1, ...]  # CHW -> 1HW
+                arr = arr[0:1, ...]
             elif arr.shape[-1] in (3, 4):
-                arr = arr.transpose(2, 0, 1)[0:1, ...]  # HWC -> CHW -> 1HW
+                arr = arr.transpose(2, 0, 1)[0:1, ...]
             else:
-                arr = arr[0:1, ...]  # ambiguous -> 1HW by slice
+                arr = arr[0:1, ...]
     elif arr.ndim == 2:
-        arr = arr[None, ...]  # 1HW
+        arr = arr[None, ...]
     else:
         raise ValueError(f"{name}: unsupported numpy dims {arr.shape} after stripping.")
 
@@ -207,7 +195,6 @@ def _alpha_from_result(result: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
         result = result.detach().float().cpu()
 
     arr = np.asarray(result)
-    # Strip to <= 3 dims, then extract
     while arr.ndim > 3:
         if arr.shape[0] > 1:
             logger.warning(f"Result has leading dim {arr.shape[0]}; taking first slice.")
@@ -216,14 +203,13 @@ def _alpha_from_result(result: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
     if arr.ndim == 2:
         alpha = arr
     elif arr.ndim == 3:
-        if arr.shape[0] in (1, 3, 4):      # CHW -> take channel 0
+        if arr.shape[0] in (1, 3, 4):
             alpha = arr[0]
-        elif arr.shape[-1] in (1, 3, 4):   # HWC -> take channel 0
+        elif arr.shape[-1] in (1, 3, 4):
             alpha = arr[..., 0]
         else:
-            alpha = arr[0]  # ambiguous
+            alpha = arr[0]
     else:
-        # 1D or 0D shouldn't happen; fallback
         alpha = np.full((512, 512), 0.5, dtype=np.float32)
 
     alpha = alpha.astype(np.float32, copy=False)
@@ -233,26 +219,18 @@ def _alpha_from_result(result: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
 
 def _hw_from_image_like(x: Union[np.ndarray, torch.Tensor]) -> Tuple[int, int]:
     """Best-effort infer (H, W) for fallback mask sizing."""
-    if torch.is_tensor(x):
-        shape = tuple(x.shape)
-    else:
-        shape = np.asarray(x).shape
-
-    # Try common orders
-    if len(shape) == 2:  # HW
+    shape = tuple(x.shape) if torch.is_tensor(x) else np.asarray(x).shape
+    if len(shape) == 2:
         return shape[0], shape[1]
     if len(shape) == 3:
-        if shape[0] in (1, 3, 4):  # CHW
+        if shape[0] in (1, 3, 4):
             return shape[1], shape[2]
-        if shape[-1] in (1, 3, 4):  # HWC
+        if shape[-1] in (1, 3, 4):
             return shape[0], shape[1]
-        # Ambiguous -> treat as CHW
         return shape[1], shape[2]
     if len(shape) >= 4:
-        # Assume leading are batch/time; try BCHW first
         if len(shape) >= 4 and (shape[1] in (1, 3, 4)):
             return shape[2], shape[3]
-        # Else BHWC-ish
         return shape[-3], shape[-2]
     return 512, 512
 
@@ -270,25 +248,16 @@ class MatAnyoneCallableWrapper:
       - Strips any extra dims from inputs before calling core.
     """
 
-    def __init__(self, inference_core, device: str = "cuda", mixed_precision: Optional[str] = "fp16"):
+    def __init__(self, inference_core, device: str = None):
         self.core = inference_core
         self.initialized = False
-        self.device = device if (device in ("cuda", "cpu")) else ("cuda" if torch.cuda.is_available() else "cpu")
-        self.mixed_precision = mixed_precision if self.device == "cuda" else None  # "fp16"|"bf16"|None
-
-    def _maybe_autocast(self):
-        if self.device == "cuda" and self.mixed_precision in ("fp16", "bf16"):
-            dtype = torch.float16 if self.mixed_precision == "fp16" else torch.bfloat16
-            return torch.autocast(device_type="cuda", dtype=dtype)
-        # no-op ctx
-        class _NullCtx:
-            def __enter__(self): return None
-            def __exit__(self, *exc): return False
-        return _NullCtx()
+        # Best-effort device selection if available
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
 
     def __call__(self, image, mask=None, **kwargs) -> np.ndarray:
         try:
-            # Preprocess (unbatched)
             img_chw = _ensure_chw_float01(image, name="image").to(self.device, non_blocking=True)
 
             if not self.initialized:
@@ -300,35 +269,32 @@ def __call__(self, image, mask=None, **kwargs) -> np.ndarray:
                 m_1hw = _ensure_1hw_float01(mask, name="mask").to(self.device, non_blocking=True)
 
                 with torch.inference_mode():
-                    with self._maybe_autocast():
-                        if hasattr(self.core, "step"):
-                            result = self.core.step(image=img_chw, mask=m_1hw, **kwargs)
-                        elif hasattr(self.core, "process_frame"):
-                            result = self.core.process_frame(img_chw, m_1hw, **kwargs)
-                        else:
-                            logger.warning("InferenceCore has no recognized frame API; echoing input mask.")
-                            return _alpha_from_result(mask)
+                    if hasattr(self.core, "step"):
+                        result = self.core.step(image=img_chw, mask=m_1hw, **kwargs)
+                    elif hasattr(self.core, "process_frame"):
+                        result = self.core.process_frame(img_chw, m_1hw, **kwargs)
+                    else:
+                        logger.warning("InferenceCore has no recognized frame API; echoing input mask.")
+                        return _alpha_from_result(mask)
 
                 self.initialized = True
                 return _alpha_from_result(result)
 
             # Subsequent frames (no mask)
             with torch.inference_mode():
-                with self._maybe_autocast():
-                    if hasattr(self.core, "step"):
-                        result = self.core.step(image=img_chw, **kwargs)
-                    elif hasattr(self.core, "process_frame"):
-                        result = self.core.process_frame(img_chw, **kwargs)
-                    else:
-                        h, w = _hw_from_image_like(image)
-                        logger.warning("InferenceCore has no recognized frame API on subsequent call; returning neutral alpha.")
-                        return np.full((h, w), 0.5, dtype=np.float32)
+                if hasattr(self.core, "step"):
+                    result = self.core.step(image=img_chw, **kwargs)
+                elif hasattr(self.core, "process_frame"):
+                    result = self.core.process_frame(img_chw, **kwargs)
+                else:
+                    h, w = _hw_from_image_like(image)
+                    logger.warning("InferenceCore has no recognized frame API on subsequent call; returning neutral alpha.")
+                    return np.full((h, w), 0.5, dtype=np.float32)
 
             return _alpha_from_result(result)
 
         except Exception as e:
             logger.error(f"MatAnyone wrapper call failed: {e}")
-            logger.debug(traceback.format_exc())
             # Fallbacks
             if mask is not None:
                 try:
@@ -353,160 +319,51 @@ def reset(self):
                 logger.debug(f"Core clear_memory() failed: {e}")
 
 
-# ------------------------------- Loader -------------------------------
-
 class MatAnyoneLoader:
-    """
-    Loads MatAnyone's InferenceCore and returns a callable wrapper.
-
-    Usage:
-        loader = MatAnyoneLoader(device="cuda")
-        session = loader.load()  # callable
-        alpha = session(frame, first_frame_mask)  # returns (H, W) float32
-    """
+    def __init__(self, device: str = "auto", model_id: str = "PeiqingYang/MatAnyone"):
+        self.device = device
+        self.model_id = model_id
+        self._processor: Optional[InferenceCore] = None  # type: ignore
+        self._wrapper: Optional[MatAnyoneCallableWrapper] = None
 
-    def __init__(self, device: str = "cuda", cache_dir: str = "./checkpoints/matanyone_cache",
-                 mixed_precision: Optional[str] = "fp16"):
-        self.device = self._select_device(device)
-        self.cache_dir = cache_dir
-        os.makedirs(self.cache_dir, exist_ok=True)
-
-        self.processor = None
-        self.wrapper = None
-        self.model_id = "PeiqingYang/MatAnyone"
-        self.load_time = 0.0
-        self.loaded = False
-        self.load_error = None
-        self.temp_dir = Path(tempfile.mkdtemp())
-        self.mixed_precision = mixed_precision if self.device == "cuda" else None
-
-    def _select_device(self, pref: str) -> str:
-        pref = (pref or "").lower()
-        if pref.startswith("cuda"):
-            return "cuda" if torch.cuda.is_available() else "cpu"
-        if pref == "cpu":
-            return "cpu"
-        return "cuda" if torch.cuda.is_available() else "cpu"
-
-    def _try_build_core(self):
+    def load(self) -> Optional[Any]:
         """
-        Try multiple constructor patterns to survive API changes.
+        Initialize and return a callable wrapper around InferenceCore.
+        Returns MatAnyoneCallableWrapper if successful, else None.
         """
-        from matanyone.inference.inference_core import InferenceCore
-
-        # 1) Preferred: from_pretrained(...)
+        global InferenceCore
         try:
-            core = InferenceCore.from_pretrained(self.model_id, device=self.device, cache_dir=self.cache_dir)
-            logger.info("Loaded MatAnyone via InferenceCore.from_pretrained(...)")
-            return core
-        except Exception as e:
-            logger.debug(f"from_pretrained failed: {e}")
-
-        # 2) Direct ctor with device/cache_dir
-        try:
-            core = InferenceCore(self.model_id, device=self.device, cache_dir=self.cache_dir)
-            logger.info("Loaded MatAnyone via InferenceCore(model_id, device, cache_dir)")
-            return core
-        except Exception as e:
-            logger.debug(f"ctor(model_id, device, cache_dir) failed: {e}")
-
-        # 3) Minimal ctor
-        core = InferenceCore(self.model_id)
-        logger.info("Loaded MatAnyone via InferenceCore(model_id) [minimal]")
-        return core
-
-    def load(self) -> Optional[MatAnyoneCallableWrapper]:
-        """Load MatAnyone and return the callable wrapper."""
-        if self.loaded and self.wrapper is not None:
-            return self.wrapper
-
-        logger.info(f"Loading MatAnyone: {self.model_id} (device={self.device})")
-        t0 = time.time()
-
-        try:
-            self.processor = self._try_build_core()
-            # Optional device move
-            try:
-                if hasattr(self.processor, "to"):
-                    self.processor.to(self.device)
-                elif hasattr(self.processor, "set_device"):
-                    self.processor.set_device(self.device)
-            except Exception as e:
-                logger.debug(f"Optional device move failed: {e}")
-
-            self.wrapper = MatAnyoneCallableWrapper(
-                self.processor, device=self.device, mixed_precision=self.mixed_precision
+            if InferenceCore is None:
+                from matanyone.inference.inference_core import InferenceCore as _IC  # type: ignore
+                InferenceCore = _IC  # type: ignore
+
+            logger.info("Loading MatAnyone InferenceCore ...")
+            self._processor = InferenceCore(self.model_id)  # type: ignore
+            logger.info("MatAnyone InferenceCore loaded successfully")
+
+            # Choose device
+            dev = (
+                "cuda" if (str(self.device).startswith("cuda") and torch.cuda.is_available()) else
+                ("cpu" if str(self.device) == "cpu" else ("cuda" if torch.cuda.is_available() else "cpu"))
             )
-            self.loaded = True
-            self.load_time = time.time() - t0
-            logger.info(f"MatAnyone loaded and wrapped in {self.load_time:.2f}s")
-            return self.wrapper
-
-        except ImportError as e:
-            self.load_error = f"MatAnyone not installed: {e}"
-            logger.error("Failed to import MatAnyone. Install with: "
-                         "pip install git+https://github.com/pq-yang/MatAnyone.git@main")
-            return None
+
+            self._wrapper = MatAnyoneCallableWrapper(self._processor, device=dev)
+            logger.info("MatAnyone wrapped with dimension-safe callable")
+            return self._wrapper
         except Exception as e:
-            self.load_error = str(e)
-            logger.error(f"Failed to load MatAnyone: {e}")
-            logger.debug(traceback.format_exc())
+            logger.error(f"Failed to load MatAnyone InferenceCore: {e}")
+            self._processor = None
+            self._wrapper = None
             return None
 
-    def cleanup(self):
-        """Cleanup temporary files and release resources."""
-        self.processor = None
-        self.wrapper = None
-
-        # Clean temp directory
-        if self.temp_dir.exists():
-            import shutil
-            shutil.rmtree(self.temp_dir, ignore_errors=True)
-
-        # Clear CUDA cache if available
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+    def get(self) -> Optional[Any]:
+        """Return the cached callable if loaded."""
+        return self._wrapper or self._processor
 
     def get_info(self) -> Dict[str, Any]:
-        """Get model information and interface flags."""
-        info = {
-            "loaded": self.loaded,
+        """Metadata for diagnostics."""
+        return {
             "model_id": self.model_id,
-            "device": str(self.device),
-            "load_time": float(self.load_time),
-            "error": self.load_error,
-            "api": "InferenceCore (wrapped)",
-            "mixed_precision": self.mixed_precision,
-        }
-        proc = self.processor
-        if proc is not None:
-            info["has_step"] = hasattr(proc, "step")
-            info["has_process_frame"] = hasattr(proc, "process_frame")
-            info["has_process_video"] = hasattr(proc, "process_video")
-        return info
-
-    def reset(self):
-        """Reset the processor for a new video."""
-        if self.wrapper:
-            self.wrapper.reset()
-        logger.info("MatAnyone session reset")
-
-    # Make the loader itself callable (direct compatibility)
-    def __call__(self, image, mask=None, **kwargs) -> np.ndarray:
-        if self.wrapper is None:
-            if self.load() is None:
-                # Fallback if loading fails
-                if mask is not None:
-                    try:
-                        return _alpha_from_result(mask)
-                    except Exception:
-                        pass
-                h, w = _hw_from_image_like(image)
-                return np.zeros((h, w), dtype=np.float32)
-        return self.wrapper(image, mask, **kwargs)
-
-
-# Backwards compatibility alias
-_MatAnyoneSession = MatAnyoneCallableWrapper
-
-__all__ = ["MatAnyoneLoader", "_MatAnyoneSession", "MatAnyoneCallableWrapper"]
+            "loaded": self._wrapper is not None or self._processor is not None,
+            "wrapped": self._wrapper is not None,
+        }