Update app.py

app.py

@@ -298,280 +298,239 @@ import os
 import re
 import time
 import asyncio
-import uvloop
+from concurrent.futures import ThreadPoolExecutor
+
 import numpy as np
-import gradio as gr
 import torch
-
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-from concurrent.futures import ThreadPoolExecutor
 import uvicorn
 
-# Official pipeline
 from kokoro import KPipeline
 
-# -------------------------
-# CPU + runtime tuning
-# -------------------------
-# Keep these conservative. HF CPU is usually 2 vCPU.
+# ----------------------------
+# CPU THREAD CAP (HF free tier is typically 2 vCPU)
+# ----------------------------
 os.environ.setdefault("OMP_NUM_THREADS", "2")
 os.environ.setdefault("MKL_NUM_THREADS", "2")
 os.environ.setdefault("NUMEXPR_NUM_THREADS", "2")
 
-torch.set_num_threads(2)
-torch.set_num_interop_threads(1)
+try:
+    torch.set_num_threads(int(os.environ.get("TORCH_NUM_THREADS", "2")))
+    torch.set_num_interop_threads(int(os.environ.get("TORCH_NUM_INTEROP_THREADS", "1")))
+except Exception:
+    pass
 
-asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+# Optional uvloop (safe to skip if not installed)
+try:
+    import uvloop  # type: ignore
+    asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+except Exception:
+    pass
 
 SAMPLE_RATE = 24000
 
-# -------------------------
-# Voices (use Kokoro voice ids)
-# -------------------------
+print("🚀 BOOTING KOKORO API ONLY (OFFICIAL PIPELINE)")
+
+# ----------------------------
+# VOICES (UI label -> kokoro voice id)
+# Client can send either label or id.
+# ----------------------------
 VOICE_CHOICES = {
-    "🇺🇸 🚺 Heart": "af_heart",
-    "🇺🇸 🚺 Bella": "af_bella",
-    "🇺🇸 🚺 Nicole": "af_nicole",
-    "🇺🇸 🚺 Aoede": "af_aoede",
-    "🇺🇸 🚺 Kore": "af_kore",
-    "🇺🇸 🚺 Sarah": "af_sarah",
-    "🇺🇸 🚺 Nova": "af_nova",
-    "🇺🇸 🚺 Sky": "af_sky",
-    "🇺🇸 🚺 Alloy": "af_alloy",
-    "🇺🇸 🚺 Jessica": "af_jessica",
-    "🇺🇸 🚺 River": "af_river",
-    "🇺🇸 🚹 Michael": "am_michael",
-    "🇺🇸 🚹 Fenrir": "am_fenrir",
-    "🇺🇸 🚹 Puck": "am_puck",
-    "🇺🇸 🚹 Echo": "am_echo",
-    "🇺🇸 🚹 Eric": "am_eric",
-    "🇺🇸 🚹 Liam": "am_liam",
-    "🇺🇸 🚹 Onyx": "am_onyx",
-    "🇺🇸 🚹 Santa": "am_santa",
-    "🇺🇸 🚹 Adam": "am_adam",
-    "🇬🇧 🚺 Emma": "bf_emma",
-    "🇬🇧 🚺 Isabella": "bf_isabella",
-    "🇬🇧 🚺 Alice": "bf_alice",
-    "🇬🇧 🚺 Lily": "bf_lily",
-    "🇬🇧 🚹 George": "bm_george",
-    "🇬🇧 🚹 Fable": "bm_fable",
-    "🇬🇧 🚹 Lewis": "bm_lewis",
+    "🇺🇸 🚺 Heart": "af_heart", "🇺🇸 🚺 Bella": "af_bella", "🇺🇸 🚺 Nicole": "af_nicole",
+    "🇺🇸 🚺 Aoede": "af_aoede", "🇺🇸 🚺 Kore": "af_kore", "🇺🇸 🚺 Sarah": "af_sarah",
+    "🇺🇸 🚺 Nova": "af_nova", "🇺🇸 🚺 Sky": "af_sky", "🇺🇸 🚺 Alloy": "af_alloy",
+    "🇺🇸 🚺 Jessica": "af_jessica", "🇺🇸 🚺 River": "af_river", "🇺🇸 🚹 Michael": "am_michael",
+    "🇺🇸 🚹 Fenrir": "am_fenrir", "🇺🇸 🚹 Puck": "am_puck", "🇺🇸 🚹 Echo": "am_echo",
+    "🇺🇸 🚹 Eric": "am_eric", "🇺🇸 🚹 Liam": "am_liam", "🇺🇸 🚹 Onyx": "am_onyx",
+    "🇺🇸 🚹 Santa": "am_santa", "🇺🇸 🚹 Adam": "am_adam", "🇬🇧 🚺 Emma": "bf_emma",
+    "🇬🇧 🚺 Isabella": "bf_isabella", "🇬🇧 🚺 Alice": "bf_alice", "🇬🇧 🚺 Lily": "bf_lily",
+    "🇬🇧 🚹 George": "bm_george", "🇬🇧 🚹 Fable": "bm_fable", "🇬🇧 🚹 Lewis": "bm_lewis",
     "🇬🇧 🚹 Daniel": "bm_daniel",
 }
