from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from mvector.predict import MVectorPredictor
import os

app = FastAPI()

# 允许跨域
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

predictor = MVectorPredictor(configs='configs/ecapa_tdnn.yml')

@app.post("/recognize")
async def recognize(file: UploadFile = File(...)):
    try:
        audio_bytes = await file.read()
        result = predictor.recognition(audio_bytes)
        return {"status": 200, "data": result}
    except Exception as e:
        return {"status": 500, "error": str(e)}

@app.post("/compare") 
async def compare(file1: UploadFile = File(...), file2: UploadFile = File(...)):
    try:
        score = predictor.contrast(
            await file1.read(),
            await file2.read()
        )
        return {"similarity": float(score)}
    except Exception as e:
        return {"status": 500, "error": str(e)}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=9001)