diffsynth-engine 0.5.1.dev4__py3-none-any.whl → 0.6.1.dev25__py3-none-any.whl

diffsynth_engine/__init__.py

@@ -11,8 +11,14 @@ from .configs import (
     FluxStateDicts,
     WanStateDicts,
     QwenImageStateDicts,
+    AttnImpl,
+    SpargeAttentionParams,
+    VideoSparseAttentionParams,
+    LoraConfig,
     ControlNetParams,
     ControlType,
+    QwenImageControlNetParams,
+    QwenImageControlType,
 )
 from .pipelines import (
     SDImagePipeline,
@@ -54,8 +60,14 @@ __all__ = [
     "FluxStateDicts",
     "WanStateDicts",
     "QwenImageStateDicts",
+    "AttnImpl",
+    "SpargeAttentionParams",
+    "VideoSparseAttentionParams",
+    "LoraConfig",
     "ControlNetParams",
     "ControlType",
+    "QwenImageControlNetParams",
+    "QwenImageControlType",
     "SDImagePipeline",
     "SDControlNet",
     "SDXLImagePipeline",

diffsynth_engine/algorithm/noise_scheduler/base_scheduler.py

@@ -6,5 +6,24 @@ def append_zero(x):
 
 
 class BaseScheduler:
+    def __init__(self):
+        self._stored_config = {}
+
+    def store_config(self):
+        self._stored_config = {
+            config_name: config_value
+            for config_name, config_value in vars(self).items()
+            if not config_name.startswith("_")
+        }
+
+    def update_config(self, config_dict):
+        for config_name, new_value in config_dict.items():
+            if hasattr(self, config_name):
+                setattr(self, config_name, new_value)
+
+    def restore_config(self):
+        for config_name, config_value in self._stored_config.items():
+            setattr(self, config_name, config_value)
+
     def schedule(self, num_inference_steps: int):
         raise NotImplementedError()

diffsynth_engine/algorithm/noise_scheduler/flow_match/recifited_flow.py

@@ -12,16 +12,23 @@ class RecifitedFlowScheduler(BaseScheduler):
     def __init__(
         self,
         shift=1.0,
-        sigma_min=0.001,
-        sigma_max=1.0,
+        sigma_min=None,
+        sigma_max=None,
         num_train_timesteps=1000,
         use_dynamic_shifting=False,
+        shift_terminal=None,
+        exponential_shift_mu=None,
     ):
+        super().__init__()
         self.shift = shift
         self.sigma_min = sigma_min
         self.sigma_max = sigma_max
         self.num_train_timesteps = num_train_timesteps
         self.use_dynamic_shifting = use_dynamic_shifting
+        self.shift_terminal = shift_terminal
+        # static mu for distill model
+        self.exponential_shift_mu = exponential_shift_mu
+        self.store_config()
 
     def _sigma_to_t(self, sigma):
         return sigma * self.num_train_timesteps
@@ -35,21 +42,30 @@ class RecifitedFlowScheduler(BaseScheduler):
     def _shift_sigma(self, sigma: torch.Tensor, shift: float):
         return shift * sigma / (1 + (shift - 1) * sigma)
 
+    def _stretch_shift_to_terminal(self, sigma: torch.Tensor):
+        one_minus_z = 1 - sigma
+        scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+        return 1 - (one_minus_z / scale_factor)
+
     def schedule(
         self,
         num_inference_steps: int,
         mu: float | None = None,
-        sigma_min: float | None = None,
-        sigma_max: float | None = None,
+        sigma_min: float = 0.001,
+        sigma_max: float = 1.0,
         append_value: float = 0,
     ):
-        sigma_min = self.sigma_min if sigma_min is None else sigma_min
-        sigma_max = self.sigma_max if sigma_max is None else sigma_max
+        sigma_min = sigma_min if self.sigma_min is None else self.sigma_min
+        sigma_max = sigma_max if self.sigma_max is None else self.sigma_max
         sigmas = torch.linspace(sigma_max, sigma_min, num_inference_steps)
+        if self.exponential_shift_mu is not None:
+            mu = self.exponential_shift_mu
         if self.use_dynamic_shifting:
             sigmas = self._time_shift(mu, 1.0, sigmas)  # FLUX
         else:
             sigmas = self._shift_sigma(sigmas, self.shift)
+        if self.shift_terminal is not None:
+            sigmas = self._stretch_shift_to_terminal(sigmas)
         timesteps = sigmas * self.num_train_timesteps
         sigmas = append(sigmas, append_value)
         return sigmas, timesteps

diffsynth_engine/conf/models/flux/flux_dit.json

@@ -101,5 +101,24 @@
             "proj_mlp": "proj_in_besides_attn",
             "proj_out": "proj_out"
         }
-    }
+    },
+    "preferred_kontext_resolutions": [
+        [672, 1568],
+        [688, 1504],
+        [720, 1456],
+        [752, 1392],
+        [800, 1328],
+        [832, 1248],
+        [880, 1184],
+        [944, 1104],
+        [1024, 1024],
+        [1104, 944],
+        [1184, 880],
+        [1248, 832],
+        [1328, 800],
+        [1392, 752],
+        [1456, 720],
+        [1504, 688],
+        [1568, 672]
+    ]
 }

diffsynth_engine/conf/models/flux/flux_vae.json

@@ -5,6 +5,8 @@
             "decoder.conv_in.weight": "decoder.conv_in.weight",
             "decoder.conv_out.bias": "decoder.conv_out.bias",
             "decoder.conv_out.weight": "decoder.conv_out.weight",
