diffsynth-engine 0.6.1.dev39__py3-none-any.whl → 0.6.1.dev41__py3-none-any.whl

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/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/utils/flag.py

@@ -6,24 +6,27 @@ 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")
+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_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")
+    if module_name:
+        if available:
+            logger.info(f"{module_name} is available")
+        else:
+            logger.info(f"{module_name} 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")
+    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:
@@ -31,37 +34,15 @@ if SDPA_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}"

{diffsynth_engine-0.6.1.dev39.dist-info → diffsynth_engine-0.6.1.dev41.dist-info}/METADATA

@@ -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.dist-info → diffsynth_engine-0.6.1.dev41.dist-info}/RECORD

@@ -86,12 +86,12 @@ diffsynth_engine/conf/tokenizers/z_image/tokenizer/tokenizer_config.json,sha256=
 diffsynth_engine/conf/tokenizers/z_image/tokenizer/vocab.json,sha256=yhDX6fs-0YV13R4neiV5wW0QjjLydDloSvoOELFECRA,2776833
 diffsynth_engine/configs/__init__.py,sha256=biluGSEw78PPwO7XFlms16iuWXDiM0Eg_qsOMMTY0NQ,1409
 diffsynth_engine/configs/controlnet.py,sha256=f3vclyP3lcAjxDGD9C1vevhqqQ7W2LL_c6Wye0uxk3Q,1180
-diffsynth_engine/configs/pipeline.py,sha256=RqhPAZOCpIMkFk-OsfiNYlqpqM-7B52ny0Zcr9Ix7wY,15310
+diffsynth_engine/configs/pipeline.py,sha256=tcnhLGdQgvEibWBZVFH3uOS1pwB6WEnHgCFSer2bT0E,15347
 diffsynth_engine/kernels/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/models/__init__.py,sha256=8Ze7cSE8InetgXWTNb0neVA2Q44K7WlE-h7O-02m2sY,119
 diffsynth_engine/models/base.py,sha256=svao__9WH8VNcyXz5o5dzywYXDcGV0YV9IfkLzDKews,2558
 diffsynth_engine/models/basic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diffsynth_engine/models/basic/attention.py,sha256=YrIxkYoekC3I7-sMTw60CL4GIKMLOTrn-eCk-iHT7E4,15701
+diffsynth_engine/models/basic/attention.py,sha256=sYoTOlloRA5KFNV18cZSkXuPY-Ck-13-kCM41AaBMu0,17225
 diffsynth_engine/models/basic/lora.py,sha256=Y6cBgrBsuDAP9FZz_fgK8vBi_EMg23saFIUSAsPIG-M,10670
 diffsynth_engine/models/basic/lora_nunchaku.py,sha256=7qhzGCzUIfDrwtWG0nspwdyZ7YUkaM4vMqzxZby2Zds,7510
 diffsynth_engine/models/basic/relative_position_emb.py,sha256=rCXOweZMcayVnNUVvBcYXMdhHS257B_PC8PZSWxvhNQ,2540
@@ -183,7 +183,7 @@ diffsynth_engine/utils/cache.py,sha256=Ivef22pCuhEq-4H00gSvkLS8ceVZoGis7OSitYL6g
 diffsynth_engine/utils/constants.py,sha256=x0-bsPRplW-KkRpLVajuC9Yv6f3QbdHgSr3XZ-eBCsQ,3745
 diffsynth_engine/utils/download.py,sha256=w9QQjllPfTUEY371UTREU7o_vvdMY-Q2DymDel3ZEZY,6792
 diffsynth_engine/utils/env.py,sha256=k749eYt_qKGq38GocDiXfkhp8nZrowFefNVTZ8R755I,363
-diffsynth_engine/utils/flag.py,sha256=KSzjnzRe7sleNCJm8IpbJQbmBY4KNV2kDrijxi27Jek,2928
+diffsynth_engine/utils/flag.py,sha256=Ubm7FF0vHG197bmJGEplp4XauBlUaQVv-zr-w6VyEIM,2493
 diffsynth_engine/utils/fp8_linear.py,sha256=k34YFWo2dc3t8aKjHaCW9CbQMOTqXxaDHk8aw8aKif4,3857
 diffsynth_engine/utils/gguf.py,sha256=ZWvw46V4g4uVyAR_oCq-4K5nPdKVrYk3u47uXMgA9lU,14092
 diffsynth_engine/utils/image.py,sha256=PiDButjv0fsRS23kpQgCLZAlBumpzQmNnolfvb5EKQ0,9626
@@ -200,8 +200,8 @@ diffsynth_engine/utils/video.py,sha256=8FCaeqIdUsWMgWI_6SO9SPynsToGcLCQAVYFTc4CD
 diffsynth_engine/utils/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/utils/memory/linear_regression.py,sha256=oW_EQEw13oPoyUrxiL8A7Ksa5AuJ2ynI2qhCbfAuZbg,3930
 diffsynth_engine/utils/memory/memory_predcit_model.py,sha256=EXprSl_zlVjgfMWNXP-iw83Ot3hyMcgYaRPv-dvyL84,3943
-diffsynth_engine-0.6.1.dev39.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.6.1.dev39.dist-info/METADATA,sha256=f_qU_vp4RcHSOgW3Agm428engf8v7TKRCt8DuxAOEi8,1164
-diffsynth_engine-0.6.1.dev39.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.6.1.dev39.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.6.1.dev39.dist-info/RECORD,,
+diffsynth_engine-0.6.1.dev41.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.6.1.dev41.dist-info/METADATA,sha256=dygz72s8iSZlS1JE54FlMvYiSmMy8uB2B0Gd_WLCSws,1164
+diffsynth_engine-0.6.1.dev41.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.6.1.dev41.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.6.1.dev41.dist-info/RECORD,,