server.py

"""
cantrell-kokoro-engine — StoryVoice TTS Backend
Docker Space, Python 3.11, FastAPI only
Supports voice blending, sentence-level silence padding, pronunciation map
"""

import io
import re
import time
import numpy as np
import soundfile as sf
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from kokoro import KPipeline
import uvicorn

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Pipeline cache ────────────────────────────────────────────────────────────
_pipelines = {}

def get_pipeline(lang_code):
    if lang_code not in _pipelines:
        _pipelines[lang_code] = KPipeline(lang_code=lang_code)
    return _pipelines[lang_code]

def lang_for_voice(voice_id: str) -> str:
    # Use first voice in a blend to determine lang
    first = voice_id.split('+')[0].strip().split(':')[0].strip()
    return "b" if first.startswith("b") else "a"

# ── Pronunciation map ─────────────────────────────────────────────────────────
PRONUNCIATION = {
    "Ybor":     "Eebore",
    "ybor":     "eebore",
    "breathed": "breethd",
    "Breathed": "Breethd",
    "Sanae":    "Suh-nay",
    "sanae":    "suh-nay",
    "Nae":      "Nay",
    "nae":      "nay",
    "Keymoni":  "Keymoney",
    "keymoni":  "keymoney",
}

def apply_pronunciation(text: str) -> str:
    for word, replacement in PRONUNCIATION.items():
        text = text.replace(word, replacement)
    return text

# ── Voice blending ────────────────────────────────────────────────────────────
# Blend format: "af_aoede:0.50 + af_sky:0.30 + af_nicole:0.20"
# Single voice: "af_heart" or "af_heart:1.0"

def parse_blend(voice_str: str):
    """Parse blend string into list of (voice_id, weight) tuples."""
    parts = [p.strip() for p in voice_str.split('+')]
    blend = []
    for part in parts:
        if ':' in part:
            vid, w = part.rsplit(':', 1)
            blend.append((vid.strip(), float(w.strip())))
        else:
            blend.append((part.strip(), 1.0))
    # Normalize weights
    total = sum(w for _, w in blend)
    return [(v, w / total) for v, w in blend]

def blend_voices(blend_list):
    """Create a blended voice tensor from a list of (voice_id, weight) tuples."""
    import torch
    blended = None
    for voice_id, weight in blend_list:
        # KPipeline loads voice pack internally; access via pipeline's voice method
        lang = "b" if voice_id.startswith("b") else "a"
        pipeline = get_pipeline(lang)
        pack = pipeline.load_voice(voice_id)
        if blended is None:
            blended = pack * weight
        else:
            blended = blended + pack * weight
    return blended

# ── Named presets ─────────────────────────────────────────────────────────────
PRESETS = {
    "male_narrator":   "am_adam:0.60 + am_michael:0.30 + am_onyx:0.10",
    "female_narrator": "af_aoede:0.50 + af_sky:0.30 + af_nicole:0.20",
}

# ── Voice registry ────────────────────────────────────────────────────────────
VOICES = [
    {"voice_id": "af_alloy",    "display_name": "Alloy",    "gender": "female", "accent": "american"},
    {"voice_id": "af_aoede",    "display_name": "Aoede",    "gender": "female", "accent": "american"},
    {"voice_id": "af_bella",    "display_name": "Bella",    "gender": "female", "accent": "american"},
    {"voice_id": "af_heart",    "display_name": "Heart",    "gender": "female", "accent": "american"},
    {"voice_id": "af_jessica",  "display_name": "Jessica",  "gender": "female", "accent": "american"},
    {"voice_id": "af_kore",     "display_name": "Kore",     "gender": "female", "accent": "american"},
    {"voice_id": "af_nicole",   "display_name": "Nicole",   "gender": "female", "accent": "american"},
    {"voice_id": "af_nova",     "display_name": "Nova",     "gender": "female", "accent": "american"},
    {"voice_id": "af_river",    "display_name": "River",    "gender": "female", "accent": "american"},
    {"voice_id": "af_sarah",    "display_name": "Sarah",    "gender": "female", "accent": "american"},
    {"voice_id": "af_sky",      "display_name": "Sky",      "gender": "female", "accent": "american"},
    {"voice_id": "am_adam",     "display_name": "Adam",     "gender": "male",   "accent": "american"},
    {"voice_id": "am_echo",     "display_name": "Echo",     "gender": "male",   "accent": "american"},
    {"voice_id": "am_eric",     "display_name": "Eric",     "gender": "male",   "accent": "american"},
    {"voice_id": "am_fenrir",   "display_name": "Fenrir",   "gender": "male",   "accent": "american"},
    {"voice_id": "am_liam",     "display_name": "Liam",     "gender": "male",   "accent": "american"},
    {"voice_id": "am_michael",  "display_name": "Michael",  "gender": "male",   "accent": "american"},
    {"voice_id": "am_onyx",     "display_name": "Onyx",     "gender": "male",   "accent": "american"},
    {"voice_id": "am_puck",     "display_name": "Puck",     "gender": "male",   "accent": "american"},
    {"voice_id": "am_santa",    "display_name": "Santa",    "gender": "male",   "accent": "american"},
    {"voice_id": "bf_alice",    "display_name": "Alice",    "gender": "female", "accent": "british"},
    {"voice_id": "bf_emma",     "display_name": "Emma",     "gender": "female", "accent": "british"},
    {"voice_id": "bf_isabella", "display_name": "Isabella", "gender": "female", "accent": "british"},
    {"voice_id": "bf_lily",     "display_name": "Lily",     "gender": "female", "accent": "british"},
    {"voice_id": "bm_daniel",   "display_name": "Daniel",   "gender": "male",   "accent": "british"},
    {"voice_id": "bm_fable",    "display_name": "Fable",    "gender": "male",   "accent": "british"},
    {"voice_id": "bm_george",   "display_name": "George",   "gender": "male",   "accent": "british"},
    {"voice_id": "bm_lewis",    "display_name": "Lewis",    "gender": "male",   "accent": "british"},
    # Named presets shown as selectable voices
    {"voice_id": "male_narrator",   "display_name": "Male Narrator (Blend)",   "gender": "male",   "accent": "american"},
    {"voice_id": "female_narrator", "display_name": "Female Narrator (Blend)", "gender": "female", "accent": "american"},
]

