|
|
import spaces |
|
|
import gradio as gr |
|
|
import os |
|
|
import torch |
|
|
from pyannote.audio import Pipeline |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
HF_TOKEN = os.environ.get('HF_TOKEN') |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
diarization_pipeline = None |
|
|
try: |
|
|
if HF_TOKEN: |
|
|
diarization_pipeline = Pipeline.from_pretrained( |
|
|
"pyannote/speaker-diarization-3.1", |
|
|
use_auth_token=HF_TOKEN |
|
|
) |
|
|
if torch.cuda.is_available(): |
|
|
diarization_pipeline.to(torch.device("cuda")) |
|
|
print("Pyannote: LOADED (GPU)") |
|
|
else: |
|
|
print("Pyannote: LOADED (CPU)") |
|
|
except Exception as e: |
|
|
print(f"Pyannote Error: {e}") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@spaces.GPU |
|
|
def diarize_audio(audio_path, min_speakers=1, max_speakers=5): |
|
|
if not diarization_pipeline: |
|
|
return {"error": "Diarization not available. Check HF_TOKEN."} |
|
|
try: |
|
|
diarization = diarization_pipeline(audio_path, min_speakers=int(min_speakers), max_speakers=int(max_speakers)) |
|
|
speakers = [] |
|
|
for turn, _, speaker in diarization.itertracks(yield_label=True): |
|
|
speakers.append({"speaker": speaker, "start": round(turn.start, 2), "end": round(turn.end, 2)}) |
|
|
return {"segments": speakers, "num_speakers": len(set(s["speaker"] for s in speakers))} |
|
|
except Exception as e: |
|
|
return {"error": str(e)} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
with gr.Blocks(title="STTR - Speaker Diarization") as demo: |
|
|
gr.Markdown("# STTR - Speaker Diarization") |
|
|
gr.Markdown("### Identify who speaks when (pyannote 3.1)") |
|
|
|
|
|
audio_in = gr.Audio(type="filepath", label="Upload Audio") |
|
|
with gr.Row(): |
|
|
min_spk = gr.Slider(1, 10, value=1, step=1, label="Min Speakers") |
|
|
max_spk = gr.Slider(1, 10, value=5, step=1, label="Max Speakers") |
|
|
btn = gr.Button("Analyze Speakers", variant="primary") |
|
|
output = gr.JSON(label="Speaker Segments") |
|
|
btn.click(diarize_audio, [audio_in, min_spk, max_spk], output, api_name="/diarize") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860))) |
