from fastapi import FastAPI, Response
from fastapi.responses import FileResponse
from kokoro import KPipeline
import soundfile as sf
import os

app = FastAPI()

# Initialize pipeline once at startup
pipeline = KPipeline(lang_code='a')

@app.post("/generate")
async def generate_audio(text: str, voice: str = "af_heart", speed: float = 1.0):
    # Generate audio
    generator = pipeline(
        text, 
        voice=voice,
        speed=speed,
        split_pattern=r'\n+'
    )
    
    # Save first segment only for demo
    for i, (gs, ps, audio) in enumerate(generator):
        sf.write(f"output_{i}.wav", audio, 24000)
        return FileResponse(
            f"output_{i}.wav",
            media_type="audio/wav",
            filename="output.wav"
        )
    
    return Response("No audio generated", status_code=400)