VOICE_MAP = {v["voice_id"]: v for v in VOICES}

SAMPLE_RATE = 24000

def resolve_voice(voice_id: str) -> str:
    voice_id = voice_id.replace(".mp3", "").strip()
    # Check presets first
    if voice_id in PRESETS:
        return PRESETS[voice_id]
    if voice_id in VOICE_MAP:
        return voice_id
    matched = next(
        (v["voice_id"] for v in VOICES if v["display_name"].lower() == voice_id.lower()),
        None
    )
    return matched or "af_heart"

def make_silence(ms: int) -> np.ndarray:
    return np.zeros(int(SAMPLE_RATE * ms / 1000), dtype=np.float32)

def split_sentences(text: str) -> list:
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    return [s.strip() for s in sentences if s.strip()]

def generate_sentence(pipeline, sentence: str, voice, speed: float) -> np.ndarray:
    chunks = []
    for _, _, audio in pipeline(sentence, voice=voice, speed=speed):
        if audio is not None and len(audio) > 0:
            chunks.append(audio)
    if not chunks:
        return np.array([], dtype=np.float32)
    return np.concatenate(chunks) if len(chunks) > 1 else chunks[0]

def generate_audio(text: str, voice_id: str, speed: float = 1.0) -> bytes:
    text = apply_pronunciation(text)
    sentences = split_sentences(text)

    # Resolve preset to blend string
    voice_str = PRESETS.get(voice_id, voice_id)
    lang = lang_for_voice(voice_str)
    pipeline = get_pipeline(lang)

    # Determine if blending needed
    is_blend = '+' in voice_str or ':' in voice_str
    if is_blend:
        blend_list = parse_blend(voice_str)
        try:
            voice = blend_voices(blend_list)
        except Exception:
            # Fallback to first voice if blending fails
            voice = blend_list[0][0]
    else:
        voice = voice_str.split(':')[0].strip()

    segments = []
    for i, sentence in enumerate(sentences):
        if not sentence:
            continue
        audio = generate_sentence(pipeline, sentence, voice, speed)
        if len(audio) > 0:
            segments.append(audio)
            if i < len(sentences) - 1:
                pause_ms = 250 if sentence.endswith(('!', '?')) else 150
                segments.append(make_silence(pause_ms))

    if not segments:
        raise ValueError("No audio generated")

    combined = np.concatenate(segments)
    buf = io.BytesIO()
    sf.write(buf, combined, SAMPLE_RATE, format="mp3")
    buf.seek(0)
    return buf.read()