+            "decoder.norm_out.bias": "decoder.conv_norm_out.bias",
+            "decoder.norm_out.weight": "decoder.conv_norm_out.weight",
             "decoder.mid.attn_1.k.bias": "decoder.blocks.1.transformer_blocks.0.to_k.bias",
             "decoder.mid.attn_1.k.weight": "decoder.blocks.1.transformer_blocks.0.to_k.weight",
             "decoder.mid.attn_1.norm.bias": "decoder.blocks.1.norm.bias",
@@ -31,8 +33,6 @@
             "decoder.mid.block_2.norm1.weight": "decoder.blocks.2.norm1.weight",
             "decoder.mid.block_2.norm2.bias": "decoder.blocks.2.norm2.bias",
             "decoder.mid.block_2.norm2.weight": "decoder.blocks.2.norm2.weight",
-            "decoder.norm_out.bias": "decoder.conv_norm_out.bias",
-            "decoder.norm_out.weight": "decoder.conv_norm_out.weight",
             "decoder.up.0.block.0.conv1.bias": "decoder.blocks.15.conv1.bias",
             "decoder.up.0.block.0.conv1.weight": "decoder.blocks.15.conv1.weight",
             "decoder.up.0.block.0.conv2.bias": "decoder.blocks.15.conv2.bias",
@@ -143,6 +143,8 @@
             "encoder.conv_in.weight": "encoder.conv_in.weight",
             "encoder.conv_out.bias": "encoder.conv_out.bias",
             "encoder.conv_out.weight": "encoder.conv_out.weight",
+            "encoder.norm_out.bias": "encoder.conv_norm_out.bias",
+            "encoder.norm_out.weight": "encoder.conv_norm_out.weight",
             "encoder.down.0.block.0.conv1.bias": "encoder.blocks.0.conv1.bias",
             "encoder.down.0.block.0.conv1.weight": "encoder.blocks.0.conv1.weight",
             "encoder.down.0.block.0.conv2.bias": "encoder.blocks.0.conv2.bias",
@@ -242,9 +244,255 @@
             "encoder.mid.block_2.norm1.bias": "encoder.blocks.13.norm1.bias",
             "encoder.mid.block_2.norm1.weight": "encoder.blocks.13.norm1.weight",
             "encoder.mid.block_2.norm2.bias": "encoder.blocks.13.norm2.bias",
