dassa

app.py

@@ -5,76 +5,87 @@ import yt_dlp as youtube_dl
 import tempfile
 import os
 import locale
+import whisper
 import datetime
 import subprocess
-import wave
-import contextlib
-import numpy as np
-from sklearn.cluster import AgglomerativeClustering
-from faster_whisper import WhisperModel
+import pyannote.audio
 from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
 from pyannote.audio import Audio
 from pyannote.core import Segment
+import wave
+import contextlib
+from sklearn.cluster import AgglomerativeClustering
+import numpy as np
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-BATCH_SIZE = 8
+BATCH_SIZE = 16
 FILE_LIMIT_MB = 1000
-YT_LENGTH_LIMIT_S = 3600  # limit to 1 hour YouTube files
+COMPUTE_TYPE = "float32" 
+YT_LENGTH_LIMIT_S = 600  # limit to 1 hour YouTube files
 
 num_speakers = 2
 language = None
-model_size = 'tiny'
-model = WhisperModel(model_size, device=device, compute_type="float32")
-embedding_model = PretrainedSpeakerEmbedding("speechbrain/spkrec-ecapa-voxceleb", device=torch.device("cpu"))
-audio = Audio()
+model_size = 'large'
+model_name = model_size
 
-def getpreferredencoding(do_setlocale=True):
+def getpreferredencoding(do_setlocale = True):
     return "UTF-8"
+
 locale.getpreferredencoding = getpreferredencoding
+embedding_model = PretrainedSpeakerEmbedding(
+    "speechbrain/spkrec-ecapa-voxceleb",
+    device=torch.device("cpu"))
+model = whisper.load_model(model_size).to(device)
+audio = Audio()
 
-def segment_embedding(segment, duration, path):
-    start = segment.start
-    end = min(duration, segment.end)
+def segment_embedding(segment,duration,path):
+    start = segment["start"]
+    # Whisper overshoots the end timestamp in the last segment
+    end = min(duration, segment["end"])
     clip = Segment(start, end)
     waveform, sample_rate = audio.crop(path, clip)
+
+    # Convert waveform to single channel
     waveform = waveform.mean(dim=0, keepdim=True)
+
     return embedding_model(waveform.unsqueeze(0))
 
 def time(secs):
-    return datetime.timedelta(seconds=round(secs))
+  return datetime.timedelta(seconds=round(secs))
 
+@spaces.GPU
 def transcribe(path, task):
     if path is None:
         raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
     
-    if not path.endswith('.wav'):
+    if path[-3:] != 'wav':
         subprocess.call(['ffmpeg', '-i', path, "audio.wav", '-y'])
         path = "audio.wav"
-    
-    segments, _ = model.transcribe(path)
-    
-    with contextlib.closing(wave.open(path, 'r')) as f:
+    result = model.transcribe(path,fp16=False)
+    segments = result["segments"]
+    print(segments)
+    with contextlib.closing(wave.open(path,'r')) as f:
         frames = f.getnframes()
         rate = f.getframerate()
         duration = frames / float(rate)
     
     embeddings = np.zeros(shape=(len(segments), 192))
     for i, segment in enumerate(segments):
-        embeddings[i] = segment_embedding(segment, duration=duration, path=path)
+        embeddings[i] = segment_embedding(segment,duration=duration,path=path)
     embeddings = np.nan_to_num(embeddings)
     clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
     labels = clustering.labels_
-    
-    output_text = ""
-    for i, segment in enumerate(segments):
-        segment.speaker = 'SPEAKER ' + str(labels[i] + 1)
-    for i, segment in enumerate(segments):
-        if i == 0 or segments[i - 1].speaker != segment.speaker:
-            output_text += "\n" + segment.speaker + ' ' + str(time(segment.start)) + '\n'
-        output_text += segment.text + ' '
-    
+    output_text=""
+    for i in range(len(segments)):
+        segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)
+    for (i, segment) in enumerate(segments):
+        if i == 0 or segments[i - 1]["speaker"] != segment["speaker"]:
+            output_text += "\n" + segment["speaker"] + ' ' + str(time(segment["start"])) + '\n'
+        output_text += segment["text"][1:] + ' '
     return output_text
 
+
+
 def _return_yt_html_embed(yt_url):
     video_id = yt_url.split("?v=")[-1]
     return f'<center><iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"></iframe></center>'
@@ -89,9 +100,11 @@ def download_yt_audio(yt_url, filename):
             "preferredquality": "192",
         }],
     }
+    
     with youtube_dl.YoutubeDL(ydl_opts) as ydl:
         ydl.download([yt_url])
 
+@spaces.GPU
 def yt_transcribe(yt_url, task):
     html_embed_str = _return_yt_html_embed(yt_url)
     
@@ -99,39 +112,45 @@ def yt_transcribe(yt_url, task):
         filepath = os.path.join(tmpdirname, "audio.wav")
         download_yt_audio(yt_url, filepath)
         
-        segments, _ = model.transcribe(filepath, batch_size=BATCH_SIZE)
+        result = model.transcribe(audio, batch_size=BATCH_SIZE)
     
-    return html_embed_str, " ".join(segment.text for segment in segments)
+    return html_embed_str, result["text"]
 
-demo = gr.Blocks(theme=gr.themes.Soft())
+demo = gr.Blocks(theme=gr.themes.Ocean())
 
 mf_transcribe = gr.Interface(
     fn=transcribe,
-    inputs=[gr.Audio(sources="microphone", type="filepath"),
-            gr.Radio(["transcribe", "translate"], label="Task", value="transcribe")],
+    inputs=[
+        gr.Audio(sources="microphone", type="filepath"),
+        gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
+    ],
     outputs="text",
     title="VerbaLens Demo 1 : Prototype",
-    description="Transcribe long-form microphone or audio inputs using Faster-Whisper.",
+    description="Transcribe long-form microphone or audio inputs using WhisperX.",
     allow_flagging="never",
 )
 
 file_transcribe = gr.Interface(
     fn=transcribe,
-    inputs=[gr.Audio(sources="upload", type="filepath", label="Audio file"),
-            gr.Radio(["transcribe", "translate"], label="Task", value="transcribe")],
+    inputs=[
+        gr.Audio(sources="upload", type="filepath", label="Audio file"),
+        gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
+    ],
     outputs="text",
     title="VerbaLens Demo 1 : Prototype",
-    description="Transcribe uploaded audio files using Faster-Whisper.",
+    description="Transcribe uploaded audio files using WhisperX.",
     allow_flagging="never",
 )
 
 yt_transcribe = gr.Interface(
     fn=yt_transcribe,
-    inputs=[gr.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
-            gr.Radio(["transcribe", "translate"], label="Task", value="transcribe")],
+    inputs=[
+        gr.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"),
+        gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
+    ],
     outputs=["html", "text"],
     title="VerbaLens Demo 1 : Prototyping",
-    description="Transcribe YouTube videos using Faster-Whisper.",
+    description="Transcribe YouTube videos using WhisperX.",
     allow_flagging="never",
 )