fix cache-dit bug2

Signed-off-by: David Chen <530634352@qq.com>

Author: david6666666

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -353,44 +353,6 @@ def enable_cache_for_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], N
     """Enable cache-dit for LTX2 pipelines (audio-video transformer blocks)."""
     transformer = pipeline.transformer
 
-    if not getattr(transformer, "_cache_dit_ltx2_patched", False):
-        def _wrap_block_forward(orig):
-            @functools.wraps(orig)
-            def _wrapped_forward(self, hidden_states, encoder_hidden_states=None, *args, **kwargs):
-                audio_hidden_states = kwargs.pop("audio_hidden_states", None)
-
-                text_encoder_hidden_states = encoder_hidden_states
-                if "encoder_hidden_states" in kwargs:
-                    text_encoder_hidden_states = kwargs.pop("encoder_hidden_states")
-                audio_encoder_hidden_states = kwargs.pop("audio_encoder_hidden_states", None)
-                temb = kwargs.pop("temb", None)
-                temb_audio = kwargs.pop("temb_audio", None)
-                temb_ca_scale_shift = kwargs.pop("temb_ca_scale_shift", None)
-                temb_ca_audio_scale_shift = kwargs.pop("temb_ca_audio_scale_shift", None)
-                temb_ca_gate = kwargs.pop("temb_ca_gate", None)
-                temb_ca_audio_gate = kwargs.pop("temb_ca_audio_gate", None)
-
-                return orig(
-                    hidden_states,
-                    audio_hidden_states,
-                    text_encoder_hidden_states,
-                    audio_encoder_hidden_states,
-                    temb,
-                    temb_audio,
-                    temb_ca_scale_shift,
-                    temb_ca_audio_scale_shift,
-                    temb_ca_gate,
-                    temb_ca_audio_gate,
-                    **kwargs,
-                )
-
-            return _wrapped_forward
-
-        for block in transformer.transformer_blocks:
-            block.forward = _wrap_block_forward(block.forward).__get__(block, block.__class__)
-
-        transformer._cache_dit_ltx2_patched = True
-
     db_cache_config = _build_db_cache_config(cache_config)
 
     calibrator_config = None

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -1359,22 +1359,22 @@ def forward(
                 )
             else:
                 hidden_states, audio_hidden_states = block(
-                    hidden_states=hidden_states,
-                    audio_hidden_states=audio_hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    audio_encoder_hidden_states=audio_encoder_hidden_states,
-                    temb=temb,
-                    temb_audio=temb_audio,
-                    temb_ca_scale_shift=video_cross_attn_scale_shift,
-                    temb_ca_audio_scale_shift=audio_cross_attn_scale_shift,
-                    temb_ca_gate=video_cross_attn_a2v_gate,
-                    temb_ca_audio_gate=audio_cross_attn_v2a_gate,
-                    video_rotary_emb=video_rotary_emb,
-                    audio_rotary_emb=audio_rotary_emb,
-                    ca_video_rotary_emb=video_cross_attn_rotary_emb,
-                    ca_audio_rotary_emb=audio_cross_attn_rotary_emb,
-                    encoder_attention_mask=encoder_attention_mask,
-                    audio_encoder_attention_mask=audio_encoder_attention_mask,
+                    hidden_states,
+                    audio_hidden_states,
+                    encoder_hidden_states,
+                    audio_encoder_hidden_states,
+                    temb,
+                    temb_audio,
+                    video_cross_attn_scale_shift,
+                    audio_cross_attn_scale_shift,
+                    video_cross_attn_a2v_gate,
+                    audio_cross_attn_v2a_gate,
+                    video_rotary_emb,
+                    audio_rotary_emb,
+                    video_cross_attn_rotary_emb,
+                    audio_cross_attn_rotary_emb,
+                    encoder_attention_mask,
+                    audio_encoder_attention_mask,
                 )
 
         # 6. Output layers (including unpatchification)