# ── Routes ────────────────────────────────────────────────────────────────────
@app.get("/")
def index():
    html = """<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<title>StoryVoice™ TTS Engine</title>
<style>
  *{box-sizing:border-box;}
  body{margin:0;background:#1a0a05;min-height:100vh;font-family:Georgia,serif;display:flex;align-items:center;justify-content:center;padding:20px;}
  .card{background:#2c1e0f;border:1px solid #c9a040;border-radius:16px;padding:40px;width:100%;max-width:560px;}
  h1{color:#c9a040;font-size:20px;letter-spacing:.08em;text-transform:uppercase;margin:0 0 4px;text-align:center;}
  .sub{color:rgba(255,255,255,.5);font-size:13px;font-style:italic;text-align:center;margin-bottom:28px;}
  label{display:block;color:#c9a040;font-size:11px;letter-spacing:.06em;text-transform:uppercase;margin-bottom:6px;}
  textarea{width:100%;background:#1a0a05;border:1px solid rgba(201,160,64,.3);border-radius:8px;color:#fff;font-family:Georgia,serif;font-size:14px;padding:12px;resize:vertical;min-height:100px;outline:none;margin-bottom:16px;}
  textarea:focus{border-color:#c9a040;}
  select,input[type=range]{width:100%;background:#1a0a05;border:1px solid rgba(201,160,64,.3);border-radius:8px;color:#fff;font-size:14px;padding:10px 12px;outline:none;margin-bottom:16px;appearance:none;}
  .row{display:flex;gap:16px;}
  .row>div{flex:1;}
  button{width:100%;background:#7a2638;border:none;color:#fff;font-family:Georgia,serif;font-size:14px;letter-spacing:.06em;text-transform:uppercase;padding:14px;border-radius:8px;cursor:pointer;margin-bottom:16px;transition:background 150ms;}
  button:hover{background:#9a3248;}
  button:disabled{opacity:.5;cursor:not-allowed;}
  audio{width:100%;margin-top:4px;}
  .status{text-align:center;color:#4caf50;font-size:12px;letter-spacing:.04em;}
</style>
</head>
<body>
<div class="card">
  <h1>StoryVoice™ TTS Engine</h1>
  <div class="sub">Cantrell Creatives — Preview</div>
  <label>Text</label>
  <textarea id="txt">Tampa Bay never slept — not really. Ybor City breathed on its own schedule, old brick and older blood soaking every corner. Marcus moved through the shadows like he belonged there, because he did. The were-animal in him recognized this territory the way most men recognized home.</textarea>
  <div class="row">
    <div>
      <label>Voice</label>
      <select id="voice"></select>
    </div>
    <div>
      <label>Speed: <span id="spd-val">0.9</span></label>
      <input type="range" id="spd" min="0.5" max="2.0" step="0.05" value="0.9" oninput="document.getElementById('spd-val').textContent=this.value">
    </div>
  </div>
  <button id="btn" onclick="generate()">Generate Preview</button>
  <audio id="player" controls style="display:none"></audio>
  <div class="status" id="status"></div>
</div>
<script>
fetch('/voices').then(r=>r.json()).then(function(data){
  var sel = document.getElementById('voice');
  data.forEach(function(v){
    var o = document.createElement('option');
    o.value = v.voice_id;
    o.textContent = v.display_name + ' (' + v.accent + ')';
    sel.appendChild(o);
  });
});
function generate(){
  var btn = document.getElementById('btn');
  var status = document.getElementById('status');
  var player = document.getElementById('player');
  btn.disabled = true;
  btn.textContent = 'Generating...';
  status.textContent = '';
  player.style.display = 'none';
  fetch('/tts-preview',{
    method:'POST',
    headers:{'Content-Type':'application/json'},
    body: JSON.stringify({
      text: document.getElementById('txt').value,
      voice_id: document.getElementById('voice').value,
      speed: parseFloat(document.getElementById('spd').value)
    })
  })
  .then(function(r){ if(!r.ok) throw new Error('Generation failed'); return r.blob(); })
  .then(function(blob){
    player.src = URL.createObjectURL(blob);
    player.style.display = 'block';
    player.play();
    status.textContent = '● Ready';
  })
  .catch(function(e){
    status.style.color='#e74c3c';
    status.textContent = 'Error: ' + e.message;
  })
  .finally(function(){
    btn.disabled = false;
    btn.textContent = 'Generate Preview';
  });
}
</script>
</body>
</html>"""
    return Response(content=html, media_type="text/html")

@app.get("/health")
def health():
    return {"status": "ok", "engine": "kokoro-82m", "voices": len(VOICES), "presets": list(PRESETS.keys()), "timestamp": int(time.time())}

@app.get("/voices")
def voices():
    return VOICES

@app.get("/presets")
def presets():
    return PRESETS

class GenerateRequest(BaseModel):
    text: str
    voice_id: str = "af_heart"
    speed: float = 1.0

@app.post("/generate")
def generate(req: GenerateRequest):
    if not req.text.strip():
        raise HTTPException(status_code=400, detail="text is required")
    vid = resolve_voice(req.voice_id)
    spd = max(0.5, min(2.0, req.speed))
    mp3 = generate_audio(req.text.strip(), vid, spd)
    return Response(content=mp3, media_type="audio/mpeg")

class PreviewRequest(BaseModel):
    text: str
    voice_id: str = "af_heart"
    speed: float = 0.9

@app.post("/tts-preview")
def tts_preview(req: PreviewRequest):
    return generate(GenerateRequest(text=req.text, voice_id=req.voice_id, speed=req.speed))

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