Update models/loaders/matanyone_loader.py

models/loaders/matanyone_loader.py

@@ -1,9 +1,14 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-MatAnyone Loader - Wrapper for Official InferenceCore API
-=========================================================
-Creates a callable wrapper around InferenceCore to maintain compatibility.
+MatAnyone Loader - Stable Callable Wrapper for InferenceCore
+===========================================================
+
+- Enforces image CHW float32 [0,1] and mask 1HW float32 [0,1]
+- Adds internal batch dim (B=1) and removes it on output
+- Works with multiple possible InferenceCore loading signatures
+- Uses torch.inference_mode() + optional autocast for speed
+- Returns a 2-D alpha mask (H,W) float32 in [0,1]
 """
 
 import os
@@ -12,133 +17,295 @@
 import tempfile
 import traceback
 from pathlib import Path
-from typing import Optional, Dict, Any, Tuple
+from typing import Optional, Dict, Any, Tuple, Union
 
 import numpy as np
 import torch
-import cv2
 
 logger = logging.getLogger(__name__)
 
 
+# ------------------------------ Helpers ------------------------------
+
+def _to_float01_np(arr: np.ndarray) -> np.ndarray:
+    """Ensure numpy array is float32 in [0,1]."""
+    if arr.dtype == np.uint8:
+        arr = arr.astype(np.float32) / 255.0
+    else:
+        arr = arr.astype(np.float32)
+    # Clamp for safety
+    np.clip(arr, 0.0, 1.0, out=arr)
+    return arr
+
+
+def _ensure_chw_float01(image: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+    """
+    Convert image to torch.FloatTensor CHW in [0,1].
+    Accepts HxWxC or CHW (numpy or tensor). Adds batch dim later.
+    """
+    if torch.is_tensor(image):
+        t = image
+        if t.ndim == 3 and t.shape[0] in (1, 3, 4):  # already CHW
+            pass
+        elif t.ndim == 3 and t.shape[-1] in (1, 3, 4):  # HWC -> CHW
+            t = t.permute(2, 0, 1)
+        elif t.ndim == 2:  # HW (grayscale)
+            t = t.unsqueeze(0)
+        else:
+            raise ValueError(f"Unsupported image tensor shape: {tuple(t.shape)}")
+        t = t.to(dtype=torch.float32)
+        # If likely 0-255, scale; otherwise clamp to [0,1]
+        if torch.max(t) > 1.5:
+            t = t / 255.0
+        t = torch.clamp(t, 0.0, 1.0)
+        return t
+    else:
+        arr = np.asarray(image)
+        if arr.ndim == 3 and arr.shape[2] in (1, 3, 4):  # HWC
+            arr = arr.transpose(2, 0, 1)  # -> CHW
+        elif arr.ndim == 2:  # HW
+            arr = arr[None, ...]  # -> 1HW
+        elif arr.ndim == 3 and arr.shape[0] in (1, 3, 4):  # already CHW
+            pass
+        else:
+            raise ValueError(f"Unsupported image numpy shape: {arr.shape}")
+        arr = _to_float01_np(arr)
+        return torch.from_numpy(arr)
+
+
+def _ensure_1hw_float01(mask: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+    """
+    Convert mask to torch.FloatTensor 1HW in [0,1].
+    Accepts HW, 1HW, CHW (C=1), HxWx1.
+    """
+    if torch.is_tensor(mask):
+        m = mask
+        if m.ndim == 2:            # HW
+            m = m.unsqueeze(0)     # 1HW
+        elif m.ndim == 3:
+            if m.shape[0] == 1:    # 1HW
+                pass
+            elif m.shape[-1] == 1: # HW1 -> 1HW
+                m = m.permute(2, 0, 1)
+            else:
+                raise ValueError(f"Mask has too many channels: {tuple(m.shape)}")
+        else:
+            raise ValueError(f"Unsupported mask tensor shape: {tuple(m.shape)}")
+        m = m.to(dtype=torch.float32)
+        if torch.max(m) > 1.5:
+            m = m / 255.0
+        m = torch.clamp(m, 0.0, 1.0)
+        return m
+    else:
+        arr = np.asarray(mask)
+        if arr.ndim == 2:               # HW
+            arr = arr[None, ...]        # 1HW
+        elif arr.ndim == 3:
+            if arr.shape[0] == 1:       # 1HW
+                pass
+            elif arr.shape[-1] == 1:    # HW1 -> 1HW
+                arr = arr.transpose(2, 0, 1)
+            else:
+                raise ValueError(f"Mask has too many channels: {arr.shape}")
+        else:
+            raise ValueError(f"Unsupported mask numpy shape: {arr.shape}")
+        arr = _to_float01_np(arr)
+        return torch.from_numpy(arr)
+
+
+def _alpha_from_result(result: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
+    """
+    Extract a 2D alpha (H,W) float32 [0,1] from a variety of possible outputs.
+    Accepts numpy/tensor with shapes: HW, 1HW, CHW(C>=1), BHWC, BCHW, etc.
+    """
+    if result is None:
+        return np.full((512, 512), 0.5, dtype=np.float32)
+
+    if torch.is_tensor(result):
+        result = result.detach().float().cpu()
+
+    arr = np.asarray(result)
+    if arr.ndim == 2:
+        alpha = arr
+    elif arr.ndim == 3:
+        # Prefer first channel for CHW/HWC
+        if arr.shape[0] in (1, 3, 4):       # CHW
+            alpha = arr[0]
+        elif arr.shape[-1] in (1, 3, 4):    # HWC
+            alpha = arr[..., 0]
+        else:
+            # Unknown 3D shape – take first slice robustly
+            alpha = arr[0]
+    elif arr.ndim == 4:
+        # Batch first: BxCxHxW or BxHxWxC
+        if arr.shape[1] in (1, 3, 4):       # BCHW
+            alpha = arr[0, 0]
+        elif arr.shape[-1] in (1, 3, 4):    # BHWC
+            alpha = arr[0, ..., 0]
+        else:
+            alpha = arr[0, 0]
+    else:
+        # Fallback
+        alpha = np.full((512, 512), 0.5, dtype=np.float32)
+
+    alpha = alpha.astype(np.float32, copy=False)
+    np.clip(alpha, 0.0, 1.0, out=alpha)
+    return alpha
+
+
+def _hw_from_image_like(x: Union[np.ndarray, torch.Tensor]) -> Tuple[int, int]:
+    """Best-effort get (H, W) from an image/mask input for neutral fallbacks."""
+    if torch.is_tensor(x):
+        shape = tuple(x.shape)
+        # Handle CHW / HWC / BCHW / BHWC / HW
+        if len(shape) == 2:                  # HW
+            return shape[0], shape[1]
+        if len(shape) == 3:
+            if shape[0] in (1, 3, 4):       # CHW
+                return shape[1], shape[2]
+            if shape[-1] in (1, 3, 4):      # HWC
+                return shape[0], shape[1]
+        if len(shape) == 4:
+            # Assume batch first
+            b, c_or_h, h_or_w, maybe_w = shape
+            # Try BCHW
+            if shape[1] in (1, 3, 4):
+                return shape[2], shape[3]
+            # Try BHWC
+            return shape[1], shape[2]
+        return 512, 512
+    else:
+        arr = np.asarray(x)
+        if arr.ndim == 2:                  # HW
+            return arr.shape[0], arr.shape[1]
+        if arr.ndim == 3:
+            if arr.shape[0] in (1, 3, 4):  # CHW
+                return arr.shape[1], arr.shape[2]
+            if arr.shape[-1] in (1, 3, 4): # HWC
+                return arr.shape[0], arr.shape[1]
+        if arr.ndim == 4:
+            # Assume batch first
+            if arr.shape[1] in (1, 3, 4):  # BCHW
+                return arr.shape[2], arr.shape[3]
+            return arr.shape[1], arr.shape[2]
+        return 512, 512
+
+
+# --------------------------- Callable Wrapper ---------------------------
+
 class MatAnyoneCallableWrapper:
     """
