Update models/matanyone_loader.py

models/matanyone_loader.py

@@ -2,18 +2,10 @@
 # -*- coding: utf-8 -*-
 """
 MatAnyone adapter — Using Official API (File-Based)
+(Enhanced logging, explicit error handling, and stage progress)
 
-Fixed to use MatAnyone's official process_video() API instead of
-bypassing it with internal tensor manipulation. This eliminates
-all 5D tensor dimension issues.
-
-Changes (2025-09-17):
-- Replaced custom tensor processing with official MatAnyone API
-- Uses file-based input/output as designed by MatAnyone authors
-- Eliminates all tensor dimension compatibility issues
-- Simplified error handling and logging
+...
 """
-
 from __future__ import annotations
 import os
 import time
@@ -23,6 +15,7 @@
 from pathlib import Path
 from typing import Optional, Callable, Tuple
 
+logging.basicConfig(level=logging.INFO)
 log = logging.getLogger(__name__)
 
 # ---------- Progress helper ----------
@@ -44,20 +37,18 @@ def _emit_progress(cb, pct: float, msg: str):
         return
     try:
         try:
-            cb(pct, msg)  # preferred (pct, msg)
+            cb(pct, msg)
         except TypeError:
-            cb(msg)       # legacy (msg)
+            cb(msg)
         _progress_last = now
         _progress_last_msg = msg
     except Exception as e:
         _progress_disabled = True
         log.warning("[progress-cb] disabled due to exception: %s", e)
 
-# ---------- Errors ----------
 class MatAnyError(RuntimeError):
     pass
 
-# ---------- CUDA helpers ----------
 def _cuda_snapshot(device: Optional[str]) -> str:
     try:
         import torch
@@ -86,30 +77,26 @@ def _safe_empty_cache():
     except Exception:
         pass
 
-# ============================================================================
-
 class MatAnyoneSession:
     """
     Simple wrapper around MatAnyone's official API.
     Uses file-based input/output as designed by the MatAnyone authors.
     """
     def __init__(self, device: Optional[str] = None, precision: str = "auto"):
+        log.info(f"[MatAnyoneSession.__init__] device={device}, precision={precision}")  # [LOG+SAFETY PATCH]
         self.device = device or ("cuda" if self._cuda_available() else "cpu")
         self.precision = precision.lower()
-        
-        # Log MatAnyone version
         try:
             version = importlib.metadata.version("matanyone")
             log.info(f"[MATANY] MatAnyone version: {version}")
         except Exception:
             log.info("[MATANY] MatAnyone version unknown")
-        
-        # Initialize MatAnyone's official API
         try:
             from matanyone import InferenceCore
             self.processor = InferenceCore("PeiqingYang/MatAnyone")
             log.info("[MATANY] MatAnyone InferenceCore initialized successfully")
         except Exception as e:
+            log.error(f"[MatAnyoneSession.__init__] Failed to initialize MatAnyone: {e}", exc_info=True)  # [LOG+SAFETY PATCH]
             raise MatAnyError(f"Failed to initialize MatAnyone: {e}")
     
     def _cuda_available(self) -> bool:
@@ -126,156 +113,93 @@ def process_stream(
         out_dir: Optional[Path] = None,
         progress_cb: Optional[Callable] = None,
     ) -> Tuple[Path, Path]:
-        """
-        Process video using MatAnyone's official API.
-        
-        Args:
-            video_path: Path to input video file
-            seed_mask_path: Path to first-frame mask PNG (white=foreground, black=background)
-            out_dir: Output directory for results
-            progress_cb: Progress callback function
-            
-        Returns:
-            Tuple of (alpha_path, foreground_path)
-        """
+        log.info(f"[MatAnyoneSession.process_stream] Start: video={video_path}, mask={seed_mask_path}, out_dir={out_dir}")  # [LOG+SAFETY PATCH]
         video_path = Path(video_path)
         if not video_path.exists():
