Add examples, better UI

app.py

@@ -5,84 +5,137 @@ import sys
 import soundfile as sf
 import numpy as np
 import torch
+import traceback
 
 repo_url = "https://huggingface.co/dangtr0408/StyleTTS2-lite-vi"
 repo_dir = "StyleTTS2-lite-vi"
-
 if not os.path.exists(repo_dir):
     subprocess.run(["git", "clone", repo_url, repo_dir])
-
-# Clone repo and load model
 sys.path.append(os.path.abspath(repo_dir))
 from inference import StyleTTS2
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
 config_path = os.path.join(repo_dir, "Models", "config.yml")
 models_path = os.path.join(repo_dir, "Models", "model.pth")
 model = StyleTTS2(config_path, models_path).eval().to(device)
+session_uploaded_files = {}
+voice_path = os.path.join(repo_dir, "reference_audio")
+eg_voices = [os.path.join(voice_path,"vn_1.wav"), os.path.join(voice_path,"vn_2.wav")]
+eg_texts = [
+    "[id_1] Với [en-us]{StyleTTS2-lite} bạn có thể sử dụng [en-us]{language tag} để mô hình chắc chắn đọc bằng tiếng Anh, [id_2]cũng như sử dụng [en-us]{speaker tag} để chuyển đổi nhanh giữa các giọng đọc.",
+    "[id_1]Chỉ với khoảng 90 triệu tham số, [id_2][en-us]{StyleTTS2-lite} có thể dễ dàng tạo giọng nói với tốc độ cao.",
+]
+
 
 # Core inference function
-def process_inputs(text_prompt, reference_audio_paths,
-                   n_merge, denoise, avg_style,stabilize):
+def main(text_prompt, denoise, avg_style, stabilize):
+    try:
+        global session_uploaded_files
+        reference_paths = [file for file in session_uploaded_files.values()]
+        speakers = {}
+        for i, path in enumerate(reference_paths, 1):
+            speaker_id = f"id_{i}"
+            speakers[speaker_id] = {
+                "path": path,
+                "lang": "vi",
+                "speed": 1.0
+        }
+
+        with torch.no_grad():
+            r = model.generate(text_prompt, speakers, avg_style, stabilize, denoise, 20, "[id_1]")
+
+        r = r / np.abs(r).max()
+        sf.write("output.wav", r, samplerate=24000)
+
+        return "output.wav", "Audio generated successfully!"
+    
+    except Exception as e:
+        error_message = traceback.format_exc()
+        return None, error_message
+
+def on_file_upload(file_list):
+    if not file_list:
+        return "No file uploaded yet."
     
-    speakers = {}
-    for i, path in enumerate(reference_audio_paths, 1):
-        speaker_id = f"id_{i}"
-        speakers[speaker_id] = {
-            "path": path,
-            "lang": "vi",
-            "speed": 1.0
-    }
-
-    with torch.no_grad():
-        r = model.generate(text_prompt, speakers, avg_style, stabilize, denoise, n_merge, "[id_1]")
-
-    r = r / np.abs(r).max()
-    sf.write("output.wav", r, samplerate=24000)
-    return "output.wav"
-
-custom_css = """
-    #custom-box {
-        min-height: 300px !important;
-        display: flex;
-        align-items: center;
-    }
-    #custom-box textarea {
-        min-height: 250px !important;
-        height: 100% !important;
-    }
-"""
+    global session_uploaded_files
+    for file_path in file_list:
+        file_name = os.path.basename(file_path)
+        session_uploaded_files[file_name] = file_path #update and remove duplicate
+
+    uploaded_infos = []
+    uploaded_file_names = list(session_uploaded_files.keys())
+    for i in range(len(session_uploaded_files)):
+        uploaded_infos.append(f"[id_{i}]: {uploaded_file_names[i]}")
+        
+    summary = "\n".join(uploaded_infos)
+    return f"Current reference audios:\n{summary}"
+
+def gen_example(text_prompt):
+    on_file_upload(eg_voices)
+    output, status = main(text_prompt, 0.6, True, True)
+    #Reset
+    on_file_upload(None)
+    global session_uploaded_files
+    session_uploaded_files = {}
+    return output, status
+
 
 # Gradio UI
-with gr.Blocks(css=custom_css) as demo:
-    gr.Markdown("## StyleTTS2-lite-vi Demo")
-    gr.Markdown("Upload a reference audio and input your text to synthesize speech with style control.")
-
-    with gr.Row():
-        text_prompt = gr.Textbox(label="Text Prompt", placeholder="Enter your text here...", elem_id="custom-box")
-        reference_audios = gr.File(label="Reference Audios", file_types=[".wav", ".mp3", ".flac"], file_count="multiple", elem_id="custom-box")
-    # Parameters
-    with gr.Accordion("Advanced Settings", open=False):
-        avg_style = gr.Checkbox(label="Use Average Styles", value=True)
-        stabilize = gr.Checkbox(label="Stabilize Speaking Speed", value=True)
-        denoise = gr.Slider(0.0, 1.0, value=0.6, label="Denoise Strength")
-        n_merge = gr.Slider(10, 30, value=16, label="Min Words to Merge")
-
-    submit_button = gr.Button("Synthesize")
-    synthesized_audio = gr.Audio(label="Synthesized Audio", type="filepath")
-
-    submit_button.click(
-        fn=process_inputs,
+with gr.Blocks() as demo:
+    gr.HTML("<h1 style='text-align: center;'>StyleTTS2‑Lite Demo</h1>")
+    gr.Markdown(
+        "Download the local inference package from Hugging Face: "
+        "[StyleTTS2‑Lite (Vietnamese)]"
+        "(https://huggingface.co/dangtr0408/StyleTTS2-lite-vi/)."
+    )
+    gr.Markdown(
+        "Please specify a language tag in your inputs if the word is not Vietnamese, e.g., [en-us]{  } for English. For more information, see "
+        "[eSpeakNG docs]"
+        "(https://github.com/espeak-ng/espeak-ng/blob/master/docs/languages.md)"
+    )
+
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1):
+            text_prompt = gr.Textbox(label="Text Prompt", placeholder="Enter your text here...", lines=4)
+        with gr.Column(scale=1):
+            avg_style = gr.Checkbox(label="Use Average Styles", value=True)
+            stabilize = gr.Checkbox(label="Stabilize Speaking Speed", value=True)
+            denoise = gr.Slider(0.0, 1.0, step=0.1, value=0.6, label="Denoise Strength")
+
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1):
+            reference_audios = gr.File(label="Reference Audios", file_types=[".wav", ".mp3"], file_count="multiple", height=150)
+            gen_button = gr.Button("Generate")
+        with gr.Column(scale=1):
+            synthesized_audio = gr.Audio(label="Generate Audio", type="filepath")
+
+    status = gr.Textbox(label="Status", interactive=False, lines=3)
+
+    reference_audios.change(
+            on_file_upload, 
+            inputs=[reference_audios], 
+            outputs=[status]
+    )
+
+    gen_button.click(
+        fn=main,
         inputs=[
             text_prompt,
-            reference_audios,
-            n_merge,
             denoise,
             avg_style,
             stabilize
         ],
-        outputs=synthesized_audio
+        outputs=[synthesized_audio, status]
+    )
+
+    gr.Examples(
+        examples=[eg_texts[0], eg_texts[1]],
+        inputs=[text_prompt],
+        outputs=[synthesized_audio, status],
+        fn=gen_example,
+        cache_examples=False,
+        label="Creation Examples",
+        run_on_click=True
     )
 
 demo.launch()