-    Callable wrapper around InferenceCore to maintain API compatibility.
-    Makes the processor work like a callable session.
+    Callable session-like wrapper around an InferenceCore instance.
+
+    Contract:
+      - First call SHOULD include a mask (1HW). If not, returns neutral 0.5 alpha.
+      - Subsequent calls do not require mask.
+      - Returns 2D alpha (H,W) float32 in [0,1].
     """
-    
-    def __init__(self, inference_core):
+
+    def __init__(self, inference_core, device: str = "cuda", mixed_precision: Optional[str] = "fp16"):
         self.core = inference_core
         self.initialized = False
-        
-    def __call__(self, image, mask=None, **kwargs):
-        """
-        Make this wrapper callable like the old session interface.
-        
-        Args:
-            image: Input image as numpy array
-            mask: Optional mask for first frame
-            
-        Returns:
-            Alpha mask as 2D numpy array
-        """
+        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 context manager
+        class _NullCtx:
+            def __enter__(self): return None
+            def __exit__(self, *exc): return False
+        return _NullCtx()
+
+    def __call__(self, image, mask=None, **kwargs) -> np.ndarray:
         try:
-            # For MatAnyone, the first frame needs initialization with a mask
+            # Preprocess → CHW/1HW tensors, then add batch
+            img_chw = _ensure_chw_float01(image).to(self.device, non_blocking=True)
+            img_bchw = img_chw.unsqueeze(0)  # B=1
+
             if not self.initialized:
                 if mask is None:
-                    # Return a default mask if no mask provided for first frame
-                    logger.warning("First frame called without mask, returning default")
-                    if isinstance(image, np.ndarray):
-                        h, w = image.shape[:2]
-                    else:
-                        h, w = 512, 512
-                    return np.ones((h, w), dtype=np.float32) * 0.5
-                
-                # Initialize with first frame and mask
-                # The exact API call depends on the InferenceCore implementation
-                # This is a placeholder - adjust based on actual API
-                if hasattr(self.core, 'step'):
-                    result = self.core.step(image=image, mask=mask)
-                elif hasattr(self.core, 'process_frame'):
-                    result = self.core.process_frame(image, mask)
-                else:
-                    # Fallback
-                    logger.warning("InferenceCore API unclear, returning input mask")
-                    return mask if isinstance(mask, np.ndarray) else np.array(mask)
-                
+                    h, w = _hw_from_image_like(image)
+                    logger.warning("MatAnyone first frame called without mask; returning neutral alpha.")
+                    return np.full((h, w), 0.5, dtype=np.float32)
+
+                m_1hw = _ensure_1hw_float01(mask).to(self.device, non_blocking=True)
+                m_b1hw = m_1hw.unsqueeze(0)  # B=1
+
+                with torch.inference_mode():
+                    with self._maybe_autocast():
+                        if hasattr(self.core, "step"):
+                            result = self.core.step(image=img_bchw, mask=m_b1hw, **kwargs)
+                        elif hasattr(self.core, "process_frame"):
+                            result = self.core.process_frame(img_bchw, m_b1hw, **kwargs)
+                        else:
+                            logger.warning("InferenceCore has no recognized frame API; echoing input mask.")
+                            return _alpha_from_result(mask)
+
                 self.initialized = True
-                return self._extract_alpha(result)
-            else:
-                # Subsequent frames - no mask needed
-                if hasattr(self.core, 'step'):
-                    result = self.core.step(image=image)
-                elif hasattr(self.core, 'process_frame'):
-                    result = self.core.process_frame(image)
-                else:
-                    # Fallback - return neutral mask
-                    if isinstance(image, np.ndarray):
-                        h, w = image.shape[:2]
+                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_bchw, **kwargs)
+                    elif hasattr(self.core, "process_frame"):
+                        result = self.core.process_frame(img_bchw, **kwargs)
                     else:
-                        h, w = 512, 512
-                    return np.ones((h, w), dtype=np.float32) * 0.5
-                
-                return self._extract_alpha(result)
-                
+                        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}")
-            # Return a fallback mask
+            logger.debug(traceback.format_exc())
+            # Fallbacks
             if mask is not None:
-                return mask if isinstance(mask, np.ndarray) else np.array(mask)
-            if isinstance(image, np.ndarray):
-                h, w = image.shape[:2]
-            else:
-                h, w = 512, 512
-            return np.ones((h, w), dtype=np.float32) * 0.5
-    
-    def _extract_alpha(self, result):
-        """Extract alpha channel from result."""
-        if result is None:
-            return np.ones((512, 512), dtype=np.float32) * 0.5
-        
-        if isinstance(result, np.ndarray):
-            if result.ndim == 2:
-                return result.astype(np.float32)
-            elif result.ndim == 3:
-                # Take first channel or average
-                return result[..., 0].astype(np.float32)
-            elif result.ndim == 4:
-                # Batch dimension - take first
-                return result[0, 0].astype(np.float32)
-        
-        # Try to convert to numpy
-        try:
-            arr = np.array(result)
-            if arr.ndim >= 2:
-                return arr[..., 0] if arr.ndim > 2 else arr
-        except:
-            pass
-        
-        return np.ones((512, 512), dtype=np.float32) * 0.5
-    
+                try:
+                    return _alpha_from_result(mask)
+                except Exception:
+                    pass
+            h, w = _hw_from_image_like(image)
+            return np.full((h, w), 0.5, dtype=np.float32)
+
     def reset(self):
-        """Reset the session state."""
+        """Reset state between videos."""
         self.initialized = False
-        if hasattr(self.core, 'reset'):
-            self.core.reset()
-        elif hasattr(self.core, 'clear_memory'):
-            self.core.clear_memory()
+        if hasattr(self.core, "reset"):
+            try:
+                self.core.reset()
+            except Exception as e:
+                logger.debug(f"Core reset() failed: {e}")
+        elif hasattr(self.core, "clear_memory"):
+            try:
+                self.core.clear_memory()
+            except Exception as e:
+                logger.debug(f"Core clear_memory() failed: {e}")
+
 
+# ------------------------------- Loader -------------------------------
 
 class MatAnyoneLoader:
     """
