diffsynth-engine 0.3.6.dev11__tar.gz → 0.3.6.dev13__tar.gz

{diffsynth_engine-0.3.6.dev11 → diffsynth_engine-0.3.6.dev13}/PKG-INFO

@@ -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 → diffsynth_engine-0.3.6.dev13}/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-0.3.6.dev11 → diffsynth_engine-0.3.6.dev13}/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-0.3.6.dev11/diffsynth_engine/pipelines/controlnet_helper.py → diffsynth_engine-0.3.6.dev13/diffsynth_engine/configs/controlnet.py

@@ -1,12 +1,30 @@
+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
-from dataclasses import dataclass
 
 ImageType = Union[Image.Image, torch.Tensor, List[Image.Image], List[torch.Tensor]]
 
 
+# FLUX ControlType
+class ControlType(Enum):
+    normal = "normal"
+    bfl_control = "bfl_control"
+    bfl_fill = "bfl_fill"
+    bfl_kontext = "bfl_kontext"
+
+    def get_in_channel(self):
+        if self in [ControlType.normal, ControlType.bfl_kontext]:
+            return 64
+        elif self == ControlType.bfl_control:
+            return 128
+        elif self == ControlType.bfl_fill:
+            return 384
+
+
 @dataclass
 class ControlNetParams:
     image: ImageType
@@ -16,11 +34,3 @@ class ControlNetParams:
     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

{diffsynth_engine-0.3.6.dev11 → diffsynth_engine-0.3.6.dev13}/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-0.3.6.dev11 → diffsynth_engine-0.3.6.dev13}/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-0.3.6.dev11 → diffsynth_engine-0.3.6.dev13}/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-0.3.6.dev11 → diffsynth_engine-0.3.6.dev13}/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-0.3.6.dev11 → diffsynth_engine-0.3.6.dev13}/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-0.3.6.dev13/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 → diffsynth_engine-0.3.6.dev13}/diffsynth_engine.egg-info/PKG-INFO

@@ -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 → diffsynth_engine-0.3.6.dev13}/diffsynth_engine.egg-info/SOURCES.txt

@@ -123,10 +123,10 @@ diffsynth_engine/models/wan/wan_text_encoder.py
 diffsynth_engine/models/wan/wan_vae.py
 diffsynth_engine/pipelines/__init__.py
 diffsynth_engine/pipelines/base.py
-diffsynth_engine/pipelines/controlnet_helper.py
 diffsynth_engine/pipelines/flux_image.py
 diffsynth_engine/pipelines/sd_image.py
 diffsynth_engine/pipelines/sdxl_image.py
+diffsynth_engine/pipelines/utils.py
 diffsynth_engine/pipelines/wan_video.py
 diffsynth_engine/processor/__init__.py
 diffsynth_engine/processor/canny_processor.py

diffsynth_engine-0.3.6.dev11/diffsynth_engine/configs/controlnet.py

@@ -1,17 +0,0 @@
-from enum import Enum
-
-
-# FLUX ControlType
-class ControlType(Enum):
-    normal = "normal"
-    bfl_control = "bfl_control"
-    bfl_fill = "bfl_fill"
-    bfl_kontext = "bfl_kontext"
-
-    def get_in_channel(self):
-        if self in [ControlType.normal, ControlType.bfl_kontext]:
-            return 64
-        elif self == ControlType.bfl_control:
-            return 128
-        elif self == ControlType.bfl_fill:
-            return 384