-            "encoder.mid.block_2.norm2.weight": "encoder.blocks.13.norm2.weight",
-            "encoder.norm_out.bias": "encoder.conv_norm_out.bias",
-            "encoder.norm_out.weight": "encoder.conv_norm_out.weight"
+            "encoder.mid.block_2.norm2.weight": "encoder.blocks.13.norm2.weight"
+        }
+    },
+    "diffusers": {
+        "rename_dict": {
+            "decoder.conv_in.bias": "decoder.conv_in.bias",
+            "decoder.conv_in.weight": "decoder.conv_in.weight",
+            "decoder.conv_out.bias": "decoder.conv_out.bias",
+            "decoder.conv_out.weight": "decoder.conv_out.weight",
+            "decoder.conv_norm_out.bias": "decoder.conv_norm_out.bias",
+            "decoder.conv_norm_out.weight": "decoder.conv_norm_out.weight",
+            "decoder.mid_block.attentions.0.to_k.bias": "decoder.blocks.1.transformer_blocks.0.to_k.bias",
+            "decoder.mid_block.attentions.0.to_k.weight": "decoder.blocks.1.transformer_blocks.0.to_k.weight",
+            "decoder.mid_block.attentions.0.group_norm.bias": "decoder.blocks.1.norm.bias",
+            "decoder.mid_block.attentions.0.group_norm.weight": "decoder.blocks.1.norm.weight",
+            "decoder.mid_block.attentions.0.to_out.0.bias": "decoder.blocks.1.transformer_blocks.0.to_out.bias",
+            "decoder.mid_block.attentions.0.to_out.0.weight": "decoder.blocks.1.transformer_blocks.0.to_out.weight",
+            "decoder.mid_block.attentions.0.to_q.bias": "decoder.blocks.1.transformer_blocks.0.to_q.bias",
+            "decoder.mid_block.attentions.0.to_q.weight": "decoder.blocks.1.transformer_blocks.0.to_q.weight",
+            "decoder.mid_block.attentions.0.to_v.bias": "decoder.blocks.1.transformer_blocks.0.to_v.bias",
+            "decoder.mid_block.attentions.0.to_v.weight": "decoder.blocks.1.transformer_blocks.0.to_v.weight",
+            "decoder.mid_block.resnets.0.conv1.bias": "decoder.blocks.0.conv1.bias",
+            "decoder.mid_block.resnets.0.conv1.weight": "decoder.blocks.0.conv1.weight",
+            "decoder.mid_block.resnets.0.conv2.bias": "decoder.blocks.0.conv2.bias",
+            "decoder.mid_block.resnets.0.conv2.weight": "decoder.blocks.0.conv2.weight",
+            "decoder.mid_block.resnets.0.norm1.bias": "decoder.blocks.0.norm1.bias",
+            "decoder.mid_block.resnets.0.norm1.weight": "decoder.blocks.0.norm1.weight",
+            "decoder.mid_block.resnets.0.norm2.bias": "decoder.blocks.0.norm2.bias",
+            "decoder.mid_block.resnets.0.norm2.weight": "decoder.blocks.0.norm2.weight",
+            "decoder.mid_block.resnets.1.conv1.bias": "decoder.blocks.2.conv1.bias",
+            "decoder.mid_block.resnets.1.conv1.weight": "decoder.blocks.2.conv1.weight",
+            "decoder.mid_block.resnets.1.conv2.bias": "decoder.blocks.2.conv2.bias",
+            "decoder.mid_block.resnets.1.conv2.weight": "decoder.blocks.2.conv2.weight",
+            "decoder.mid_block.resnets.1.norm1.bias": "decoder.blocks.2.norm1.bias",
+            "decoder.mid_block.resnets.1.norm1.weight": "decoder.blocks.2.norm1.weight",
+            "decoder.mid_block.resnets.1.norm2.bias": "decoder.blocks.2.norm2.bias",
+            "decoder.mid_block.resnets.1.norm2.weight": "decoder.blocks.2.norm2.weight",
+            "decoder.up_blocks.0.resnets.0.conv1.bias": "decoder.blocks.3.conv1.bias",
+            "decoder.up_blocks.0.resnets.0.conv1.weight": "decoder.blocks.3.conv1.weight",
+            "decoder.up_blocks.0.resnets.0.conv2.bias": "decoder.blocks.3.conv2.bias",
+            "decoder.up_blocks.0.resnets.0.conv2.weight": "decoder.blocks.3.conv2.weight",
+            "decoder.up_blocks.0.resnets.0.norm1.bias": "decoder.blocks.3.norm1.bias",
+            "decoder.up_blocks.0.resnets.0.norm1.weight": "decoder.blocks.3.norm1.weight",
+            "decoder.up_blocks.0.resnets.0.norm2.bias": "decoder.blocks.3.norm2.bias",
+            "decoder.up_blocks.0.resnets.0.norm2.weight": "decoder.blocks.3.norm2.weight",
+            "decoder.up_blocks.0.resnets.1.conv1.bias": "decoder.blocks.4.conv1.bias",
+            "decoder.up_blocks.0.resnets.1.conv1.weight": "decoder.blocks.4.conv1.weight",
+            "decoder.up_blocks.0.resnets.1.conv2.bias": "decoder.blocks.4.conv2.bias",
+            "decoder.up_blocks.0.resnets.1.conv2.weight": "decoder.blocks.4.conv2.weight",
+            "decoder.up_blocks.0.resnets.1.norm1.bias": "decoder.blocks.4.norm1.bias",
+            "decoder.up_blocks.0.resnets.1.norm1.weight": "decoder.blocks.4.norm1.weight",
+            "decoder.up_blocks.0.resnets.1.norm2.bias": "decoder.blocks.4.norm2.bias",
+            "decoder.up_blocks.0.resnets.1.norm2.weight": "decoder.blocks.4.norm2.weight",
+            "decoder.up_blocks.0.resnets.2.conv1.bias": "decoder.blocks.5.conv1.bias",
+            "decoder.up_blocks.0.resnets.2.conv1.weight": "decoder.blocks.5.conv1.weight",
+            "decoder.up_blocks.0.resnets.2.conv2.bias": "decoder.blocks.5.conv2.bias",
+            "decoder.up_blocks.0.resnets.2.conv2.weight": "decoder.blocks.5.conv2.weight",
+            "decoder.up_blocks.0.resnets.2.norm1.bias": "decoder.blocks.5.norm1.bias",
+            "decoder.up_blocks.0.resnets.2.norm1.weight": "decoder.blocks.5.norm1.weight",
+            "decoder.up_blocks.0.resnets.2.norm2.bias": "decoder.blocks.5.norm2.bias",
+            "decoder.up_blocks.0.resnets.2.norm2.weight": "decoder.blocks.5.norm2.weight",
+            "decoder.up_blocks.0.upsamplers.0.conv.bias": "decoder.blocks.6.conv.bias",
+            "decoder.up_blocks.0.upsamplers.0.conv.weight": "decoder.blocks.6.conv.weight",
+            "decoder.up_blocks.1.resnets.0.conv1.bias": "decoder.blocks.7.conv1.bias",
+            "decoder.up_blocks.1.resnets.0.conv1.weight": "decoder.blocks.7.conv1.weight",
+            "decoder.up_blocks.1.resnets.0.conv2.bias": "decoder.blocks.7.conv2.bias",
+            "decoder.up_blocks.1.resnets.0.conv2.weight": "decoder.blocks.7.conv2.weight",
+            "decoder.up_blocks.1.resnets.0.norm1.bias": "decoder.blocks.7.norm1.bias",
+            "decoder.up_blocks.1.resnets.0.norm1.weight": "decoder.blocks.7.norm1.weight",
+            "decoder.up_blocks.1.resnets.0.norm2.bias": "decoder.blocks.7.norm2.bias",
+            "decoder.up_blocks.1.resnets.0.norm2.weight": "decoder.blocks.7.norm2.weight",
+            "decoder.up_blocks.1.resnets.1.conv1.bias": "decoder.blocks.8.conv1.bias",
+            "decoder.up_blocks.1.resnets.1.conv1.weight": "decoder.blocks.8.conv1.weight",
+            "decoder.up_blocks.1.resnets.1.conv2.bias": "decoder.blocks.8.conv2.bias",
+            "decoder.up_blocks.1.resnets.1.conv2.weight": "decoder.blocks.8.conv2.weight",
+            "decoder.up_blocks.1.resnets.1.norm1.bias": "decoder.blocks.8.norm1.bias",
+            "decoder.up_blocks.1.resnets.1.norm1.weight": "decoder.blocks.8.norm1.weight",
+            "decoder.up_blocks.1.resnets.1.norm2.bias": "decoder.blocks.8.norm2.bias",
+            "decoder.up_blocks.1.resnets.1.norm2.weight": "decoder.blocks.8.norm2.weight",
+            "decoder.up_blocks.1.resnets.2.conv1.bias": "decoder.blocks.9.conv1.bias",
+            "decoder.up_blocks.1.resnets.2.conv1.weight": "decoder.blocks.9.conv1.weight",
+            "decoder.up_blocks.1.resnets.2.conv2.bias": "decoder.blocks.9.conv2.bias",
+            "decoder.up_blocks.1.resnets.2.conv2.weight": "decoder.blocks.9.conv2.weight",
+            "decoder.up_blocks.1.resnets.2.norm1.bias": "decoder.blocks.9.norm1.bias",
+            "decoder.up_blocks.1.resnets.2.norm1.weight": "decoder.blocks.9.norm1.weight",
+            "decoder.up_blocks.1.resnets.2.norm2.bias": "decoder.blocks.9.norm2.bias",
+            "decoder.up_blocks.1.resnets.2.norm2.weight": "decoder.blocks.9.norm2.weight",
+            "decoder.up_blocks.1.upsamplers.0.conv.bias": "decoder.blocks.10.conv.bias",
+            "decoder.up_blocks.1.upsamplers.0.conv.weight": "decoder.blocks.10.conv.weight",
+            "decoder.up_blocks.2.resnets.0.conv1.bias": "decoder.blocks.11.conv1.bias",
+            "decoder.up_blocks.2.resnets.0.conv1.weight": "decoder.blocks.11.conv1.weight",
+            "decoder.up_blocks.2.resnets.0.conv2.bias": "decoder.blocks.11.conv2.bias",
+            "decoder.up_blocks.2.resnets.0.conv2.weight": "decoder.blocks.11.conv2.weight",
+            "decoder.up_blocks.2.resnets.0.conv_shortcut.bias": "decoder.blocks.11.conv_shortcut.bias",
+            "decoder.up_blocks.2.resnets.0.conv_shortcut.weight": "decoder.blocks.11.conv_shortcut.weight",
+            "decoder.up_blocks.2.resnets.0.norm1.bias": "decoder.blocks.11.norm1.bias",
+            "decoder.up_blocks.2.resnets.0.norm1.weight": "decoder.blocks.11.norm1.weight",
+            "decoder.up_blocks.2.resnets.0.norm2.bias": "decoder.blocks.11.norm2.bias",
+            "decoder.up_blocks.2.resnets.0.norm2.weight": "decoder.blocks.11.norm2.weight",
+            "decoder.up_blocks.2.resnets.1.conv1.bias": "decoder.blocks.12.conv1.bias",
+            "decoder.up_blocks.2.resnets.1.conv1.weight": "decoder.blocks.12.conv1.weight",
+            "decoder.up_blocks.2.resnets.1.conv2.bias": "decoder.blocks.12.conv2.bias",
+            "decoder.up_blocks.2.resnets.1.conv2.weight": "decoder.blocks.12.conv2.weight",
+            "decoder.up_blocks.2.resnets.1.norm1.bias": "decoder.blocks.12.norm1.bias",
+            "decoder.up_blocks.2.resnets.1.norm1.weight": "decoder.blocks.12.norm1.weight",
+            "decoder.up_blocks.2.resnets.1.norm2.bias": "decoder.blocks.12.norm2.bias",
+            "decoder.up_blocks.2.resnets.1.norm2.weight": "decoder.blocks.12.norm2.weight",
+            "decoder.up_blocks.2.resnets.2.conv1.bias": "decoder.blocks.13.conv1.bias",
+            "decoder.up_blocks.2.resnets.2.conv1.weight": "decoder.blocks.13.conv1.weight",
+            "decoder.up_blocks.2.resnets.2.conv2.bias": "decoder.blocks.13.conv2.bias",
+            "decoder.up_blocks.2.resnets.2.conv2.weight": "decoder.blocks.13.conv2.weight",
+            "decoder.up_blocks.2.resnets.2.norm1.bias": "decoder.blocks.13.norm1.bias",
+            "decoder.up_blocks.2.resnets.2.norm1.weight": "decoder.blocks.13.norm1.weight",
+            "decoder.up_blocks.2.resnets.2.norm2.bias": "decoder.blocks.13.norm2.bias",
+            "decoder.up_blocks.2.resnets.2.norm2.weight": "decoder.blocks.13.norm2.weight",
+            "decoder.up_blocks.2.upsamplers.0.conv.bias": "decoder.blocks.14.conv.bias",
+            "decoder.up_blocks.2.upsamplers.0.conv.weight": "decoder.blocks.14.conv.weight",
+            "decoder.up_blocks.3.resnets.0.conv1.bias": "decoder.blocks.15.conv1.bias",
+            "decoder.up_blocks.3.resnets.0.conv1.weight": "decoder.blocks.15.conv1.weight",
+            "decoder.up_blocks.3.resnets.0.conv2.bias": "decoder.blocks.15.conv2.bias",
+            "decoder.up_blocks.3.resnets.0.conv2.weight": "decoder.blocks.15.conv2.weight",
+            "decoder.up_blocks.3.resnets.0.conv_shortcut.bias": "decoder.blocks.15.conv_shortcut.bias",
+            "decoder.up_blocks.3.resnets.0.conv_shortcut.weight": "decoder.blocks.15.conv_shortcut.weight",
+            "decoder.up_blocks.3.resnets.0.norm1.bias": "decoder.blocks.15.norm1.bias",
+            "decoder.up_blocks.3.resnets.0.norm1.weight": "decoder.blocks.15.norm1.weight",
+            "decoder.up_blocks.3.resnets.0.norm2.bias": "decoder.blocks.15.norm2.bias",
+            "decoder.up_blocks.3.resnets.0.norm2.weight": "decoder.blocks.15.norm2.weight",
+            "decoder.up_blocks.3.resnets.1.conv1.bias": "decoder.blocks.16.conv1.bias",
+            "decoder.up_blocks.3.resnets.1.conv1.weight": "decoder.blocks.16.conv1.weight",
+            "decoder.up_blocks.3.resnets.1.conv2.bias": "decoder.blocks.16.conv2.bias",
+            "decoder.up_blocks.3.resnets.1.conv2.weight": "decoder.blocks.16.conv2.weight",
+            "decoder.up_blocks.3.resnets.1.norm1.bias": "decoder.blocks.16.norm1.bias",
+            "decoder.up_blocks.3.resnets.1.norm1.weight": "decoder.blocks.16.norm1.weight",
+            "decoder.up_blocks.3.resnets.1.norm2.bias": "decoder.blocks.16.norm2.bias",
+            "decoder.up_blocks.3.resnets.1.norm2.weight": "decoder.blocks.16.norm2.weight",
+            "decoder.up_blocks.3.resnets.2.conv1.bias": "decoder.blocks.17.conv1.bias",
+            "decoder.up_blocks.3.resnets.2.conv1.weight": "decoder.blocks.17.conv1.weight",
+            "decoder.up_blocks.3.resnets.2.conv2.bias": "decoder.blocks.17.conv2.bias",
+            "decoder.up_blocks.3.resnets.2.conv2.weight": "decoder.blocks.17.conv2.weight",
+            "decoder.up_blocks.3.resnets.2.norm1.bias": "decoder.blocks.17.norm1.bias",
+            "decoder.up_blocks.3.resnets.2.norm1.weight": "decoder.blocks.17.norm1.weight",
+            "decoder.up_blocks.3.resnets.2.norm2.bias": "decoder.blocks.17.norm2.bias",
+            "decoder.up_blocks.3.resnets.2.norm2.weight": "decoder.blocks.17.norm2.weight",
+            "encoder.conv_in.bias": "encoder.conv_in.bias",
+            "encoder.conv_in.weight": "encoder.conv_in.weight",
+            "encoder.conv_out.bias": "encoder.conv_out.bias",
+            "encoder.conv_out.weight": "encoder.conv_out.weight",
+            "encoder.conv_norm_out.bias": "encoder.conv_norm_out.bias",
+            "encoder.conv_norm_out.weight": "encoder.conv_norm_out.weight",
+            "encoder.down_blocks.0.resnets.0.conv1.bias": "encoder.blocks.0.conv1.bias",
+            "encoder.down_blocks.0.resnets.0.conv1.weight": "encoder.blocks.0.conv1.weight",
+            "encoder.down_blocks.0.resnets.0.conv2.bias": "encoder.blocks.0.conv2.bias",
+            "encoder.down_blocks.0.resnets.0.conv2.weight": "encoder.blocks.0.conv2.weight",
+            "encoder.down_blocks.0.resnets.0.norm1.bias": "encoder.blocks.0.norm1.bias",
+            "encoder.down_blocks.0.resnets.0.norm1.weight": "encoder.blocks.0.norm1.weight",
+            "encoder.down_blocks.0.resnets.0.norm2.bias": "encoder.blocks.0.norm2.bias",
+            "encoder.down_blocks.0.resnets.0.norm2.weight": "encoder.blocks.0.norm2.weight",
+            "encoder.down_blocks.0.resnets.1.conv1.bias": "encoder.blocks.1.conv1.bias",
+            "encoder.down_blocks.0.resnets.1.conv1.weight": "encoder.blocks.1.conv1.weight",
+            "encoder.down_blocks.0.resnets.1.conv2.bias": "encoder.blocks.1.conv2.bias",
+            "encoder.down_blocks.0.resnets.1.conv2.weight": "encoder.blocks.1.conv2.weight",
+            "encoder.down_blocks.0.resnets.1.norm1.bias": "encoder.blocks.1.norm1.bias",
+            "encoder.down_blocks.0.resnets.1.norm1.weight": "encoder.blocks.1.norm1.weight",
+            "encoder.down_blocks.0.resnets.1.norm2.bias": "encoder.blocks.1.norm2.bias",
+            "encoder.down_blocks.0.resnets.1.norm2.weight": "encoder.blocks.1.norm2.weight",
+            "encoder.down_blocks.0.downsamplers.0.conv.bias": "encoder.blocks.2.conv.bias",
+            "encoder.down_blocks.0.downsamplers.0.conv.weight": "encoder.blocks.2.conv.weight",
+            "encoder.down_blocks.1.resnets.0.conv1.bias": "encoder.blocks.3.conv1.bias",
+            "encoder.down_blocks.1.resnets.0.conv1.weight": "encoder.blocks.3.conv1.weight",
+            "encoder.down_blocks.1.resnets.0.conv2.bias": "encoder.blocks.3.conv2.bias",
+            "encoder.down_blocks.1.resnets.0.conv2.weight": "encoder.blocks.3.conv2.weight",
+            "encoder.down_blocks.1.resnets.0.conv_shortcut.bias": "encoder.blocks.3.conv_shortcut.bias",
+            "encoder.down_blocks.1.resnets.0.conv_shortcut.weight": "encoder.blocks.3.conv_shortcut.weight",
+            "encoder.down_blocks.1.resnets.0.norm1.bias": "encoder.blocks.3.norm1.bias",
+            "encoder.down_blocks.1.resnets.0.norm1.weight": "encoder.blocks.3.norm1.weight",
+            "encoder.down_blocks.1.resnets.0.norm2.bias": "encoder.blocks.3.norm2.bias",
+            "encoder.down_blocks.1.resnets.0.norm2.weight": "encoder.blocks.3.norm2.weight",
+            "encoder.down_blocks.1.resnets.1.conv1.bias": "encoder.blocks.4.conv1.bias",
+            "encoder.down_blocks.1.resnets.1.conv1.weight": "encoder.blocks.4.conv1.weight",
+            "encoder.down_blocks.1.resnets.1.conv2.bias": "encoder.blocks.4.conv2.bias",
+            "encoder.down_blocks.1.resnets.1.conv2.weight": "encoder.blocks.4.conv2.weight",
+            "encoder.down_blocks.1.resnets.1.norm1.bias": "encoder.blocks.4.norm1.bias",
+            "encoder.down_blocks.1.resnets.1.norm1.weight": "encoder.blocks.4.norm1.weight",
+            "encoder.down_blocks.1.resnets.1.norm2.bias": "encoder.blocks.4.norm2.bias",
+            "encoder.down_blocks.1.resnets.1.norm2.weight": "encoder.blocks.4.norm2.weight",
+            "encoder.down_blocks.1.downsamplers.0.conv.bias": "encoder.blocks.5.conv.bias",
+            "encoder.down_blocks.1.downsamplers.0.conv.weight": "encoder.blocks.5.conv.weight",
+            "encoder.down_blocks.2.resnets.0.conv1.bias": "encoder.blocks.6.conv1.bias",
+            "encoder.down_blocks.2.resnets.0.conv1.weight": "encoder.blocks.6.conv1.weight",
+            "encoder.down_blocks.2.resnets.0.conv2.bias": "encoder.blocks.6.conv2.bias",
+            "encoder.down_blocks.2.resnets.0.conv2.weight": "encoder.blocks.6.conv2.weight",
+            "encoder.down_blocks.2.resnets.0.conv_shortcut.bias": "encoder.blocks.6.conv_shortcut.bias",
+            "encoder.down_blocks.2.resnets.0.conv_shortcut.weight": "encoder.blocks.6.conv_shortcut.weight",
+            "encoder.down_blocks.2.resnets.0.norm1.bias": "encoder.blocks.6.norm1.bias",
+            "encoder.down_blocks.2.resnets.0.norm1.weight": "encoder.blocks.6.norm1.weight",
+            "encoder.down_blocks.2.resnets.0.norm2.bias": "encoder.blocks.6.norm2.bias",
+            "encoder.down_blocks.2.resnets.0.norm2.weight": "encoder.blocks.6.norm2.weight",
+            "encoder.down_blocks.2.resnets.1.conv1.bias": "encoder.blocks.7.conv1.bias",
+            "encoder.down_blocks.2.resnets.1.conv1.weight": "encoder.blocks.7.conv1.weight",
+            "encoder.down_blocks.2.resnets.1.conv2.bias": "encoder.blocks.7.conv2.bias",
+            "encoder.down_blocks.2.resnets.1.conv2.weight": "encoder.blocks.7.conv2.weight",
+            "encoder.down_blocks.2.resnets.1.norm1.bias": "encoder.blocks.7.norm1.bias",
+            "encoder.down_blocks.2.resnets.1.norm1.weight": "encoder.blocks.7.norm1.weight",
+            "encoder.down_blocks.2.resnets.1.norm2.bias": "encoder.blocks.7.norm2.bias",
+            "encoder.down_blocks.2.resnets.1.norm2.weight": "encoder.blocks.7.norm2.weight",
+            "encoder.down_blocks.2.downsamplers.0.conv.bias": "encoder.blocks.8.conv.bias",
+            "encoder.down_blocks.2.downsamplers.0.conv.weight": "encoder.blocks.8.conv.weight",
+            "encoder.down_blocks.3.resnets.0.conv1.bias": "encoder.blocks.9.conv1.bias",
+            "encoder.down_blocks.3.resnets.0.conv1.weight": "encoder.blocks.9.conv1.weight",
+            "encoder.down_blocks.3.resnets.0.conv2.bias": "encoder.blocks.9.conv2.bias",
+            "encoder.down_blocks.3.resnets.0.conv2.weight": "encoder.blocks.9.conv2.weight",
+            "encoder.down_blocks.3.resnets.0.norm1.bias": "encoder.blocks.9.norm1.bias",
+            "encoder.down_blocks.3.resnets.0.norm1.weight": "encoder.blocks.9.norm1.weight",
+            "encoder.down_blocks.3.resnets.0.norm2.bias": "encoder.blocks.9.norm2.bias",
+            "encoder.down_blocks.3.resnets.0.norm2.weight": "encoder.blocks.9.norm2.weight",
+            "encoder.down_blocks.3.resnets.1.conv1.bias": "encoder.blocks.10.conv1.bias",
+            "encoder.down_blocks.3.resnets.1.conv1.weight": "encoder.blocks.10.conv1.weight",
+            "encoder.down_blocks.3.resnets.1.conv2.bias": "encoder.blocks.10.conv2.bias",
+            "encoder.down_blocks.3.resnets.1.conv2.weight": "encoder.blocks.10.conv2.weight",
+            "encoder.down_blocks.3.resnets.1.norm1.bias": "encoder.blocks.10.norm1.bias",
+            "encoder.down_blocks.3.resnets.1.norm1.weight": "encoder.blocks.10.norm1.weight",
+            "encoder.down_blocks.3.resnets.1.norm2.bias": "encoder.blocks.10.norm2.bias",
+            "encoder.down_blocks.3.resnets.1.norm2.weight": "encoder.blocks.10.norm2.weight",
+            "encoder.mid_block.attentions.0.to_k.bias": "encoder.blocks.12.transformer_blocks.0.to_k.bias",
+            "encoder.mid_block.attentions.0.to_k.weight": "encoder.blocks.12.transformer_blocks.0.to_k.weight",
+            "encoder.mid_block.attentions.0.group_norm.bias": "encoder.blocks.12.norm.bias",
+            "encoder.mid_block.attentions.0.group_norm.weight": "encoder.blocks.12.norm.weight",
+            "encoder.mid_block.attentions.0.to_out.0.bias": "encoder.blocks.12.transformer_blocks.0.to_out.bias",
+            "encoder.mid_block.attentions.0.to_out.0.weight": "encoder.blocks.12.transformer_blocks.0.to_out.weight",
+            "encoder.mid_block.attentions.0.to_q.bias": "encoder.blocks.12.transformer_blocks.0.to_q.bias",
+            "encoder.mid_block.attentions.0.to_q.weight": "encoder.blocks.12.transformer_blocks.0.to_q.weight",
+            "encoder.mid_block.attentions.0.to_v.bias": "encoder.blocks.12.transformer_blocks.0.to_v.bias",
+            "encoder.mid_block.attentions.0.to_v.weight": "encoder.blocks.12.transformer_blocks.0.to_v.weight",
+            "encoder.mid_block.resnets.0.conv1.bias": "encoder.blocks.11.conv1.bias",
+            "encoder.mid_block.resnets.0.conv1.weight": "encoder.blocks.11.conv1.weight",
+            "encoder.mid_block.resnets.0.conv2.bias": "encoder.blocks.11.conv2.bias",
+            "encoder.mid_block.resnets.0.conv2.weight": "encoder.blocks.11.conv2.weight",
+            "encoder.mid_block.resnets.0.norm1.bias": "encoder.blocks.11.norm1.bias",
+            "encoder.mid_block.resnets.0.norm1.weight": "encoder.blocks.11.norm1.weight",
+            "encoder.mid_block.resnets.0.norm2.bias": "encoder.blocks.11.norm2.bias",
+            "encoder.mid_block.resnets.0.norm2.weight": "encoder.blocks.11.norm2.weight",
+            "encoder.mid_block.resnets.1.conv1.bias": "encoder.blocks.13.conv1.bias",
+            "encoder.mid_block.resnets.1.conv1.weight": "encoder.blocks.13.conv1.weight",
+            "encoder.mid_block.resnets.1.conv2.bias": "encoder.blocks.13.conv2.bias",
+            "encoder.mid_block.resnets.1.conv2.weight": "encoder.blocks.13.conv2.weight",
+            "encoder.mid_block.resnets.1.norm1.bias": "encoder.blocks.13.norm1.bias",
+            "encoder.mid_block.resnets.1.norm1.weight": "encoder.blocks.13.norm1.weight",
+            "encoder.mid_block.resnets.1.norm2.bias": "encoder.blocks.13.norm2.bias",
+            "encoder.mid_block.resnets.1.norm2.weight": "encoder.blocks.13.norm2.weight"
         }
     }
 }

