Update app.py

app.py

@@ -26,8 +26,14 @@ 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="cpu")
+    return spkr_model
 
-spkr_model = SpeakerEmbedding(device="cpu")
+spkr_model = spkr_model_init()
+spkr_model.to("cuda")
 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"
@@ -45,9 +51,7 @@ 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):
@@ -69,8 +73,10 @@ def call_bot(history, ref_embs, request: gr.Request):
     }
     formated_history["conversations"].append({"from": "assistant"})
     print(formated_history)
-    ref_embs = torch.tensor(ref_embs, dtype=torch.float32, device="cuda")
-    ref_embs_mask = torch.tensor([1], device="cuda")
+    ref_embs = torch.tensor(ref_embs, dtype=torch.float32, device="cpu")
+    ref_embs_mask = torch.tensor([1], device="cpu")
+    ref_embs.to("cuda")
+    ref_embs_mask.to("cuda")
     out = eval_model(model, tokenizer, tokenizer_voila, model_type, "chat_aiao", formated_history, ref_embs, ref_embs_mask, max_new_tokens=512)
     if 'audio' in out:
         wav, sr = out['audio']
@@ -93,8 +99,10 @@ def run_tts(text, ref_embs):
         "conversations": [{'from': "user", 'text': text}],
     }
     formated_history["conversations"].append({"from": "assistant"})
-    ref_embs = torch.tensor(ref_embs, dtype=torch.float32, device="cuda")
-    ref_embs_mask = torch.tensor([1], device="cuda")
+    ref_embs = torch.tensor(ref_embs, dtype=torch.float32, device="cpu")
+    ref_embs_mask = torch.tensor([1], device="cpu")
+    ref_embs.to("cuda")
+    ref_embs_mask.to("cuda")
     out = eval_model(model, tokenizer, tokenizer_voila, model_type, "chat_tts", formated_history, ref_embs, ref_embs_mask, max_new_tokens=512)
     if 'audio' in out:
         wav, sr = out['audio']