aesda

app.py

@@ -5,87 +5,76 @@ import yt_dlp as youtube_dl
 import tempfile
 import os
 import locale
-import whisper
 import datetime
 import subprocess
-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
+from sklearn.cluster import AgglomerativeClustering
+from faster_whisper import WhisperModel
+from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
+from pyannote.audio import Audio
+from pyannote.core import Segment
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 BATCH_SIZE = 8
 FILE_LIMIT_MB = 1000
-COMPUTE_TYPE = "float32" 
 YT_LENGTH_LIMIT_S = 3600  # limit to 1 hour YouTube files
 
 num_speakers = 2
 language = None
 model_size = 'tiny'
-model_name = model_size
+model = WhisperModel(model_size, device=device, compute_type="float32")
+embedding_model = PretrainedSpeakerEmbedding("speechbrain/spkrec-ecapa-voxceleb", device=torch.device("cpu"))
+audio = Audio()
 
-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"]
-    # Whisper overshoots the end timestamp in the last segment
-    end = min(duration, segment["end"])
+def segment_embedding(segment, duration, path):
+    start = segment.start
+    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 path[-3:] != 'wav':
+    if not path.endswith('.wav'):
         subprocess.call(['ffmpeg', '-i', path, "audio.wav", '-y'])
         path = "audio.wav"
-    result = model.transcribe(path,fp16=False)
-    segments = result["segments"]
-    print(segments)
-    with contextlib.closing(wave.open(path,'r')) as f:
+    
+    segments, _ = model.transcribe(path)
+    
+    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 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:] + ' '
+    
+    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 + ' '
+    
     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>'
@@ -100,11 +89,9 @@ 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)
     
@@ -112,45 +99,39 @@ def yt_transcribe(yt_url, task):
         filepath = os.path.join(tmpdirname, "audio.wav")
         download_yt_audio(yt_url, filepath)
         
-        result = model.transcribe(audio, batch_size=BATCH_SIZE)
+        segments, _ = model.transcribe(filepath, batch_size=BATCH_SIZE)
     
-    return html_embed_str, result["text"]
+    return html_embed_str, " ".join(segment.text for segment in segments)
 
-demo = gr.Blocks(theme=gr.themes.Ocean())
+demo = gr.Blocks(theme=gr.themes.Soft())
 
 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 WhisperX.",
+    description="Transcribe long-form microphone or audio inputs using Faster-Whisper.",
     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 WhisperX.",
+    description="Transcribe uploaded audio files using Faster-Whisper.",
     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 WhisperX.",
+    description="Transcribe YouTube videos using Faster-Whisper.",
     allow_flagging="never",
 )