import io
import base64
import numpy as np
import soundfile as sf
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from kokoro import KPipeline
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
from pydantic import BaseModel
import uvicorn
# ── Model loading ──────────────────────────────────────────────────────────────
MODEL_NAME = "HuggingFaceTB/SmolLM2-360M-Instruct"
print("Loading LLM...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)
model.eval()
print("✅ LLM loaded")
print("Loading TTS...")
tts = KPipeline(lang_code="a")
print("✅ TTS loaded")
# ── FastAPI ────────────────────────────────────────────────────────────────────
app = FastAPI(title="AI Text + TTS API")
class QuestionRequest(BaseModel):
question: str
voice: str = "af_heart"
max_new_tokens: int = 120
def run_llm(question: str, max_new_tokens: int = 120) -> str:
messages = [
{"role": "system", "content": "You are a helpful NLP chatbot. Reply simply and shortly."},
{"role": "user", "content": question}
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
output_ids = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.eos_token_id
)
new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
def run_tts(text: str, voice: str = "af_heart") -> np.ndarray:
parts = []
for _, _, audio in tts(text, voice=voice):
parts.append(audio)
return np.concatenate(parts) if parts else np.array([])
# ── API routes ─────────────────────────────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
@app.get("/ui", response_class=HTMLResponse)
def serve_ui():
return HTMLResponse(content=UI_HTML)
@app.get("/health")
def health():
return {"status": "healthy", "llm": MODEL_NAME, "tts": "kokoro"}
@app.post("/api/text")
def generate_text(req: QuestionRequest):
if not req.question.strip():
raise HTTPException(400, "Question cannot be empty")
try:
answer = run_llm(req.question, req.max_new_tokens)
return JSONResponse({"question": req.question, "answer": answer})
except Exception as e:
raise HTTPException(500, str(e))
@app.post("/api/tts")
def generate_tts(req: QuestionRequest):
if not req.question.strip():
raise HTTPException(400, "Question cannot be empty")
try:
answer = run_llm(req.question, req.max_new_tokens)
audio = run_tts(answer, req.voice)
if audio.size == 0:
raise HTTPException(500, "TTS produced no audio")
buf = io.BytesIO()
sf.write(buf, audio, 24000, format="WAV")
buf.seek(0)
return StreamingResponse(
buf,
media_type="audio/wav",
headers={"X-Answer-Text": answer}
)
except HTTPException:
raise
except Exception as e:
raise HTTPException(500, str(e))
@app.post("/api/both")
def generate_both(req: QuestionRequest):
if not req.question.strip():
raise HTTPException(400, "Question cannot be empty")
try:
answer = run_llm(req.question, req.max_new_tokens)
audio = run_tts(answer, req.voice)
audio_b64 = ""
if audio.size > 0:
buf = io.BytesIO()
sf.write(buf, audio, 24000, format="WAV")
audio_b64 = base64.b64encode(buf.getvalue()).decode()
return JSONResponse({
"question": req.question,
"answer": answer,
"audio_base64": audio_b64,
"sample_rate": 24000
})
except Exception as e:
raise HTTPException(500, str(e))
# ── Embedded UI ────────────────────────────────────────────────────────────────
UI_HTML = """
AI Text + TTS
🎙️ AI Text + Text-to-Speech
SmolLM2 generates answers · Kokoro converts them to speech
API Endpoints
GET/healthModel status
POST/api/textJSON → text answer
POST/api/ttsJSON → WAV audio stream
POST/api/bothJSON → text + base64 audio
"""
# ── Entry point ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)