Update app.py

app.py

@@ -4,20 +4,27 @@ from TTS.utils.synthesizer import Synthesizer
 import gradio as gr
 import tempfile
 
-# Variables
 MAX_TXT_LEN = 800
-MODEL_DIR = "kbd-vits-tts-male"
-MODEL_URL = "https://huggingface.co/anzorq/kbd-vits-tts-male/resolve/main/checkpoint_56000.pth"
-CONFIG_URL = "https://huggingface.co/anzorq/kbd-vits-tts-male/resolve/main/config_35000.json"
+BASE_DIR = "kbd-vits-tts-{}"
+MALE_MODEL_URL = "https://huggingface.co/anzorq/kbd-vits-tts-male/resolve/main/checkpoint_56000.pth"
+MALE_CONFIG_URL = "https://huggingface.co/anzorq/kbd-vits-tts-male/resolve/main/config_35000.json"
+FEMALE_MODEL_URL = "https://huggingface.co/anzorq/kbd-vits-tts-female/resolve/main/best_model_56351.pth"
+FEMALE_CONFIG_URL = "https://huggingface.co/anzorq/kbd-vits-tts-female/resolve/main/config.json"
 
-# Downloading model and config
-if not os.path.exists(MODEL_DIR):
-    os.makedirs(MODEL_DIR)
-download_url(MODEL_URL, MODEL_DIR, "model.pth")
-download_url(CONFIG_URL, MODEL_DIR, "config.json")
+def download_model_and_config(gender):
+    dir_path = BASE_DIR.format(gender)
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+    model_url = MALE_MODEL_URL if gender == "male" else FEMALE_MODEL_URL
+    config_url = MALE_CONFIG_URL if gender == "male" else FEMALE_CONFIG_URL
+    download_url(model_url, dir_path, "model.pth")
+    download_url(config_url, dir_path, "config.json")
+    return dir_path
 
+download_model_and_config("male")
+download_model_and_config("female")
 
-def tts(text: str):
+def tts(text: str, voice: str="Male"):
     if len(text) > MAX_TXT_LEN:
         text = text[:MAX_TXT_LEN]
         print(f"Input text was cutoff since it went over the {MAX_TXT_LEN} character limit.")
@@ -25,8 +32,10 @@ def tts(text: str):
 
     text = text.replace("I", "ӏ") #replace capital is with "Palochka" symbol
 
+    model_dir = BASE_DIR.format("male" if voice == "Male" else "female")
+
     # synthesize
-    synthesizer = Synthesizer(f"{MODEL_DIR}/model.pth", f"{MODEL_DIR}/config.json")
+    synthesizer = Synthesizer(f"{model_dir}/model.pth", f"{model_dir}/config.json")
     wavs = synthesizer.tts(text)
 
     # return output
@@ -34,16 +43,22 @@ def tts(text: str):
         synthesizer.save_wav(wavs, fp)
         return fp.name
 
-# Gradio interface
 iface = gr.Interface(
     fn=tts,
-    inputs=gr.Textbox(
-        label="Text",
-        value="Default text here if you need it.",
-    ),
+    inputs=[
+        gr.Textbox(
+            label="Text",
+            value="Default text here if you need it.",
+        ),
+        gr.Radio(
+            choices=["Male", "Female"],
+            value="Male",  # Set Male as the default choice
+            label="Voice"
+        )
+    ],
     outputs=gr.Audio(label="Output", type='filepath'),
     title="KBD TTS",
     live=False
 )
 
-iface.launch(share=False)
+iface.launch(share=False)