Update app.py

app.py

@@ -27,33 +27,32 @@ from huggingface_hub import hf_hub_download
 from infer import load_model, eval_model
 from spkr import SpeakerEmbedding
 
-spkr_model = None
-
 @spaces.GPU
 def spkr_model_init():
-    global spkr_model
-    if spkr_model is None:
-        spkr_model = SpeakerEmbedding(device="cuda")
+    spkr_model = SpeakerEmbedding(device="cuda")
     return spkr_model
 
-model, tokenizer, tokenizer_voila, model_type = load_model("maitrix-org/Voila-chat", "maitrix-org/Voila-Tokenizer")
-default_ref_file = "examples/character_ref_emb_demo.pkl"
-default_ref_name = "Homer Simpson"
-million_voice_ref_file = hf_hub_download(repo_id="maitrix-org/Voila-million-voice", filename="character_ref_emb_chunk0.pkl", repo_type="dataset")
+def init_models():
+    global spkr_model, model, tokenizer, tokenizer_voila, model_type
+    global default_ref_emb_mask_list, million_voice_ref_emb_mask_list
+    
+    spkr_model = spkr_model_init()
+    model, tokenizer, tokenizer_voila, model_type = load_model("maitrix-org/Voila-chat", "maitrix-org/Voila-Tokenizer")
+    default_ref_file = "examples/character_ref_emb_demo.pkl"
+    million_voice_ref_file = hf_hub_download(repo_id="maitrix-org/Voila-million-voice", filename="character_ref_emb_chunk0.pkl", repo_type="dataset")
+    
+    default_ref_emb_mask_list = pickle.load(open(default_ref_file, "rb"))
+    million_voice_ref_emb_mask_list = pickle.load(open(million_voice_ref_file, "rb"))
 
 instruction = "You are a smart AI agent created by Maitrix.org."
 save_path = os.environ.get("GRADIO_TEMP_DIR", tempfile.gettempdir())
+default_ref_name = "Homer Simpson"
 
 intro = """**Voila**
 
 For more demos, please goto [https://voila.maitrix.org](https://voila.maitrix.org)."""
 
-default_ref_emb_mask_list = pickle.load(open(default_ref_file, "rb"))
-million_voice_ref_emb_mask_list = pickle.load(open(million_voice_ref_file, "rb"))
-
 def get_ref_embs(ref_audio):
-    global spkr_model
-    spkr_model = spkr_model_init()
     wav, sr = torchaudio.load(ref_audio)
     ref_embs = spkr_model(wav, sr).cpu()
     return ref_embs
@@ -259,14 +258,18 @@ def get_asr_tab():
     )
     return asr_tab
 
-with gr.Blocks(fill_height=True) as demo:
-    with gr.Tab("Chat"):
-        chat_tab = get_chat_tab()
-    with gr.Tab("TTS"):
-        tts_tab = get_tts_tab()
-    with gr.Tab("ASR"):
-        asr_tab = get_asr_tab()
-    demo.unload(delete_directory)
+def create_demo():
+    with gr.Blocks(fill_height=True) as demo:
+        with gr.Tab("Chat"):
+            chat_tab = get_chat_tab()
+        with gr.Tab("TTS"):
+            tts_tab = get_tts_tab()
+        with gr.Tab("ASR"):
+            asr_tab = get_asr_tab()
+        demo.unload(delete_directory)
+    return demo
 
 if __name__ == "__main__":
+    init_models()
+    demo = create_demo()
     demo.launch()