Update app.py

app.py

@@ -22,18 +22,21 @@ def quantize_model(model):
     )
     return quantized_model
 
-model = quantize_model(model)
+# Only quantize the vocoder, as the main model might not be compatible
 vocoder = quantize_model(vocoder)
 
-# JIT compile the models for faster inference
-model = torch.jit.script(model)
-vocoder = torch.jit.script(vocoder)
+# Move models to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+vocoder = vocoder.to(device)
+speaker_embeddings = speaker_embeddings.to(device)
 
 # Use inference mode for faster computation
 @torch.inference_mode()
 def text_to_speech(text):
-    inputs = processor(text=text, return_tensors="pt")
+    inputs = processor(text=text, return_tensors="pt").to(device)
     speech = model.generate_speech(inputs["input_ids"], speaker_embeddings, vocoder=vocoder)
+    speech = speech.cpu()  # Move back to CPU for saving
     output_path = "output.wav"
     sf.write(output_path, speech.numpy(), samplerate=16000)
     return output_path