diffsynth-engine 0.6.1.dev39__tar.gz → 0.6.1.dev41__tar.gz

{diffsynth_engine-0.6.1.dev39 → diffsynth_engine-0.6.1.dev41}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.6.1.dev39
+Version: 0.6.1.dev41
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

{diffsynth_engine-0.6.1.dev39 → diffsynth_engine-0.6.1.dev41}/diffsynth_engine/configs/pipeline.py

@@ -26,6 +26,7 @@ class AttnImpl(Enum):
     FA2 = "fa2"  # Flash Attention 2
     FA3 = "fa3"  # Flash Attention 3
     FA3_FP8 = "fa3_fp8"  # Flash Attention 3 with FP8
+    FA4 = "fa4"  # Flash Attention 4
     AITER = "aiter"  # Aiter Flash Attention
     AITER_FP8 = "aiter_fp8"  # Aiter Flash Attention with FP8
     XFORMERS = "xformers"  # XFormers

{diffsynth_engine-0.6.1.dev39 → diffsynth_engine-0.6.1.dev41}/diffsynth_engine/models/basic/attention.py

@@ -6,6 +6,7 @@ from typing import Optional
 
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.flag import (
+    FLASH_ATTN_4_AVAILABLE,
     FLASH_ATTN_3_AVAILABLE,
     FLASH_ATTN_2_AVAILABLE,
     XFORMERS_AVAILABLE,
@@ -21,7 +22,8 @@ FA3_MAX_HEADDIM = 256
 
 logger = logging.get_logger(__name__)
 
-
+if FLASH_ATTN_4_AVAILABLE:
+    from flash_attn.cute.interface import flash_attn_func as flash_attn4
 if FLASH_ATTN_3_AVAILABLE:
     from flash_attn_interface import flash_attn_func as flash_attn3
 if FLASH_ATTN_2_AVAILABLE:
@@ -142,6 +144,7 @@ def attention(
         "fa2",
         "fa3",
         "fa3_fp8",
+        "fa4",
         "aiter",
         "aiter_fp8",
         "xformers",
@@ -152,6 +155,22 @@ def attention(
     ]
     flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM
     if attn_impl is None or attn_impl == "auto":
+        if FLASH_ATTN_4_AVAILABLE:
+            # FA4 also has the same max-head-256 limitation as FA3
+            if flash_attn3_compatible and attn_mask is None:
+                attn_out = flash_attn4(q, k, v, softmax_scale=scale)
+                if isinstance(attn_out, tuple):
+                    attn_out = attn_out[0]
+                return attn_out
+            else:
+                if not flash_attn3_compatible:
+                    logger.warning(
+                        f"head_dim={q.shape[-1]}, but flash_attn_4 only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation"
+                    )
+                else:
+                    logger.debug(
+                        "flash_attn_4 does not support attention mask, will use fallback attention implementation"
+                    )
         if FLASH_ATTN_3_AVAILABLE:
             if flash_attn3_compatible and attn_mask is None:
                 return flash_attn3(q, k, v, softmax_scale=scale)
@@ -213,6 +232,17 @@ def attention(
                 v = v.to(dtype=DTYPE_FP8)
                 out = aiter_flash_attn_fp8(q, k, v, softmax_scale=scale)
                 return out.to(dtype=origin_dtype)
+        if attn_impl == "fa4":
+            if not flash_attn3_compatible:
+                raise RuntimeError(
+                    f"head_dim={q.shape[-1]}, but flash_attn_4 only supports head dimension at most {FA3_MAX_HEADDIM}"
+                )
+            if attn_mask is not None:
+                raise RuntimeError("flash_attn_4 does not support attention mask")
+            attn_out = flash_attn4(q, k, v, softmax_scale=scale)
+            if isinstance(attn_out, tuple):
+                attn_out = attn_out[0]
+            return attn_out
         if attn_impl == "fa2":
             return flash_attn2(q, k, v, softmax_scale=scale)
         if attn_impl == "xformers":

diffsynth_engine-0.6.1.dev41/diffsynth_engine/utils/flag.py

@@ -0,0 +1,62 @@
+import importlib
+import torch
+
+from diffsynth_engine.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+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
+
+    if module_name:
+        if available:
+            logger.info(f"{module_name} is available")
+        else:
+            logger.info(f"{module_name} is not available")
+
+    return 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")
+
+
+# 有损
+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")
+
+NUNCHAKU_AVAILABLE = check_module_available("nunchaku", "Nunchaku")
+NUNCHAKU_IMPORT_ERROR = None
+if not NUNCHAKU_AVAILABLE:
+    import sys
+    torch_version = getattr(torch, "__version__", "unknown")
+    python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+    NUNCHAKU_IMPORT_ERROR = (
+        "\n\n"
+        "ERROR: This model requires the 'nunchaku' library for quantized inference, but it is not installed.\n"
+        "'nunchaku' is not available on PyPI and must be installed manually.\n\n"
+        "Please follow these steps:\n"
+        "1. Visit the nunchaku releases page: https://github.com/nunchaku-tech/nunchaku/releases\n"
+        "2. Find the wheel (.whl) file that matches your environment:\n"
+        f"   - PyTorch version: {torch_version}\n"
+        f"   - Python version: {python_version}\n"
+        f"   - Operating System: {sys.platform}\n"
+        "3. Copy the URL of the correct wheel file.\n"
+        "4. Install it using pip, for example:\n"
+        "   pip install nunchaku @ https://.../your_specific_nunchaku_file.whl\n"
+    )

{diffsynth_engine-0.6.1.dev39 → diffsynth_engine-0.6.1.dev41}/diffsynth_engine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth_engine
-Version: 0.6.1.dev39
+Version: 0.6.1.dev41
 Author: MuseAI x ModelScope
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent

diffsynth_engine-0.6.1.dev39/diffsynth_engine/utils/flag.py

@@ -1,81 +0,0 @@
-import importlib
-import torch
-
-from diffsynth_engine.utils import logging
-
-logger = logging.get_logger(__name__)
-
-
-# 无损
-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")
-
-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")
-
-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")
-
-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")
-
-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_IMPORT_ERROR = None
-if NUNCHAKU_AVAILABLE:
-    logger.info("Nunchaku is available")
-else:
-    logger.info("Nunchaku is not available")
-    import sys
-    torch_version = getattr(torch, "__version__", "unknown")
-    python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
-    NUNCHAKU_IMPORT_ERROR = (
-        "\n\n"
-        "ERROR: This model requires the 'nunchaku' library for quantized inference, but it is not installed.\n"
-        "'nunchaku' is not available on PyPI and must be installed manually.\n\n"
-        "Please follow these steps:\n"
-        "1. Visit the nunchaku releases page: https://github.com/nunchaku-tech/nunchaku/releases\n"
-        "2. Find the wheel (.whl) file that matches your environment:\n"
-        f"   - PyTorch version: {torch_version}\n"
-        f"   - Python version: {python_version}\n"
-        f"   - Operating System: {sys.platform}\n"
-        "3. Copy the URL of the correct wheel file.\n"
-        "4. Install it using pip, for example:\n"
-        "   pip install nunchaku @ https://.../your_specific_nunchaku_file.whl\n"
-    )