Use flash_attention_2 if available

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import spaces
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
 from qwen_vl_utils import process_vision_info
-
+from transformers.utils import is_flash_attn_2_available
 
 subprocess.run(
     "pip install flash-attn --no-build-isolation",
@@ -29,7 +29,8 @@ def run_inference(model_key, input_type, text, image, video, fps):
     model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
         model_id,
         torch_dtype="auto",
-        device_map="auto"
+        device_map="auto",
+        attn_implementation="flash_attention_2" if is_flash_attn_2_available() else None,
     )
     processor = AutoProcessor.from_pretrained(model_id)