safety check module available (#225)

* safety check module available

* 删除重复日志

Author: qzzz95

diffsynth_engine/utils/flag.py

@@ -6,68 +6,43 @@
 logger = logging.get_logger(__name__)
 
 
-# 无损
-FLASH_ATTN_4_AVAILABLE = importlib.util.find_spec("flash_attn.cute.interface") is not None
-if FLASH_ATTN_4_AVAILABLE:
-    logger.info("Flash attention 4 is available")
-else:
-    logger.info("Flash attention 4 is not available")
+def check_module_available(module_path: str, module_name: str = None) -> bool:
+    try:
+        available = importlib.util.find_spec(module_path) is not None
+    except (ModuleNotFoundError, AttributeError, ValueError):
+        available = False
 
-FLASH_ATTN_3_AVAILABLE = importlib.util.find_spec("flash_attn_interface") is not None
-if FLASH_ATTN_3_AVAILABLE:
-    logger.info("Flash attention 3 is available")
-else:
-    logger.info("Flash attention 3 is not available")
+    if module_name:
+        if available:
+            logger.info(f"{module_name} is available")
+        else:
+            logger.info(f"{module_name} is not available")
 
-FLASH_ATTN_2_AVAILABLE = importlib.util.find_spec("flash_attn") is not None
-if FLASH_ATTN_2_AVAILABLE:
-    logger.info("Flash attention 2 is available")
-else:
-    logger.info("Flash attention 2 is not available")
+    return available
 
-XFORMERS_AVAILABLE = importlib.util.find_spec("xformers") is not None
-if XFORMERS_AVAILABLE:
-    logger.info("XFormers is available")
-else:
-    logger.info("XFormers is not available")
+
+# 无损
+FLASH_ATTN_4_AVAILABLE = check_module_available("flash_attn.cute.interface", "Flash attention 4")
+FLASH_ATTN_3_AVAILABLE = check_module_available("flash_attn_interface", "Flash attention 3")
+FLASH_ATTN_2_AVAILABLE = check_module_available("flash_attn", "Flash attention 2")
+XFORMERS_AVAILABLE = check_module_available("xformers", "XFormers")
+AITER_AVAILABLE = check_module_available("aiter", "Aiter")
 
 SDPA_AVAILABLE = hasattr(torch.nn.functional, "scaled_dot_product_attention")
 if SDPA_AVAILABLE:
     logger.info("Torch SDPA is available")
 else:
     logger.info("Torch SDPA is not available")
 
-AITER_AVAILABLE = importlib.util.find_spec("aiter") is not None
-if AITER_AVAILABLE:
-    logger.info("Aiter is available")
-else:
-    logger.info("Aiter is not available")
 
 # 有损
-SAGE_ATTN_AVAILABLE = importlib.util.find_spec("sageattention") is not None
-if SAGE_ATTN_AVAILABLE:
-    logger.info("Sage attention is available")
-else:
-    logger.info("Sage attention is not available")
-
-SPARGE_ATTN_AVAILABLE = importlib.util.find_spec("spas_sage_attn") is not None
-if SPARGE_ATTN_AVAILABLE:
-    logger.info("Sparge attention is available")
-else:
-    logger.info("Sparge attention is not available")
+SAGE_ATTN_AVAILABLE = check_module_available("sageattention", "Sage attention")
+SPARGE_ATTN_AVAILABLE = check_module_available("spas_sage_attn", "Sparge attention")
+VIDEO_SPARSE_ATTN_AVAILABLE = check_module_available("vsa", "Video sparse attention")
 
-VIDEO_SPARSE_ATTN_AVAILABLE = importlib.util.find_spec("vsa") is not None
-if VIDEO_SPARSE_ATTN_AVAILABLE:
-    logger.info("Video sparse attention is available")
-else:
-    logger.info("Video sparse attention is not available")
-
-NUNCHAKU_AVAILABLE = importlib.util.find_spec("nunchaku") is not None
+NUNCHAKU_AVAILABLE = check_module_available("nunchaku", "Nunchaku")
 NUNCHAKU_IMPORT_ERROR = None
-if NUNCHAKU_AVAILABLE:
-    logger.info("Nunchaku is available")
-else:
-    logger.info("Nunchaku is not available")
+if not NUNCHAKU_AVAILABLE:
     import sys
     torch_version = getattr(torch, "__version__", "unknown")
     python_version = f"{sys.version_info.major}.{sys.version_info.minor}"