Upload server.py with huggingface_hub

server.py

@@ -0,0 +1,267 @@
+import logging
+import os
+import io
+import numpy as np
+import torch
+import soundfile as sf
+import uuid
+import time
+from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional
+
+from omnivoice import OmniVoice, OmniVoiceGenerationConfig
+from text_preprocessor import chunk_text
+
+logging.basicConfig(
+    level=logging.WARNING,
+    format="%(asctime)s %(name)s %(levelname)s: %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+# FastAPI app
+app = FastAPI(title="Arabic TTS Server (OmniVoice)")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[
+        "https://arabic-tts-frontend.web.app",
+        "https://arabic-tts-frontend.firebaseapp.com",
+        "http://localhost:3000",
+        "http://localhost:8000"
+    ],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global variables for model
+CHECKPOINT = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice")
+model = None
+sampling_rate = 24000
+
+# Simple In-Memory Database for State Tracking (Step 4 preview)
+tasks_db = {}
+
+@app.on_event("startup")
+async def startup_event():
+    global model, sampling_rate
+    print(f"Loading OmniVoice model from {CHECKPOINT} ...")
+    model = OmniVoice.from_pretrained(
+        CHECKPOINT,
+        load_asr=True,
+        device_map="cpu",  # Using CPU by default, adjust if a GPU is available.
+    )
+    sampling_rate = model.sampling_rate
+    print("Model loaded successfully!")
+
+class SynthesizeRequest(BaseModel):
+    text: str
+    voice: Optional[str] = "Auto"
+    speed: Optional[float] = 1.0
+
+import os
+import shutil
+
+def process_audio_task(task_id: str, chunks: list[str], speed: float, voice_id: str):
+    """
+    Background worker that iterivately generates audio, saves chunks to disk,
+    concatenates them, and handles cleanup. Validates consistent voice!
+    """
+    try:
+        total_chunks = len(chunks)
+        tasks_db[task_id]["status"] = "processing"
+        tasks_db[task_id]["total_chunks"] = total_chunks
+        
+        chunk_dir = os.path.join("audio_chunks", task_id)
+        os.makedirs(chunk_dir, exist_ok=True)
+        
+        gen_config = OmniVoiceGenerationConfig(
+            num_step=32,
+            guidance_scale=2.0,
+            denoise=True,
+            preprocess_prompt=False,
+            postprocess_output=False,
+        )
+
+        master_voice_prompt = None
+
+        # Check if user requested a specific built-in voice
+        if voice_id and voice_id != "Auto":
+            voice_path = os.path.join("voices", f"{voice_id}.wav")
+            text_path = os.path.join("voices", f"{voice_id}.txt")
+            if os.path.exists(voice_path):
+                ref_text = None
+                if os.path.exists(text_path):
+                    with open(text_path, "r", encoding="utf-8") as f:
+                        ref_text = f.read().strip()
+                try:
+                    master_voice_prompt = model.create_voice_clone_prompt(ref_audio=voice_path, ref_text=ref_text)
+                except Exception as e:
+                    logger.warning(f"Voice clone setup failed: {e}")
+
+        for i, chunk in enumerate(chunks):
+            # Update state tracker
+            tasks_db[task_id]["current_chunk"] = i + 1
+            tasks_db[task_id]["progress"] = int(((i) / total_chunks) * 100)
+            
+            chunk_path = os.path.join(chunk_dir, f"chunk_{i}.wav")
+            
+            # Check if already generated for resume ability
+            if not os.path.exists(chunk_path):
+                kw = dict(
+                    text=chunk,
+                    language="Auto",
+                    generation_config=gen_config
+                )
+                if speed is not None and speed != 1.0:
+                    kw["speed"] = speed
+
+                # Apply consistent voice cloning (prevents mid-book gender switching)
+                if master_voice_prompt is not None:
+                    kw["voice_clone_prompt"] = master_voice_prompt
+
+                # Generate Audio via OmniVoice
+                audio = model.generate(**kw)
+                waveform = audio[0].squeeze(0).numpy()
+                
+                # Save chunk to disk incrementally
+                sf.write(chunk_path, waveform, sampling_rate, format='wav', subtype='PCM_16')
+
+                # If Auto mode, use the FIRST successfully generated chunk as the reference 
+                # voice for ALL subsequent chunks. This locks the randomly chosen voice!
+                if master_voice_prompt is None and i == 0:
+                    try:
+                        master_voice_prompt = model.create_voice_clone_prompt(ref_audio=chunk_path, ref_text=chunk)
+                    except Exception as e:
+                        logger.warning(f"Could not extract voice clone from chunk 0: {e}")
+            
+        # All chunks generated, now concatenate
+        tasks_db[task_id]["status"] = "stitching"
+        
+        all_data = []
+        sr = sampling_rate
+        for i in range(total_chunks):
+            chunk_path = os.path.join(chunk_dir, f"chunk_{i}.wav")
+            data, sr = sf.read(chunk_path)
+            all_data.append(data)
+            
+        final_waveform = np.concatenate(all_data)
+        
+        final_dir = os.path.join("static", "audio")
+        os.makedirs(final_dir, exist_ok=True)
+        final_path = os.path.join(final_dir, f"{task_id}.wav")
+        
+        sf.write(final_path, final_waveform, sr, format='wav', subtype='PCM_16')
+        
+        # Cleanup temporary chunks
+        shutil.rmtree(chunk_dir)
+            
+        tasks_db[task_id]["status"] = "completed"
+        tasks_db[task_id]["progress"] = 100
+        tasks_db[task_id]["download_url"] = f"audio/{task_id}.wav"
+        
+    except Exception as e:
+        logger.error(f"Background task failed: {str(e)}")
+        tasks_db[task_id]["status"] = "failed"
+        tasks_db[task_id]["error"] = str(e)
+
+@app.post("/synthesize")
+async def synthesize(req: SynthesizeRequest, background_tasks: BackgroundTasks):
+    if not model:
+        raise HTTPException(status_code=500, detail="Model not loaded yet.")
+
+    try:
+        # Step 1 Integration: chunk the text
+        chunks = chunk_text(req.text.strip())
+        if not chunks:
+             raise HTTPException(status_code=400, detail="Text is empty or invalid.")
+
+        task_id = str(uuid.uuid4())
+        
+        # Step 4 preview: Initialize state tracking
+        tasks_db[task_id] = {
+            "task_id": task_id,
+            "status": "pending",
+            "progress": 0,
+            "current_chunk": 0,
+            "total_chunks": len(chunks)
+        }
+
+        # Step 2 Integration: Start the background process instead of blocking
+        background_tasks.add_task(process_audio_task, task_id, chunks, req.speed, req.voice)
+        
+        return JSONResponse(content={"task_id": task_id, "message": "Audio generation started in the background."})
+
+    except Exception as e:
+        logger.error(f"Synthesis failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/status/{task_id}")
+async def get_status(task_id: str):
+    if task_id not in tasks_db:
+         raise HTTPException(status_code=404, detail="Task not found.")
+    return JSONResponse(content=tasks_db[task_id])
+# ── Distributed Worker: /synthesize_chunk endpoint ───────────────────────────
+class ChunkRequest(BaseModel):
+    text: str
+    voice: Optional[str] = "Auto"
+    speed: Optional[float] = 1.0
+    chunk_index: int = 0
+
+@app.post("/synthesize_chunk")
+async def synthesize_chunk(req: ChunkRequest):
+    """
+    Lightweight endpoint for distributed workers.
+    Receives a single text chunk, generates audio, and returns it as bytes.
+    The Orchestrator calls this on multiple Workers simultaneously.
+    """
+    if not model:
+        raise HTTPException(status_code=500, detail="Model not loaded.")
+    
+    gen_config = OmniVoiceGenerationConfig(
+        num_step=32,
+        guidance_scale=2.0,
+        denoise=True,
+        preprocess_prompt=False,
+        postprocess_output=False,
+    )
+    
+    kw = dict(text=req.text.strip(), language="Auto", generation_config=gen_config)
+    if req.speed != 1.0:
+        kw["speed"] = req.speed
+    
+    # Voice selection
+    if req.voice and req.voice != "Auto":
+        voice_path = os.path.join("voices", f"{req.voice}.wav")
+        if os.path.exists(voice_path):
+            try:
+                kw["voice_clone_prompt"] = model.create_voice_clone_prompt(ref_audio=voice_path)
+            except Exception as e:
+                logger.warning(f"Voice clone failed: {e}")
+    
+    try:
+        audio = model.generate(**kw)
+        waveform = audio[0].squeeze(0).numpy()
+        
+        # Return the waveform as bytes in a WAV container
+        buffer = io.BytesIO()
+        sf.write(buffer, waveform, sampling_rate, format='wav', subtype='PCM_16')
+        buffer.seek(0)
+        
+        headers = {"X-Chunk-Index": str(req.chunk_index)}
+        return StreamingResponse(buffer, media_type="audio/wav", headers=headers)
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Ensure voices directory is explicitly available
+os.makedirs("voices", exist_ok=True)
+app.mount("/voices", StaticFiles(directory="voices"), name="voices")
+
+# Mount static files directly on root
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
+