diffsynth-engine 0.6.1.dev26__tar.gz → 0.6.1.dev27__tar.gz

{diffsynth_engine-0.6.1.dev26 → diffsynth_engine-0.6.1.dev27}/PKG-INFO

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

{diffsynth_engine-0.6.1.dev26 → diffsynth_engine-0.6.1.dev27}/diffsynth_engine/pipelines/wan_s2v.py

@@ -656,6 +656,8 @@ class WanSpeech2VideoPipeline(WanVideoPipeline):
         )
 
         with LoRAContext():
+            cls._auto_enable_vsa(state_dicts.model, config)
+
             dit = WanS2VDiT.from_state_dict(
                 state_dicts.model,
                 config=model_config,

{diffsynth_engine-0.6.1.dev26 → diffsynth_engine-0.6.1.dev27}/diffsynth_engine/pipelines/wan_video.py

@@ -4,7 +4,7 @@ from typing import Callable, List, Dict, Tuple, Optional
 from tqdm import tqdm
 from PIL import Image
 
-from diffsynth_engine.configs import WanPipelineConfig, WanStateDicts
+from diffsynth_engine.configs import WanPipelineConfig, WanStateDicts, AttnImpl, VideoSparseAttentionParams
 from diffsynth_engine.algorithm.noise_scheduler.flow_match import RecifitedFlowScheduler
 from diffsynth_engine.algorithm.sampler import FlowMatchEulerSampler
 from diffsynth_engine.models.wan.wan_dit import WanDiT
@@ -584,6 +584,8 @@ class WanVideoPipeline(BasePipeline):
             dit_state_dict = state_dicts.model
 
         with LoRAContext():
+            cls._auto_enable_vsa(dit_state_dict, config)
+
             dit = WanDiT.from_state_dict(
                 dit_state_dict,
                 config=dit_config,
@@ -668,6 +670,16 @@ class WanVideoPipeline(BasePipeline):
             vae_type = "wan2.2-vae"
         return vae_type
 
+    @staticmethod
+    def _auto_enable_vsa(state_dict: Dict[str, torch.Tensor], config: WanPipelineConfig):
+        def has_any_key(*xs):
+            return any(x in state_dict for x in xs)
+
+        if has_any_key("blocks.0.to_gate_compress.weight", "blocks.0.self_attn.gate_compress.weight"):
+            config.dit_attn_impl = AttnImpl.VSA
+            if config.attn_params is None:
+                config.attn_params = VideoSparseAttentionParams(sparsity=0.9)
+
     def compile(self):
         self.dit.compile_repeated_blocks()
         if self.dit2 is not None:

{diffsynth_engine-0.6.1.dev26 → diffsynth_engine-0.6.1.dev27}/diffsynth_engine.egg-info/PKG-INFO

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