clean up mem

comic_panel_extractor/annorator_server.py

@@ -63,13 +63,13 @@ def load_yolo_boxes(image_path: str, label_path: str, detect: bool = False):
         if detect and not os.path.exists(label_path):
             from .yolo_manager import YOLOManager
             from .utils import Config
-            yolo_manager = YOLOManager()
-            weights_path = f'{current_path}/{Config.YOLO_MODEL_NAME}.pt'
+            with YOLOManager() as yolo_manager:
+                weights_path = f'{current_path}/{Config.YOLO_MODEL_NAME}.pt'
 
-            yolo_manager.load_model(weights_path)
+                yolo_manager.load_model(weights_path)
 
-            # Run inference
-            _, label_path = yolo_manager.annotate_images(image_paths=[image_path], output_dir=IMAGE_LABEL_ROOT, save_image=False, label_path=label_path)
+                # Run inference
+                _, label_path = yolo_manager.annotate_images(image_paths=[image_path], output_dir=IMAGE_LABEL_ROOT, save_image=False, label_path=label_path)
 
         if os.path.exists(label_path):
             with open(label_path, "r") as f:

comic_panel_extractor/yolo_manager.py

@@ -228,4 +228,34 @@ class YOLOManager:
                 
             except Exception as e:
                 print(f"❌ Error processing {image_path}: {str(e)}")
-                return None, None
+                return None, None
+
+    def __enter__(self):
+        # When entering context, just return self
+        return self
+
+    def __del__(self):
+        # On exit, unload model and clear cache
+        self.unload_model()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # On exit, unload model and clear cache
+        self.unload_model()
+
+    def unload_model(self):
+        if self.model is not None:
+            print("🧹 Unloading YOLO model and clearing CUDA cache...")
+            try:
+                import torch
+                import gc
+                del self.model
+                self.model = None
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                    torch.cuda.ipc_collect()
+                print("✅ Model unloaded and GPU cache cleared.")
+            except Exception as e:
+                print(f"❌ Error unloading model: {e}")
+        else:
+            print("⚠️ No model loaded to unload.")