baddarbukhari's picture
Update api.py
23d028a verified
Raw
History Blame Contribute Delete
4.52 kB
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
import io
import soundfile as sf
from kittentts import KittenTTS
app = FastAPI(
title="KittenTTS API",
description="Full KittenTTS API with all models and voices",
version="1.0.0"
)
# =========================
# AVAILABLE MODELS
# =========================
AVAILABLE_MODELS = {
"mini": {
"id": "KittenML/kitten-tts-mini-0.8",
"params": "80M",
"size": "80MB",
"description": "Highest quality model"
},
"micro": {
"id": "KittenML/kitten-tts-micro-0.8",
"params": "40M",
"size": "41MB",
"description": "Balanced performance"
},
"nano": {
"id": "KittenML/kitten-tts-nano-0.8",
"params": "15M",
"size": "56MB",
"description": "Lightweight model"
},
"nano-int8": {
"id": "KittenML/kitten-tts-nano-0.8-int8",
"params": "15M",
"size": "25MB",
"description": "Ultra lightweight INT8 model"
}
}
# =========================
# LOAD DEFAULT MODEL
# =========================
current_model_name = "nano-int8"
tts_model = KittenTTS(AVAILABLE_MODELS[current_model_name]["id"])
AVAILABLE_VOICES = [
"Bella",
"Jasper",
"Luna",
"Bruno",
"Rosie",
"Hugo",
"Kiki",
"Leo"
]
# =========================
# REQUEST MODEL
# =========================
class TTSRequest(BaseModel):
text: str
voice: str = "Bella"
model: str = current_model_name
speed: float = 1.0
# =========================
# ROOT ENDPOINT
# =========================
@app.get("/")
def root():
return {
"status": "running",
"engine": "KittenTTS",
"current_model": current_model_name,
"voices": AVAILABLE_VOICES,
"models": list(AVAILABLE_MODELS.keys())
}
# =========================
# GET MODELS
# =========================
@app.get("/models")
def get_models():
return AVAILABLE_MODELS
# =========================
# GET VOICES
# =========================
@app.get("/voices")
def get_voices():
return {
"voices": AVAILABLE_VOICES,
"count": len(AVAILABLE_VOICES)
}
# =========================
# SWITCH MODEL
# =========================
@app.post("/switch-model")
def switch_model(model_name: str):
global tts_model, current_model_name
if model_name not in AVAILABLE_MODELS:
raise HTTPException(
status_code=400,
detail="Invalid model name"
)
tts_model = KittenTTS(
AVAILABLE_MODELS[model_name]["id"]
)
current_model_name = model_name
return {
"status": "success",
"current_model": model_name
}
# =========================
# GENERATE SPEECH
# =========================
@app.post("/tts")
def generate_tts(req: TTSRequest):
global tts_model, current_model_name
# switch model automatically if requested
if req.model != current_model_name:
if req.model not in AVAILABLE_MODELS:
raise HTTPException(
status_code=400,
detail="Invalid model"
)
tts_model = KittenTTS(
AVAILABLE_MODELS[req.model]["id"]
)
current_model_name = req.model
# validate voice
if req.voice not in AVAILABLE_VOICES:
raise HTTPException(
status_code=400,
detail="Invalid voice"
)
try:
audio = tts_model.generate(
req.text,
voice=req.voice
)
buffer = io.BytesIO()
sf.write(
buffer,
audio,
24000,
format="WAV"
)
buffer.seek(0)
return StreamingResponse(
buffer,
media_type="audio/wav",
headers={
"Content-Disposition":
f"attachment; filename=kitten-{req.voice}.wav"
}
)
except Exception as e:
raise HTTPException(
status_code=500,
detail=str(e)
)
# =========================
# API INFO
# =========================
@app.get("/info")
def info():
return {
"engine": "KittenTTS",
"version": "0.8",
"models_available": len(AVAILABLE_MODELS),
"voices_available": len(AVAILABLE_VOICES),
"current_model": current_model_name,
"voices": AVAILABLE_VOICES,
"models": AVAILABLE_MODELS
}