diffsynth-engine 0.3.6.dev11__py3-none-any.whl → 0.3.6.dev13__py3-none-any.whl

diffsynth_engine/__init__.py

@@ -3,13 +3,14 @@ from .configs import (
     SDXLPipelineConfig,
     FluxPipelineConfig,
     WanPipelineConfig,
+    ControlNetParams,
+    ControlType,
 )
 from .pipelines import (
     FluxImagePipeline,
     SDXLImagePipeline,
     SDImagePipeline,
     WanVideoPipeline,
-    ControlNetParams,
 )
 from .models.flux import FluxControlNet, FluxIPAdapter, FluxRedux
 from .models.sd import SDControlNet
@@ -44,6 +45,7 @@ __all__ = [
     "FluxReplaceByControlTool",
     "FluxReduxRefTool",
     "ControlNetParams",
+    "ControlType",
     "fetch_model",
     "fetch_modelscope_model",
     "fetch_civitai_model",

diffsynth_engine/configs/__init__.py

@@ -8,7 +8,7 @@ from .pipeline import (
     FluxPipelineConfig,
     WanPipelineConfig,
 )
-from .controlnet import ControlType
+from .controlnet import ControlType, ControlNetParams
 
 __all__ = [
     "BaseConfig",
@@ -20,4 +20,5 @@ __all__ = [
     "FluxPipelineConfig",
     "WanPipelineConfig",
     "ControlType",
+    "ControlNetParams",
 ]

diffsynth_engine/configs/controlnet.py

@@ -1,5 +1,13 @@
+from dataclasses import dataclass
 from enum import Enum
 
+import torch
+import torch.nn as nn
+from typing import List, Union, Optional
+from PIL import Image
+
+ImageType = Union[Image.Image, torch.Tensor, List[Image.Image], List[torch.Tensor]]
+
 
 # FLUX ControlType
 class ControlType(Enum):
@@ -15,3 +23,14 @@ class ControlType(Enum):
             return 128
         elif self == ControlType.bfl_fill:
             return 384
+
+
+@dataclass
+class ControlNetParams:
+    image: ImageType
+    scale: float = 1.0
+    model: Optional[nn.Module] = None
+    mask: Optional[ImageType] = None
+    control_start: float = 0
+    control_end: float = 1
+    processor_name: Optional[str] = None  # only used for sdxl controlnet union now

diffsynth_engine/models/basic/attention.py

@@ -14,6 +14,8 @@ from diffsynth_engine.utils.flag import (
     SPARGE_ATTN_AVAILABLE,
 )
 
+FA3_MAX_HEADDIM = 256
+
 logger = logging.get_logger(__name__)
 
 
@@ -130,31 +132,40 @@ def attention(
         "sage_attn",
         "sparge_attn",
     ]
+    flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM
     if attn_impl is None or attn_impl == "auto":
         if FLASH_ATTN_3_AVAILABLE:
-            return flash_attn3(q, k, v, softmax_scale=scale)
-        elif XFORMERS_AVAILABLE:
+            if flash_attn3_compatible:
+                return flash_attn3(q, k, v, softmax_scale=scale)
+            else:
+                logger.warning(
+                    f"head_dim={q.shape[-1]}, but flash_attn_3 only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation"
+                )
+        if XFORMERS_AVAILABLE:
             return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale)
-        elif SDPA_AVAILABLE:
+        if SDPA_AVAILABLE:
             return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale)
-        elif FLASH_ATTN_2_AVAILABLE:
+        if FLASH_ATTN_2_AVAILABLE:
             return flash_attn2(q, k, v, softmax_scale=scale)
-        else:
-            return eager_attn(q, k, v, attn_mask=attn_mask, scale=scale)
+        return eager_attn(q, k, v, attn_mask=attn_mask, scale=scale)
     else:
         if attn_impl == "eager":
             return eager_attn(q, k, v, attn_mask=attn_mask, scale=scale)
-        elif attn_impl == "flash_attn_3":
+        if attn_impl == "flash_attn_3":
+            if not flash_attn3_compatible:
+                raise RuntimeError(
+                    f"head_dim={q.shape[-1]}, but flash_attn_3 only supports head dimension at most {FA3_MAX_HEADDIM}"
+                )
             return flash_attn3(q, k, v, softmax_scale=scale)
-        elif attn_impl == "flash_attn_2":
+        if attn_impl == "flash_attn_2":
             return flash_attn2(q, k, v, softmax_scale=scale)
-        elif attn_impl == "xformers":
+        if attn_impl == "xformers":
             return xformers_attn(q, k, v, attn_mask=attn_mask, scale=scale)
