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]) # 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")