Update app.py

app.py

@@ -5,13 +5,20 @@ import shutil
 import tempfile
 
 # Set the PATH and LD_LIBRARY_PATH for CUDA 12.3
-cuda_bin_path = "/usr/local/cuda-12.3/bin"
-cuda_lib_path = "/usr/local/cuda-12.3/lib64"
+cuda_bin_path = "/usr/local/cuda/bin"
+cuda_lib_path = "/usr/local/cuda/lib64"
 
 # Update the environment variables
 os.environ['PATH'] = f"{cuda_bin_path}:{os.environ.get('PATH', '')}"
 os.environ['LD_LIBRARY_PATH'] = f"{cuda_lib_path}:{os.environ.get('LD_LIBRARY_PATH', '')}"
 
+# Install flash attention
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
+
 from huggingface_hub import snapshot_download 
 
 # Create xcodec_mini_infer folder