[Feat] support cache-dit for DreamID-Omni (#3265)

Signed-off-by: hanzheli <hanzheli@kuaishou.com>
Signed-off-by: fywc <hanzheli@kuaishou.com>

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -101,3 +101,4 @@ Key arguments:
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
+- `--cache-backend`: enable `cache_dit` for acceleration.

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -63,6 +63,7 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.")
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -146,12 +147,29 @@ def main() -> None:
         hsdp_replicate_size=args.hsdp_replicate_size,
     )
 
+    cache_config = None
+    if args.cache_backend == "cache_dit":
+        cache_config = {
+            "Fn_compute_blocks": 1,
+            "Bn_compute_blocks": 0,
+            "max_warmup_steps": 4,
+            "max_cached_steps": 20,
+            "residual_diff_threshold": 0.24,
+            "max_continuous_cached_steps": 3,
+            "enable_taylorseer": False,
+            "taylorseer_order": 1,
+            "scm_steps_mask_policy": None,
+            "scm_steps_policy": "dynamic",
+        }
+
     omni = Omni(
         model=args.model,
         parallel_config=parallel_config,
         model_type=args.model_type,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_layerwise_offload=args.enable_layerwise_offload,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
     )
     start = time.perf_counter()
     outputs = omni.generate(prompt, sampling_params)

tests/diffusion/cache/test_cache_backends.py

@@ -127,6 +127,10 @@ class TestCacheDiTBackend:
         """Test HunyuanImage3 custom cache-dit enabler is registered."""
         assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS
 
