Update models/loaders/matanyone_loader.py

models/loaders/matanyone_loader.py

@@ -319,7 +319,16 @@ def reset(self):
                 logger.debug(f"Core clear_memory() failed: {e}")
 
 
+# --------------------------- Main Loader Class ---------------------------
+
 class MatAnyoneLoader:
+    """
+    Loader for MatAnyone InferenceCore with cleanup support.
+    
+    Provides a consistent interface with other model loaders,
+    including proper resource cleanup.
+    """
+    
     def __init__(self, device: str = "auto", model_id: str = "PeiqingYang/MatAnyone"):
         self.device = device
         self.model_id = model_id
@@ -366,4 +375,49 @@ def get_info(self) -> Dict[str, Any]:
             "model_id": self.model_id,
             "loaded": self._wrapper is not None or self._processor is not None,
             "wrapped": self._wrapper is not None,
-        }
+        }
+    
+    def cleanup(self):
+        """
+        Clean up all resources associated with MatAnyone.
+        
+        This method ensures proper cleanup of:
+        - The wrapper's state and memory
+        - The InferenceCore processor
+        - Any CUDA tensors in memory
+        """
+        logger.debug("Starting MatAnyone cleanup...")
+        
+        # Clean up wrapper first
+        if self._wrapper:
+            try:
+                self._wrapper.reset()
+                logger.debug("MatAnyone wrapper reset completed")
+            except Exception as e:
+                logger.debug(f"Wrapper reset failed (non-critical): {e}")
+            self._wrapper = None
+        
+        # Clean up processor
+        if self._processor:
+            try:
+                # Try various cleanup methods that might exist
+                if hasattr(self._processor, 'cleanup'):
+                    self._processor.cleanup()
+                elif hasattr(self._processor, 'clear'):
+                    self._processor.clear()
+                elif hasattr(self._processor, 'reset'):
+                    self._processor.reset()
+                logger.debug("MatAnyone processor cleanup attempted")
+            except Exception as e:
+                logger.debug(f"Processor cleanup failed (non-critical): {e}")
+            self._processor = None
+        
+        # Clear any CUDA cache if using GPU
+        if self.device != "cpu" and torch.cuda.is_available():
+            try:
+                torch.cuda.empty_cache()
+                logger.debug("CUDA cache cleared for MatAnyone")
+            except Exception as e:
+                logger.debug(f"CUDA cache clear failed: {e}")
+        
+        logger.info("MatAnyone resources cleaned up")