+ALLOWED_VOICE_IDS = set(VOICE_CHOICES.values())
+
+# ✅ DEFAULT VOICE = ONYX
+DEFAULT_VOICE_LABEL = "🇺🇸 🚹 Onyx"
+DEFAULT_VOICE_ID = VOICE_CHOICES[DEFAULT_VOICE_LABEL]
+DEFAULT_SPEED = 1.0
+
+def voice_to_lang_code(voice_id: str) -> str:
+    if voice_id.startswith("bf_") or voice_id.startswith("bm_"):
+        return "b"  # British
+    return "a"      # American
+
+# ----------------------------
+# PIPELINES (keep hot in RAM)
+# ----------------------------
+PIPELINES = {
+    "a": KPipeline(lang_code="a"),
+    "b": KPipeline(lang_code="b"),
+}
 
-DEFAULT_VOICE_UI = "🇺🇸 🚺 Bella"
-DEFAULT_VOICE = VOICE_CHOICES[DEFAULT_VOICE_UI]
+# ----------------------------
+# TEXT NORMALIZATION (from your provided docs)
+# ----------------------------
+_SENT_BOUNDARY = re.compile(r"([.!?;:])\s+")
+_MULTI_NL = re.compile(r"\n{3,}")
+_CAMEL = re.compile(r"([a-z])([A-Z])")
+_ALLCAPS = re.compile(r"\b([A-Z]{2,})\b")
+
+def normalize_text(text: str) -> str:
+    if not text:
+        return ""
+    return text.replace("Kokoro", "[Kokoro](/kˈOkəɹO/)")
+
+def reduce_name_skips(text: str) -> str:
+    if not text:
+        return ""
+    text = _ALLCAPS.sub(lambda m: " ".join(list(m.group(1))), text)
+    text = _CAMEL.sub(r"\1 \2", text)
+    return text
 
-# -------------------------
-# Kokoro pipeline (global)
-# -------------------------
-print("🚀 BOOTING KOKORO (OFFICIAL PIPELINE)")
-PIPELINE = KPipeline(lang_code="a")
+def inject_newlines_for_fast_stream(text: str) -> str:
+    text = normalize_text(text).strip()
+    if not text:
+        return ""
+    text = _SENT_BOUNDARY.sub(r"\1\n", text)
+    text = _MULTI_NL.sub("\n\n", text)
+
+    # Ensure a small first segment for faster first audio
+    if "\n" not in text and len(text) > 90:
+        cut = text.rfind(" ", 0, 70)
+        if cut < 35:
+            cut = 70
+        text = text[:cut].strip() + "\n" + text[cut:].strip()
 
