install flash attention

main.py

@@ -3,6 +3,7 @@ import logging
 import json
 import torch
 import asyncio
+import subprocess
 
 import gradio as gr
 import numpy as np
@@ -41,6 +42,15 @@ MODEL_ID = os.getenv("MODEL_ID", "openai/whisper-large-v3-turbo")
 LANGUAGE = os.getenv("LANGUAGE", "english").lower()
 
 device = get_device(force_cpu=False)
+
+# Install Flash Attention 2 if device is "cuda"
+if device == "cuda":
+    subprocess.run(
+        ["pip", "install", "flash-attn", "--no-build-isolation"],
+        env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
+        shell=True,
+    )
+
 torch_dtype, np_dtype = get_torch_and_np_dtypes(device, use_bfloat16=False)
 logger.info(f"Using device: {device}, torch_dtype: {torch_dtype}, np_dtype: {np_dtype}")