from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
import whisper
import io
import torch
import ffmpeg
import numpy as np
import soundfile as sf

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load Whisper model
model = whisper.load_model("tiny")

@app.post("/transcribe")
async def transcribe(audio: UploadFile = File(...)):
    # Đọc dữ liệu âm thanh vào memory
    audio_bytes = await audio.read()

    # Convert bytes thành numpy array
    with io.BytesIO(audio_bytes) as audio_buffer:
        audio_array, sample_rate = sf.read(audio_buffer)

    # Whisper yêu cầu sample rate 16000Hz
    if sample_rate != 16000:
        audio_array = whisper.pad_or_trim(whisper.resample(audio_array, sample_rate, 16000))

    # Chuyển numpy array thành tensor
    audio_tensor = torch.from_numpy(audio_array).float()

    # Chuyển thành Mel Spectrogram
    mel = whisper.log_mel_spectrogram(audio_tensor).unsqueeze(0)

    # Chạy mô hình để nhận diện giọng nói
    result = model.decode(mel)

    return {"text": result.text}