-# -------------------------
-# Helpers
-# -------------------------
-def _to_numpy_audio(audio):
-    # Kokoro may return a torch.Tensor or numpy array
-    if isinstance(audio, torch.Tensor):
-        return audio.detach().cpu().numpy()
-    return np.asarray(audio)
-
-def _float_to_int16(audio_f32):
-    audio_f32 = np.clip(audio_f32, -1.0, 1.0).astype(np.float32)
-    return (audio_f32 * 32767.0).astype(np.int16)
-
-def trim_silence(audio_f32, threshold=0.01, pad=240):
-    # audio_f32 is float32, shape [N]
-    if audio_f32.size == 0:
-        return audio_f32
-    mask = np.abs(audio_f32) > threshold
-    if not np.any(mask):
-        return audio_f32
-    start = int(np.argmax(mask))
-    end = int(len(mask) - np.argmax(mask[::-1]))
-    start = max(0, start - pad)
-    end = min(len(audio_f32), end + pad)
-    return audio_f32[start:end]
-
-def crossfade_concat(a, b, overlap=1200):
-    # overlap ~ 50ms at 24k
-    if a is None:
-        return b
-    if b is None:
-        return a
-    if len(a) < overlap or len(b) < overlap:
-        return np.concatenate([a, b])
-
-    fade_out = np.linspace(1.0, 0.0, overlap, dtype=np.float32)
-    fade_in = 1.0 - fade_out
-
-    a_tail = a[-overlap:] * fade_out
-    b_head = b[:overlap] * fade_in
-
-    mixed = a_tail + b_head
-    return np.concatenate([a[:-overlap], mixed, b[overlap:]])
-
-def tuned_splitter(text):
-    # First chunk small for fast first packet, later chunks larger for efficiency
-    parts = re.split(r"([.,!?;:\n]+)", text)
-    buf = ""
-    chunk_idx = 0
-    for p in parts:
-        buf += p
-        if chunk_idx == 0:
-            threshold = 80
-        elif chunk_idx == 1:
-            threshold = 140
-        elif chunk_idx == 2:
-            threshold = 220
-        else:
-            threshold = 320
-
-        if re.search(r"[.,!?;:\n]$", buf) and len(buf) >= threshold:
-            s = buf.strip()
-            if s:
-                yield s
-                chunk_idx += 1
-            buf = ""
-
-    s = buf.strip()
-    if s:
-        yield s
-
-def normalize_names_minimally(text):
-    # Cheap heuristics to reduce skipped acronyms and CamelCase
-    # 1) Split ALLCAPS as letters: "AI" -> "A I"
-    text = re.sub(r"\b([A-Z]{2,})\b", lambda m: " ".join(list(m.group(1))), text)
-    # 2) Split CamelCase boundaries: "OpenAI" -> "Open AI"
-    text = re.sub(r"([a-z])([A-Z])", r"\1 \2", text)
-    # Keep your Kokoro IPA hint example
-    text = text.replace("Kokoro", "Kokoro")  # keep as-is unless you inject IPA tags in client
     return text
 
-def synthesize_one_chunk(chunk, voice_id, speed):
-    # Make sure no nested splitting happens inside a chunk
+# ----------------------------
+# AUDIO CONVERSION
+# ----------------------------
+def audio_to_int16_np(audio):
+    if isinstance(audio, torch.Tensor):
+        a = audio.detach().cpu()
+        a = torch.clamp(a, -1.0, 1.0)
+        return (a * 32767.0).to(torch.int16).numpy()
+
+    a = np.asarray(audio)
+    a = np.clip(a, -1.0, 1.0)
+    return (a * 32767.0).astype(np.int16)
+
+def audio_to_pcm_bytes(audio) -> bytes:
+    return audio_to_int16_np(audio).tobytes()
+
+# ----------------------------
+# OFFICIAL GENERATION PATH (single pipeline call per request)
+# ----------------------------
+def kokoro_audio_iter(text: str, voice_id: str, speed: float):
+    lang_code = voice_to_lang_code(voice_id)
+    pipeline = PIPELINES[lang_code]
+    prepared = inject_newlines_for_fast_stream(text)
+    if not prepared:
+        return
+
     with torch.inference_mode():
