From d01eaca39a747a8ef1273767f0f809c09dc494c4 Mon Sep 17 00:00:00 2001 From: fywc Date: Sat, 23 May 2026 08:09:33 +0800 Subject: [PATCH] [Feat] support cache-dit for DreamID-Omni (#3265) Signed-off-by: hanzheli Signed-off-by: fywc --- .../x_to_video_audio/x_to_video_audio.md | 1 + .../x_to_video_audio/x_to_video_audio.py | 18 ++ tests/diffusion/cache/test_cache_backends.py | 45 +++ tests/diffusion/cache/test_cache_dit.py | 1 + .../test_dreamid_omni_cache_dit.py | 290 ++++++++++++++++++ .../diffusion/cache/cache_dit_backend.py | 69 +++++ .../diffusion/models/dreamid_omni/fusion.py | 6 +- 7 files changed, 428 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md index 7621a6058..921fc1bd9 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ b/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. diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index 544c285ab..96e5b3592 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/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) diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index ecb993945..b68944c6c 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/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.""" diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 570c4f238..40a564485 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/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, diff --git a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py b/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py new file mode 100644 index 000000000..a3f57064f --- /dev/null +++ b/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 "" + 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" diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 3f8229722..7c27c31eb 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/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, diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py index a7ba745e9..d0b00bc5b 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py +++ b/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