[Refactor][1/N]diffusion/cache: cache-dit (#5226)

Signed-off-by: yangjianjuan <66548561+yangjianjuan@users.noreply.github.com>

tests/diffusion/cache/test_cache_dit.py

@@ -5,15 +5,18 @@
 Model specific tests for CacheDiT enablement.
 """
 
+import ast
 import sys
+from pathlib import Path
 from unittest.mock import Mock, patch
 
 import pytest
 import torch
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
 
-import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
+import vllm_omni.diffusion.cache.cachedit as cd_backend
+import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific
+from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
@@ -42,8 +45,73 @@ SEPARATE_CFG_TRANSFORMERS = [
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
 
 
+def test_custom_cache_dit_enablers_are_registered_explicitly():
+    expected_enablers = {
+        "Wan22Pipeline": cd_model_specific.enable_cache_for_wan22,
+        "Wan22I2VPipeline": cd_model_specific.enable_cache_for_wan22,
+        "Wan22TI2VPipeline": cd_model_specific.enable_cache_for_wan22,
+        "Wan22VACEPipeline": cd_model_specific.enable_cache_for_wan22,
+        "Wan22S2VPipeline": cd_model_specific.enable_cache_for_wan22_s2v,
+        "Cosmos3OmniDiffusersPipeline": cd_model_specific.enable_cache_for_cosmos3,
+        "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3,
+        "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2,
+    }
+
+    with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True):
+        cd_model_specific.register_custom_dit_enablers()
+        assert cd_backend.CUSTOM_DIT_ENABLERS == expected_enablers
+
+
 def test_wan22_vace_uses_wan22_custom_cache_dit_enabler():
-    assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_backend.enable_cache_for_wan22
+    assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_model_specific.enable_cache_for_wan22
+
+
+@pytest.mark.parametrize(
+    "pipeline_name",
+    ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"],
+)
+def test_cosmos3_aliases_use_cosmos3_custom_cache_dit_enabler(pipeline_name: str):
+    assert cd_backend.CUSTOM_DIT_ENABLERS[pipeline_name] is cd_model_specific.enable_cache_for_cosmos3
+
+
+def test_cachedit_public_api_is_explicit():
+    assert set(cd_backend.__all__) == {
+        "BagelCachedAdapter",
+        "CUSTOM_DIT_ENABLERS",
+        "CacheDiTAdapterConfig",
+        "CacheDiTBackend",
+        "CacheDiTConfig",
+        "SensenovaCachedAdapter",
+        "cache_summary",
+        "enable_cache_for_dit",
+    }
+    assert not hasattr(cd_backend, "enable_cache_for_wan22")
+    assert not hasattr(cd_backend, "enable_cache_for_wan22_s2v")
+
+
+def test_cachedit_consumers_use_package_api():
+    cache_dir = Path(cd_backend.__file__).resolve().parents[1]
+    package_root = cache_dir.parents[1]
+    legacy_module = "vllm_omni.diffusion.cache.cache_dit_backend"
+    internal_prefix = "vllm_omni.diffusion.cache.cachedit."
+
+    invalid_imports = []
+    for source_path in package_root.rglob("*.py"):
+        if source_path.is_relative_to(cache_dir):
+            continue
+        tree = ast.parse(source_path.read_text(encoding="utf-8"))
+        for node in ast.walk(tree):
+            modules = []
+            if isinstance(node, ast.ImportFrom) and node.module is not None:
+                modules.append(node.module)
+            elif isinstance(node, ast.Import):
+                modules.extend(alias.name for alias in node.names)
+
+            for module in modules:
+                if module == legacy_module or module.startswith(internal_prefix):
+                    invalid_imports.append(f"{source_path.relative_to(package_root)}:{node.lineno}: {module}")
+
+    assert not invalid_imports, "Cache-DiT consumers must use the package API:\n" + "\n".join(invalid_imports)
 
 
 @pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS)
@@ -54,12 +122,12 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
     assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
 
 
-@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
-@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
 def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
     """Ensure that Wan22, which has a custom enabler, setts custom CFG correctly."""
     mock_pipeline = Mock()
-    cd_backend.enable_cache_for_wan22(mock_pipeline, SAMPLE_CACHE_CONFIG)
+    cd_model_specific.enable_cache_for_wan22(mock_pipeline, SAMPLE_CACHE_CONFIG)
 
     mock_cache_dit.enable_cache.assert_called_once()
     adapter_kwargs = mock_block_adapter.call_args.kwargs
@@ -67,8 +135,8 @@ def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_blo
 
 
 @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")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
 def test_wan22_custom_enabler_passes_taylorseer_calibrator(
     mock_cache_dit,
     mock_block_adapter,
@@ -82,7 +150,7 @@ def test_wan22_custom_enabler_passes_taylorseer_calibrator(
         mock_pipeline.transformer_2 = None
     cache_config = DiffusionCacheConfig(enable_taylorseer=True, taylorseer_order=1)
 
-    cd_backend.enable_cache_for_wan22(mock_pipeline, cache_config)
+    cd_model_specific.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"]
@@ -94,8 +162,8 @@ def test_wan22_custom_enabler_passes_taylorseer_calibrator(
         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")
+@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
 def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
     """Cosmos3 should cache only the repeated GEN pathway blocks."""
     mock_pipeline = Mock()
@@ -103,7 +171,7 @@ def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
     mock_pipeline.transformer.gen_layers = gen_layers
     mock_pipeline.transformer._cache_dit_adapter_config = Cosmos3VFMTransformer._cache_dit_adapter_config
 
-    cd_backend.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG)
+    cd_model_specific.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG)
 
     mock_cache_dit.enable_cache.assert_called_once()
     adapter_kwargs = mock_block_adapter.call_args.kwargs