+            log.error(f"[MatAnyoneSession.process_stream] Video file not found: {video_path}")  # [LOG+SAFETY PATCH]
             raise MatAnyError(f"Video file not found: {video_path}")
-        
         if seed_mask_path and not Path(seed_mask_path).exists():
+            log.error(f"[MatAnyoneSession.process_stream] Seed mask not found: {seed_mask_path}")  # [LOG+SAFETY PATCH]
             raise MatAnyError(f"Seed mask not found: {seed_mask_path}")
-        
         out_dir = Path(out_dir) if out_dir else video_path.parent / "matanyone_output"
         out_dir.mkdir(parents=True, exist_ok=True)
-        
         log.info(f"[MATANY] Processing video: {video_path}")
         log.info(f"[MATANY] Using mask: {seed_mask_path}")
         log.info(f"[MATANY] Output directory: {out_dir}")
-        
         _emit_progress(progress_cb, 0.0, "Initializing MatAnyone processing...")
-        
         try:
-            # Use MatAnyone's official API
             start_time = time.time()
-            
             _emit_progress(progress_cb, 0.1, "Running MatAnyone video matting...")
-            
-            # Call the official process_video method
             foreground_path, alpha_path = self.processor.process_video(
                 input_path=str(video_path),
                 mask_path=str(seed_mask_path) if seed_mask_path else None,
                 output_path=str(out_dir)
             )
-            
             processing_time = time.time() - start_time
             log.info(f"[MATANY] Processing completed in {processing_time:.1f}s")
             log.info(f"[MATANY] Foreground output: {foreground_path}")
             log.info(f"[MATANY] Alpha output: {alpha_path}")
-            
-            # Convert to Path objects
             fg_path = Path(foreground_path) if foreground_path else None
             al_path = Path(alpha_path) if alpha_path else None
-            
-            # Verify outputs exist
             if not fg_path or not fg_path.exists():
+                log.error(f"[MatAnyoneSession.process_stream] Foreground output not created: {fg_path}")  # [LOG+SAFETY PATCH]
                 raise MatAnyError(f"Foreground output not created: {fg_path}")
             if not al_path or not al_path.exists():
+                log.error(f"[MatAnyoneSession.process_stream] Alpha output not created: {al_path}")  # [LOG+SAFETY PATCH]
                 raise MatAnyError(f"Alpha output not created: {al_path}")
-            
             _emit_progress(progress_cb, 1.0, "MatAnyone processing complete")
-            
-            return al_path, fg_path  # Return (alpha, foreground) to match expected order
-            
+            log.info(f"[MatAnyoneSession.process_stream] Success, returning paths.")  # [LOG+SAFETY PATCH]
+            return al_path, fg_path  # (alpha, foreground)
         except Exception as e:
-            log.error(f"[MATANY] Processing failed: {e}")
+            log.error(f"[MatAnyoneSession.process_stream] Processing failed: {e}", exc_info=True)  # [LOG+SAFETY PATCH]
             raise MatAnyError(f"MatAnyone processing failed: {e}")
-        
         finally:
             _safe_empty_cache()
 
-# ============================================================================
-# MatAnyoneModel Wrapper Class for app_hf.py compatibility
-# ============================================================================
-
 class MatAnyoneModel:
     """Wrapper class for MatAnyone to match app_hf.py interface"""
-    
     def __init__(self, device="cuda"):
+        log.info(f"[MatAnyoneModel.__init__] device={device}")  # [LOG+SAFETY PATCH]
         self.device = device
         self.session = None
         self.loaded = False
-        log.info(f"Initializing MatAnyoneModel on device: {device}")
-        
-        # Initialize the session
         self._load_model()
-    
     def _load_model(self):
-        """Load the MatAnyone session"""
         try:
             self.session = MatAnyoneSession(device=self.device, precision="auto")
             self.loaded = True
