[BugFix] Fix Wan2.2-VACE-Fun cache and serving lifecycle (#4667)

Signed-off-by: Ianniu <985683179yzy@gmail.com>
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: Ianniu <985683179yzy@gmail.com>

tests/diffusion/cache/test_cache_dit.py

@@ -42,6 +42,10 @@ SEPARATE_CFG_TRANSFORMERS = [
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
 
 
+def test_wan22_vace_uses_wan22_custom_cache_dit_enabler():
+    assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_backend.enable_cache_for_wan22
+
+
 @pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS)
 def test_cache_dit_configs_have_separate_cfg(transformer_model):
     """Check models with separate CFG set has_separate_cfg=True in their configs."""
@@ -62,6 +66,34 @@ def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_blo
     assert adapter_kwargs["has_separate_cfg"] is True
 
 
+@pytest.mark.parametrize("has_transformer_2", [False, True])
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+def test_wan22_custom_enabler_passes_taylorseer_calibrator(
+    mock_cache_dit,
+    mock_block_adapter,
+    has_transformer_2,
+):
+    mock_pipeline = Mock()
+    mock_pipeline.transformer.blocks = [Mock()]
+    if has_transformer_2:
+        mock_pipeline.transformer_2.blocks = [Mock()]
+    else:
+        mock_pipeline.transformer_2 = None
+    cache_config = DiffusionCacheConfig(enable_taylorseer=True, taylorseer_order=1)
+
+    cd_backend.enable_cache_for_wan22(mock_pipeline, cache_config)
+
+    enable_cache_kwargs = mock_cache_dit.enable_cache.call_args.kwargs
+    calibrator_config = enable_cache_kwargs["calibrator_config"]
+    assert calibrator_config is not None
+    assert calibrator_config.taylorseer_order == 1
+
+    adapter_kwargs = mock_block_adapter.call_args.kwargs
+    for modifier in adapter_kwargs["params_modifiers"]:
+        assert modifier._context_kwargs["calibrator_config"] is calibrator_config
+
+
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
 def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -14,6 +14,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
     create_vace_transformer_from_config,
     get_wan22_vace_pre_process_func,
 )
+from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -93,6 +94,10 @@ def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch
     }
 
 
+def test_vace_transformer_layerwise_offloads_conditioning_blocks() -> None:
+    assert WanVACETransformer3DModel._layerwise_offload_blocks_attrs == ["vace_blocks", "blocks"]
+
+
 def test_vace_prepare_masks_encodes_spatial_stride_and_reference_padding() -> None:
     pipeline = _make_vace_pipeline()
     mask = torch.ones(1, 3, 5, 16, 16)