v3 of test

app.py

@@ -5,35 +5,89 @@ import yt_dlp as youtube_dl
 import whisperx
 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
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-BATCH_SIZE = 4
+BATCH_SIZE = 8
 FILE_LIMIT_MB = 1000
-COMPUTE_TYPE = "float32" 
+COMPUTE_TYPE = "float16" 
 YT_LENGTH_LIMIT_S = 3600  # limit to 1 hour YouTube files
 
-model = whisperx.load_model("large-v2", device,compute_type=COMPUTE_TYPE)
+num_speakers = 2
+language = 'French'
+model_size = 'large'
+model_name = model_size
+
+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)
+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"])
+    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))
 
 @spaces.GPU
-def transcribe(inputs, task):
-    if inputs is None:
+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.")
     
-    audio = whisperx.load_audio(inputs)
-    result = model.transcribe(audio, batch_size=BATCH_SIZE)
-    model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
-    result = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False)
-    diarize_model = whisperx.DiarizationPipeline(use_auth_token=os.getenv("HF_TOKEN"), device=device)
-    diarize_segments = diarize_model(audio)
-    result = whisperx.assign_word_speakers(diarize_segments, result)
-    output_text = ""
-    for segment in result['segments']:
-        speaker = segment.get('speaker', 'Unknown Speaker')
-        text = segment['text']
-        output_text += f"{speaker}: {text}\n"
-        
+    if path[-3:] != 'wav':
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            filepath = os.path.join(tmpdirname, "audio.wav")
+            subprocess.call(['ffmpeg', '-i', path, filepath, '-y'])
+            path = filepath
+    result = model.transcribe(path)
+    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 = 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:] + ' '
     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>'

requirements.txt

@@ -11,4 +11,4 @@ transformers>=4.19.0
 yt-dlp
 more_itertools
 faster-whisper
-git+https://github.com/m-bain/whisperx.git
+git+https://github.com/openai/whisper.git