auto enable vsa (#203)

Author: akaitsuki-ii

diffsynth_engine/pipelines/wan_s2v.py

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

diffsynth_engine/pipelines/wan_video.py

@@ -4,7 +4,7 @@
 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 @@ def _from_state_dict(cls, state_dicts: WanStateDicts, config: WanPipelineConfig)
             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 @@ def _get_vae_type(vae_state_dict: Dict[str, torch.Tensor]) -> str:
             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: