Update coqui.py

coqui.py

@@ -1,11 +1,12 @@
 import sys
 import io, os, stat
+import torch
+
 import subprocess
 import random
 from zipfile import ZipFile
 import uuid
 import time
-import torch
 import torchaudio
 import numpy as np
 
@@ -78,9 +79,13 @@ if not os.path.exists(checkpoint_path):
 if not os.path.exists(vocab_path):
     raise FileNotFoundError(f"Vocab file not found at {vocab_path}")
 
+if not os.environ.get('CUDA_HOME'):
+    print(f"ENV var CUDA_HOME is not set, defaulting to: '/usr/local/cuda'")
+    os.environ['CUDA_HOME'] = f"/usr/local/cuda"
+
 model.load_checkpoint(
     config,
-    checkpoint_path=checkpoint_path,
+    checkpoint_dir=model_path,
     vocab_path=vocab_path,
     eval=True,
     use_deepspeed=True,
@@ -410,4 +415,4 @@ def predict(
         gr.Warning("Please accept the Terms & Condition!")
         return (
             None,
-        )
+        )