import torch, os
import soundfile as sf
import numpy as np
from fastapi import FastAPI
from fastapi.responses import FileResponse
from pydantic import BaseModel
from transformers import VitsModel, AutoTokenizer

print("Loading models...")

# ── Luhya — Benjamin's fine-tuned Swahili MMS ─────────────
luhya_model     = VitsModel.from_pretrained("Benjamin-png/swahili-mms-tts-finetuned")
luhya_tokenizer = AutoTokenizer.from_pretrained("Benjamin-png/swahili-mms-tts-finetuned")
luhya_model.eval()
print("✅ Luhya (Benjamin MMS) loaded")

# ── Kikuyu — VITS ─────────────────────────────────────────
kikuyu_model     = VitsModel.from_pretrained("gateremark/kikuyu-tts-v1")
kikuyu_tokenizer = AutoTokenizer.from_pretrained("gateremark/kikuyu-tts-v1")
kikuyu_model.eval()
print("✅ Kikuyu loaded")

# ── Swahili — Facebook MMS base ───────────────────────────
swahili_model     = VitsModel.from_pretrained("facebook/mms-tts-swh")
swahili_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-swh")
swahili_model.eval()
print("✅ Swahili (Facebook MMS) loaded")

# ── FastAPI ────────────────────────────────────────────────
app = FastAPI()

class TTSRequest(BaseModel):
    text:     str
    language: str = "luhya"   # "luhya", "kikuyu", or "swahili"

@app.get("/")
def root():
    return {"status": "ok", "languages": ["luhya", "kikuyu", "swahili"]}

@app.post("/predict")
def predict(req: TTSRequest):
    lang = req.language.lower().strip()

    if lang == "kikuyu":
        model, tokenizer = kikuyu_model, kikuyu_tokenizer
    elif lang == "swahili":
        model, tokenizer = swahili_model, swahili_tokenizer
    else:
        # luhya — use Benjamin fine-tuned Swahili MMS directly
        model, tokenizer = luhya_model, luhya_tokenizer

    inputs = tokenizer(text=req.text.strip(), return_tensors="pt")
    with torch.no_grad():
        output = model(**inputs)

    wav = output.waveform.squeeze().cpu().numpy()
    sr  = model.config.sampling_rate
    sf.write("/tmp/output.wav", wav, sr)
    return FileResponse("/tmp/output.wav", media_type="audio/wav")

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