ttasd

app.py

@@ -20,7 +20,18 @@ def transcribe(inputs, task):
     
     audio = whisperx.load_audio(inputs)
     result = model.transcribe(audio, batch_size=BATCH_SIZE)
-    return result["text"]
+    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"
+        
+    return output_text
 
 def _return_yt_html_embed(yt_url):
     video_id = yt_url.split("?v=")[-1]