Update app/main.py

app/main.py

@@ -1,274 +1,129 @@
-import io
-import re
-import os
-import gc
-import time
-import logging
-import asyncio
-import numpy as np
-import torch
-import soundfile as sf
-from typing import Optional, List, AsyncGenerator
-from contextlib import asynccontextmanager
+# app/main.py
+
+import json
+import httpx
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse, Response, JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, Response
 from pydantic import BaseModel, Field
-from transformers import AutoModelForCausalLM
-from yarngpt.audiotokenizer import AudioTokenizerV2
-from concurrent.futures import ThreadPoolExecutor
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
-log = logging.getLogger("yarngpt-tts")
-
-BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-if os.path.exists("/data"):
-    CACHE_DIR = "/data/.cache"
-    log.info("Using persistent storage at /data for cache")
-else:
-    CACHE_DIR = os.path.join(BASE_DIR, ".cache")
-    log.info(f"Using local cache directory: {CACHE_DIR}")
-
-os.environ['HF_HOME'] = CACHE_DIR
-os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'huggingface')
-os.environ['TORCH_HOME'] = os.path.join(CACHE_DIR, 'torch')
-os.environ['HUGGINGFACE_HUB_CACHE'] = os.path.join(CACHE_DIR, 'huggingface')
-
-MODELS_DIR = "/app/models"
-os.makedirs(MODELS_DIR, exist_ok=True)
-os.makedirs(CACHE_DIR, exist_ok=True)
-
-MODEL_ID = os.getenv("MODEL_ID", "saheedniyi/YarnGPT2b")
-WAV_TOKENIZER_CONFIG = os.getenv("WAV_TOKENIZER_CONFIG", "/app/models/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml")
-WAV_TOKENIZER_CKPT = os.getenv("WAV_TOKENIZER_CKPT", "/app/models/wavtokenizer_large_speech_320_24k.ckpt")
-SAMPLE_RATE = int(os.getenv("SAMPLE_RATE", "24000"))
-WORD_LIMIT = int(os.getenv("WORD_LIMIT", "25"))
-MAX_TEXT_LENGTH = int(os.getenv("MAX_TEXT_LENGTH", "30000"))
-GENERATION_TEMP = float(os.getenv("GENERATION_TEMP", "0.1"))
-REPEAT_PENALTY = float(os.getenv("REPEAT_PENALTY", "1.1"))
-MAX_GEN_LENGTH = int(os.getenv("MAX_GEN_LENGTH", "4000"))
-SILENCE_TOKEN = int(os.getenv("SILENCE_TOKEN", "453"))
-SILENCE_FRAMES = int(os.getenv("SILENCE_FRAMES", "38"))
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE = torch.bfloat16 if DEVICE == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
-log.info(f"Device: {DEVICE} | dtype: {DTYPE}")
-
-DEFAULT_LANGUAGE = "english"
-DEFAULT_SPEAKER = "jude"
-SUPPORTED_LANGUAGES = ["english", "yoruba", "igbo", "hausa", "pidgin"]
-SUPPORTED_SPEAKERS = ["idera", "jude", "tayo", "zainab", "chisom", "regina", "umar", "emma", "osagie", "amara"]
 
-audio_tokenizer: Optional[AudioTokenizerV2] = None
-model_lm: Optional[AutoModelForCausalLM] = None
-executor = ThreadPoolExecutor(max_workers=4)
-semaphore = asyncio.Semaphore(600)
-request_queue = asyncio.Queue(maxsize=1000)
+CONFIG_PATH = "app/config.json"
 
-class TTSRequest(BaseModel):
-    text: str = Field(..., min_length=1, max_length=MAX_TEXT_LENGTH)
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    global audio_tokenizer, model_lm
-    log.info("Loading YarnGPT2b …")
-    t0 = time.time()
-    
-    if not os.path.exists(WAV_TOKENIZER_CONFIG):
-        raise RuntimeError(f"Model config not found at: {WAV_TOKENIZER_CONFIG}")
-    if not os.path.exists(WAV_TOKENIZER_CKPT):
-        raise RuntimeError(f"Model checkpoint not found at: {WAV_TOKENIZER_CKPT}")
-    
-    audio_tokenizer = AudioTokenizerV2(MODEL_ID, WAV_TOKENIZER_CKPT, WAV_TOKENIZER_CONFIG)
-    
-    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-    torch.cuda.set_device(0)
-    
-    model_lm = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID, 
-        torch_dtype=DTYPE, 
-        device_map=None,
-        low_cpu_mem_usage=True
-    ).to(audio_tokenizer.device)
-    
-    model_lm.eval()
-    
-    if DEVICE == "cuda":
-        torch.backends.cudnn.benchmark = True
-    
-    log.info(f"✓ Model loaded in {time.time()-t0:.1f}s")
-    asyncio.create_task(queue_processor())
-    yield
-    del model_lm, audio_tokenizer
-    gc.collect()
-    if DEVICE == "cuda":
-        torch.cuda.empty_cache()
-    log.info("Model unloaded.")
 
