Update app.py

app.py

@@ -9,38 +9,49 @@ 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
+from transformers.utils import (
+    is_torch_available,
+    _is_package_available,
+    is_torch_mlu_available
+)
+
+def diagnose_flash_attn_2_availability():
+    if not is_torch_available():
+        return "PyTorch is not available."
+
+    if not _is_package_available("flash_attn"):
+        return "flash_attn package is not installed."
+
+    import torch
+
+    if not (torch.cuda.is_available() or is_torch_mlu_available()):
+        return "Neither CUDA nor MLU is available."
+
+    flash_attn_version = importlib.metadata.version("flash_attn")
+
+    if torch.version.cuda:
+        required_version = "2.1.0"
+        if version.parse(flash_attn_version) < version.parse(required_version):
+            return f"CUDA is available, but flash_attn version {flash_attn_version} is installed. Version >= {required_version} is required."
+    elif torch.version.hip:
+        required_version = "2.0.4"
+        if version.parse(flash_attn_version) < version.parse(required_version):
+            return f"HIP is available, but flash_attn version {flash_attn_version} is installed. Version >= {required_version} is required."
+    elif is_torch_mlu_available():
+        required_version = "2.3.3"
+        if version.parse(flash_attn_version) < version.parse(required_version):
+            return f"MLU is available, but flash_attn version {flash_attn_version} is installed. Version >= {required_version} is required."
+    else:
+        return "Unknown PyTorch backend."
+
+    return "All requirements for Flash Attention 2 are met."
 
-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.")
+# 使用诊断函数
+result = diagnose_flash_attn_2_availability()
+if result != "All requirements for Flash Attention 2 are met.":
+    print(f"Flash Attention 2 is not available: {result}")
+    print("Using `flash_attention_2` requires having the correct version of `flash_attn` installed.")
 else:
     print("Flash Attention 2 can be used.")