Update chunkedTranscriber.py

chunkedTranscriber.py

@@ -6,15 +6,18 @@ import torch
 import spaces
 import torchaudio
 import numpy as np
-from scipy.signal import resample
-from pyannote.audio import Pipeline
+
+
+from df.enhance import enhance, init_df
 from dotenv import load_dotenv
 load_dotenv()
+from scipy.signal import resample
+from pyannote.audio import Pipeline
 from difflib import SequenceMatcher
 from transformers import Wav2Vec2ForSequenceClassification, AutoFeatureExtractor, Wav2Vec2ForCTC, AutoProcessor, AutoTokenizer, AutoModelForSeq2SeqLM
-from difflib import SequenceMatcher
-import logging
 
+
+import logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
@@ -38,7 +41,7 @@ class ChunkedTranscriber:
         pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=os.getenv("HF_TOKEN"))
         return pipeline
 
-    @spaces.GPU(duration=60)
+    @spaces.GPU(duration=180)
     def diarize_audio(self, audio_path):
         """
         Perform speaker diarization on the input audio.
@@ -53,7 +56,7 @@ class ChunkedTranscriber:
         return processor, model
 
     
-    @spaces.GPU(duration=60)
+    @spaces.GPU(duration=180)
     def language_identification(self, model, processor, chunk, device="cuda"):
         inputs = processor(chunk, sampling_rate=16_000, return_tensors="pt")
         model.to(device)
@@ -77,7 +80,7 @@ class ChunkedTranscriber:
         return model, processor
 
 
-    @spaces.GPU(duration=60)
+    @spaces.GPU(duration=180)
     def mms_transcription(self, model, processor, chunk, device="cuda"):
 
         inputs = processor(chunk, sampling_rate=16_000, return_tensors="pt")
@@ -102,7 +105,7 @@ class ChunkedTranscriber:
         return model, tokenizer
 
     
-    @spaces.GPU(duration=60)
+    @spaces.GPU(duration=180)
     def text2text_translation(self, translation_model, translation_tokenizer, transcript, device="cuda"):
         # model, tokenizer = load_translation_model()
 
@@ -336,6 +339,10 @@ class ChunkedTranscriber:
         return translation
 
 
+    def audio_denoising(): 
+        model, df_state = init_df()
+        enhanced_audio = enhance(model, df_state, noisy_audio)
+        return enhanced_audio
 
     def transcribe_audio(self, audio_path, translate=False):
         """
@@ -353,7 +360,7 @@ class ChunkedTranscriber:
                 'end_time': turn.end,
                 'speaker': speaker
             })
-
+        
         audio = self.load_audio(audio_path)
         chunks = self.preprocess_audio(audio)