-app = FastAPI(title="YarnGPT2b TTS Service", description="Nigerian-accented Text-to-Speech via YarnGPT2b", version="2.0.0", lifespan=lifespan)
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+def load_config() -> dict:
+    with open(CONFIG_PATH, "r") as f:
+        return json.load(f)
 
-def split_text_into_chunks(text: str, word_limit: int = WORD_LIMIT) -> List[str]:
-    text = re.sub(r"\s+", " ", text.strip())
-    text = text.replace("...", ".")
-    raw_sentences = re.split(r'(?<=[.!?])\s+', text)
-    chunks: List[str] = []
-    for sentence in raw_sentences:
-        sentence = sentence.strip()
-        if not sentence:
-            continue
-        chunks.append(".")
-        words = sentence.split()
-        for i in range(0, len(words), word_limit):
-            chunks.append(" ".join(words[i:i + word_limit]))
-    return chunks
 
-def speed_change_np(audio: np.ndarray, speed: float) -> np.ndarray:
-    if speed == 1.0:
-        return audio
-    original_len = len(audio)
-    target_len = int(original_len / speed)
-    indices = np.linspace(0, original_len - 1, target_len)
-    return np.interp(indices, np.arange(original_len), audio).astype(np.float32)
+config = load_config()
 
-async def generate_codes_for_chunk_async(chunk: str, language: str, speaker: str, temperature: float, repetition_penalty: float) -> List[int]:
-    loop = asyncio.get_event_loop()
-    def _generate():
-        prompt = audio_tokenizer.create_prompt(chunk, lang=language, speaker_name=speaker)
-        input_ids = audio_tokenizer.tokenize_prompt(prompt)
-        if isinstance(input_ids, torch.Tensor):
-            input_ids = input_ids.to(audio_tokenizer.device)
-        with torch.inference_mode():
-            output = model_lm.generate(
-                input_ids=input_ids, 
-                temperature=temperature, 
-                repetition_penalty=repetition_penalty, 
-                max_length=MAX_GEN_LENGTH, 
-                do_sample=temperature > 0,
-                pad_token_id=model_lm.config.eos_token_id,
-                attention_mask=torch.ones_like(input_ids).to(audio_tokenizer.device)
-            )
-        return audio_tokenizer.get_codes(output)
-    return await loop.run_in_executor(executor, _generate)
+EL_API_KEY = config["elevenlabs"]["api_key"]
+VOICE_ID = config["elevenlabs"]["voice_id"]
+MODEL_ID = config["elevenlabs"]["model_id"]
 
-def codes_to_audio_np(codes: List[int]) -> np.ndarray:
-    audio_tensor = audio_tokenizer.get_audio(codes)
-    return audio_tensor.squeeze().cpu().float().numpy()
+ELEVENLABS_STREAM_URL = (
+    f"https://api.elevenlabs.io/v1/text-to-speech/{VOICE_ID}/stream"
+)
 
-def make_wav_header(sample_rate: int, channels: int = 1, bits: int = 16) -> bytes:
-    byte_rate = sample_rate * channels * bits // 8
-    block_align = channels * bits // 8
-    return (b"RIFF" + b"\xff\xff\xff\xff" + b"WAVEfmt " + (16).to_bytes(4, "little") + (1).to_bytes(2, "little") + channels.to_bytes(2, "little") + sample_rate.to_bytes(4, "little") + byte_rate.to_bytes(4, "little") + block_align.to_bytes(2, "little") + bits.to_bytes(2, "little") + b"data" + b"\xff\xff\xff\xff")
+HEADERS = {
+    "xi-api-key": EL_API_KEY,
+    "Content-Type": "application/json",
+    "Accept": "audio/mpeg",
+}
 
