Env Setup

app.py

@@ -11,19 +11,32 @@ import spaces
 import time
 import subprocess
 
+print("\n=== Environment Setup ===")
+
 if torch.cuda.is_available():
+    print(f"GPU detected: {torch.cuda.get_device_name(0)}")
     try:
         subprocess.run(
             "pip install flash-attn --no-build-isolation",
-            env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
             shell=True,
             check=True,
         )
-        print("✅ flash-attn installed (GPU detected)")
+        print("✅ flash-attn installed successfully")
     except subprocess.CalledProcessError as e:
         print("⚠️ flash-attn installation failed:", e)
 else:
     print("⚙️ CPU detected — skipping flash-attn installation")
+    # Disable flash-attn references safely
+    os.environ["DISABLE_FLASH_ATTN"] = "1"
+    os.environ["FLASH_ATTENTION_SKIP_CUDA_BUILD"] = "TRUE"
+try:
+    from transformers.utils import import_utils
+    if "flash_attn" not in import_utils.PACKAGE_DISTRIBUTION_MAPPING:
+        import_utils.PACKAGE_DISTRIBUTION_MAPPING["flash_attn"] = "flash-attn"
+except Exception as e:
+    print("⚠️ Patch skipped:", e)
+
+    
 token = os.environ["HF_TOKEN"]