-        elif attn_impl == "sdpa":
+        if attn_impl == "sdpa":
             return sdpa_attn(q, k, v, attn_mask=attn_mask, scale=scale)
-        elif attn_impl == "sage_attn":
+        if attn_impl == "sage_attn":
             return sage_attn(q, k, v, attn_mask=attn_mask, scale=scale)
-        elif attn_impl == "sparge_attn":
+        if attn_impl == "sparge_attn":
             return sparge_attn(
                 q,
                 k,
@@ -166,8 +177,7 @@ def attention(
                 cdfthreshd=kwargs.get("sparge_cdfthreshd", 0.98),
                 pvthreshd=kwargs.get("sparge_pvthreshd", 50),
             )
-        else:
-            raise ValueError(f"Invalid attention implementation: {attn_impl}")
+        raise ValueError(f"Invalid attention implementation: {attn_impl}")
 
 
 class Attention(nn.Module):
@@ -240,32 +250,42 @@ def long_context_attention(
         "sage_attn",
         "sparge_attn",
     ]
+    flash_attn3_compatible = q.shape[-1] <= FA3_MAX_HEADDIM
     if attn_impl is None or attn_impl == "auto":
         if FLASH_ATTN_3_AVAILABLE:
-            attn_func = LongContextAttention(attn_type=AttnType.FA3)
-        elif SDPA_AVAILABLE:
-            attn_func = LongContextAttention(attn_type=AttnType.TORCH)
-        elif FLASH_ATTN_2_AVAILABLE:
-            attn_func = LongContextAttention(attn_type=AttnType.FA)
-        else:
-            raise ValueError("No available long context attention implementation")
+            if flash_attn3_compatible:
+                return LongContextAttention(attn_type=AttnType.FA3)(q, k, v, softmax_scale=scale)
+            else:
+                logger.warning(
+                    f"head_dim={q.shape[-1]}, but flash_attn_3 only supports head dimension at most {FA3_MAX_HEADDIM}, will use fallback attention implementation"
+                )
+        if SDPA_AVAILABLE:
+            return LongContextAttention(attn_type=AttnType.TORCH)(q, k, v, softmax_scale=scale)
+        if FLASH_ATTN_2_AVAILABLE:
+            return LongContextAttention(attn_type=AttnType.FA)(q, k, v, softmax_scale=scale)
+        raise ValueError("No available long context attention implementation")
     else:
         if attn_impl == "flash_attn_3":
-            attn_func = LongContextAttention(attn_type=AttnType.FA3)
-        elif attn_impl == "flash_attn_2":
-            attn_func = LongContextAttention(attn_type=AttnType.FA)
-        elif attn_impl == "sdpa":
-            attn_func = LongContextAttention(attn_type=AttnType.TORCH)
-        elif attn_impl == "sage_attn":
-            attn_func = LongContextAttention(attn_type=AttnType.SAGE_FP8)
-        elif attn_impl == "sparge_attn":
+            if flash_attn3_compatible:
+                return LongContextAttention(attn_type=AttnType.FA3)(q, k, v, softmax_scale=scale)
+            else:
+                raise RuntimeError(
+                    f"head_dim={q.shape[-1]}, but flash_attn_3 only supports head dimension at most {FA3_MAX_HEADDIM}"
+                )
+        if attn_impl == "flash_attn_2":
+            return LongContextAttention(attn_type=AttnType.FA)(q, k, v, softmax_scale=scale)
+        if attn_impl == "sdpa":
+            return LongContextAttention(attn_type=AttnType.TORCH)(q, k, v, softmax_scale=scale)
+        if attn_impl == "sage_attn":
+            return LongContextAttention(attn_type=AttnType.SAGE_FP8)(q, k, v, softmax_scale=scale)
+        if attn_impl == "sparge_attn":
             attn_processor = SparseAttentionMeansim()
             # default args from spas_sage2_attn_meansim_cuda
             attn_processor.smooth_k = torch.tensor(kwargs.get("sparge_smooth_k", True))
             attn_processor.simthreshd1 = torch.tensor(kwargs.get("sparge_simthreshd1", 0.6))
             attn_processor.cdfthreshd = torch.tensor(kwargs.get("sparge_cdfthreshd", 0.98))
             attn_processor.pvthreshd = torch.tensor(kwargs.get("sparge_pvthreshd", 50))
-            attn_func = LongContextAttention(attn_type=AttnType.SPARSE_SAGE, attn_processor=attn_processor)
-        else:
-            raise ValueError(f"Invalid long context attention implementation: {attn_impl}")
-    return attn_func(q, k, v, softmax_scale=scale)
+            return LongContextAttention(attn_type=AttnType.SPARSE_SAGE, attn_processor=attn_processor)(
+                q, k, v, softmax_scale=scale
+            )
+        raise ValueError(f"Invalid long context attention implementation: {attn_impl}")

diffsynth_engine/pipelines/__init__.py

@@ -1,5 +1,4 @@
 from .base import BasePipeline, LoRAStateDictConverter
-from .controlnet_helper import ControlNetParams
 from .flux_image import FluxImagePipeline
 from .sdxl_image import SDXLImagePipeline
 from .sd_image import SDImagePipeline
@@ -13,5 +12,4 @@ __all__ = [
     "SDXLImagePipeline",
     "SDImagePipeline",
     "WanVideoPipeline",
-    "ControlNetParams",
 ]

diffsynth_engine/pipelines/flux_image.py

@@ -17,10 +17,10 @@ from diffsynth_engine.models.flux import (
     flux_dit_config,
     flux_text_encoder_config,
 )
-from diffsynth_engine.configs import FluxPipelineConfig, ControlType
+from diffsynth_engine.configs import FluxPipelineConfig, ControlType, ControlNetParams
 from diffsynth_engine.models.basic.lora import LoRAContext
 from diffsynth_engine.pipelines import BasePipeline, LoRAStateDictConverter
-from diffsynth_engine.pipelines.controlnet_helper import ControlNetParams, accumulate
+from diffsynth_engine.pipelines.utils import accumulate
 from diffsynth_engine.tokenizers import CLIPTokenizer, T5TokenizerFast
 from diffsynth_engine.algorithm.noise_scheduler import RecifitedFlowScheduler
 from diffsynth_engine.algorithm.sampler import FlowMatchEulerSampler

diffsynth_engine/pipelines/sd_image.py

@@ -6,12 +6,12 @@ from typing import Callable, Dict, Optional, List
 from tqdm import tqdm
 from PIL import Image, ImageOps
 
-from diffsynth_engine.configs import SDPipelineConfig
+from diffsynth_engine.configs import SDPipelineConfig, ControlNetParams
 from diffsynth_engine.models.base import split_suffix
 from diffsynth_engine.models.basic.lora import LoRAContext
 from diffsynth_engine.models.sd import SDTextEncoder, SDVAEDecoder, SDVAEEncoder, SDUNet, sd_unet_config
 from diffsynth_engine.pipelines import BasePipeline, LoRAStateDictConverter
-from diffsynth_engine.pipelines.controlnet_helper import ControlNetParams, accumulate
+from diffsynth_engine.pipelines.utils import accumulate
 from diffsynth_engine.tokenizers import CLIPTokenizer
 from diffsynth_engine.algorithm.noise_scheduler import ScaledLinearScheduler
 from diffsynth_engine.algorithm.sampler import EulerSampler

diffsynth_engine/pipelines/sdxl_image.py

@@ -6,7 +6,7 @@ from typing import Callable, Dict, Optional, List
 from tqdm import tqdm
 from PIL import Image, ImageOps
 
-from diffsynth_engine.configs import SDXLPipelineConfig
+from diffsynth_engine.configs import SDXLPipelineConfig, ControlNetParams
 from diffsynth_engine.models.base import split_suffix
 from diffsynth_engine.models.basic.lora import LoRAContext
 from diffsynth_engine.models.basic.timestep import TemporalTimesteps
@@ -19,7 +19,7 @@ from diffsynth_engine.models.sdxl import (
     sdxl_unet_config,
 )
 from diffsynth_engine.pipelines import BasePipeline, LoRAStateDictConverter
-from diffsynth_engine.pipelines.controlnet_helper import ControlNetParams, accumulate
+from diffsynth_engine.pipelines.utils import accumulate
 from diffsynth_engine.tokenizers import CLIPTokenizer
 from diffsynth_engine.algorithm.noise_scheduler import ScaledLinearScheduler
 from diffsynth_engine.algorithm.sampler import EulerSampler

diffsynth_engine/pipelines/utils.py

@@ -0,0 +1,6 @@
+def accumulate(result, new_item):
+    if result is None:
+        return new_item
+    for i, item in enumerate(new_item):
+        result[i] += item
+    return result

{diffsynth_engine-0.3.6.dev11.dist-info → diffsynth_engine-0.3.6.dev13.dist-info}/METADATA

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

{diffsynth_engine-0.3.6.dev11.dist-info → diffsynth_engine-0.3.6.dev13.dist-info}/RECORD

@@ -1,4 +1,4 @@
-diffsynth_engine/__init__.py,sha256=PnsxBE7qAW_5yDsrl1S-I3UraXMOQKTHWxAfKHbwIYQ,1279
+diffsynth_engine/__init__.py,sha256=ysgNUqKZwce7rt_JdytIOPAJH5KYiH_LQqh-JQ51ZY8,1315
 diffsynth_engine/algorithm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/algorithm/noise_scheduler/__init__.py,sha256=YvcwE2tCNua-OAX9GEPm0EXsINNWH4XvJMNZb-uaZMM,745
 diffsynth_engine/algorithm/noise_scheduler/base_scheduler.py,sha256=WICrLEh7b2TdZMMEN14NqiYydj7dxXT6RolXymKiMk8,188
@@ -60,15 +60,15 @@ diffsynth_engine/conf/tokenizers/wan/umt5-xxl/special_tokens_map.json,sha256=e4q
 diffsynth_engine/conf/tokenizers/wan/umt5-xxl/spiece.model,sha256=45CaZ7eAZQs1z1Kax4KtK2sm5tH4SdP7tqhykF9FJFg,4548313
 diffsynth_engine/conf/tokenizers/wan/umt5-xxl/tokenizer.json,sha256=bhl7TT29cdoUtOslX0-pHJwfIGiyCi3iRylnyj0iYCs,16837417
 diffsynth_engine/conf/tokenizers/wan/umt5-xxl/tokenizer_config.json,sha256=7Zo6iw-qcacKMoR-BDX-A25uES1N9O23u0ipIeNE3AU,61728
-diffsynth_engine/configs/__init__.py,sha256=dZ80g2GB3B2YdmoGMp9yvwK3FRJI5j8vShhB9L95j1U,460
-diffsynth_engine/configs/controlnet.py,sha256=OF_cznEw-NpGTM9vP_mIApr4MAJCywWoDWkcUWCz-bs,434
+diffsynth_engine/configs/__init__.py,sha256=qvfbnHf3wK9THPU_mFr1Qx_lU80BaUp5HpxUmjoNy60,502
+diffsynth_engine/configs/controlnet.py,sha256=EpUkCdRNk2G5uo56syaOzPFdR9g0sDHRXckagmMsgaQ,948
 diffsynth_engine/configs/pipeline.py,sha256=NPQlNz-AOpi8qFzRob0RNnOqSc8C-vCdHbstLyUugeo,7703
 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=sbyyGP-ENnqicr6cxjEmXRf6dWrmKjCu6k5yamuJ518,2665
 diffsynth_engine/models/utils.py,sha256=r5xLSEog1_ODaFrpqzJvAj3r23PQiEpgivzErClTZTg,1561
 diffsynth_engine/models/basic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-diffsynth_engine/models/basic/attention.py,sha256=Hquc9H57N37hz5JbCJHpuPGX1smG403DNiOZTnlCYYA,9103
+diffsynth_engine/models/basic/attention.py,sha256=vwQZi4MPTE5AV3Tv3KDAQ4TIGOy4UlD5n93zsbtBMWg,10321
 diffsynth_engine/models/basic/lora.py,sha256=qEh44zfh7ZBblLpjmKzwzAxmTlVyY0wu9IkGsnr7Ih8,10614
 diffsynth_engine/models/basic/relative_position_emb.py,sha256=rCXOweZMcayVnNUVvBcYXMdhHS257B_PC8PZSWxvhNQ,2540
 diffsynth_engine/models/basic/timestep.py,sha256=WJODYqkSXEM0wcS42YkkfrGwxWt0e60zMTkDdUBQqBw,2810
@@ -107,12 +107,12 @@ diffsynth_engine/models/wan/wan_dit.py,sha256=gUd9KeMl7y_VPLntGoGtT2Io94opPiKlrr
 diffsynth_engine/models/wan/wan_image_encoder.py,sha256=LYwcfCcQmXf9FP08DGaU2bfaPgFfdpJ23OpJP8UCggo,14397
 diffsynth_engine/models/wan/wan_text_encoder.py,sha256=bkphxtqNNwXcEA_OaUrwV9CvICV-s16awu5Z9gjjzsM,10912
 diffsynth_engine/models/wan/wan_vae.py,sha256=RxyuHExQmRjGBAqhZdIbtwZFdCibTzh__U4-Sa00zdI,29004
-diffsynth_engine/pipelines/__init__.py,sha256=Ewarnhf4K-sYFfSG4mghDoJh5FZKG9Xiz2DFZizNZ-I,452
+diffsynth_engine/pipelines/__init__.py,sha256=kTvANqHcMPrHqiJVg-XohfqRdW6Cj4aElfItTb1B7Vs,380
 diffsynth_engine/pipelines/base.py,sha256=yVp4hSPCqk98azzy3ykKBfPAufvq_ncTFOURN95z7d0,12178
-diffsynth_engine/pipelines/controlnet_helper.py,sha256=b6HnJFJfMKZq9s5DQ-9Se8OTSDeHVk4AskONSwcRShg,680
-diffsynth_engine/pipelines/flux_image.py,sha256=CPI7WwXJz60rpQN_ZfVV5kcWxg8WlrYNNe6QI1E5EPk,48851
-diffsynth_engine/pipelines/sd_image.py,sha256=hiL2gvQcPgniRm8TlzUKhoo5bGUmmVlYDS__E_WFDiE,17834
-diffsynth_engine/pipelines/sdxl_image.py,sha256=qWpE5q0CeDrsKZVIxYHpYLyZAkKMqACDzf4RPPvdn7A,21587
+diffsynth_engine/pipelines/flux_image.py,sha256=MtQqTnCqQjIFovhA3lzBXpnkS4DkZH2PtFUwNZdl42M,48839
+diffsynth_engine/pipelines/sd_image.py,sha256=5dGIa6crtklO7xPd1eeBVkqj54Pe89Uo3bMyXVEaXxM,17822
+diffsynth_engine/pipelines/sdxl_image.py,sha256=Ns4bCSO3BtCXdjGJEQ0s5oY0S3jrp5yE5lhfon-iNiw,21575
+diffsynth_engine/pipelines/utils.py,sha256=VfSTwRejSVSKXIa7w0VhObmvaBFRvDP-uiYsHHkPAgs,165
 diffsynth_engine/pipelines/wan_video.py,sha256=vi_xW-jU4PeMtZzjkfQbnj8eOymJrTZMrOQau6tx6ks,20187
 diffsynth_engine/processor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 diffsynth_engine/processor/canny_processor.py,sha256=hV30NlblTkEFUAmF_O-LJrNlGVM2SFrqq6okfF8VpOo,602
@@ -144,8 +144,8 @@ diffsynth_engine/utils/parallel.py,sha256=gbIeilfOYsqeDcgkaP68TfLjIXxvD0KfLiAsR_
 diffsynth_engine/utils/platform.py,sha256=2lXdw6YkqcRONCeT98n4cyg1Ii8Ybbyj2Ns72Se9tlk,496
 diffsynth_engine/utils/prompt.py,sha256=YItMchoVzsG6y-LB4vzzDUWrkhKRVlt1HfVhxZjSxMQ,280
 diffsynth_engine/utils/video.py,sha256=Ne0rd2lb59UT1q5EotpjlY7OT8F9oTCFDyo1ST77uoQ,1004
-diffsynth_engine-0.3.6.dev11.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
-diffsynth_engine-0.3.6.dev11.dist-info/METADATA,sha256=pFhxaPrL9JrwA9ZnbtL0VLrS6gVC52wFwRu2Yn_vqQc,1069
-diffsynth_engine-0.3.6.dev11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-diffsynth_engine-0.3.6.dev11.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
-diffsynth_engine-0.3.6.dev11.dist-info/RECORD,,
+diffsynth_engine-0.3.6.dev13.dist-info/licenses/LICENSE,sha256=x7aBqQuVI0IYnftgoTPI_A0I_rjdjPPQkjnU6N2nikM,11346
+diffsynth_engine-0.3.6.dev13.dist-info/METADATA,sha256=2jH1jlJdbUga4JOoDHfRyKEn6E4xQ1w9wRhLKVYaqRk,1069
+diffsynth_engine-0.3.6.dev13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+diffsynth_engine-0.3.6.dev13.dist-info/top_level.txt,sha256=6zgbiIzEHLbhgDKRyX0uBJOV3F6VnGGBRIQvSiYYn6w,17
+diffsynth_engine-0.3.6.dev13.dist-info/RECORD,,

diffsynth_engine/pipelines/controlnet_helper.py

@@ -1,26 +0,0 @@
-import torch
-import torch.nn as nn
-from typing import List, Union, Optional
-from PIL import Image
-from dataclasses import dataclass
-
-ImageType = Union[Image.Image, torch.Tensor, List[Image.Image], List[torch.Tensor]]
-
-
-@dataclass
-class ControlNetParams:
-    image: ImageType
-    scale: float = 1.0
-    model: Optional[nn.Module] = None
-    mask: Optional[ImageType] = None
-    control_start: float = 0
-    control_end: float = 1
-    processor_name: Optional[str] = None  # only used for sdxl controlnet union now
-
-
-def accumulate(result, new_item):
-    if result is None:
-        return new_item
-    for i, item in enumerate(new_item):
-        result[i] += item
-    return result