diffsynth_engine/conf/models/wan/dit/wan_dit_keymap.json

@@ -0,0 +1,41 @@
+{
+    "diffusers": {
+        "global_rename_dict": {
+            "patch_embedding": "patch_embedding",
+            "condition_embedder.text_embedder.linear_1": "text_embedding.0",
+            "condition_embedder.text_embedder.linear_2": "text_embedding.2",
+            "condition_embedder.time_embedder.linear_1": "time_embedding.0",
+            "condition_embedder.time_embedder.linear_2": "time_embedding.2",
+            "condition_embedder.time_proj": "time_projection.1",
+            "condition_embedder.image_embedder.norm1": "img_emb.proj.0",
+            "condition_embedder.image_embedder.ff.net.0.proj": "img_emb.proj.1",
+            "condition_embedder.image_embedder.ff.net.2": "img_emb.proj.3",
+            "condition_embedder.image_embedder.norm2": "img_emb.proj.4",
+            "condition_embedder.image_embedder.pos_embed": "img_emb.emb_pos",
+            "proj_out": "head.head",
+            "scale_shift_table": "head.modulation"
+        },
+        "rename_dict": {
+            "attn1.to_q": "self_attn.q",
+            "attn1.to_k": "self_attn.k",
+            "attn1.to_v": "self_attn.v",
+            "attn1.to_out.0": "self_attn.o",
+            "attn1.norm_q": "self_attn.norm_q",
+            "attn1.norm_k": "self_attn.norm_k",
+            "to_gate_compress": "self_attn.gate_compress",
+            "attn2.to_q": "cross_attn.q",
+            "attn2.to_k": "cross_attn.k",
+            "attn2.to_v": "cross_attn.v",
+            "attn2.to_out.0": "cross_attn.o",
+            "attn2.norm_q": "cross_attn.norm_q",
+            "attn2.norm_k": "cross_attn.norm_k",
+            "attn2.add_k_proj": "cross_attn.k_img",
+            "attn2.add_v_proj": "cross_attn.v_img",
+            "attn2.norm_added_k": "cross_attn.norm_k_img",
+            "norm2": "norm3",
+            "ffn.net.0.proj": "ffn.0",
+            "ffn.net.2": "ffn.2",
+            "scale_shift_table": "modulation"
+        }
+    }
+}

