Update modules/diarize/diarizer.py

modules/diarize/diarizer.py

@@ -4,6 +4,7 @@ from typing import List, Union, BinaryIO, Optional
 import numpy as np
 import time
 import logging
+import gc
 
 from modules.utils.paths import DIARIZATION_MODELS_DIR
 from modules.diarize.diarize_pipeline import DiarizationPipeline, assign_word_speakers
@@ -115,6 +116,16 @@ class Diarizer:
         )
         logger.disabled = False
 
+    def offload(self):
+        """Offload the model and free up the memory"""
+        if self.pipe is not None:
+            del self.pipe
+            self.pipe = None
+        if self.device == "cuda":
+            torch.cuda.empty_cache()
+            torch.cuda.reset_max_memory_allocated()
+        gc.collect()
+    
     @staticmethod
     def get_device():
         if torch.cuda.is_available():