Update cfg skip (#337)

videox_fun/models/__init__.py

@@ -70,6 +70,9 @@ if importlib.util.find_spec("paifuser") is not None:
     WanTransformer3DModel.enable_cfg_skip = enable_cfg_skip()(WanTransformer3DModel.enable_cfg_skip)
     WanTransformer3DModel.disable_cfg_skip = disable_cfg_skip()(WanTransformer3DModel.disable_cfg_skip)
     WanTransformer3DModel.share_cfg_skip = share_cfg_skip()(WanTransformer3DModel.share_cfg_skip)
+
+    QwenImageTransformer2DModel.enable_cfg_skip = enable_cfg_skip()(QwenImageTransformer2DModel.enable_cfg_skip)
+    QwenImageTransformer2DModel.disable_cfg_skip = disable_cfg_skip()(QwenImageTransformer2DModel.disable_cfg_skip)
     print("Import CFG Skip Turbo")
 
     # --------------------------------------------------------------- #

videox_fun/models/wan_transformer3d_s2v.py

@@ -23,6 +23,7 @@ from .wan_audio_injector import (AudioInjector_WAN, CausalAudioEncoder,
 from .wan_transformer3d import (Wan2_2Transformer3DModel, WanAttentionBlock,
                                 WanLayerNorm, WanSelfAttention,
                                 sinusoidal_embedding_1d)
+from ..utils import cfg_skip
 
 
 def zero_module(module):
@@ -542,6 +543,7 @@ class Wan2_2Transformer3DModel_S2V(Wan2_2Transformer3DModel):
 
         return hidden_states
 
+    @cfg_skip()
     def forward(
         self,
         x,

videox_fun/models/wan_transformer3d_vace.py

@@ -13,6 +13,7 @@ from diffusers.utils import is_torch_version
 
 from .wan_transformer3d import (WanAttentionBlock, WanTransformer3DModel,
                                 sinusoidal_embedding_1d)
+from ..utils import cfg_skip
 
 
 VIDEOX_OFFLOAD_VACE_LATENTS = os.environ.get("VIDEOX_OFFLOAD_VACE_LATENTS", False)
@@ -178,6 +179,7 @@ class VaceWanTransformer3DModel(WanTransformer3DModel):
         hints = torch.unbind(c)[:-1]
         return hints
 
+    @cfg_skip()
     def forward(
         self,
         x,