diffsynth_engine/configs/__init__.py

@@ -17,8 +17,17 @@ from .pipeline import (
     WanStateDicts,
     WanS2VStateDicts,
     QwenImageStateDicts,
+    AttnImpl,
+    SpargeAttentionParams,
+    VideoSparseAttentionParams,
+    LoraConfig,
+)
+from .controlnet import (
+    ControlType,
+    ControlNetParams,
+    QwenImageControlType,
+    QwenImageControlNetParams,
 )
-from .controlnet import ControlType, ControlNetParams
 
 __all__ = [
     "BaseConfig",
@@ -39,6 +48,12 @@ __all__ = [
     "WanStateDicts",
     "WanS2VStateDicts",
     "QwenImageStateDicts",
+    "AttnImpl",
+    "SpargeAttentionParams",
+    "VideoSparseAttentionParams",
+    "LoraConfig",
     "ControlType",
     "ControlNetParams",
+    "QwenImageControlType",
+    "QwenImageControlNetParams",
 ]

diffsynth_engine/configs/controlnet.py

@@ -34,3 +34,16 @@ class ControlNetParams:
     control_start: float = 0
     control_end: float = 1
     processor_name: Optional[str] = None  # only used for sdxl controlnet union now
+
+
+class QwenImageControlType(Enum):
+    eligen = "eligen"
+    in_context = "in_context"
+
+
+@dataclass
+class QwenImageControlNetParams:
+    image: ImageType
+    model: str
+    control_type: QwenImageControlType
+    scale: float = 1.0

diffsynth_engine/configs/pipeline.py

@@ -1,5 +1,6 @@
 import os
 import torch
+from enum import Enum
 from dataclasses import dataclass, field
 from typing import List, Dict, Tuple, Optional
 
@@ -19,14 +20,36 @@ class BaseConfig:
     offload_to_disk: bool = False
 
 
+class AttnImpl(Enum):
+    AUTO = "auto"
+    EAGER = "eager"  # Native Attention
+    FA2 = "fa2"  # Flash Attention 2
+    FA3 = "fa3"  # Flash Attention 3
+    FA3_FP8 = "fa3_fp8"  # Flash Attention 3 with FP8
+    XFORMERS = "xformers"  # XFormers
+    SDPA = "sdpa"  # Scaled Dot Product Attention
+    SAGE = "sage"  # Sage Attention
+    SPARGE = "sparge"  # Sparge Attention
+    VSA = "vsa"  # Video Sparse Attention
+
+
+@dataclass
+class SpargeAttentionParams:
+    smooth_k: bool = True
+    cdfthreshd: float = 0.6
+    simthreshd1: float = 0.98
+    pvthreshd: float = 50.0
+
+
+@dataclass
+class VideoSparseAttentionParams:
+    sparsity: float = 0.9
+
+
 @dataclass
 class AttentionConfig:
-    dit_attn_impl: str = "auto"
-    # Sparge Attention
-    sparge_smooth_k: bool = True
-    sparge_cdfthreshd: float = 0.6
-    sparge_simthreshd1: float = 0.98
-    sparge_pvthreshd: float = 50.0
+    dit_attn_impl: AttnImpl = AttnImpl.AUTO
+    attn_params: Optional[SpargeAttentionParams | VideoSparseAttentionParams] = None
 
 
 @dataclass
@@ -221,14 +244,11 @@ class QwenImagePipelineConfig(AttentionConfig, OptimizationConfig, ParallelConfi
     encoder_dtype: torch.dtype = torch.bfloat16
     vae_dtype: torch.dtype = torch.float32
 
+    load_encoder: bool = True
+
     # override OptimizationConfig
     fbcache_relative_l1_threshold = 0.009
 
-    # override BaseConfig
-    vae_tiled: bool = True
-    vae_tile_size: Tuple[int, int] = (34, 34)
-    vae_tile_stride: Tuple[int, int] = (18, 16)
-
     @classmethod
     def basic_config(
         cls,
@@ -352,3 +372,9 @@ def init_parallel_config(config: FluxPipelineConfig | QwenImagePipelineConfig |
         config.tp_degree = 1
     else:
         raise ValueError("sp_ulysses_degree and sp_ring_degree must be specified together")
+
+
+@dataclass
+class LoraConfig:
+    scale: float
+    scheduler_config: Optional[Dict] = None

diffsynth_engine/models/base.py

@@ -57,7 +57,7 @@ class PreTrainedModel(nn.Module):
     def get_tp_plan(self):
         raise NotImplementedError(f"{self.__class__.__name__} does not support TP")
 
-    def get_fsdp_modules(self):
+    def get_fsdp_module_cls(self):
         raise NotImplementedError(f"{self.__class__.__name__} does not support FSDP")