Update app.py

app.py

@@ -1,10 +1,12 @@
+!pip install nvidia-ml-py3
 import gradio as gr
 import torch
 import soundfile as sf
 from snac import SNAC
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-device = torch.device("cpu")  # Changed to CPU mode
+# Ensure the code uses NVIDIA GPUs
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def find_last_instance_of_separator(lst, element=50258):
     reversed_list = lst[::-1]
@@ -55,7 +57,7 @@ def reconstruct_tensors(flattened_output):
             tensor2.append(flattened_output[i+5])
             tensor3.append(flattened_output[i+6])
             tensor3.append(flattened_output[i+7])
-            codes = [list_to_torch_tensor(tensor1), list_to_torch_tensor(tensor2), list_to_torch_tensor(tensor3)]
+            codes = [list_to_torch_tensor(tensor1).to(device), list_to_torch_tensor(tensor2).to(device), list_to_torch_tensor(tensor3).to(device)]
 
     if n_tensors == 15:
         for i in range(0, len(flattened_output), 16):
@@ -74,20 +76,20 @@ def reconstruct_tensors(flattened_output):
             tensor3.append(flattened_output[i+13])
             tensor4.append(flattened_output[i+14])
             tensor4.append(flattened_output[i+15])
-            codes = [list_to_torch_tensor(tensor1), list_to_torch_tensor(tensor2), list_to_torch_tensor(tensor3), list_to_torch_tensor(tensor4)]
+            codes = [list_to_torch_tensor(tensor1).to(device), list_to_torch_tensor(tensor2).to(device), list_to_torch_tensor(tensor3).to(device), list_to_torch_tensor(tensor4).to(device)]
 
     return codes
 
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained("Lwasinam/voicera-jenny-finetune")
     model = AutoModelForCausalLM.from_pretrained("Lwasinam/voicera-jenny-finetune").to(device)
-    snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval()
+    snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to(device)
     return model, tokenizer, snac_model
 
 def SpeechDecoder(codes, snac_model):
     codes = codes.squeeze(0).tolist()
     reconstructed_codes = reconstruct_tensors(codes)
-    audio_hat = snac_model.to(device).decode(reconstructed_codes)
+    audio_hat = snac_model.decode(reconstructed_codes)
     audio_path = "reconstructed_audio.wav"
     sf.write(audio_path, audio_hat.squeeze().cpu().detach().numpy(), 24000)
     return audio_path
@@ -117,4 +119,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch()