[Diffusion] Add Minimax-H3 modular pipeline support (#5720)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: David Chen <530634352@qq.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: David Chen <530634352@qq.com>

vllm_omni/diffusion/cache/cachedit/backend.py

@@ -3,10 +3,13 @@
 
 """Generic Cache-DiT backend lifecycle and integration."""
 
+from __future__ import annotations
+
 from collections.abc import Callable
-from typing import Any, TypeAlias
+from typing import TYPE_CHECKING, Any, TypeAlias, cast
 
 import cache_dit
+import torch.nn as nn
 from cache_dit import BlockAdapter, DBCacheConfig
 from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter
 from vllm.logger import init_logger
@@ -18,6 +21,9 @@ from vllm_omni.diffusion.cache.cachedit.config import (
 )
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+
 logger = init_logger(__name__)
 
 RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None]
@@ -27,14 +33,18 @@ CacheDiTEnabler: TypeAlias = Callable[[Any, DiffusionCacheConfig], RefreshCacheC
 CUSTOM_DIT_ENABLERS: dict[str, CacheDiTEnabler] = {}
 
 
-def cache_summary(pipeline: Any, details: bool = True) -> None:
+def _dit_module_names(pipeline: SupportsComponentDiscovery) -> tuple[str, ...]:
+    """Return the pipeline DiT attributes that are present at runtime."""
+    names = getattr(pipeline, "_dit_modules", None)
+    if not isinstance(names, (list, tuple)):
+        names = ("transformer",)
+    return tuple(name for name in names if isinstance(name, str) and getattr(pipeline, name, None) is not None)
+
+
+def cache_summary(pipeline: SupportsComponentDiscovery, details: bool = True) -> None:
     """Log Cache-DiT statistics for every transformer on the pipeline."""
 
-    transformers = [
-        transformer
-        for attribute in ("transformer", "transformer_2")
-        if (transformer := getattr(pipeline, attribute, None)) is not None
-    ]
+    transformers = [getattr(pipeline, name) for name in _dit_module_names(pipeline)]
     for transformer in transformers:
         cache_dit.summary(transformer, details=details)
 
@@ -42,19 +52,34 @@ def cache_summary(pipeline: Any, details: bool = True) -> None:
         logger.warning("CacheDiT summary failed; this pipeline has no defined transformer attribute")
 
 
-def _default_get_pipeline_transformer(pipeline: Any) -> Any:
-    return pipeline.transformer
+def _default_get_pipeline_transformer(
+    pipeline: SupportsComponentDiscovery,
+) -> nn.Module:
+    return cast(nn.Module, getattr(pipeline, "transformer"))
+
+
+def _make_pipeline_transformer_getter(
+    name: str,
+) -> Callable[[SupportsComponentDiscovery], nn.Module]:
+    def get_pipeline_transformer(
+        pipeline: SupportsComponentDiscovery,
+    ) -> nn.Module:
+        return cast(nn.Module, getattr(pipeline, name))
+
+    return get_pipeline_transformer
 
 
 def _build_cache_context_refresh(
     cache_config: DiffusionCacheConfig,
-    get_pipeline_transformer: Callable[[Any], Any] = _default_get_pipeline_transformer,
+    get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer,
 ) -> RefreshCacheContextFunc:
     """Build the cache context refresh callback for one transformer."""
 
     projected_config = CacheDiTConfig.from_diffusion_config(cache_config)
 
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+    def refresh_cache_context(
+        pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True
+    ) -> None:
         transformer = get_pipeline_transformer(pipeline)
 
         # Cache-DiT has no predefined SCM mask for these small step counts.
@@ -80,10 +105,11 @@ def _build_cache_context_refresh(
 
 
 def enable_cache_for_dit(
-    pipeline: Any,
+    pipeline: SupportsComponentDiscovery,
     cache_config: DiffusionCacheConfig,
     block_adapter: BlockAdapter | None = None,
     adapter_cls: type[CachedAdapter] | None = None,
+    get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer,
 ) -> RefreshCacheContextFunc:
     """Enable Cache-DiT for a standard single-transformer DiT pipeline."""
 
@@ -98,7 +124,7 @@ def enable_cache_for_dit(
         db_cache_config.max_warmup_steps,
     )
 
-    transformer = _default_get_pipeline_transformer(pipeline)
+    transformer = get_pipeline_transformer(pipeline)
     cache_target = transformer if block_adapter is None else block_adapter
     if adapter_cls is not None:
         adapter_cls.apply(
@@ -128,13 +154,16 @@ def enable_cache_for_dit(
             calibrator_config=calibrator_config,
         )
 
-    return _build_cache_context_refresh(cache_config)
+    return _build_cache_context_refresh(cache_config, get_pipeline_transformer)
 
 
-def _maybe_build_block_adapter(pipeline: Any) -> BlockAdapter | None:
+def _maybe_build_block_adapter(
+    pipeline: SupportsComponentDiscovery,
+    get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer,
+) -> BlockAdapter | None:
     """Build the model-declared block adapter, when one is configured."""
 
-    transformer = _default_get_pipeline_transformer(pipeline)
+    transformer = get_pipeline_transformer(pipeline)
     adapter_config: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None)
     if adapter_config is None:
         logger.info(
@@ -160,10 +189,13 @@ def _maybe_build_block_adapter(pipeline: Any) -> BlockAdapter | None:
     )
 
 
-def _maybe_get_cached_adapter_cls(pipeline: Any) -> type[CachedAdapter] | None:
+def _maybe_get_cached_adapter_cls(
+    pipeline: SupportsComponentDiscovery,
+    get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer,
+) -> type[CachedAdapter] | None:
     """Return the custom cached adapter declared by the transformer."""
 
-    transformer = _default_get_pipeline_transformer(pipeline)
+    transformer = get_pipeline_transformer(pipeline)
     adapter_config: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None)
     return None if adapter_config is None else adapter_config.cached_adapter_cls
 
@@ -180,29 +212,37 @@ class CacheDiTBackend(CacheBackend):
             config = cache_config
 
         super().__init__(config)
-        self._refresh_func: RefreshCacheContextFunc | None = None
+        self._refresh_funcs: list[RefreshCacheContextFunc] = []
 
-    def enable(self, pipeline: Any) -> None:
+    def enable(self, pipeline: SupportsComponentDiscovery) -> None:
         pipeline_name = type(pipeline).__name__
         custom_enabler = CUSTOM_DIT_ENABLERS.get(pipeline_name)
         if custom_enabler is not None:
             logger.info("Using custom cache-dit enabler for model: %s", pipeline_name)
-            self._refresh_func = custom_enabler(pipeline, self.config)
+            self._refresh_funcs = [custom_enabler(pipeline, self.config)]
         else:
-            block_adapter = _maybe_build_block_adapter(pipeline)
-            adapter_cls = _maybe_get_cached_adapter_cls(pipeline)
-            self._refresh_func = enable_cache_for_dit(
-                pipeline,
-                self.config,
-                block_adapter,
-                adapter_cls,
-            )
+            self._refresh_funcs = []
+            for name in _dit_module_names(pipeline):
+                get_transformer = _make_pipeline_transformer_getter(name)
+                block_adapter = _maybe_build_block_adapter(pipeline, get_transformer)
+                adapter_cls = _maybe_get_cached_adapter_cls(pipeline, get_transformer)
+                self._refresh_funcs.append(
+                    enable_cache_for_dit(
+                        pipeline,
+                        self.config,
+                        block_adapter,
+                        adapter_cls,
+                        get_transformer,
+                    )
+                )
+            if not self._refresh_funcs:
+                raise ValueError(f"Pipeline {pipeline_name} has no declared DiT modules for Cache-DiT")
 
         self.enabled = True
         logger.info("Cache-dit enabled successfully on %s", pipeline_name)
 
-    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        if not self.enabled or self._refresh_func is None:
+    def refresh(self, pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True) -> None:
+        if not self.enabled or not self._refresh_funcs:
             logger.warning("Cache-dit is not enabled. Cannot refresh cache context.")
             return
 
@@ -211,7 +251,8 @@ class CacheDiTBackend(CacheBackend):
                 "Refreshing cache context for transformer with num_inference_steps: %s",
                 num_inference_steps,
             )
-        self._refresh_func(pipeline, num_inference_steps, verbose)
+        for refresh_func in self._refresh_funcs:
+            refresh_func(pipeline, num_inference_steps, verbose)
 
 
 __all__ = [