-SILENCE_AUDIO = np.zeros(int(SAMPLE_RATE * 0.5), dtype=np.float32)
+app = FastAPI(
+    title="Production TTS Service",
+    version="1.0.2",
+    docs_url="/docs",
+    redoc_url="/redoc",
+)
 
-async def queue_processor():
-    while True:
-        try:
-            future, req_data = await asyncio.wait_for(request_queue.get(), timeout=1.0)
-            try:
-                result = await process_request_async(req_data)
-                future.set_result(result)
-            except Exception as e:
-                future.set_exception(e)
-            finally:
-                request_queue.task_done()
-        except asyncio.TimeoutError:
-            continue
-        except Exception as e:
-            log.error(f"Queue processor error: {e}")
 
-async def process_request_async(req: TTSRequest) -> bytes:
-    chunks = split_text_into_chunks(req.text, WORD_LIMIT)
-    all_codes = []
-    for chunk in chunks:
-        if chunk == ".":
-            all_codes.extend([SILENCE_TOKEN] * SILENCE_FRAMES)
-        else:
-            try:
-                codes = await generate_codes_for_chunk_async(chunk, DEFAULT_LANGUAGE, DEFAULT_SPEAKER, GENERATION_TEMP, REPEAT_PENALTY)
-                all_codes.extend(codes)
-            except Exception as e:
-                log.error(f"Chunk error: {e}")
-                all_codes.extend([SILENCE_TOKEN] * SILENCE_FRAMES)
-    audio_np = codes_to_audio_np(all_codes)
-    audio_np = speed_change_np(audio_np, 1.0)
-    buf = io.BytesIO()
-    sf.write(buf, audio_np, SAMPLE_RATE, format="WAV", subtype="PCM_16")
-    return buf.getvalue()
+class TTSRequest(BaseModel):
+    text: str = Field(..., min_length=1, max_length=5000)
+    stability: float = Field(0.5, ge=0.0, le=1.0)
+    similarity_boost: float = Field(0.5, ge=0.0, le=1.0)
 
-async def audio_stream_generator(req: TTSRequest) -> AsyncGenerator[bytes, None]:
-    chunks = split_text_into_chunks(req.text, WORD_LIMIT)
-    log.info(f"Streaming {len(chunks)} chunks")
-    yield make_wav_header(SAMPLE_RATE)
-    silence_bytes = (SILENCE_AUDIO * 32767).astype(np.int16).tobytes()
-    for chunk in chunks:
-        if chunk == ".":
-            yield silence_bytes
-            continue
-        try:
-            codes = await generate_codes_for_chunk_async(chunk, DEFAULT_LANGUAGE, DEFAULT_SPEAKER, GENERATION_TEMP, REPEAT_PENALTY)
-            audio_np = codes_to_audio_np(codes)
-            audio_np = speed_change_np(audio_np, 1.0)
-            audio_i16 = (audio_np * 32767).astype(np.int16)
-            yield audio_i16.tobytes()
-        except Exception as e:
-            log.error(f"Stream chunk error: {e}")
-            yield silence_bytes
 
-@app.get("/health", tags=["Meta"])
+@app.get("/health")
 async def health():
