Update App.py

App.py

@@ -5,59 +5,34 @@ from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
+from typing import Optional
 
 torch.set_num_threads(int(os.environ.get("TORCH_THREADS", "2")))
 
-HERE = os.path.dirname(os.path.abspath(__file__))
-SECRET = os.environ.get("TTS_SECRET_KEY", "")           # Space Settings -> Secrets me daalo
-MODEL_REPO = "saryps-labs/pocket-tts-hindi"
-LOCAL_DIR = "/home/user/hindi_model"
-VOICE_FILE = os.path.join(HERE, "voice.wav")
-FALLBACK_VOICE = "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
-QUANTIZE = os.environ.get("TTS_QUANTIZE", "0") == "1"
+SECRET = os.environ.get("TTS_SECRET_KEY", "")          # Space Settings -> Secrets me daalo
+DEFAULT_VOICE = os.environ.get("VOICE", "hf_alpha")    # hf_alpha, hf_beta (female) | hm_omega, hm_psi (male)
+ALLOWED_VOICES = {"hf_alpha", "hf_beta", "hm_omega", "hm_psi"}
+DEFAULT_SPEED = float(os.environ.get("SPEED", "1.0"))
 MAX_CHARS = int(os.environ.get("MAX_CHARS", "6000"))
-GAP_SECONDS = 0.22
-
-model = None
-voice_state = None
 SR = 24000
+GAP_SECONDS = 0.15
+
+pipeline = None
 state_msg = "loading"
 last_error = ""
 gen_lock = threading.Lock()
 
 
-def get_config_path() -> str:
-    cfg = os.environ.get("TTS_CONFIG") or os.path.join(LOCAL_DIR, "config.yaml")
-    if not os.path.exists(cfg):
-        from huggingface_hub import snapshot_download
-        snapshot_download(MODEL_REPO, local_dir=LOCAL_DIR)
-        cfg = os.path.join(LOCAL_DIR, "config.yaml")
-    return cfg
-
-
 def load_everything():
-    global model, voice_state, SR, state_msg, last_error
+    global pipeline, state_msg, last_error
     try:
-        from pocket_tts import TTSModel
-        cfg = get_config_path()
-        print("Config:", cfg)
-        try:
-            m = TTSModel.load_model(cfg, temp=0.3, lsd_decode_steps=1, eos_threshold=-1.0)
-        except TypeError:
-            m = TTSModel.load_model(cfg)
-        if QUANTIZE:
-            try:
-                m = torch.quantization.quantize_dynamic(m, {torch.nn.Linear}, dtype=torch.qint8)
-                print("int8 quantization ON")
-            except Exception as e:
-                print("quantize fail, fp32 chal raha he:", e)
-        voice_path = VOICE_FILE if os.path.exists(VOICE_FILE) else FALLBACK_VOICE
-        print("Voice prompt:", voice_path)
-        voice_state = m.get_state_for_audio_prompt(voice_path)
-        SR = int(getattr(m, "sample_rate", 24000))
-        model = m
+        from kokoro import KPipeline
+        p = KPipeline(lang_code="h", repo_id="hexgrad/Kokoro-82M")
+        for _ in p("नमस्ते दोस्तों", voice=DEFAULT_VOICE):   # warm-up
+            break
+        pipeline = p
         state_msg = "ready"
-        print("Model ready, sample rate", SR)
+        print("Kokoro ready")
     except Exception as e:
         state_msg = "error"
         last_error = f"{type(e).__name__}: {e}"
@@ -90,31 +65,32 @@ def authorized(request: Request) -> bool:
 
 
 def split_clean(text: str):
-    """Text ko vakyon me todo aur punctuation/English hata do (model punctuation nahi samajhta)."""
-    text = text.replace("\r", "\n")
-    parts = re.split(r"[\u0964\u0965\.\?\!\n;:]+", text)
-    out = []
-    for p in parts:
-        p = re.sub(r"[^\u0900-\u0963\u0966-\u097F0-9\s]", " ", p)
+    """Text saaf karo, vakyon me todo, phir ~200 chars ke tukdon me jodo (Kokoro ko medium length pasand he)."""
+    text = re.sub(r"[\*#_`~>\[\]\{\}\(\)<>|\\/]", " ", text)
+    text = text.replace("\u0964", ".").replace("\u0965", ".").replace("\r", "\n")
+    pieces = re.split(r"(?<=[\.\?\!])\s+|\n+", text)
+    sents = []
+    for p in pieces:
         p = re.sub(r"\s+", " ", p).strip()
