Update transcription_diarization.py

transcription_diarization.py

@@ -1,5 +1,6 @@
 import os
 import torch
+import gc
 import math
 from moviepy.editor import VideoFileClip
 from pyannote.audio import Pipeline
@@ -141,14 +142,26 @@ def process_video(video_path, diarization_access_token, language):
     diarization = pipeline(audio_path)
     print("Diarization complete.")
 
+    # Clear GPU memory after diarization
+    torch.cuda.empty_cache()
+    gc.collect()
+
     print("Performing transcription...")
     transcription, chunks = transcribe_audio(audio_path, language)
     print("Transcription complete.")
 
+    # Clear GPU memory after transcription
+    torch.cuda.empty_cache()
+    gc.collect()
+
     combined_srt_path = f"{base_name}_combined.srt"
     create_combined_srt(chunks, diarization, combined_srt_path)
     print(f"Combined SRT file created and saved to {combined_srt_path}")
 
     os.remove(audio_path)
 
+    # Final GPU memory clear
+    torch.cuda.empty_cache()
+    gc.collect()
+
     return combined_srt_path