Add fallback for LlamaFlashAttention2 in ensure_llama_flash_attn2 function

app.py

@@ -13,6 +13,18 @@ import numpy as np
 import base64
 from io import StringIO, BytesIO
 
+def ensure_llama_flash_attn2():
+    try:
+        from transformers.models.llama import modeling_llama as llama_mod
+    except Exception:
+        return
+    if not hasattr(llama_mod, "LlamaFlashAttention2"):
+        class LlamaFlashAttention2:  # fallback shim; not used when attn impl is SDPA
+            pass
+        llama_mod.LlamaFlashAttention2 = LlamaFlashAttention2
+
+ensure_llama_flash_attn2()
+
 MODEL_NAME = "deepseek-ai/DeepSeek-OCR-2"
 MODEL_REVISION = "e6322a289fe5b5218278d276d4e7c58e8103f46a"
 DOTS_OCR_MODEL = "rednote-hilab/dots.ocr"
@@ -405,4 +417,4 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
                              [text_out, md_out, raw_out, img_out, gallery])
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(theme=gr.themes.Soft())
+    demo.queue(max_size=20).launch(theme=gr.themes.Soft())