Update diarization.py

diarization.py

@@ -118,6 +118,7 @@ def create_combined_srt(transcription_chunks, diarization, output_path):
 
 @spaces.GPU(duration=600)
 def process_video(video_path, diarization_access_token, language):
+    import torch
     base_name = os.path.splitext(video_path)[0]
     audio_path = f"{base_name}.wav"
     extract_audio(video_path, audio_path)
@@ -125,7 +126,7 @@ def process_video(video_path, diarization_access_token, language):
     # Diarization
     print("Performing diarization...")
     pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=diarization_access_token)
-    pipeline = pipeline.to(torch.device("cpu"))
+    pipeline = pipeline.to("cuda")
     diarization = pipeline(audio_path)
     print("Diarization complete.")
 
@@ -142,6 +143,8 @@ def process_video(video_path, diarization_access_token, language):
     # Clean up
     os.remove(audio_path)
 
+    return combined_srt_path
+
 if __name__ == "__main__":
     video_path = r"C:\Users\reab5\Downloads\MediaHuman\Music\test1.mp4"
     # Get Hugging Face token from Space secret