Update app.py

app.py

@@ -31,6 +31,16 @@ GPU_ID = args.gpu
 DEFAULT_MODEL_PATH = args.model_path
 OUTPUT_DIR = args.output_dir
 
+if 'CUDA_HOME' not in os.environ:
+    for cuda_path in ['/usr/local/cuda', '/usr/cuda', '/opt/cuda']:
+        if os.path.exists(cuda_path):
+            os.environ['CUDA_HOME'] = cuda_path
+            print(cuda_path)
+            break
+    if 'CUDA_HOME' not in os.environ:
+        os.environ['CUDA_HOME'] = '/usr/local/cuda'
+        print("set default cuda path in: /usr/local/cuda")
+
 # Create necessary directories
 os.makedirs("outputs", exist_ok=True)
 # Create project tmp directory instead of using system temp