-    gpu_info = {}
-    if DEVICE == "cuda":
-        gpu_info = {"gpu_name": torch.cuda.get_device_name(0), "vram_total_gb": round(torch.cuda.get_device_properties(0).total_memory / 1e9, 2), "vram_reserved_gb": round(torch.cuda.memory_reserved(0) / 1e9, 2), "vram_used_gb": round(torch.cuda.memory_allocated(0) / 1e9, 2)}
-    models_status = {"config_exists": os.path.exists(WAV_TOKENIZER_CONFIG), "ckpt_exists": os.path.exists(WAV_TOKENIZER_CKPT), "models_dir": MODELS_DIR, "cache_dir": CACHE_DIR}
-    return {"status": "ok", "model": MODEL_ID, "device": DEVICE, "dtype": str(DTYPE), "sample_rate": SAMPLE_RATE, "queue_size": request_queue.qsize(), "models": models_status, **gpu_info}
-
-@app.get("/voices", tags=["Meta"])
-async def list_voices():
-    return {"languages": SUPPORTED_LANGUAGES, "speakers": SUPPORTED_SPEAKERS}
-
-@app.post("/tts/stream", tags=["TTS"])
-async def stream_tts(payload: TTSRequest):
-    if model_lm is None:
-        raise HTTPException(503, "Model not loaded yet.")
-    async with semaphore:
-        return StreamingResponse(audio_stream_generator(payload), media_type="audio/wav", headers={"Content-Disposition": "inline; filename=stream.wav"})
-
-@app.post("/tts/buffered", tags=["TTS"])
-async def buffered_tts(payload: TTSRequest):
-    if model_lm is None:
-        raise HTTPException(503, "Model not loaded yet.")
-    async with semaphore:
-        loop = asyncio.get_event_loop()
-        future = loop.create_future()
-        try:
-            await asyncio.wait_for(request_queue.put((future, payload)), timeout=10000.0)
-            audio_data = await asyncio.wait_for(future, timeout=120000.0)
-            duration_s = len(audio_data) / (SAMPLE_RATE * 2)
-            return Response(content=audio_data, media_type="audio/wav", headers={"Content-Disposition": "attachment; filename=tts.wav", "X-Audio-Duration-Sec": str(round(duration_s, 2))})
-        except asyncio.TimeoutError:
-            raise HTTPException(503, "Service busy, please try again later")
-
-@app.post("/tts/codes", tags=["TTS (Advanced)"])
-async def get_codes(payload: TTSRequest):
-    if model_lm is None:
-        raise HTTPException(503, "Model not loaded yet.")
-    async with semaphore:
-        chunks = split_text_into_chunks(payload.text, WORD_LIMIT)
-        all_codes = []
-        for chunk in chunks:
-            if chunk == ".":
-                all_codes.extend([SILENCE_TOKEN] * SILENCE_FRAMES)
-            else:
-                try:
-                    codes = await generate_codes_for_chunk_async(chunk, DEFAULT_LANGUAGE, DEFAULT_SPEAKER, GENERATION_TEMP, REPEAT_PENALTY)
-                    all_codes.extend(codes)
-                except Exception as e:
-                    log.error(f"Code generation failed: {e}")
-                    all_codes.extend([SILENCE_TOKEN] * SILENCE_FRAMES)
-        return JSONResponse({"codes": all_codes, "total_tokens": len(all_codes)})
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("app.main:app", host="0.0.0.0", port=7860, reload=False, workers=1, limit_concurrency=1000, backlog=2048, timeout_keep_alive=5)
+    return {"status": "ok"}
+
+
+@app.post("/tts")
+async def text_to_speech(payload: TTSRequest):
+    body = {
+        "text": payload.text,
+        "model_id": MODEL_ID,
+        "voice_settings": {
+            "stability": payload.stability,
+            "similarity_boost": payload.similarity_boost,
+        },
+    }
+
+    async def audio_stream():
+        async with httpx.AsyncClient(timeout=None) as client:
+            async with client.stream(
+                method="POST",
+                url=ELEVENLABS_STREAM_URL,
+                headers=HEADERS,
+                json=body,
+            ) as response:
+
+                if response.status_code != 200:
+                    error = await response.aread()
+                    raise HTTPException(
+                        status_code=502,
+                        detail=error.decode(),
+                    )
+
+                async for chunk in response.aiter_bytes():
+                    yield chunk
+
+    return StreamingResponse(
+        audio_stream(),
+        media_type="audio/mpeg",
+        headers={
+            "Content-Disposition": "inline; filename=tts.mp3",
+        },
+    )
+
+
+
+@app.post("/tts/buffered")
+async def text_to_speech_buffered(payload: TTSRequest):
+    body = {
+        "text": payload.text,
+        "model_id": MODEL_ID,
+        "voice_settings": {
+            "stability": payload.stability,
+            "similarity_boost": payload.similarity_boost,
+        },
+    }
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        response = await client.post(
+            ELEVENLABS_STREAM_URL,
+            headers=HEADERS,
+            json=body,
+        )
+
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=502,
+            detail=response.text,
+        )
+
+    if not response.content:
+        raise HTTPException(
+            status_code=500,
+            detail="Received empty audio buffer",
+        )
+
+    return Response(
+        content=response.content,
+        media_type="audio/mpeg",
+        headers={
+            "Content-Disposition": "attachment; filename=tts.mp3",
+            "Content-Length": str(len(response.content)),
+        },
+    )