from flask import Flask, request, jsonify
from pyannote.audio.pipelines import SpeakerDiarization
from pyannote.audio import Audio
import torch
import os
import wave

app = Flask(__name__)

# Cargar el modelo de diarización de Hugging Face
model_name = "pyannote/speaker-diarization-3.0"
pipeline = SpeakerDiarization.from_pretrained(model_name)

@app.route("/diarize", methods=["POST"])
def diarize_audio():
    if "audio" not in request.files:
        return jsonify({"error": "No file uploaded"}), 400
    
    audio_file = request.files["audio"]
    audio_path = "temp_audio.wav"
    audio_file.save(audio_path)

    # Realizar la diarización
    diarization = pipeline({"uri": "file", "audio": audio_path})

    # Formatear los resultados
    result = []
    for segment, _, speaker in diarization.itertracks(yield_label=True):
        result.append({
            "speaker": speaker,
            "start_time": segment.start,
            "end_time": segment.end
        })
    
    return jsonify(result)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)