update tts

src/server/main.py

@@ -69,10 +69,29 @@ class Settings(BaseSettings):
 
 settings = Settings()
 
-# TTS Setup
-tts_repo_id = "ai4bharat/IndicF5"
-tts_model = AutoModel.from_pretrained(tts_repo_id, trust_remote_code=True).to(device)
+# TTS Manager
+class TTSManager:
+    def __init__(self, device_type=device):
+        self.device_type = device_type
+        self.model = None
+        self.repo_id = "ai4bharat/IndicF5"
+
+    async def load(self):
+        logger.info("Loading TTS model IndicF5...")
+        self.model = await asyncio.to_thread(
+            AutoModel.from_pretrained,
+            self.repo_id,
+            trust_remote_code=True
+        )
+        self.model = self.model.to(self.device_type)
+        logger.info("TTS model IndicF5 loaded")
+
+    def synthesize(self, text, ref_audio_path, ref_text):
+        if not self.model:
+            raise ValueError("TTS model not loaded")
+        return self.model(text, ref_audio_path=ref_audio_path, ref_text=ref_text)
 
+# TTS Constants
 EXAMPLES = [
     {
         "audio_name": "KAN_F (Happy)",
@@ -99,7 +118,7 @@ def load_audio_from_url(url: str):
         return sample_rate, audio_data
     raise HTTPException(status_code=500, detail="Failed to load reference audio from URL.")
 
-def synthesize_speech(text: str, ref_audio_name: str, ref_text: str):
+def synthesize_speech(tts_manager: TTSManager, text: str, ref_audio_name: str, ref_text: str):
     ref_audio_url = None
     for example in EXAMPLES:
         if example["audio_name"] == ref_audio_name:
@@ -119,7 +138,7 @@ def synthesize_speech(text: str, ref_audio_name: str, ref_text: str):
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
         sf.write(temp_audio.name, audio_data, samplerate=sample_rate, format='WAV')
         temp_audio.flush()
-        audio = tts_model(text, ref_audio_path=temp_audio.name, ref_text=ref_text)
+        audio = tts_manager.synthesize(text, ref_audio_path=temp_audio.name, ref_text=ref_text)
 
     if audio.dtype == np.int16:
         audio = audio.astype(np.float32) / 32768.0
@@ -233,6 +252,7 @@ class ASRModelManager:
 llm_manager = LLMManager(settings.llm_model_name)
 model_manager = ModelManager()
 asr_manager = ASRModelManager()
+tts_manager = TTSManager()
 ip = IndicProcessor(inference=True)
 
 # Pydantic Models
@@ -278,6 +298,7 @@ async def lifespan(app: FastAPI):
         tasks = [
             asyncio.create_task(llm_manager.load()),
             asyncio.create_task(asr_manager.load()),
+            asyncio.create_task(tts_manager.load()),
             asyncio.create_task(model_manager.load_model('eng_Latn', 'kan_Knda', 'eng_indic')),
             asyncio.create_task(model_manager.load_model('kan_Knda', 'eng_Latn', 'indic_eng')),
             asyncio.create_task(model_manager.load_model('kan_Knda', 'hin_Deva', 'indic_indic')),
@@ -314,11 +335,14 @@ app.state.limiter = limiter
 # API Endpoints
 @app.post("/audio/speech", response_class=StreamingResponse)
 async def synthesize_kannada(request: KannadaSynthesizeRequest):
+    if not tts_manager.model:
+        raise HTTPException(status_code=503, detail="TTS model still loading, please try again later")
     kannada_example = next(ex for ex in EXAMPLES if ex["audio_name"] == "KAN_F (Happy)")
     if not request.text.strip():
         raise HTTPException(status_code=400, detail="Text to synthesize cannot be empty.")
     
     audio_buffer = synthesize_speech(
+        tts_manager,
         text=request.text,
         ref_audio_name="KAN_F (Happy)",
         ref_text=kannada_example["ref_text"]
@@ -610,6 +634,8 @@ async def speech_to_speech(
     file: UploadFile = File(...),
     language: str = Query(..., enum=list(asr_manager.model_language.keys())),
 ) -> StreamingResponse:
+    if not tts_manager.model:
+        raise HTTPException(status_code=503, detail="TTS model still loading, please try again later")
     transcription = await transcribe_audio(file, language)
     logger.info(f"Transcribed text: {transcription.text}")