from fastapi import FastAPI, UploadFile, File
import soundfile as sf
from faster_whisper import WhisperModel
import os
import tempfile

# 设置 Hugging Face 缓存目录
os.environ["HF_HOME"] = "/tmp/hf_cache"

# 初始化 FastAPI 应用
app = FastAPI()

# 加载 Faster-Whisper 模型
print("✅ THIS IS NEW APP.PY VERSION")
model = WhisperModel("ct2_model", compute_type="int8", device="cpu")

@app.get("/")
def root():
    return {
        "message": "CTranslate2 Whisper API is running.",
        "usage": "POST /transcribe with an audio file (.wav, .mp3, etc.)"
    }

@app.post("/transcribe")
async def transcribe(file: UploadFile = File(...)):
    # 保存上传音频到临时文件
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
        temp_path = temp_file.name
        temp_file.write(await file.read())

    try:
        # 使用 Faster-Whisper 进行推理
        segments, info = model.transcribe(temp_path)
        transcription = " ".join([segment.text for segment in segments])

        return {"text": transcription}
    finally:
        # 删除临时文件
        os.remove(temp_path)