-        if len(p) < 3:
+        if len(p) < 2:
             continue
-        while len(p) > 220:
-            cut = p.rfind(" ", 0, 200)
-            cut = cut if cut > 50 else 200
-            out.append(p[:cut].strip())
+        while len(p) > 300:
+            cut = p.rfind(" ", 0, 260)
+            cut = cut if cut > 80 else 260
+            sents.append(p[:cut].strip())
             p = p[cut:].strip()
         if p:
-            out.append(p)
-    return out
-
-
-def iter_audio(text: str):
-    if hasattr(model, "generate_audio_stream"):
-        for a in model.generate_audio_stream(voice_state, text):
-            yield a
-    else:
-        yield model.generate_audio(voice_state, text)
+            sents.append(p)
+    chunks, cur = [], ""
+    for s in sents:
+        if cur and len(cur) + len(s) + 1 > 200:
+            chunks.append(cur)
+            cur = s
+        else:
+            cur = (cur + " " + s).strip()
+    if cur:
+        chunks.append(cur)
+    return chunks
 
 
 def to_pcm16(a) -> bytes:
@@ -123,21 +99,26 @@ def to_pcm16(a) -> bytes:
     return (arr * 32767.0).astype("<i2").tobytes()
 
 
-def pcm_stream(sentences):
+def pcm_stream(chunks, voice, speed):
     silence = np.zeros(int(GAP_SECONDS * SR), dtype="<i2").tobytes()
-    with gen_lock:
-        for s in sentences:
+    with gen_lock:                        # ek samay me ek hi generation
+        for c in chunks:
             try:
-                for a in iter_audio(s):
-                    yield to_pcm16(a)
+                for item in pipeline(c, voice=voice, speed=speed, split_pattern=r"\n+"):
+                    audio = item.audio if hasattr(item, "audio") else item[2]
+                    if audio is None:
+                        continue
+                    yield to_pcm16(audio)
                 yield silence
             except Exception:
-                print("Sentence fail:", s[:60])
+                print("Chunk fail:", c[:60])
                 traceback.print_exc()
 
 
 class TTSReq(BaseModel):
     text: str
+    voice: Optional[str] = None
+    speed: Optional[float] = None
 
 
 @app.get("/")
@@ -156,15 +137,17 @@ def tts(req: TTSReq, request: Request):
         return JSONResponse({"error": "unauthorized"}, status_code=401)
     if state_msg == "loading":
         return JSONResponse({"error": "model abhi load ho raha he, 1 minute baad try karo"}, status_code=503)
-    if state_msg == "error" or model is None:
+    if state_msg == "error" or pipeline is None:
         return JSONResponse({"error": "model load fail hua: " + last_error}, status_code=500)
     if len(req.text) > MAX_CHARS:
         return JSONResponse({"error": f"text bahut lamba he (max {MAX_CHARS} characters)"}, status_code=413)
-    sentences = split_clean(req.text)
-    if not sentences:
-        return JSONResponse({"error": "text me Hindi (देवनागरी) shabd nahi mile"}, status_code=400)
+    chunks = split_clean(req.text)
+    if not chunks:
+        return JSONResponse({"error": "text me padhne layak kuch nahi mila"}, status_code=400)
+    voice = req.voice if req.voice in ALLOWED_VOICES else DEFAULT_VOICE
+    speed = min(max(req.speed or DEFAULT_SPEED, 0.7), 1.3)
     return StreamingResponse(
-        pcm_stream(sentences),
+        pcm_stream(chunks, voice, speed),
         media_type="application/octet-stream",
         headers={"X-Sample-Rate": str(SR), "Cache-Control": "no-store"},
     )