#8 Move synthesis to package

app.py

@@ -1,15 +1,8 @@
 import tempfile
-
 import gradio as gr
-
-from TTS.utils.synthesizer import Synthesizer
-import requests
-from os.path import exists
-from ukrainian_tts.formatter import preprocess_text
 from datetime import datetime
 from enum import Enum
-import torch
-
+from ukrainian_tts.tts import TTS
 
 class StressOption(Enum):
     AutomaticStress = "Автоматичні наголоси (за словником) 📖"
@@ -24,44 +17,11 @@ class VoiceOption(Enum):
     Olga = "Ольга (жіночий) 👩"
 
 
-def download(url, file_name):
-    if not exists(file_name):
-        print(f"Downloading {file_name}")
-        r = requests.get(url, allow_redirects=True)
-        with open(file_name, "wb") as file:
-            file.write(r.content)
-    else:
-        print(f"Found {file_name}. Skipping download...")
-
-
-print("downloading uk/mykyta/vits-tts")
-release_number = "v3.0.0"
-model_link = f"https://github.com/robinhad/ukrainian-tts/releases/download/{release_number}/model-inference.pth"
-config_link = f"https://github.com/robinhad/ukrainian-tts/releases/download/{release_number}/config.json"
-speakers_link = f"https://github.com/robinhad/ukrainian-tts/releases/download/{release_number}/speakers.pth"
-
-model_path = "model.pth"
-config_path = "config.json"
-speakers_path = "speakers.pth"
-
-download(model_link, model_path)
-download(config_link, config_path)
-download(speakers_link, speakers_path)
-
 badge = (
     "https://visitor-badge-reloaded.herokuapp.com/badge?page_id=robinhad.ukrainian-tts"
 )
 
-synthesizer = Synthesizer(
-    model_path,
-    config_path,
-    speakers_path,
-    None,
-    None,
-)
-
-if synthesizer is None:
-    raise NameError("model not found")
+ukr_tts = TTS()
 
 
 def tts(text: str, voice: str, stress: str):
@@ -81,17 +41,15 @@ def tts(text: str, voice: str, stress: str):
         VoiceOption.Olga.value: "olga",
     }
     speaker_name = voice_mapping[voice]
-    text = preprocess_text(text, autostress_with_model)
+
     text_limit = 7200
     text = (
         text if len(text) < text_limit else text[0:text_limit]
     )  # mitigate crashes on hf space
-    print("Converted:", text)
+    
 
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
-        with torch.no_grad():
-            wavs = synthesizer.tts(text, speaker_name=speaker_name)
-            synthesizer.save_wav(wavs, fp)
+        ukr_tts.tts(text, speaker_name, autostress_with_model, fp)
         return fp.name, text
 
 

ukrainian_tts/tts.py

@@ -0,0 +1,82 @@
+from io import BytesIO
+import requests
+from os.path import exists
+from TTS.utils.synthesizer import Synthesizer
+from enum import Enum
+from .formatter import preprocess_text
+from torch import no_grad
+
+class Voices(Enum):
+    Olena = "olena"
+    Mykyta = "mykyta"
+    Lada = "lada"
+    Dmytro = "dmytro"
+    Olga = "olga"
+
+
+class StressOption(Enum):
+    Dictionary = "dictionary"
+    Model = "model"
+
+
+class TTS:
+    def __init__(self, cache_folder=None) -> None:
+        self.__setup_cache(cache_folder)
+
+
+    def tts(self, text: str, voice: str, stress: str, output_fp=BytesIO()):
+        autostress_with_model = (
+            True if stress == StressOption.Model.value else False
+        )
+
+        if voice not in [option.value for option in Voices]:
+            raise ValueError("Invalid value for voice selected! Please use one of the following values: {', '.join([option.value for option in Voices])}.")
+
+        text = preprocess_text(text, autostress_with_model)
+
+        with no_grad():
+            wavs = self.synthesizer.tts(text, speaker_name=voice)
+            self.synthesizer.save_wav(wavs, output_fp)
+
+        output_fp.seek(0)
+
+        return output_fp
+
+
+    def __setup_cache(self, cache_folder=None):
+        print("downloading uk/mykyta/vits-tts")
+        release_number = "v3.0.0"
+        model_link = f"https://github.com/robinhad/ukrainian-tts/releases/download/{release_number}/model-inference.pth"
+        config_link = f"https://github.com/robinhad/ukrainian-tts/releases/download/{release_number}/config.json"
+        speakers_link = f"https://github.com/robinhad/ukrainian-tts/releases/download/{release_number}/speakers.pth"
+
+        model_path = "model.pth"
+        config_path = "config.json"
+        speakers_path = "speakers.pth"
+
+        self.__download(model_link, model_path)
+        self.__download(config_link, config_path)
+        self.__download(speakers_link, speakers_path)
+
+        self.synthesizer = Synthesizer(
+            model_path,
+            config_path,
+            speakers_path,
+            None,
+            None,
+        )
+
+        if self.synthesizer is None:
+            raise NameError("model not found")
+
+
+    def __download(self, url, file_name):
+        if not exists(file_name):
+            print(f"Downloading {file_name}")
+            r = requests.get(url, allow_redirects=True)
+            with open(file_name, "wb") as file:
+                file.write(r.content)
+        else:
+            print(f"Found {file_name}. Skipping download...")
+
+