Update app.py

app.py

@@ -7,14 +7,6 @@ import spaces
 import torch
 
 
-# Install flash-attn for GPU only (after spaces import)
-if torch.cuda.is_available():
-    print("CUDA detected - installing flash-attn for optimal GPU performance...")
-    subprocess.run(
-        "pip install flash-attn --no-build-isolation",
-        env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-        shell=True,
-    )
 
 import gradio as gr
 from PIL import Image
@@ -29,9 +21,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Choose best attention implementation based on device
 if device == "cuda":
-    attn_implementation = "flash_attention_2"  # Best for GPU
+    attn_implementation = "sdpa"  
     dtype = torch.bfloat16
-    print("Using flash_attention_2 for GPU")
+    print("Using sdpa for GPU")
 else:
     attn_implementation = "eager"  # Best for CPU
     dtype = torch.float32