Update models/loaders/model_loader.py
Browse files
models/loaders/model_loader.py
CHANGED
|
@@ -108,7 +108,8 @@ def load_all_models(
|
|
| 108 |
progress_callback(0.1, "Loading SAM2 model...")
|
| 109 |
|
| 110 |
sam2_start = time.time()
|
| 111 |
-
|
|
|
|
| 112 |
sam2_time = time.time() - sam2_start
|
| 113 |
|
| 114 |
if sam2_model:
|
|
@@ -306,7 +307,14 @@ def _cleanup_models(self):
|
|
| 306 |
"""Internal cleanup of loaded models"""
|
| 307 |
# Clean up SAM2
|
| 308 |
if self.sam2_loader:
|
| 309 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 310 |
if self.sam2_predictor:
|
| 311 |
try:
|
| 312 |
del self.sam2_predictor
|
|
@@ -316,7 +324,20 @@ def _cleanup_models(self):
|
|
| 316 |
|
| 317 |
# Clean up MatAnyone
|
| 318 |
if self.matanyone_loader:
|
| 319 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 320 |
if self.matanyone_model:
|
| 321 |
try:
|
| 322 |
del self.matanyone_model
|
|
@@ -331,4 +352,4 @@ def _cleanup_models(self):
|
|
| 331 |
# Garbage collection
|
| 332 |
gc.collect()
|
| 333 |
|
| 334 |
-
logger.debug("Model cleanup completed")
|
|
|
|
| 108 |
progress_callback(0.1, "Loading SAM2 model...")
|
| 109 |
|
| 110 |
sam2_start = time.time()
|
| 111 |
+
# CHANGED: Force tiny model instead of auto-detection
|
| 112 |
+
sam2_model = self.sam2_loader.load("tiny") # Force tiny model for faster loading and less memory usage
|
| 113 |
sam2_time = time.time() - sam2_start
|
| 114 |
|
| 115 |
if sam2_model:
|
|
|
|
| 307 |
"""Internal cleanup of loaded models"""
|
| 308 |
# Clean up SAM2
|
| 309 |
if self.sam2_loader:
|
| 310 |
+
try:
|
| 311 |
+
if hasattr(self.sam2_loader, 'cleanup'):
|
| 312 |
+
self.sam2_loader.cleanup()
|
| 313 |
+
else:
|
| 314 |
+
logger.debug("SAM2 loader has no cleanup method")
|
| 315 |
+
except Exception as e:
|
| 316 |
+
logger.debug(f"SAM2 loader cleanup error: {e}")
|
| 317 |
+
|
| 318 |
if self.sam2_predictor:
|
| 319 |
try:
|
| 320 |
del self.sam2_predictor
|
|
|
|
| 324 |
|
| 325 |
# Clean up MatAnyone
|
| 326 |
if self.matanyone_loader:
|
| 327 |
+
try:
|
| 328 |
+
if hasattr(self.matanyone_loader, 'cleanup'):
|
| 329 |
+
self.matanyone_loader.cleanup()
|
| 330 |
+
else:
|
| 331 |
+
# MatAnyone doesn't have cleanup, but we can clean the wrapper
|
| 332 |
+
if hasattr(self.matanyone_loader, '_wrapper') and self.matanyone_loader._wrapper:
|
| 333 |
+
if hasattr(self.matanyone_loader._wrapper, 'reset'):
|
| 334 |
+
self.matanyone_loader._wrapper.reset()
|
| 335 |
+
self.matanyone_loader._processor = None
|
| 336 |
+
self.matanyone_loader._wrapper = None
|
| 337 |
+
logger.debug("MatAnyone loader cleaned up manually")
|
| 338 |
+
except Exception as e:
|
| 339 |
+
logger.debug(f"MatAnyone loader cleanup error: {e}")
|
| 340 |
+
|
| 341 |
if self.matanyone_model:
|
| 342 |
try:
|
| 343 |
del self.matanyone_model
|
|
|
|
| 352 |
# Garbage collection
|
| 353 |
gc.collect()
|
| 354 |
|
| 355 |
+
logger.debug("Model cleanup completed")
|