-        gen = PIPELINE(
-            chunk,
+        gen = pipeline(
+            prepared,
             voice=voice_id,
             speed=float(speed),
-            split_pattern=r"\n+",  # chunk text has no newlines in practice
+            split_pattern=r"\n+",
         )
-        # gen yields (gs, ps, audio)
-        out_audio = None
         for _, _, audio in gen:
-            audio_np = _to_numpy_audio(audio).astype(np.float32)
-            audio_np = trim_silence(audio_np)
-            out_audio = crossfade_concat(out_audio, audio_np, overlap=1200)
-        return out_audio
-
-# -------------------------
-# Warmup to remove cold start latency
-# -------------------------
+            yield audio
+
 def warmup():
     try:
         t0 = time.time()
-        _ = synthesize_one_chunk("Warmup.", DEFAULT_VOICE, 1.0)
-        dt = time.time() - t0
-        print(f"✅ Warmup done in {dt:.2f}s")
+        for _ in kokoro_audio_iter("Hello.", DEFAULT_VOICE_ID, 1.0):
+            break
+        print(f"✅ WARMUP DONE in {time.time() - t0:.2f}s")
     except Exception as e:
-        print(f"⚠️ Warmup failed: {e}")
-
-# Run warmup in background thread once
-WARMUP_EXECUTOR = ThreadPoolExecutor(max_workers=1)
-WARMUP_EXECUTOR.submit(warmup)
-
-# -------------------------
-# Streaming strategy
-# -------------------------
-def stream_generator(text, voice_ui, speed):
-    voice_id = VOICE_CHOICES.get(voice_ui, DEFAULT_VOICE)
-    text = normalize_names_minimally(text)
-
-    print("--- START UI STREAM ---")
-    first = True
-
-    # Buffer audio after the first packet to reduce gaps from too many tiny yields
-    buffer_audio = None
-    buffer_min_seconds = 0.9
-
-    for chunk_idx, chunk in enumerate(tuned_splitter(text)):
-        t0 = time.time()
-        audio_f32 = synthesize_one_chunk(chunk, voice_id, speed)
-        if audio_f32 is None or len(audio_f32) == 0:
-            continue
-
-        dt = time.time() - t0
-        print(f"⚡ UI chunk {chunk_idx}: {len(chunk)} chars in {dt:.2f}s")
-
-        if first:
-            # First packet: yield immediately for low perceived latency
-            first = False
-            yield (SAMPLE_RATE, _float_to_int16(audio_f32))
-            continue
-
-        buffer_audio = crossfade_concat(buffer_audio, audio_f32, overlap=1200)
-        if buffer_audio is not None:
-            if len(buffer_audio) >= int(buffer_min_seconds * SAMPLE_RATE):
-                yield (SAMPLE_RATE, _float_to_int16(buffer_audio))
-                buffer_audio = None
-
-    if buffer_audio is not None and len(buffer_audio) > 0:
-        yield (SAMPLE_RATE, _float_to_int16(buffer_audio))
+        print(f"⚠️ WARMUP FAILED: {e}")
 
-    print("--- END UI STREAM ---")
-
-# -------------------------
-# API (FastAPI + WS)
-# -------------------------
+# ----------------------------
+# FASTAPI APP (API ONLY)
+# ----------------------------
 api = FastAPI()
 
-# One inference worker is the right call on 2 vCPU
 INFERENCE_EXECUTOR = ThreadPoolExecutor(max_workers=1)
-INFERENCE_QUEUE = asyncio.Queue()
+INFERENCE_QUEUE: asyncio.Queue = asyncio.Queue(maxsize=64)
+
+@api.get("/health")
+async def health():
+    return {"ok": True, "model": "kokoro", "sample_rate": SAMPLE_RATE, "default_voice": DEFAULT_VOICE_ID}
 
 async def audio_engine_loop():
     print("⚡ API AUDIO PIPELINE STARTED")
     loop = asyncio.get_running_loop()
 
     while True:
-        job = await INFERENCE_QUEUE.get()
-        text, voice_id, speed, ws = job
+        ws, voice_id, speed, text = await INFERENCE_QUEUE.get()
 
-        try:
-            if ws.client_state.value > 1:
-                continue
+        if ws.client_state.value > 1:
+            continue
+
+        frame_q: asyncio.Queue = asyncio.Queue(maxsize=8)
+
+        def _worker():
+            try:
+                first = True
+                started = time.time()
+                for audio in kokoro_audio_iter(text, voice_id, speed):
+                    b = audio_to_pcm_bytes(audio)
+                    loop.call_soon_threadsafe(frame_q.put_nowait, b)
+                    if first:
+                        first = False
+                        dt = time.time() - started
+                        print(f"⚡ first audio ready in {dt:.2f}s")
+                loop.call_soon_threadsafe(frame_q.put_nowait, None)
+            except Exception as e:
+                print(f"API Worker Error: {e}")
+                try:
+                    loop.call_soon_threadsafe(frame_q.put_nowait, None)
+                except Exception:
+                    pass
+
+        INFERENCE_EXECUTOR.submit(_worker)
 
-            # Run synthesis in the single worker thread
-            audio_f32 = await loop.run_in_executor(
-                INFERENCE_EXECUTOR,
-                lambda: synthesize_one_chunk(text, voice_id, speed),
-            )
+        while True:
+            frame = await frame_q.get()
+            if frame is None:
+                break
 
-            if audio_f32 is None or len(audio_f32) == 0:
-                continue
+            if ws.client_state.value > 1:
+                break
 
-            pcm = _float_to_int16(audio_f32).tobytes()
             try:
-                await ws.send_bytes(pcm)
+                await ws.send_bytes(frame)
             except Exception:
-                pass
-
-        except Exception as e:
-            print(f"API Engine Error: {e}")
+                break
 
 @api.on_event("startup")
 async def startup():
+    loop = asyncio.get_running_loop()
+    await loop.run_in_executor(INFERENCE_EXECUTOR, warmup)
     asyncio.create_task(audio_engine_loop())
 
+def resolve_voice(value: str) -> str:
+    if not value:
+        return DEFAULT_VOICE_ID
+
+    if value in VOICE_CHOICES:
+        vid = VOICE_CHOICES[value]
+    else:
+        vid = value.strip()
+
+    if vid not in ALLOWED_VOICE_IDS:
+        return DEFAULT_VOICE_ID
+    return vid
+
 @api.websocket("/ws/audio")
 async def websocket_endpoint(ws: WebSocket):
     await ws.accept()
 
-    voice_id = DEFAULT_VOICE
-    speed = 1.0
-    loop = asyncio.get_running_loop()
+    voice_id = DEFAULT_VOICE_ID  # ✅ default Onyx
+    speed = DEFAULT_SPEED
 
     print(f"✅ Client connected: {ws.client}")
 
@@ -590,54 +549,42 @@ async def websocket_endpoint(ws: WebSocket):
             try:
                 data = await ws.receive_json()
             except WebSocketDisconnect:
+                print("❌ Client disconnected")
                 break
             except Exception:
                 break
 
-            if "config" in data:
-                voice_ui = data.get("voice", DEFAULT_VOICE_UI)
-                voice_id = VOICE_CHOICES.get(voice_ui, DEFAULT_VOICE)
-                speed = float(data.get("speed", speed))
+            is_config = ("config" in data) or (data.get("type") == "config")
+            if is_config:
+                voice_id = resolve_voice(str(data.get("voice", voice_id)))
+                try:
+                    speed = float(data.get("speed", speed))
+                except Exception:
+                    speed = DEFAULT_SPEED
+
+            has_text = ("text" in data) or (data.get("type") == "text")
+            if has_text:
+                raw = data.get("text", "")
+                raw = reduce_name_skips(raw)
+                raw = normalize_text(raw)
+
+                if raw and raw.strip():
+                    try:
+                        INFERENCE_QUEUE.put_nowait((ws, voice_id, speed, raw))
+                    except asyncio.QueueFull:
+                        try:
+                            await ws.send_json({"type": "error", "message": "server_busy"})
+                        except Exception:
+                            pass
+
+            if "flush" in data or data.get("type") == "flush":
+                try:
+                    await ws.send_json({"type": "flushed"})
+                except Exception:
+                    pass
 
-            if "text" in data:
-                raw = data["text"]
-                raw = normalize_names_minimally(raw)
-
-                # First chunk tiny, rest larger, same as UI
-                for chunk in tuned_splitter(raw):
-                    if not chunk.strip():
-                        continue
-                    await INFERENCE_QUEUE.put((chunk, voice_id, speed, ws))
-
-            if "flush" in data:
-                pass
-
-    except Exception as e:
-        print(f"🔥 Critical WS Error: {e}")
     finally:
         heartbeat_task.cancel()
 
-# -------------------------
-# Gradio UI
-# -------------------------
-with gr.Blocks(title="Kokoro TTS") as app:
-    gr.Markdown("## ⚡ Kokoro-82M (Official Pipeline, Low Latency)")
-    with gr.Row():
-        with gr.Column():
-            text_in = gr.Textbox(
-                label="Input Text",
-                lines=4,
-                value="The system is live. Use the UI or connect to /ws/audio.",
-            )
-            voice_in = gr.Dropdown(list(VOICE_CHOICES.keys()), value=DEFAULT_VOICE_UI, label="Voice")
-            speed_in = gr.Slider(0.5, 2.0, value=1.0, label="Speed")
-            btn = gr.Button("Generate", variant="primary")
-        with gr.Column():
-            audio_out = gr.Audio(streaming=True, autoplay=True, label="Audio Stream")
-
-    btn.click(stream_generator, inputs=[text_in, voice_in, speed_in], outputs=[audio_out])
-
-final_app = gr.mount_gradio_app(api, app, path="/")
-
 if __name__ == "__main__":
-    uvicorn.run(final_app, host="0.0.0.0", port=7860)
+    uvicorn.run(api, host="0.0.0.0", port=7860)