Handle multiple languages

.gitignore

@@ -2,9 +2,10 @@ venv
 *.wav
 .vscode
 __pycache__
-model.pt
+*.pt
 dist
 *egg-info
 build
 sessions
-samples
+samples
+langs.json

silero_api_server/server.py

@@ -3,7 +3,7 @@
 import os
 from fastapi import FastAPI, Response, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 import uvicorn
@@ -48,6 +48,8 @@ class SampleText(BaseModel):
 class SessionPayload(BaseModel):
     path: Optional[str]
 
+class Language(BaseModel):
+    id: str
 @app.get("/tts/speakers")
 def speakers(request: Request):
     voices = [
@@ -75,18 +77,27 @@ def generate(voice: Voice):
 
 @app.get("/tts/sample")
 def play_sample(speaker: str):
-    return FileResponse(f"{SAMPLE_PATH}/{speaker}.wav")
+    return FileResponse(f"{SAMPLE_PATH}/{speaker}.wav",status_code=200)
 
 @app.post("/tts/generate-samples")
 def generate_samples(sample_text: Optional[str] = ""):
     tts_service.update_sample_text(sample_text)
     tts_service.generate_samples()
     return Response("Generated samples",status_code=200)
 
-@app.post("/tts/init_session")
+@app.post("/tts/session")
 def init_session(sessionPayload: SessionPayload):
     tts_service.init_sessions_path(sessionPayload.path)
     return Response(f"Session path created at {sessionPayload.path}")
 
+@app.get("/tts/language")
+def get_languages():
+    return JSONResponse(list(tts_service.langs.keys()),status_code=200)
+
+@app.post("/tts/language")
+def set_language(language: Language):
+    tts_service.load_model(language.id)
+    return Response(status_code=200)
+
 if __name__ == "__main__":
     uvicorn.run(app,host="0.0.0.0",port=8001)

silero_api_server/tts.py

@@ -1,56 +1,71 @@
 # V3
 import os, time
 import shutil
+import requests
 import torch
 import torch.package
 import torchaudio
 from hashlib import md5
 from loguru import logger
 from pydub import AudioSegment
 from pathlib import Path
+import json
 
 class SileroTtsService:
     """
     Generate TTS wav files using Silero
     """
-    def __init__(self, sample_path) -> None:
+    def __init__(self, sample_path, lang="v3_en.pt") -> None:
         self.sample_text = "The fallowed fallen swindle auspacious goats in portable power stations."
         self.sample_path = Path(sample_path)
         self.sessions_path = None
+
         # Silero works fine on CPU
         self.device = torch.device('cpu')
         torch.set_num_threads(4)
         torchaudio.set_audio_backend("soundfile")
 
-        # Make sure we  have the model
-        self.local_file = Path('model.pt')
-        if not Path.is_file(self.local_file):
-            logger.warning(f"First run, downloading Silero model. This could take some time...") 
-            torch.hub.download_url_to_file('https://models.silero.ai/models/tts/en/v3_en.pt',
-                                        self.local_file)  
-            logger.info(f"Model download completed.")
-
-
         # Make sure we have the sample path
         if not self.sample_path.exists():
-            self.sample_path.mkdir()        
-
-        self.model = torch.package.PackageImporter(self.local_file).load_pickle("tts_models", "model")
-        self.model.to(self.device)
+            self.sample_path.mkdir()   
 
         self.sample_rate = 48000 
         logger.info(f"TTS Service loaded successfully")
 
         # Prevent generation failure due to too long input
         self.max_char_length = 600
 
+        # Get language model URLs
+        self.langs = self.list_languages()
+
+        # Load model
+        self.load_model(lang)
+
     def init_sessions_path(self, sessions_path="sessions"):
         self.sessions_path = Path(sessions_path)
         if not self.sessions_path.exists():
             self.sessions_path.mkdir()
+
+    def load_model(self, lang_model="v3_en.pt"):
+        # Download the model. Default to en.
+        if lang_model not in self.langs:
+            raise Exception(f"{lang_model} not in {list(self.langs.values())}")
+
+        model_url = self.langs[lang_model]
+        self.model_file = Path(lang_model)
+
+        if not Path.is_file(self.model_file):
+            logger.warning(f"Downloading Silero {lang_model} model...") 
+            torch.hub.download_url_to_file(model_url,
+                                        self.model_file)  
+            logger.info(f"Model download completed.")
+
+        self.model = torch.package.PackageImporter(self.model_file).load_pickle("tts_models", "model")
+        self.model.to(self.device)
 
     def generate(self, speaker, text, session=""):
         if len(text) > self.max_char_length:
+            # Handle long text input
             text_chunks = self.split_text(text)
             combined_wav = AudioSegment.empty()
 
@@ -124,3 +139,32 @@ def update_sample_text(self,text: str):
         self.sample_text = text
         logger.info(f"Sample text updated to {self.sample_text}")  
 
+    def list_languages(self):
+        'Grab all v3 model links from https://models.silero.ai/models/tts'
+        lang_file = Path('langs.json')
+        if lang_file.exists():
+            with lang_file.open('r') as fh:
+                logger.info('Loading cached language index')
+                return json.load(fh)
+        logger.info('Loading remote language index')
+        lang_base_url = 'https://models.silero.ai/models/tts'
+        lang_urls = {}
+
+        # Parse initial web directory for languages
+        response = requests.get(lang_base_url)
+        langs = [lang.split('/')[0] for lang in response.text.split('<a href="')][1:]
+
+        # Enter each web directory and grab v3 model file links
+        for lang in langs:
+            response = requests.get(f"{lang_base_url}/{lang}")
+            if not response.ok:
+                raise f"Failed to get languages: {response.status_code}"
+            lang_files = [f.split('"')[0] for f in response.text.split('<a href="')][1:]
+
+            # If a valid v3 file, add to list
+            for lang_file in lang_files:
+                if lang_file.startswith('v3'):
+                    lang_urls[lang_file]=f"{lang_base_url}/{lang}/{lang_file}"
+        with open('langs.json','w') as fh:
+            json.dump(lang_urls,fh)
+        return lang_urls

test/silero_test.py

@@ -0,0 +1,15 @@
+import unittest
+from silero_api_server.tts import SileroTtsService
+
+class SileroApiServerTest(unittest.TestCase):
+    def test_list_languages(self):
+        svc = SileroTtsService('samples')
+        lang = svc.list_languages()
+        self.assertIn('https://models.silero.ai/models/tts/es/v3_es.pt',lang.values())
+
+    def test_init_model(self):
+        svc = SileroTtsService('samples')
+        lang = svc.list_languages()
+        svc.load_model(list(lang.keys())[0])
+        model_file = svc.model_file
+        self.assertTrue(model_file.exists())