Update app.py

app.py

@@ -12,12 +12,21 @@ cuda_lib_path = "/usr/local/cuda/lib64"
 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,
-)
+# Install required package
+def install_flash_attn():
+    try:
+        print("Installing flash-attn...")
+        subprocess.run(
+            ["pip", "install", "flash-attn", "--no-build-isolation"], 
+            check=True
+        )
+        print("flash-attn installed successfully!")
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to install flash-attn: {e}")
+        exit(1)
+
+# Install flash-attn
+install_flash_attn()
 
 from huggingface_hub import snapshot_download