+    def test_dreamid_custom_enabler_registered(self):
+        """Test DreamIDOmni custom cache-dit enabler is registered."""
+        assert "DreamIDOmniPipeline" in CUSTOM_DIT_ENABLERS
+
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
@@ -169,6 +173,47 @@ class TestCacheDiTBackend:
         with pytest.raises(ValueError, match="HunyuanImage3Pipeline"):
             backend.enable(mock_pipeline)
 
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.transformer = Mock()
+        mock_pipeline.transformer.fused_blocks = Mock()
+
+        mock_cache_dit.enable_cache = Mock()
+        mock_cache_dit.refresh_context = Mock()
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        assert backend._refresh_func is not None
+        mock_block_adapter.assert_called_once()
+        adapter_kwargs = mock_block_adapter.call_args.kwargs
+        assert adapter_kwargs["transformer"] is mock_pipeline.transformer
+        assert adapter_kwargs["blocks"] is mock_pipeline.transformer.fused_blocks
+        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_0
+        assert adapter_kwargs["has_separate_cfg"] is True
+        assert len(adapter_kwargs["params_modifiers"]) == 1
+        mock_cache_dit.enable_cache.assert_called_once()
+
+        backend.refresh(mock_pipeline, num_inference_steps=12)
+        mock_cache_dit.refresh_context.assert_called_once()
+        call_args = mock_cache_dit.refresh_context.call_args
+        assert call_args[0][0] is mock_pipeline.transformer
+        assert call_args[1]["num_inference_steps"] == 12
+
+    def test_enable_dreamid_pipeline_requires_fused_blocks(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.transformer = Mock(spec=[])
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+
+        with pytest.raises(ValueError, match="DreamIDOmniPipeline"):
+            backend.enable(mock_pipeline)
+
 
 class TestTeaCacheBackend:
     """Test TeaCacheBackend implementation."""

tests/diffusion/cache/test_cache_dit.py

@@ -15,6 +15,7 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 SEPARATE_CFG_ENABLERS = [
+    cd_backend.enable_cache_for_dreamid_omni,
     cd_backend.enable_cache_for_ltx2,
     cd_backend.enable_cache_for_helios,
     cd_backend.enable_cache_for_wan22,

tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py

@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E test for DreamID-Omni offline inference with Cache-Dit.
+
+Uses the official DreamID-Omni ``test_case/twoip`` example assets instead of
+synthetic inputs so the coverage matches the upstream demo structure.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import re
+import time
+from functools import lru_cache
+from pathlib import Path
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import soundfile as sf
+from PIL import Image
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow]
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni"
+MODEL_PATH_ENV = "DREAMID_OMNI_MODEL_PATH"
+EXAMPLE_ROOT_ENV = "DREAMID_OMNI_EXAMPLE_ROOT"
+NUM_INFERENCE_STEPS = 45
+
+CACHE_CONFIG = {
+    "Fn_compute_blocks": 1,
+    "Bn_compute_blocks": 0,
+    "max_warmup_steps": 4,
+    "max_cached_steps": 20,
+    "residual_diff_threshold": 0.24,
+    "max_continuous_cached_steps": 3,
+    "enable_taylorseer": False,
+    "taylorseer_order": 1,
+    "scm_steps_mask_policy": None,
+    "scm_steps_policy": "dynamic",
+}
+
+
+class DreamIDGenerationResult(NamedTuple):
+    video: np.ndarray
+    audio: np.ndarray
+    fps: int
+    duration_s: float
+    elapsed_s: float
+
+
+class DreamIDExampleCase(NamedTuple):
+    prompt: str
+    image_paths: tuple[Path, ...]
+    audio_paths: tuple[Path, ...]
+
+
+def _require_model_path() -> str:
+    model_path = os.environ.get(MODEL_PATH_ENV, str(DEFAULT_MODEL_PATH))
+    if not Path(model_path).exists():
+        pytest.skip(f"DreamID-Omni model bundle not found at {model_path}. Set {MODEL_PATH_ENV}.")
+
+    if importlib.util.find_spec("dreamid_omni") is None:
+        pytest.skip(
+            "dreamid_omni dependency package is not importable. "
+            "Run examples/offline_inference/x_to_video_audio/download_dreamid_omni.py first."
+        )
+
+    return model_path
+
+
+def _dreamid_repo_root_candidates() -> list[Path]:
+    candidates: list[Path] = []
+    example_root = os.environ.get(EXAMPLE_ROOT_ENV)
+    if example_root:
+        candidates.append(Path(example_root).expanduser().resolve())
+
+    spec = importlib.util.find_spec("dreamid_omni")
+    if spec is not None and spec.origin:
+        candidates.append(Path(spec.origin).resolve().parents[1] / "test_case" / "twoip")
+
+    return candidates
+
+
+@lru_cache(maxsize=1)
+def _load_example_case() -> DreamIDExampleCase:
+    for candidate_root in _dreamid_repo_root_candidates():
+        caption_path = candidate_root / "captions" / "20.json"
+        image_paths = (
+            candidate_root / "imgs" / "20" / "0.png",
+            candidate_root / "imgs" / "20" / "1.png",
+        )
+        audio_paths = (
+            candidate_root / "audios" / "20" / "0.wav",
+            candidate_root / "audios" / "20" / "1.wav",
+        )
+        if caption_path.exists() and all(path.exists() for path in image_paths + audio_paths):
+            prompt = _load_prompt_from_json(caption_path)
+            return DreamIDExampleCase(
+                prompt=prompt,
+                image_paths=image_paths,
+                audio_paths=audio_paths,
+            )
+
+    searched = ", ".join(str(path) for path in _dreamid_repo_root_candidates()) or "<none>"
+    pytest.skip(
+        "DreamID-Omni official example assets not found. "
+        f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. "
+        f"Searched: {searched}"
+    )
+
+
+def _load_prompt_from_json(path: Path) -> str:
+    raw_text = path.read_text(encoding="utf-8").strip()
+    parsed = json.loads(raw_text)
+    if isinstance(parsed, str):
+        prompt = parsed
+    elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str):
+        prompt = parsed["content"]
+    else:
+        raise AssertionError(f"Unexpected DreamID caption format in {path}")
+
+    prompt = re.sub(
+        r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
+        "",
+        prompt,
+        flags=re.DOTALL,
+    ).strip()
+    prompt = re.sub(
+        r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
+        "",
+        prompt,
+        flags=re.DOTALL,
+    ).strip()
+    prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
+    prompt = re.sub(r"\n\s*\n", "\n", prompt).strip()
+    return prompt
+
+
+def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]:
+    images: list[Image.Image] = []
+    for path in example_case.image_paths:
+        with Image.open(path) as image:
+            images.append(image.convert("RGB"))
+    return images
+
+
+def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]:
+    audios: list[np.ndarray] = []
+    for path in example_case.audio_paths:
+        audio_array, sample_rate = sf.read(str(path), dtype="float32")
+        audio_np = np.asarray(audio_array, dtype=np.float32)
+        if audio_np.ndim == 2:
+            audio_np = np.mean(audio_np, axis=1)
+        audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)]
+        audios.append(audio_np.reshape(-1))
+    return audios
+
+
+def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        height=704,
+        width=1280,
+        num_inference_steps=num_inference_steps,
+        seed=42,
+        extra_args={
+            "solver_name": "unipc",
+            "shift": 5.0,
+        },
+    )
+
+
+def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]:
+    return {
+        "prompt": example_case.prompt,
+        "video_negative_prompt": "jitter, bad hands, blur, distortion",
+        "audio_negative_prompt": "robotic, muffled, echo, distorted",
+        "multi_modal_data": {
+            "image": _load_example_images(example_case),
+            "audio": _load_example_audios(example_case),
+        },
+    }
+
+
+def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult:
+    """Extract video and audio from raw Omni outputs"""
+    assert outputs, "DreamID-Omni returned no outputs"
+
+    result = outputs[0]
+    assert hasattr(result, "images") and result.images, "not return video frames"
+
+    video = result.images[0]
+    assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}"
+    assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}"
+
+    multimodal_output = getattr(result, "multimodal_output", None) or {}
+    assert "audio" in multimodal_output, "not return audio"
+    audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1)
+    assert audio.size > 0, "returned empty audio"
+
+    audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000))
+    fps = multimodal_output.get("fps", 24)
+    duration_s = float(audio.size) / float(audio_sample_rate)
+
+    generation = DreamIDGenerationResult(
+        video=video.astype(np.float32, copy=False),
+        audio=audio,
+        fps=fps,
+        duration_s=duration_s,
+        elapsed_s=elapsed_s,
+    )
+    return generation
+
+
+def _run_generation(
+    *,
+    model_path: str,
+    example_case: DreamIDExampleCase,
+    use_cache_dit: bool,
+) -> DreamIDGenerationResult:
+    runner_kwargs = {
+        "model_type": "dreamid-omni",
+        "parallel_config": DiffusionParallelConfig(cfg_parallel_size=1),
+    }
+    if use_cache_dit:
+        runner_kwargs["cache_backend"] = "cache_dit"
+        runner_kwargs["cache_config"] = CACHE_CONFIG
+
+    with OmniRunner(model_path, **runner_kwargs) as runner:
+        warmup_payload = _request_payload(example_case)
+        _ = runner.omni.generate(warmup_payload, _sampling_params(2))
+
+        elapsed_samples: list[float] = []
+        extracted_result: DreamIDGenerationResult | None = None
+
+        request_payload = _request_payload(example_case)
+        start = time.perf_counter()
+        outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS))
+        elapsed_s = time.perf_counter() - start
+        elapsed_samples.append(elapsed_s)
+        if extracted_result is None:
+            extracted_result = _extract_generation(outputs, elapsed_s)
+
+        assert extracted_result is not None
+        median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64)))
+        return extracted_result._replace(elapsed_s=median_elapsed_s)
+
+
+def _video_mae(a: np.ndarray, b: np.ndarray) -> float:
+    assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}"
+    return float(np.mean(np.abs(a - b)))
+
+
+def _audio_rms(audio: np.ndarray) -> float:
+    audio = np.asarray(audio, dtype=np.float32).reshape(-1)
+    return float(np.sqrt(np.mean(audio**2)))
+
+
+def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float:
+    rms_a = _audio_rms(a.audio)
+    rms_b = _audio_rms(b.audio)
+    return rms_b / max(rms_a, 1e-6)
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_dreamid_omni_cache_dit_e2e() -> None:
+    model_path = _require_model_path()
+    example_case = _load_example_case()
+
+    baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False)
+    cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True)
+
+    video_mae = _video_mae(baseline.video, cached.video)
+    audio_rms_ratio = _audio_rms_ratio(baseline, cached)
+    speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf")
+
+    assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}"
+    assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}"
+    assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x"

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -1320,6 +1320,74 @@ def enable_cache_for_glm_image(pipeline: Any, cache_config: Any) -> Callable[[in
     return refresh_cache_context
 
 
+def enable_cache_for_dreamid_omni(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for DreamID-Omni fused pipeline.
+
+    Args:
+        pipeline: The DreamIDOmni pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
+    """
+    transformer = getattr(pipeline, "transformer", None)
+    fused_blocks = getattr(transformer, "fused_blocks", None)
+    if transformer is None or fused_blocks is None:
+        raise ValueError("DreamIDOmniPipeline cache-dit expects pipeline.transformer.fused_blocks.")
+
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on DreamID-Omni fused blocks: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=fused_blocks,
+            forward_pattern=ForwardPattern.Pattern_0,
+            params_modifiers=[modifier],
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = pipeline.transformer
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for ERNIE-Image pipeline.
 
@@ -1521,6 +1589,7 @@ CUSTOM_DIT_ENABLERS.update(
         "BagelPipeline": enable_cache_for_bagel,
         "GlmImagePipeline": enable_cache_for_glm_image,
         "Flux2Pipeline": enable_cache_for_flux2,
+        "DreamIDOmniPipeline": enable_cache_for_dreamid_omni,
         "ErnieImagePipeline": enable_cache_for_ernie_image,
         "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15,
         "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15,

vllm_omni/diffusion/models/dreamid_omni/fusion.py

@@ -128,8 +128,8 @@ class FusedBlock(nn.Module):
 
     def forward(
         self,
-        vid,
-        audio,
+        hidden_states,
+        encoder_hidden_states,
         attn: Attention,
         vid_e,
         vid_seq_lens,
@@ -148,6 +148,8 @@ class FusedBlock(nn.Module):
         audio_ref_lengths,
         audio_freqs_scaling,
     ):
+        vid = hidden_states
+        audio = encoder_hidden_states
         vid_block = self.vid_block
         audio_block = self.audio_block