Optimize gradio

app.py

@@ -18,20 +18,17 @@ 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 = [
+    "Chỉ với khoảng 90 triệu tham số, [en-us]{StyleTTS2-lite} có thể dễ dàng tạo giọng nói với tốc độ cao.",
     "[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 main(text_prompt, denoise, avg_style, stabilize):
+def main(reference_paths, 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}"
@@ -42,11 +39,10 @@ def main(text_prompt, denoise, avg_style, stabilize):
         }
 
         with torch.no_grad():
-            r = model.generate(text_prompt, speakers, avg_style, stabilize, denoise, 20, "[id_1]")
-
-        r = r / np.abs(r).max()
+            r = model.generate(text_prompt, speakers, avg_style, stabilize, denoise, 18, "[id_1]") #Should seperate style computation process to style caching.
+            r = r / np.abs(r).max()
+            
         sf.write("output.wav", r, samplerate=24000)
-
         return "output.wav", "Audio generated successfully!"
     
     except Exception as e:
@@ -55,29 +51,24 @@ def main(text_prompt, denoise, avg_style, stabilize):
 
 def on_file_upload(file_list):
     if not file_list:
-        return "No file uploaded yet."
+        return None, "No file uploaded yet."
     
-    global session_uploaded_files
+    unique_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
+        unique_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]}")
+    uploaded_file_names = list(unique_files.keys())
+    for i in range(len(uploaded_file_names)):
+        uploaded_infos.append(f"[id_{i+1}]: {uploaded_file_names[i]}")
         
     summary = "\n".join(uploaded_infos)
-    return f"Current reference audios:\n{summary}"
+    return list(unique_files.values()), 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
+def gen_example(reference_paths, text_prompt):
+    output, status = main(reference_paths, text_prompt, 0.6, True, True)
+    return output, eg_voices, status
 
 
 # Gradio UI
@@ -112,14 +103,15 @@ with gr.Blocks() as demo:
     status = gr.Textbox(label="Status", interactive=False, lines=3)
 
     reference_audios.change(
-            on_file_upload, 
-            inputs=[reference_audios], 
-            outputs=[status]
+        on_file_upload, 
+        inputs=[reference_audios], 
+        outputs=[reference_audios, status]
     )
 
     gen_button.click(
         fn=main,
         inputs=[
+            reference_audios,
             text_prompt,
             denoise,
             avg_style,
@@ -129,12 +121,12 @@ with gr.Blocks() as demo:
     )
 
     gr.Examples(
-        examples=[eg_texts[0], eg_texts[1]],
-        inputs=[text_prompt],
-        outputs=[synthesized_audio, status],
+        examples=[[eg_voices, eg_texts[0]], [eg_voices, eg_texts[1]]],
+        inputs=[reference_audios, text_prompt],
+        outputs=[synthesized_audio, reference_audios, status],
         fn=gen_example,
         cache_examples=False,
-        label="Creation Examples",
+        label="Examples",
         run_on_click=True
     )