Update app_exp.py

app_exp.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import torch
 import os
@@ -19,6 +20,25 @@ subprocess.run(["pip3", "install", "-U", "cache-dit"], check=True)
 
 import cache_dit
 
+enable_fa3 = False  # default if FA3 cannot be loaded
+
+try:
+    print("Installing FlashAttention 3...")
+    flash_attention_wheel = hf_hub_download(
+        repo_id="rahul7star/flash-attn-3",
+        repo_type="model",
+        filename="128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl",
+    )
+    subprocess.run(["pip", "install", flash_attention_wheel], check=True)
+    site.addsitedir(site.getsitepackages()[0])
+    importlib.invalidate_caches()
+    enable_fa3 = True
+    print("✅ FlashAttention 3 installed and enabled")
+except Exception as e:
+    print(f"⚠️ Could not install FlashAttention 3: {e}")
+    # enable_fa3 remains False
+
+
 # ============================================================
 # 1️⃣ Repository & Weights
 # ============================================================