Update chunkedTranscriber.py

chunkedTranscriber.py

@@ -3,18 +3,16 @@ import gc
 import sys
 import time
 import torch
+import spaces
 import torchaudio
 import numpy as np
 from scipy.signal import resample
 from pyannote.audio import Pipeline
 from dotenv import load_dotenv
 load_dotenv()
-import logging
-import time
 from difflib import SequenceMatcher
 from transformers import Wav2Vec2ForSequenceClassification, AutoFeatureExtractor, Wav2Vec2ForCTC, AutoProcessor, AutoTokenizer, AutoModelForSeq2SeqLM
 from difflib import SequenceMatcher
-import gc
 
 class ChunkedTranscriber:
     def __init__(self, chunk_size=5, overlap=1, sample_rate=16000):
@@ -32,6 +30,7 @@ class ChunkedTranscriber:
         pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
         return pipeline
 
+    @spaces.GPU(duration=60)
     def diarize_audio(self, audio_path):
         """
         Perform speaker diarization on the input audio.
@@ -45,7 +44,8 @@ class ChunkedTranscriber:
         model = Wav2Vec2ForSequenceClassification.from_pretrained(model_id)
         return processor, model
 
-
+    
+    @spaces.GPU(duration=60)
     def language_identification(self, model, processor, chunk, device="cuda"):
         inputs = processor(chunk, sampling_rate=16_000, return_tensors="pt")
         model.to(device)
@@ -69,6 +69,7 @@ class ChunkedTranscriber:
         return model, processor
 
 
+    @spaces.GPU(duration=60)
     def mms_transcription(self, model, processor, chunk, device="cuda"):
 
         inputs = processor(chunk, sampling_rate=16_000, return_tensors="pt")
@@ -92,7 +93,8 @@ class ChunkedTranscriber:
         model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
         return model, tokenizer
 
-
+    
+    @spaces.GPU(duration=60)
     def text2text_translation(self, translation_model, translation_tokenizer, transcript, device="cuda"):
         # model, tokenizer = load_translation_model()
 
@@ -108,6 +110,7 @@ class ChunkedTranscriber:
         gc.collect()
         return translation_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
 
+        
     def preprocess_audio(self, audio):
         """
         Create overlapping chunks with improved timing logic