from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
import torchaudio
import ChatTTS

# Configure Torch
torch._dynamo.config.cache_size_limit = 64
torch._dynamo.config.suppress_errors = True
torch.set_float32_matmul_precision('high')

# Initialize FastAPI app and ChatTTS model
app = FastAPI()
chat = ChatTTS.Chat()
chat.load(compile=False)

# Define the request model
class TextRequest(BaseModel):
    text: str

@app.post("/synthesize/")
async def synthesize_speech(request: TextRequest):
    try:
        # Perform inference
        wavs = chat.infer([request.text])
        
        # Save the generated audio
        output_file = "output.wav"
        torchaudio.save(output_file, torch.from_numpy(wavs[0]), 24000)

        return {"message": "Speech synthesized successfully", "audio_file": output_file}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Root endpoint
@app.get("/")
def read_root():
    return {"message": "Welcome to the ChatTTS API. Use the /synthesize/ endpoint to generate speech."}