Update diarization.py

diarization.py

@@ -1,6 +1,5 @@
 import os
 import torch
-import torchvision
 import math
 from moviepy.editor import VideoFileClip, AudioFileClip
 from pyannote.audio import Pipeline
@@ -11,6 +10,51 @@ from collections import defaultdict
 import numpy as np
 import spaces
 
+class LazyDiarizationPipeline:
+    def __init__(self):
+        self.pipeline = None
+
+    @spaces.GPU(duration=120)
+    def get_pipeline(self, diarization_access_token):
+        if self.pipeline is None:
+            self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=diarization_access_token)
+            self.pipeline = self.pipeline.to("cuda")
+        return self.pipeline
+
+lazy_diarization_pipeline = LazyDiarizationPipeline()
+
+class LazyTranscriptionPipeline:
+    def __init__(self):
+        self.model = None
+        self.processor = None
+        self.pipe = None
+
+    @spaces.GPU(duration=120)
+    def get_pipeline(self, language):
+        if self.pipe is None:
+            model_id = "openai/whisper-large-v3"
+            self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                model_id, torch_dtype=torch.float16, low_cpu_mem_usage=True, use_safetensors=True
+            )
+            self.model.to("cuda")
+            self.processor = AutoProcessor.from_pretrained(model_id)
+            self.pipe = pipeline(
+                "automatic-speech-recognition",
+                model=self.model,
+                tokenizer=self.processor.tokenizer,
+                feature_extractor=self.processor.feature_extractor,
+                max_new_tokens=128,
+                chunk_length_s=30,
+                batch_size=1,
+                return_timestamps=True,
+                torch_dtype=torch.float16,
+                device="cuda",
+                generate_kwargs={"language": language}
+            )
+        return self.pipe
+
+lazy_transcription_pipeline = LazyTranscriptionPipeline()
+
 def extract_audio(video_path, audio_path):
     video = VideoFileClip(video_path)
     audio = video.audio
@@ -21,30 +65,7 @@ def format_timestamp(seconds):
 
 @spaces.GPU(duration=300)
 def transcribe_audio(audio_path, language):
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-    model_id = "openai/whisper-large-v3"
-
-    model = AutoModelForSpeechSeq2Seq.from_pretrained(
-        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
-    )
-    model.to(device)
-
-    processor = AutoProcessor.from_pretrained(model_id)
-
-    pipe = pipeline(
-        "automatic-speech-recognition",
-        model=model,
-        tokenizer=processor.tokenizer,
-        feature_extractor=processor.feature_extractor,
-        max_new_tokens=128,
-        chunk_length_s=30,
-        batch_size=1,
-        return_timestamps=True,
-        torch_dtype=torch_dtype,
-        device=device,
-        generate_kwargs={"language": language}
-    )
+    pipe = lazy_transcription_pipeline.get_pipeline(language)
 
     audio, sr = librosa.load(audio_path, sr=16000)
     duration = len(audio) / sr
@@ -118,15 +139,13 @@ 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)
 
     # Diarization
     print("Performing diarization...")
-    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=diarization_access_token)
-    pipeline = pipeline.to("cuda")
+    pipeline = lazy_diarization_pipeline.get_pipeline(diarization_access_token)
     diarization = pipeline(audio_path)
     print("Diarization complete.")
 
@@ -144,13 +163,3 @@ def process_video(video_path, diarization_access_token, language):
     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
-    access_token = os.environ.get('hf_secret')
-    if not access_token:
-        raise ValueError("HF_TOKEN not found in environment variables. Please set it in the Space secrets.")
-
-    language = "en"
-    process_video(video_path, access_token, language)