Update app.py

app.py

@@ -27,34 +27,27 @@ from huggingface_hub import hf_hub_download
 from infer import load_model, eval_model
 from spkr import SpeakerEmbedding
 
-@spaces.GPU
-def spkr_model_init():
-    spkr_model = SpeakerEmbedding(device="cuda")
-    return spkr_model
-
-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"))
+spkr_model = SpeakerEmbedding(device="cpu")
+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")
 
 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):
     wav, sr = torchaudio.load(ref_audio)
+    spkr_model.to("cuda")
     ref_embs = spkr_model(wav, sr).cpu()
+    spkr_model.to("cpu")
     return ref_embs
 
 def delete_directory(request: gr.Request):
@@ -258,18 +251,14 @@ def get_asr_tab():
     )
     return asr_tab
 
-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
+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)
 
 if __name__ == "__main__":
-    init_models()
-    demo = create_demo()
     demo.launch()