Update app.py

app.py

@@ -8,6 +8,42 @@ from transformers import AutoModelForCausalLM
 from transformers import TextIteratorStreamer
 from threading import Thread
 
+import importlib.metadata
+from importlib import import_module
+from transformers.utils import is_flash_attn_2_available
+from packaging import version
+
+def check_flash_attention_2_requirements():
+    # 检查 Flash Attention 2 是否可用
+    flash_attn_2_available = is_flash_attn_2_available()
+    if not flash_attn_2_available:
+        raise ImportError("Flash Attention 2 is not available.")
+
+    # 获取已安装的 flash_attn 版本
+    try:
+        installed_version = importlib.metadata.version("flash_attn")
+    except importlib.metadata.PackageNotFoundError:
+        raise ImportError("flash_attn package is not installed.")
+
+    # 解析已安装的版本和所需的最低版本
+    parsed_installed_version = version.parse(installed_version)
+    required_version = version.parse("2.6.3")
+
+    # 检查版本是否满足要求
+    if parsed_installed_version < required_version:
+        raise ImportError(f"flash_attn version {installed_version} is installed, but version >= 2.6.3 is required.")
+
+    print("All requirements for Flash Attention 2 are met.")
+
+# 使用 try-except 块来捕获和显示具体的错误
+try:
+    check_flash_attention_2_requirements()
+except ImportError as e:
+    print(f"Error: {e}")
+    print("Using `flash_attention_2` requires having `flash_attn>=2.6.3` installed.")
+else:
+    print("Flash Attention 2 can be used.")
+
 model_name = 'AIDC-AI/Ovis2-16B'
 
 # load model