-            log.info("MatAnyoneModel loaded successfully")
+            log.info("[MatAnyoneModel._load_model] Loaded successfully")  # [LOG+SAFETY PATCH]
         except Exception as e:
-            log.error(f"Error loading MatAnyoneModel: {e}")
+            log.error(f"[MatAnyoneModel._load_model] Error loading: {e}", exc_info=True)  # [LOG+SAFETY PATCH]
             self.loaded = False
-    
     def replace_background(self, video_path, masks, background_path):
-        """Replace background in video using MatAnyone"""
+        log.info(f"[MatAnyoneModel.replace_background] Start")  # [LOG+SAFETY PATCH]
         if not self.loaded:
+            log.error("[MatAnyoneModel.replace_background] Model not loaded")  # [LOG+SAFETY PATCH]
             raise MatAnyError("MatAnyoneModel not loaded")
-        
         try:
-            from pathlib import Path
-            import tempfile
-            
-            # Convert paths to Path objects
             video_path = Path(video_path)
-            
-            # For now, we expect masks to be a path to the first-frame mask
             mask_path = Path(masks) if isinstance(masks, (str, Path)) else None
-            
-            # Create output directory
             with tempfile.TemporaryDirectory() as temp_dir:
                 output_dir = Path(temp_dir)
-                
-                # Process the video stream
                 alpha_path, fg_path = self.session.process_stream(
                     video_path=video_path,
                     seed_mask_path=mask_path,
                     out_dir=output_dir,
                     progress_cb=None
                 )
-                
-                # Return the foreground video path
-                # In a full implementation, you'd composite with the background_path
+                log.info(f"[MatAnyoneModel.replace_background] Success, returning fg_path: {fg_path}")  # [LOG+SAFETY PATCH]
                 return str(fg_path)
-                
         except Exception as e:
-            log.error(f"Error in replace_background: {e}")
+            log.error(f"[MatAnyoneModel.replace_background] Error: {e}", exc_info=True)  # [LOG+SAFETY PATCH]
             raise MatAnyError(f"Background replacement failed: {e}")
 
-# ============================================================================
-# Helper function for pipeline integration
-# ============================================================================
-
 def create_matanyone_session(device=None):
-    """Create a MatAnyone session for use in pipeline"""
+    log.info(f"[create_matanyone_session] device={device}")  # [LOG+SAFETY PATCH]
     return MatAnyoneSession(device=device)
 
 def run_matanyone_on_files(video_path, mask_path, output_dir, device="cuda", progress_callback=None):
-    """
-    Run MatAnyone on video and mask files.
-    
-    Args:
-        video_path: Path to input video
-        mask_path: Path to first-frame mask PNG
-        output_dir: Directory for outputs
-        device: Device to use (cuda/cpu)
-        progress_callback: Progress callback function
-        
-    Returns:
-        Tuple of (alpha_path, foreground_path) or (None, None) on failure
-    """
+    log.info(f"[run_matanyone_on_files] Start: video={video_path}, mask={mask_path}, out={output_dir}, device={device}")  # [LOG+SAFETY PATCH]
     try:
         session = MatAnyoneSession(device=device)
         alpha_path, fg_path = session.process_stream(
@@ -284,7 +208,8 @@ def run_matanyone_on_files(video_path, mask_path, output_dir, device="cuda", pro
             out_dir=Path(output_dir),
             progress_cb=progress_callback
         )
+        log.info(f"[run_matanyone_on_files] Success, returning (alpha, fg): {alpha_path}, {fg_path}")  # [LOG+SAFETY PATCH]
         return str(alpha_path), str(fg_path)
     except Exception as e:
-        log.error(f"MatAnyone processing failed: {e}")
-        return None, None
+        log.error(f"[run_matanyone_on_files] MatAnyone processing failed: {e}", exc_info=True)  # [LOG+SAFETY PATCH]
+        return None, None