-    Official MatAnyone loader using InferenceCore API with callable wrapper.
+    Loads MatAnyone's InferenceCore and returns a callable wrapper.
+
+    Usage:
+        loader = MatAnyoneLoader(device="cuda")
+        session = loader.load()  # callable
+        alpha = session(frame, first_frame_mask)  # 2-D float32 [0,1]
     """
-    
-    def __init__(self, device: str = "cuda", cache_dir: str = "./checkpoints/matanyone_cache"):
+
+    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"
@@ -146,103 +313,139 @@ def __init__(self, device: str = "cuda", cache_dir: str = "./checkpoints/matanyo
         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:
-        """Select best available device."""
         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 load(self):
-        """Load MatAnyone using official InferenceCore API and wrap it."""
-        if self.loaded and self.wrapper:
+
+    def _try_build_core(self):
+        """
+        Try multiple constructor patterns to survive API changes.
+        """
+        from matanyone.inference.inference_core import InferenceCore
+
+        # 1) Preferred: from_pretrained(...)
+        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
+        try:
+            core = InferenceCore(self.model_id)
+            logger.info("Loaded MatAnyone via InferenceCore(model_id) [minimal]")
+            return core
+        except Exception as e:
+            logger.debug(f"ctor(model_id) failed: {e}")
+            raise  # Propagate last error
+
+    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 from HF: {self.model_id} (device={self.device})")
+
+        logger.info(f"Loading MatAnyone: {self.model_id} (device={self.device})")
         t0 = time.time()
-        
+
         try:
-            # Import the official API
-            from matanyone.inference.inference_core import InferenceCore
-            
-            # Create the InferenceCore processor
-            self.processor = InferenceCore(self.model_id)
-            
-            # Wrap it to make it callable
-            self.wrapper = MatAnyoneCallableWrapper(self.processor)
-            
+            self.processor = self._try_build_core()
+            # If the core has an explicit to(device) or set_device, try to use it
+            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
+            )
             self.loaded = True
             self.load_time = time.time() - t0
-            logger.info(f"MatAnyone loaded and wrapped successfully in {self.load_time:.2f}s")
+            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(f"Failed to import MatAnyone. Install with: pip install git+https://github.com/pq-yang/MatAnyone.git@main")
+            logger.error("Failed to import MatAnyone. Install with: "
+                         "pip install git+https://github.com/pq-yang/MatAnyone.git@main")
             return None
-            
         except Exception as e:
             self.load_error = str(e)
             logger.error(f"Failed to load MatAnyone: {e}")
             logger.debug(traceback.format_exc())
             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_info(self) -> Dict[str, Any]:
-        """Get model information."""
+        """Get model information and interface flags."""
         info = {
             "loaded": self.loaded,
             "model_id": self.model_id,
             "device": str(self.device),
-            "load_time": self.load_time,
+            "load_time": float(self.load_time),
             "error": self.load_error,
-            "api": "InferenceCore (wrapped)"
+            "api": "InferenceCore (wrapped)",
+            "mixed_precision": self.mixed_precision,
         }
-        
-        # Add interface info
-        if self.processor:
-            info["has_step"] = hasattr(self.processor, 'step')
-            info["has_process_frame"] = hasattr(self.processor, 'process_frame')
-            info["has_process_video"] = hasattr(self.processor, 'process_video')
-        
+        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")
-    
-    # Compatibility - make the loader itself callable
-    def __call__(self, image, mask=None, **kwargs):
-        """Direct call compatibility."""
-        if not self.wrapper:
-            if not self.load():
+
+    # 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:
-                    return mask if isinstance(mask, np.ndarray) else np.array(mask)
-                return np.zeros(image.shape[:2], dtype=np.float32)
-        
+                    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)
 
 
-# For backwards compatibility
+# Backwards compatibility alias (legacy session naming)
 _MatAnyoneSession = MatAnyoneCallableWrapper
 
-__all__ = ["MatAnyoneLoader", "_MatAnyoneSession", "MatAnyoneCallableWrapper"]
+__all__ = ["MatAnyoneLoader", "_MatAnyoneSession", "MatAnyoneCallableWrapper"]