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,
- )
+ )