[Diffusion] [Model] Dreamid omni (#1855)

Signed-off-by: dengyunyang <584797741@qq.com>

docs/.nav.yml

@@ -28,6 +28,7 @@ nav:
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
+      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
       - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md

docs/models/supported_models.md

@@ -51,6 +51,7 @@ th {
 |`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
 |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
 |`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
+|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
 
 
 ## List of Supported Models for NPU

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -39,6 +39,7 @@ The following table shows which models are currently supported by parallelism me
 | **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |        N/A        |
 | **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        N/A        |
 | **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        ✅        |
+| **DreamID-Omni**           | `XuGuo699/DreamID-Omni`       |     ❌      |    ❌    |      ✅       |        ❌        |         ❌          |        N/A        |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.

docs/user_guide/diffusion_acceleration.md

@@ -77,6 +77,7 @@ The following table shows which models are currently supported by each accelerat
 |-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
 
 ### Quantization
 

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -0,0 +1,58 @@
+# X-To-Video-Audio
+
+The `DreamID-Omni` pipeline generates short videos from text, image and video.
+
+## Local CLI Usage
+### Download the Model locally
+Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
+
+```bash
+python download_dreamid_omni.py --output-dir ./dreamid_omni
+```
+After download, the model directory will look like this:
+
+```
+dreamid_omni/
+├── DreamID-Omni/
+│   ├── dreamid_omni.safetensors
+├── MMAudio/
+│   ├── ext_weights/
+│   │   ├── best_netG.pt
+│   │   ├── v1-16.pth
+├── Wan2.2-TI2V-5B/
+│   ├── google/*
+│   ├── models_t5_umt5-xxl-enc-bf16.pth
+│   ├── Wan2.2_VAE.pth
+│
+├── model_index.json # create by download_dreamid_omni.py
+```
+
+### Run the Inference
+```
+python x_to_video_audio.py \
+  --model /xx/dreamid_omni \
+  --prompt "Two people walking together and singing happily" \
+  --image-path ./example0.png ./example1.png \
+  --audio-path ./example0.wav ./example1.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --output dreamid_omni.mp4
+```
+In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
+The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
+
+Key arguments:
+- `--prompt`: text description (string).
+- `--model`: path to the model local directory.
+- `--height/--width`: output resolution (defaults 704 * 1024).
+- `--image-path`: path to the input image list.
+- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
+- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
+- `--num-inference-steps`: number of denoising steps (defaults 45).
+- `--video-negative-prompt`: negative prompt for video generation.
+- `--audio-negative-prompt`: negative prompt for audio generation.
+- `--enable-cpu-offload`: enable CPU offload (defaults False).

examples/offline_inference/x_to_video_audio/download_dreamid_omni.py

@@ -0,0 +1,103 @@
+import argparse
+import fcntl
+import os
+import site
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
+DEPENDENCY_BRANCH = "omni"
+CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
+LOCK_FILE = CACHE_DIR / ".install.lock"
+DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
+
+
+def download_dependency():
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+    with open(LOCK_FILE, "w") as f:
+        fcntl.flock(f, fcntl.LOCK_EX)
+        if not DEPENDENCY_DIR.exists():
+            print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
+            subprocess.run(
+                ["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
+                check=True,
+            )
+            print("Download finished.")
+        fcntl.flock(f, fcntl.LOCK_UN)
+
+    # write .pth to site-packages
+    site_packages = Path(site.getsitepackages()[0])
+    pth_file = site_packages / "vllm_omni_dependency.pth"
+    pth_file.write_text(str(DEPENDENCY_DIR))
+    print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
+
+
+def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
+    """Download files from HF repo and log time + destination."""
+    if os.path.exists(local_dir):
+        print(f"Directory {local_dir} already exists. Skipping download.")
+        return
+    print(f"Starting download from {repo_id} into {local_dir}")
+    start_time = time.time()
+
+    snapshot_download(
+        repo_id=repo_id,
+        local_dir=local_dir,
+        local_dir_use_symlinks=False,
+        allow_patterns=allow_patterns,
+    )
+
+    elapsed = time.time() - start_time
+    print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
+
+
+def main(output_dir: str):
+    # Wan2.2
+    wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
+    timed_download(
+        repo_id="Wan-AI/Wan2.2-TI2V-5B",
+        local_dir=wan_dir,
+        allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
+    )
+
+    # MMAudio
+    mm_audio_dir = os.path.join(output_dir, "MMAudio")
+    timed_download(
+        repo_id="hkchengrex/MMAudio",
+        local_dir=mm_audio_dir,
+        allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
+    )
+
+    dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
+
+    timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
+
+    # Now we construct the config file
+    import json
+
+    data = {
+        "_class_name": "DreamIDOmniPipeline",
+        "fusion": "DreamID-Omni/dreamid_omni.safetensors",
+    }
+
+    with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+
+    print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
+
+    # now we download the dependency code
+    download_dependency()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download models from Hugging Face")
+    parser.add_argument(
+        "--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
+    )
+    args = parser.parse_args()
+    main(args.output_dir)

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -0,0 +1,58 @@
+# X-To-Video-Audio
+
+The `DreamID-Omni` pipeline generates short videos from text, image and video.
+
+## Local CLI Usage
+### Download the Model locally
+Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
+
+```bash
+python download_dreamid_omni.py --output-dir ./dreamid_omni
+```
+After download, the model directory will look like this:
+
+```
+dreamid_omni/
+├── DreamID-Omni/
+│   ├── dreamid_omni.safetensors
+├── MMAudio/
+│   ├── ext_weights/
+│   │   ├── best_netG.pt
+│   │   ├── v1-16.pth
+├── Wan2.2-TI2V-5B/
+│   ├── google/*
+│   ├── models_t5_umt5-xxl-enc-bf16.pth
+│   ├── Wan2.2_VAE.pth
+│
+├── model_index.json # create by download_dreamid_omni.py
+```
+
+### Run the Inference
+```
+python x_to_video_audio.py \
+  --model /xx/dreamid_omni \
+  --prompt "Two people walking together and singing happily" \
+  --image-path ./example0.png ./example1.png \
+  --audio-path ./example0.wav ./example1.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --output dreamid_omni.mp4
+```
+In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
+The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
+
+Key arguments:
+- `--prompt`: text description (string).
+- `--model`: path to the model local directory.
+- `--height/--width`: output resolution (defaults 704 * 1024).
+- `--image-path`: path to the input image list.
+- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
+- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
+- `--num-inference-steps`: number of denoising steps (defaults 45).
+- `--video-negative-prompt`: negative prompt for video generation.
+- `--audio-negative-prompt`: negative prompt for audio generation.
+- `--enable-cpu-offload`: enable CPU offload (defaults False).

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import re
+import time
+
+import librosa
+from PIL import Image
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
+AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).")
+    parser.add_argument("--model", required=True, help="DreamID ckpt root directory.")
+    parser.add_argument("--model-type", default="dreamid-omni", help="Model type.")
+    parser.add_argument("--prompt", default=None, help="Text prompt.")
+
+    parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
+    parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
+    parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
+
+    parser.add_argument("--height", type=int, default=704, help="Video height.")
+    parser.add_argument("--width", type=int, default=1280, help="Video width.")
+    parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
+    parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
+    parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
+    parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
+    parser.add_argument(
+        "--cfg-parallel-size",
+        type=int,
+        default=1,
+        choices=[1, 2],
+        help="Number of GPUs used for classifier free guidance parallel size.",
+    )
+    parser.add_argument(
+        "--video-negative-prompt",
+        default="jitter, bad hands, blur, distortion",
+        help="Negative prompt for video.",
+    )
+    parser.add_argument(
+        "--audio-negative-prompt",
+        default="robotic, muffled, echo, distorted",
+        help="Negative prompt for audio.",
+    )
+    parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
+    parser.add_argument(
+        "--enable-cpu-offload",
+        action="store_true",
+        default=False,
+        help="Enable CPU offloading for diffusion models.",
+    )
+    return parser.parse_args()
+
+
+def load_image_and_audio(image_paths, audio_paths):
+    image = []
+    audio = []
+
+    for path in image_paths:
+        with Image.open(path) as img:
+            img = img.convert("RGB")
+            image.append(img)
+
+    for path in audio_paths:
+        audio_array, sr = librosa.load(path, sr=16000)
+        audio_array = audio_array[int(sr * 1) : int(sr * 3)]
+        audio.append(audio_array)
+    return image, audio
+
+
+def main() -> None:
+    args = parse_args()
+    if args.prompt is None and args.prompt_file is None:
+        raise ValueError("Either --prompt or --prompt-file must be provided.")
+
+    text_prompt = args.prompt
+    if args.prompt_file:
+        import json
+
+        with open(args.prompt_file) as f:
+            text_prompt = json.load(f)
+            text_prompt = re.sub(
+                r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
+            ).strip()
+            text_prompt = re.sub(
+                r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
+            ).strip()
+            text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
+            text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
+
+    image, audio = load_image_and_audio(args.image_path, args.audio_path)
+
+    prompt = {
+        "prompt": text_prompt,
+        "video_negative_prompt": args.video_negative_prompt,
+        "audio_negative_prompt": args.audio_negative_prompt,
+        "multi_modal_data": {"image": image, "audio": audio},
+    }
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+        seed=args.seed,
+        extra_args={
+            "solver_name": args.solver_name,
+            "shift": args.shift,
+        },
+    )
+
+    parallel_config = DiffusionParallelConfig(
+        cfg_parallel_size=args.cfg_parallel_size,
+    )
+
+    omni = Omni(
+        model=args.model,
+        parallel_config=parallel_config,
+        model_type=args.model_type,
+        enable_cpu_offload=args.enable_cpu_offload,
+    )
+    start = time.perf_counter()
+    outputs = omni.generate(prompt, sampling_params)
+    elapsed = time.perf_counter() - start
+
+    if not outputs:
+        raise RuntimeError("No output returned from DreamID-Omni.")
+    output = outputs[0].request_output
+    generated_video = output[0].images[0][0]
+    generated_audio = output[0].images[0][1]
+    try:
+        from dreamid_omni.utils.io_utils import save_video
+    except Exception as e:
+        raise RuntimeError(f"Failed to extract video and audio from DreamID-Omni output. Error: {e}")
+    output_path = args.output
+    save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000)
+    print(f"Saved generated video to {output_path}")
+    print(f"Total time: {elapsed:.2f}s")
+
+
+if __name__ == "__main__":
+    main()

vllm_omni/diffusion/data.py

@@ -322,6 +322,7 @@ class OmniDiffusionConfig:
 
     dtype: torch.dtype = torch.bfloat16
 
+    model_config: dict[str, Any] = field(default_factory=dict)
     tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
 
     # Attention

vllm_omni/diffusion/diffusion_engine.py

@@ -6,6 +6,7 @@ import time
 from collections.abc import Iterable
 from typing import Any
 
+import numpy as np
 import PIL.Image
 from vllm.logger import init_logger
 
@@ -30,6 +31,13 @@ def supports_image_input(model_class_name: str) -> bool:
     return bool(getattr(model_cls, "support_image_input", False))
 
 
+def supports_audio_input(model_class_name: str) -> bool:
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    if model_cls is None:
+        return False
+    return bool(getattr(model_cls, "support_audio_input", False))
+
+
 def image_color_format(model_class_name: str) -> str:
     model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
     return getattr(model_cls, "color_format", "RGB")
@@ -369,7 +377,19 @@ class DiffusionEngine:
             dummy_image = PIL.Image.new(color_format, (width, height))
         else:
             dummy_image = None
-        prompt: OmniTextPrompt = {"prompt": "dummy run", "multi_modal_data": {"image": dummy_image}}
+
+        if supports_audio_input(self.od_config.model_class_name):
+            audio_sr = 16000
+            audio_duration_sec = 4
+            audio_array = np.random.randn(audio_sr * audio_duration_sec).astype(np.float32)
+            dummy_audio = audio_array[audio_sr * 1 : audio_sr * 3]
+        else:
+            dummy_audio = None
+
+        prompt: OmniTextPrompt = {
+            "prompt": "dummy run",
+            "multi_modal_data": {"image": dummy_image, "audio": dummy_audio},
+        }
         req = OmniDiffusionRequest(
             prompts=[prompt],
             sampling_params=OmniDiffusionSamplingParams(

vllm_omni/diffusion/models/dreamid_omni/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .pipeline_dreamid_omni import DreamIDOmniPipeline
+
+__all__ = [
+    "DreamIDOmniPipeline",
+]

vllm_omni/diffusion/models/dreamid_omni/fusion.py

@@ -0,0 +1,335 @@
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.attention.layer import Attention
+
+try:
+    from dreamid_omni.modules.model import WanLayerNorm, WanRMSNorm
+except ImportError:
+    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
+
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import WanModel, rope_apply
+
+
+class FusionModel(nn.Module):
+    def __init__(self, video_config=None, audio_config=None):
+        super().__init__()
+        has_video = True
+        has_audio = True
+        if video_config is not None:
+            self.video_model = WanModel(**video_config)
+        else:
+            has_video = False
+            self.video_model = None
+            print("Warning: No video model is provided!")
+
+        if audio_config is not None:
+            self.audio_model = WanModel(**audio_config)
+        else:
+            has_audio = False
+            self.audio_model = None
+            print("Warning: No audio model is provided!")
+
+        if has_video and has_audio:
+            assert len(self.video_model.blocks) == len(self.audio_model.blocks)
+            self.num_blocks = len(self.video_model.blocks)
+
+            self.inject_cross_attention_kv_projections()
+        self.device = get_local_device()
+
+        self.num_heads = self.video_model.num_heads
+        self.head_dim = self.video_model.dim // self.video_model.num_heads
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            num_kv_heads=self.num_heads,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+        )
+
+    def inject_cross_attention_kv_projections(self):
+        for vid_block in self.video_model.blocks:
+            vid_block.cross_attn.k_fusion = nn.Linear(vid_block.dim, vid_block.dim)
+            vid_block.cross_attn.v_fusion = nn.Linear(vid_block.dim, vid_block.dim)
+            vid_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
+            vid_block.cross_attn.norm_k_fusion = (
+                WanRMSNorm(vid_block.dim, eps=1e-6) if vid_block.qk_norm else nn.Identity()
+            )
+
+        for audio_block in self.audio_model.blocks:
+            audio_block.cross_attn.k_fusion = nn.Linear(audio_block.dim, audio_block.dim)
+            audio_block.cross_attn.v_fusion = nn.Linear(audio_block.dim, audio_block.dim)
+            audio_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
+            audio_block.cross_attn.norm_k_fusion = (
+                WanRMSNorm(audio_block.dim, eps=1e-6) if audio_block.qk_norm else nn.Identity()
+            )
+
+    def merge_kwargs(self, vid_kwargs, audio_kwargs):
+        """
+        keys in each kwarg:
+        e
+        seq_lens
+        grid_sizes
+        freqs
+        context
+        context_lens
+        """
+        merged_kwargs = {}
+        for key in vid_kwargs:
+            merged_kwargs[f"vid_{key}"] = vid_kwargs[key]
+        for key in audio_kwargs:
+            merged_kwargs[f"audio_{key}"] = audio_kwargs[key]
+        return merged_kwargs
+
+    def single_fusion_cross_attention_forward(
+        self,
+        cross_attn_block,
+        src_seq,
+        src_grid_sizes,
+        src_freqs,
+        target_seq,
+        target_seq_lens,
+        target_grid_sizes,
+        target_freqs,
+        context,
+        context_lens,
+        src_ref_lengths=None,
+        src_freqs_scaling=None,
+        target_ref_lengths=None,
+        target_freqs_scaling=None,
+    ):
+        b, n, d = src_seq.size(0), cross_attn_block.num_heads, cross_attn_block.head_dim
+        if hasattr(cross_attn_block, "k_img"):
+            ## means is i2v block
+            q, k, v, k_img, v_img = cross_attn_block.qkv_fn(src_seq, context)
+        else:
+            ## means is t2v block
+            q, k, v = cross_attn_block.qkv_fn(src_seq, context)
+            k_img = v_img = None
+
+        x = self.attn(q, k, v)
+
+        if k_img is not None:
+            img_x = self.attn(q, k_img, v_img)
+            x = x + img_x
+
+        # is_vid = src_grid_sizes.shape[1] > 1
+        # compute target attention
+        target_seq = cross_attn_block.pre_attn_norm_fusion(target_seq)
+        k_target = cross_attn_block.norm_k_fusion(cross_attn_block.k_fusion(target_seq)).view(b, -1, n, d)
+        v_target = cross_attn_block.v_fusion(target_seq).view(b, -1, n, d)
+
+        q = rope_apply(q, src_grid_sizes, src_freqs, ref_lengths=src_ref_lengths, freqs_scaling=src_freqs_scaling)
+        k_target = rope_apply(
+            k_target,
+            target_grid_sizes,
+            target_freqs,
+            ref_lengths=target_ref_lengths,
+            freqs_scaling=target_freqs_scaling,
+        )
+
+        target_x = self.attn(q, k_target, v_target)
+
+        x = x + target_x
+
+        x = x.flatten(2)  # [B, L/P, C]
+
+        x = cross_attn_block.o(x)
+        return x
+
+    def single_fusion_cross_attention_ffn_forward(
+        self,
+        attn_block,
+        src_seq,
+        src_grid_sizes,
+        src_freqs,
+        target_seq,
+        target_seq_lens,
+        target_grid_sizes,
+        target_freqs,
+        context,
+        context_lens,
+        src_e,
+        src_ref_lengths=None,
+        src_freqs_scaling=None,
+        target_ref_lengths=None,
+        target_freqs_scaling=None,
+    ):
+        src_seq = src_seq + self.single_fusion_cross_attention_forward(
+            attn_block.cross_attn,
+            attn_block.norm3(src_seq),
+            src_grid_sizes=src_grid_sizes,
+            src_freqs=src_freqs,
+            target_seq=target_seq,
+            target_seq_lens=target_seq_lens,
+            target_grid_sizes=target_grid_sizes,
+            target_freqs=target_freqs,
+            context=context,
+            context_lens=context_lens,
+            src_ref_lengths=src_ref_lengths,
+            src_freqs_scaling=src_freqs_scaling,
+            target_ref_lengths=target_ref_lengths,
+            target_freqs_scaling=target_freqs_scaling,
+        )
+        y = attn_block.ffn(attn_block.norm2(src_seq).bfloat16() * (1 + src_e[4].squeeze(2)) + src_e[3].squeeze(2))
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            src_seq = src_seq + y * src_e[5].squeeze(2)
+        return src_seq
+
+    def single_fusion_block_forward(
+        self,
+        vid_block,
+        audio_block,
+        vid,
+        audio,
+        vid_e,
+        vid_seq_lens,
+        vid_grid_sizes,
+        vid_freqs,
+        vid_context,
+        vid_context_lens,
+        vid_ref_lengths,
+        vid_freqs_scaling,
+        audio_e,
+        audio_seq_lens,
+        audio_grid_sizes,
+        audio_freqs,
+        audio_context,
+        audio_context_lens,
+        audio_ref_lengths,
+        audio_freqs_scaling,
+    ):
+        ## audio modulation
+        assert audio_e.dtype == torch.bfloat16
+        assert len(audio_e.shape) == 4 and audio_e.size(2) == 6 and audio_e.shape[1] == audio.shape[1], (
+            f"{audio_e.shape}, {audio.shape}"
+        )
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            audio_e = audio_block.modulation(audio_e).chunk(6, dim=2)
+        assert audio_e[0].dtype == torch.bfloat16
+
+        # audio self-attention
+        audio_y = audio_block.self_attn(
+            audio_block.norm1(audio).bfloat16() * (1 + audio_e[1].squeeze(2)) + audio_e[0].squeeze(2),
+            audio_seq_lens,
+            audio_grid_sizes,
+            audio_freqs,
+            ref_lengths=audio_ref_lengths,
+            freqs_scaling=audio_freqs_scaling,
+        )
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            audio = audio + audio_y * audio_e[2].squeeze(2)
+
+        ## video modulation
+        assert len(vid_e.shape) == 4 and vid_e.size(2) == 6 and vid_e.shape[1] == vid.shape[1], (
+            f"{vid_e.shape}, {vid.shape}"
+        )
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            vid_e = vid_block.modulation(vid_e).chunk(6, dim=2)
+
+        # video self-attention
+        vid_y = vid_block.self_attn(
+            vid_block.norm1(vid).bfloat16() * (1 + vid_e[1].squeeze(2)) + vid_e[0].squeeze(2),
+            vid_seq_lens,
+            vid_grid_sizes,
+            vid_freqs,
+            ref_lengths=vid_ref_lengths,
+        )
+
+        with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+            vid = vid + vid_y * vid_e[2].squeeze(2)
+
+        og_audio = audio
+
+        # audio cross-attention
+        audio = self.single_fusion_cross_attention_ffn_forward(
+            audio_block,
+            audio,
+            audio_grid_sizes,
+            audio_freqs,
+            vid,
+            vid_seq_lens,
+            vid_grid_sizes,
+            vid_freqs,
+            audio_context,
+            audio_context_lens,
+            audio_e,
+            src_ref_lengths=audio_ref_lengths,
+            src_freqs_scaling=audio_freqs_scaling,
+            target_ref_lengths=vid_ref_lengths,
+            target_freqs_scaling=None,
+        )
+
+        assert not torch.equal(og_audio, audio), "Audio should be changed after cross-attention!"
+
+        # video cross-attention
+        vid = self.single_fusion_cross_attention_ffn_forward(
+            vid_block,
+            vid,
+            vid_grid_sizes,
+            vid_freqs,
+            og_audio,
+            audio_seq_lens,
+            audio_grid_sizes,
+            audio_freqs,
+            vid_context,
+            vid_context_lens,
+            vid_e,
+            src_ref_lengths=vid_ref_lengths,
+            src_freqs_scaling=None,
+            target_ref_lengths=audio_ref_lengths,
+            target_freqs_scaling=audio_freqs_scaling,
+        )
+
+        return vid, audio
+
+    def forward(
+        self,
+        vid,
+        audio,
+        t,
+        vid_context,
+        audio_context,
+        vid_seq_len,
+        audio_seq_len,
+        ref_ip_lengths=None,
+        ref_audio_lengths=None,
+        slg_layer=False,
+        freqs_scaling=None,
+    ):
+        vid, vid_e, vid_kwargs = self.video_model.prepare_transformer_block_kwargs(
+            x=vid, t=t, context=vid_context, seq_len=vid_seq_len, ref_lengths=ref_ip_lengths
+        )
+
+        audio, audio_e, audio_kwargs = self.audio_model.prepare_transformer_block_kwargs(
+            x=audio,
+            t=t,
+            context=audio_context,
+            seq_len=audio_seq_len,
+            ref_lengths=ref_audio_lengths,
+            freqs_scaling=freqs_scaling,
+        )
+
+        kwargs = self.merge_kwargs(vid_kwargs, audio_kwargs)
+
+        for i in range(self.num_blocks):
+            """
+            1 fusion block refers to 1 audio block with 1 video block.
+            """
+
+            vid_block = self.video_model.blocks[i]
+            audio_block = self.audio_model.blocks[i]
+
+            vid, audio = self.single_fusion_block_forward(
+                vid_block=vid_block, audio_block=audio_block, vid=vid, audio=audio, **kwargs
+            )
+
+        vid = self.video_model.post_transformer_block_out(vid, vid_kwargs["grid_sizes"], vid_e)
+        audio = self.audio_model.post_transformer_block_out(audio, audio_kwargs["grid_sizes"], audio_e)
+
+        return vid, audio
+
+    def set_rope_params(self):
+        self.video_model.set_rope_params()
+        self.audio_model.set_rope_params()

vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py

@@ -0,0 +1,508 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+import math
+import os
+
+import torch
+import torch.distributed
+from diffusers import FlowMatchEulerDiscreteScheduler
+from PIL import Image, ImageOps
+from torch import nn
+from torchvision.transforms import Compose, Normalize
+from tqdm import tqdm
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+try:
+    from dreamid_omni.utils.divisible_crop import DivisibleCrop
+    from dreamid_omni.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps
+    from dreamid_omni.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+    from dreamid_omni.utils.model_loading_utils import (
+        init_mmaudio_vae,
+        init_text_model,
+        init_wan_vae_2_2,
+        load_fusion_checkpoint,
+    )
+    from dreamid_omni.utils.rearrange import Rearrange
+    from dreamid_omni.utils.resize import NaResize
+except ImportError:
+    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
+from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
+
+logger = logging.getLogger(__name__)
+
+
+AUDIO_CONFIG = {
+    "patch_size": [1],
+    "model_type": "t2a",
+    "dim": 3072,
+    "ffn_dim": 14336,
+    "freq_dim": 256,
+    "num_heads": 24,
+    "num_layers": 30,
+    "in_dim": 20,
+    "out_dim": 20,
+    "text_len": 512,
+    "window_size": [-1, -1],
+    "qk_norm": True,
+    "cross_attn_norm": True,
+    "eps": 1e-6,
+    "temporal_rope_scaling_factor": 0.19676,
+}
+
+VIDEO_CONFIG = {
+    "patch_size": [1, 2, 2],
+    "model_type": "ti2v",
+    "dim": 3072,
+    "ffn_dim": 14336,
+    "freq_dim": 256,
+    "num_heads": 24,
+    "num_layers": 30,
+    "in_dim": 48,
+    "out_dim": 48,
+    "text_len": 512,
+    "window_size": [-1, -1],
+    "qk_norm": True,
+    "cross_attn_norm": True,
+    "eps": 1e-6,
+}
+
+
+class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput):
+    """DreamID-Omni pipeline for vLLM-Omni."""
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+
+        self.device = get_local_device()
+        model = od_config.model
+        self.target_dtype = od_config.dtype
+
+        # Init Models
+        ## Load VAEs
+        vae_model_video = init_wan_vae_2_2(model, rank=self.device)
+        vae_model_video.model.requires_grad_(False).eval()
+        vae_model_video.model = vae_model_video.model.bfloat16()
+        self.vae_model_video = vae_model_video
+        self.vae = self.vae_model_video
+
+        vae_model_audio = init_mmaudio_vae(model, rank=self.device)
+        vae_model_audio.requires_grad_(False).eval()
+        self.vae_model_audio = vae_model_audio.bfloat16()
+
+        # Load T5 text model
+        self.text_model = init_text_model(model, rank=self.device)
+        self.text_encoder = self.text_model.model
+
+        # Fusion model
+        ## load audio/video model config
+        Fusion_model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG)
+
+        checkpoint_path = self.od_config.model_config.get("fusion", None)
+        assert checkpoint_path is not None, "fusion checkpoint path is None"
+        load_fusion_checkpoint(Fusion_model, checkpoint_path=os.path.join(model, checkpoint_path))
+        self.model = Fusion_model
+        self.transformer = self.model
+
+        # Fixed attributes, non-configurable
+        self.audio_latent_channel = AUDIO_CONFIG.get("in_dim")
+        self.video_latent_channel = VIDEO_CONFIG.get("in_dim")
+        self.video_latent_length = 31
+        self.audio_latent_length = 157
+        self.target_area = 960 * 960
+
+        # Fixed attributes, CFG scales
+        self.video_cfg_scale = 3.0
+        self.video_ref_cfg_scale = 1.5
+        self.audio_cfg_scale = 4.0
+        self.audio_ref_cfg_scale = 2.0
+
+        # Schedulers will be set in forward
+        self.scheduler_video = None
+        self.scheduler_audio = None
+
+    def load_image_latent_ref_ip_video(
+        self,
+        images,
+        audios,
+        video_frame_height_width,
+    ):
+        # Load size.
+        patch_size = self.model.video_model.patch_size
+        vae_stride = [4, 16, 16]
+
+        # Load image and video.
+        ref_vae_latents = {
+            "image": [],
+            "audio": [],
+        }
+        video_h = video_frame_height_width[0]
+        video_w = video_frame_height_width[1]
+        ref_audio_lengths = []
+        for img in images:
+            # Calculate the required size to keep aspect ratio and fill the rest with padding.
+            img_ratio = img.width / img.height
+            target_ratio = video_w / video_h
+
+            if img_ratio > target_ratio:  # Image is wider than target
+                new_width = video_w
+                new_height = int(new_width / img_ratio)
+            else:  # Image is taller than target
+                new_height = video_h
+                new_width = int(new_height * img_ratio)
+
+            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+            # Create a new image with the target size and place the resized image in the center
+            delta_w = video_w - img.size[0]
+            delta_h = video_h - img.size[1]
+            padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+            new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))
+
+            # Transform to tensor and normalize.
+            image_transform = Compose(
+                [
+                    NaResize(
+                        resolution=math.sqrt(
+                            video_frame_height_width[0] * video_frame_height_width[1]
+                        ),  # 256*448, 480*832
+                        mode="area",
+                        downsample_only=True,
+                    ),
+                    DivisibleCrop((vae_stride[1] * patch_size[1], vae_stride[2] * patch_size[2])),
+                    Normalize(0.5, 0.5),
+                    Rearrange("t c h w -> c t h w"),
+                ]
+            )
+            new_img = image_transform([new_img])
+            new_img = new_img.transpose(0, 1)
+            new_img = new_img.to(self.device)
+            new_img = new_img.to(self.target_dtype)
+
+            img_vae_latent = self.vae_model_video.wrapped_encode(new_img[:, :, None]).to(self.target_dtype).squeeze(0)
+            ref_vae_latents["image"].append(img_vae_latent)
+
+        for audio_array in audios:
+            audio_tensor = torch.from_numpy(audio_array).float().unsqueeze(0)
+            # print(audio_tensor.shape)#torch.Size([1, 81280])
+            audio_vae_latent = self.vae_model_audio.wrapped_encode(audio_tensor)
+            # print(audio_vae_latent.shape)#torch.Size([1, 20, length])
+            audio_length = audio_vae_latent.shape[2]
+            ref_audio_lengths.append(audio_length)
+            audio_vae_latent = audio_vae_latent.squeeze(0).transpose(0, 1)
+            ref_vae_latents["audio"].append(audio_vae_latent)
+
+        ref_vae_latents["image"] = torch.cat(ref_vae_latents["image"], dim=1)
+
+        ref_vae_latents["audio"] = torch.cat(ref_vae_latents["audio"], dim=0)
+
+        return ref_vae_latents, ref_audio_lengths
+
+    def load_weights(self, weights):
+        pass
+
+    def get_scheduler_time_steps(self, sampling_steps, solver_name="unipc", device=0, shift=5.0):
+        torch.manual_seed(4)
+
+        if solver_name == "unipc":
+            sample_scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=1000, shift=1, use_dynamic_shifting=False
+            )
+            sample_scheduler.set_timesteps(sampling_steps, device=device, shift=shift)
+            timesteps = sample_scheduler.timesteps
+
+        elif solver_name == "dpm++":
+            sample_scheduler = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=1000, shift=1, use_dynamic_shifting=False
+            )
+            sampling_sigmas = get_sampling_sigmas(sampling_steps, shift=shift)
+            timesteps, _ = retrieve_timesteps(sample_scheduler, device=device, sigmas=sampling_sigmas)
+
+        elif solver_name == "euler":
+            sample_scheduler = FlowMatchEulerDiscreteScheduler(shift=shift)
+            timesteps, sampling_steps = retrieve_timesteps(
+                sample_scheduler,
+                sampling_steps,
+                device=device,
+            )
+
+        else:
+            raise NotImplementedError("Unsupported solver.")
+
+        return sample_scheduler, timesteps
+
+    def diffuse(
+        self,
+        video_noise: torch.Tensor,
+        audio_noise: torch.Tensor,
+        latents_ref_image: torch.Tensor,
+        latents_ref_audio: torch.Tensor,
+        timesteps_video: torch.Tensor,
+        timesteps_audio: torch.Tensor,
+        text_embeddings_video_pos: torch.Tensor,
+        text_embeddings_video_neg: torch.Tensor,
+        text_embeddings_audio_pos: torch.Tensor,
+        text_embeddings_audio_neg: torch.Tensor,
+        max_seq_len_video: int,
+        max_seq_len_audio: int,
+        freqs_scaling_tensor: torch.Tensor,
+        ref_ip_num: int,
+        ref_audio_length: int,
+        ref_audio_lengths: list,
+        scheduler_video,
+        scheduler_audio,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Diffusion loop with CFG parallel support for DreamID-Omni.
+        """
+        for i, (t_v, t_a) in tqdm(enumerate(zip(timesteps_video, timesteps_audio))):
+            timestep_input = torch.full((1,), t_v, device=self.device)
+
+            model_input_video = torch.cat([video_noise[:, :-ref_ip_num], latents_ref_image], dim=1)
+            model_input_video_neg = torch.cat(
+                [video_noise[:, :-ref_ip_num], torch.zeros_like(latents_ref_image)], dim=1
+            )
+
+            model_input_audio = torch.cat([audio_noise[:-ref_audio_length, :], latents_ref_audio], dim=0)
+
+            model_input_audio_neg = torch.cat(
+                [audio_noise[:-ref_audio_length, :], torch.zeros_like(latents_ref_audio)], dim=0
+            )
+            ref_ip_lengths = [ref_ip_num]
+
+            common_args = {
+                "vid_seq_len": max_seq_len_video,
+                "audio_seq_len": max_seq_len_audio,
+                "freqs_scaling": freqs_scaling_tensor,
+                "ref_ip_lengths": [ref_ip_lengths],
+                "ref_audio_lengths": [ref_audio_lengths],
+            }
+            pos_args = {
+                **common_args,
+                "audio_context": [text_embeddings_audio_pos],
+                "vid_context": [text_embeddings_video_pos],
+            }
+            neg_args = {
+                **common_args,
+                "audio_context": [text_embeddings_audio_neg],
+                "vid_context": [text_embeddings_video_neg],
+            }
+
+            if get_classifier_free_guidance_world_size() > 1:
+                # Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
+                cfg_group = get_cfg_group()
+                cfg_rank = get_classifier_free_guidance_rank()
+
+                if cfg_rank == 0:
+                    pred_vid, pred_audio = self.model(
+                        vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **pos_args
+                    )
+                    pre_vid_ip_neg, _ = self.model(
+                        vid=[model_input_video_neg], audio=[model_input_audio], t=timestep_input, **pos_args
+                    )
+                    pred_vid_0 = pred_vid[0]
+                    pred_audio_0 = pred_audio[0]
+                    pre_vid_ip_0 = pre_vid_ip_neg[0]
+                    pred_refaudio_0 = torch.zeros_like(pred_audio_0)  # dummy tensor
+                else:
+                    pred_vid, pred_audio = self.model(
+                        vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **neg_args
+                    )
+                    _, pred_refaudio_neg = self.model(
+                        vid=[model_input_video], audio=[model_input_audio_neg], t=timestep_input, **pos_args
+                    )
+                    pred_vid_0 = pred_vid[0]
+                    pred_audio_0 = pred_audio[0]
+                    pre_vid_ip_0 = torch.zeros_like(pred_vid_0)  # dummy tensor
+                    pred_refaudio_0 = pred_refaudio_neg[0]
+
+                pred_vid_gathered = cfg_group.all_gather(pred_vid_0, separate_tensors=True)
+                pred_audio_gathered = cfg_group.all_gather(pred_audio_0, separate_tensors=True)
+                pre_vid_ip_gathered = cfg_group.all_gather(pre_vid_ip_0, separate_tensors=True)
+                pred_refaudio_gathered = cfg_group.all_gather(pred_refaudio_0, separate_tensors=True)
+
+                pred_vid_pos = [pred_vid_gathered[0]]
+                pred_vid_neg = [pred_vid_gathered[1]]
+                pred_audio_pos = [pred_audio_gathered[0]]
+                pred_audio_neg = [pred_audio_gathered[1]]
+                pre_vid_ip_neg = [pre_vid_ip_gathered[0]]
+                pred_refaudio_neg = [pred_refaudio_gathered[1]]
+            else:
+                pred_vid_pos, pred_audio_pos = self.model(
+                    vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **pos_args
+                )
+
+                pred_vid_neg, pred_audio_neg = self.model(
+                    vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **neg_args
+                )
+
+                pre_vid_ip_neg, _ = self.model(
+                    vid=[model_input_video_neg], audio=[model_input_audio], t=timestep_input, **pos_args
+                )
+
+                _, pred_refaudio_neg = self.model(
+                    vid=[model_input_video], audio=[model_input_audio_neg], t=timestep_input, **pos_args
+                )
+
+            pred_video_guided = (
+                pred_vid_neg[0]
+                + self.video_cfg_scale * (pred_vid_pos[0] - pred_vid_neg[0])
+                + self.video_ref_cfg_scale * (pred_vid_pos[0] - pre_vid_ip_neg[0])
+            )
+
+            pred_audio_guided = (
+                pred_audio_neg[0]
+                + self.audio_cfg_scale * (pred_audio_pos[0] - pred_audio_neg[0])
+                + self.audio_ref_cfg_scale * (pred_audio_pos[0] - pred_refaudio_neg[0])
+            )
+            video_noise = scheduler_video.step(
+                pred_video_guided.unsqueeze(0), t_v, video_noise.unsqueeze(0), return_dict=False
+            )[0].squeeze(0)
+            audio_noise = scheduler_audio.step(
+                pred_audio_guided.unsqueeze(0), t_a, audio_noise.unsqueeze(0), return_dict=False
+            )[0].squeeze(0)
+        return video_noise, audio_noise
+
+    def forward(
+        self,
+        request: OmniDiffusionRequest,
+        **kwargs,
+    ) -> DiffusionOutput:
+        """Main forward pass for DreamID-Omni pipeline for R2AV task."""
+        # Extract parameters from request
+        r_prompts = request.prompts[0]
+        if isinstance(r_prompts, str):
+            prompt = r_prompts
+            video_negative_prompt = ""
+            audio_negative_prompt = ""
+        else:
+            prompt = r_prompts.get("prompt", "")
+            video_negative_prompt = r_prompts.get("video_negative_prompt", "")
+            audio_negative_prompt = r_prompts.get("audio_negative_prompt", "")
+
+        multi_modal_data = r_prompts.get("multi_modal_data", {}) if not isinstance(r_prompts, str) else None
+        raw_images = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+        raw_audios = multi_modal_data.get("audio", None) if multi_modal_data is not None else None
+        if raw_images is None or raw_audios is None:
+            raise ValueError("This model requires image and audio to run.")
+        if not isinstance(raw_images, list):
+            raw_images = [raw_images]
+        if not isinstance(raw_audios, list):
+            raw_audios = [raw_audios]
+        if len(raw_images) != len(raw_audios):
+            raise ValueError("'image' and 'audio' lists must have the same length.")
+
+        height = request.sampling_params.height
+        width = request.sampling_params.width
+        num_inference_steps = request.sampling_params.num_inference_steps
+        shift = request.sampling_params.extra_args.get("shift", 5.0)
+        solver_name = request.sampling_params.extra_args.get("solver_name", "unipc")
+        seed = request.sampling_params.seed if request.sampling_params.seed is not None else 42
+
+        # 1. Prepare reference latents
+        ref_vae_latents, ref_audio_lengths = self.load_image_latent_ref_ip_video(
+            images=raw_images,
+            audios=raw_audios,
+            video_frame_height_width=(height, width),
+        )
+
+        latents_ref_image = ref_vae_latents["image"]
+        latents_ref_audio = ref_vae_latents["audio"]
+        ref_ip_num = latents_ref_image.shape[1]
+        ref_audio_length = latents_ref_audio.shape[0]
+
+        # 2. scheduler
+        scheduler_video, timesteps_video = self.get_scheduler_time_steps(
+            sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift
+        )
+        scheduler_audio, timesteps_audio = self.get_scheduler_time_steps(
+            sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift
+        )
+
+        # 3. text embedding
+        text_embeddings = self.text_model([prompt, video_negative_prompt, audio_negative_prompt], device=self.device)
+        text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings]
+        text_embeddings_audio_pos = text_embeddings[0]
+        text_embeddings_video_pos = text_embeddings[0]
+        text_embeddings_video_neg = text_embeddings[1]
+        text_embeddings_audio_neg = text_embeddings[2]
+
+        video_latent_h, video_latent_w = height // 16, width // 16
+
+        video_noise_len = self.video_latent_length + ref_ip_num
+        audio_noise_len = self.audio_latent_length + ref_audio_length
+        freqs_scaling_tensor = torch.tensor(
+            self.video_latent_length / self.audio_latent_length, device=self.device, dtype=self.target_dtype
+        )
+
+        video_noise = torch.randn(
+            (self.video_latent_channel, video_noise_len, video_latent_h, video_latent_w),
+            device=self.device,
+            dtype=self.target_dtype,
+            generator=torch.Generator(device=self.device).manual_seed(seed),
+        )
+        audio_noise = torch.randn(
+            (audio_noise_len, self.audio_latent_channel),
+            device=self.device,
+            dtype=self.target_dtype,
+            generator=torch.Generator(device=self.device).manual_seed(seed),
+        )
+
+        _patch_size_h, _patch_size_w = self.model.video_model.patch_size[1], self.model.video_model.patch_size[2]
+
+        max_seq_len_video = (
+            video_noise.shape[1] * video_noise.shape[2] * video_noise.shape[3] // (_patch_size_h * _patch_size_w)
+        )
+
+        max_seq_len_audio = audio_noise_len
+
+        with torch.amp.autocast("cuda", enabled=self.target_dtype != torch.float32, dtype=self.target_dtype):
+            video_noise, audio_noise = self.diffuse(
+                video_noise,
+                audio_noise,
+                latents_ref_image,
+                latents_ref_audio,
+                timesteps_video,
+                timesteps_audio,
+                text_embeddings_video_pos,
+                text_embeddings_video_neg,
+                text_embeddings_audio_pos,
+                text_embeddings_audio_neg,
+                max_seq_len_video,
+                max_seq_len_audio,
+                freqs_scaling_tensor,
+                ref_ip_num,
+                ref_audio_length,
+                ref_audio_lengths,
+                scheduler_video,
+                scheduler_audio,
+            )
+
+        video_noise_for_decode = video_noise[:, :-ref_ip_num]
+        audio_noise_for_decode = audio_noise[:-ref_audio_length, :]
+
+        audio_latents_for_vae = audio_noise_for_decode.unsqueeze(0).transpose(1, 2)
+        generated_audio = self.vae_model_audio.wrapped_decode(audio_latents_for_vae).squeeze().cpu().float().numpy()
+
+        video_latents_for_vae = video_noise_for_decode.unsqueeze(0)
+        generated_video = self.vae_model_video.wrapped_decode(video_latents_for_vae).squeeze(0).cpu().float().numpy()
+
+        return DiffusionOutput(output=(generated_video, generated_audio))

vllm_omni/diffusion/models/dreamid_omni/wan2_2.py

@@ -0,0 +1,506 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+
+import torch
+import torch.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+from vllm_omni.diffusion.attention.layer import Attention
+
+try:
+    from dreamid_omni.modules.model import (
+        ChannelLastConv1d,
+        ConvMLP,
+        Head,
+        MLPProj,
+        ModulationAdd,
+        WanLayerNorm,
+        WanRMSNorm,
+        rope_apply,  # diff from  wan2.2
+        rope_params,
+        sinusoidal_embedding_1d,
+    )
+except ImportError:
+    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
+
+
+class WanSelfAttention(nn.Module):
+    """Optimized self-attention module using vLLM layers."""
+
+    def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        # Unified attention layer
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            num_kv_heads=self.num_heads,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+        )
+
+    # query, key, value function
+    def qkv_fn(self, x):
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x).view(b, s, n, d)
+        return q, k, v
+
+    def forward(self, x, seq_lens, grid_sizes, freqs, ref_lengths=None, freqs_scaling=None):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+            ref_lengths(Tensor): Shape [B]
+        """
+        q, k, v = self.qkv_fn(x)
+
+        x = self.attn(
+            rope_apply(q, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling),
+            rope_apply(k, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling),
+            v,
+        )
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+    def qkv_fn(self, x, context):
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+
+        return q, k, v
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        q, k, v = self.qkv_fn(x, context)
+
+        # compute attention
+        x = self.attn(q, k, v)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanI2VCrossAttention(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, *args, **kwargs):
+        raise RuntimeError("WanT2VCrossAttention is currently disabled and should not be used.")
+
+
+WAN_CROSSATTENTION_CLASSES = {
+    "t2v_cross_attn": WanT2VCrossAttention,
+    "i2v_cross_attn": WanI2VCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        cross_attn_type,
+        dim,
+        ffn_dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=False,
+        eps=1e-6,
+        additional_emb_length=None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # layers
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps)
+        self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        if cross_attn_type == "i2v_cross_attn":
+            assert additional_emb_length is not None, "additional_emb_length should be specified for i2v_cross_attn"
+            self.cross_attn = WanI2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, additional_emb_length)
+        else:
+            assert additional_emb_length is None, "additional_emb_length should be None for t2v_cross_attn"
+            self.cross_attn = WanT2VCrossAttention(
+                dim,
+                num_heads,
+                (-1, -1),
+                qk_norm,
+                eps,
+            )
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim))
+
+        # modulation
+        self.modulation = ModulationAdd(dim, 6)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, L1, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        assert e.dtype == torch.bfloat16
+        assert len(e.shape) == 4 and e.size(2) == 6 and e.shape[1] == x.shape[1], f"{e.shape}, {x.shape}"
+        with amp.autocast("cuda", dtype=torch.bfloat16):
+            e = self.modulation(e).chunk(6, dim=2)
+        assert e[0].dtype == torch.bfloat16
+
+        # self-attention
+        y = self.self_attn(
+            self.norm1(x).bfloat16() * (1 + e[1].squeeze(2)) + e[0].squeeze(2), seq_lens, grid_sizes, freqs
+        )
+        with amp.autocast("cuda", dtype=torch.bfloat16):
+            x = x + y * e[2].squeeze(2)
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e):
+            x = x + self.cross_attn(self.norm3(x), context, context_lens)
+            y = self.ffn(self.norm2(x).bfloat16() * (1 + e[4].squeeze(2)) + e[3].squeeze(2))
+            with amp.autocast("cuda", dtype=torch.bfloat16):
+                x = x + y * e[5].squeeze(2)
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+        return x
+
+
+class WanModel(ModelMixin, ConfigMixin):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video, text-to-audio.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        model_type="t2v",
+        patch_size=(1, 2, 2),
+        text_len=512,
+        in_dim=16,
+        dim=2048,
+        ffn_dim=8192,
+        freq_dim=256,
+        text_dim=4096,
+        additional_emb_dim=None,
+        additional_emb_length=None,
+        out_dim=16,
+        num_heads=16,
+        num_layers=32,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=True,
+        temporal_rope_scaling_factor=1.0,
+        eps=1e-6,
+    ):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        assert model_type in [
+            "t2v",
+            "i2v",
+            "t2a",
+            "tt2a",
+            "ti2v",
+        ]  ## tt2a means text transcript + text description to audio (to support both TTS and T2A
+        self.model_type = model_type
+        is_audio_type = "a" in self.model_type
+        is_video_type = "v" in self.model_type
+        assert is_audio_type ^ is_video_type, "Either audio or video model should be specified"
+        if is_audio_type:
+            ## audio model
+            assert len(patch_size) == 1 and patch_size[0] == 1, (
+                "Audio model should only accept 1 dimensional input, and we dont do patchify"
+            )
+
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        self.temporal_rope_scaling_factor = temporal_rope_scaling_factor
+        self.is_audio_type = is_audio_type
+        self.is_video_type = is_video_type
+        # embeddings
+        if is_audio_type:
+            ## hardcoded to MMAudio
+            self.patch_embedding = nn.Sequential(
+                ChannelLastConv1d(in_dim, dim, kernel_size=7, padding=3),
+                nn.SiLU(),
+                ConvMLP(dim, dim * 4, kernel_size=7, padding=3),
+            )
+        else:
+            self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
+
+        self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim))
+
+        self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks
+        ## so i2v and tt2a share the same cross attention while t2v and t2a share the same cross attention
+        cross_attn_type = "t2v_cross_attn" if model_type in ["t2v", "t2a", "ti2v"] else "i2v_cross_attn"
+
+        if cross_attn_type == "t2v_cross_attn":
+            assert additional_emb_dim is None and additional_emb_length is None, (
+                "additional_emb_length should be None for t2v and t2a model"
+            )
+        else:
+            assert additional_emb_dim is not None and additional_emb_length is not None, (
+                "additional_emb_length should be specified for i2v and tt2a model"
+            )
+
+        self.blocks = nn.ModuleList(
+            [
+                WanAttentionBlock(
+                    cross_attn_type,
+                    dim,
+                    ffn_dim,
+                    num_heads,
+                    window_size,
+                    qk_norm,
+                    cross_attn_norm,
+                    eps,
+                    additional_emb_length,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # head
+        self.head = Head(dim, out_dim, patch_size, eps)
+
+        self.set_rope_params()
+
+        if model_type in ["i2v", "tt2a"]:
+            self.img_emb = MLPProj(additional_emb_dim, dim)
+
+    def set_rope_params(self):
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        dim = self.dim
+        num_heads = self.num_heads
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+
+        if self.is_audio_type:
+            ## to be determined
+            # self.freqs = rope_params(1024, d, freqs_scaling=temporal_rope_scaling_factor)
+            self.freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=self.temporal_rope_scaling_factor)
+        else:
+            self.freqs = torch.cat(
+                [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))],
+                dim=1,
+            )
+
+    def prepare_transformer_block_kwargs(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        ref_lengths=None,
+        freqs_scaling=None,
+    ):
+        # params
+        ## need to change!
+        device = next(self.patch_embedding.parameters()).device
+
+        if self.is_audio_type and freqs_scaling is not None:
+            if isinstance(freqs_scaling, torch.Tensor):
+                scale_val = freqs_scaling.item()
+            else:
+                scale_val = freqs_scaling
+            d = self.dim // self.num_heads
+
+            current_freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=scale_val).to(device)
+        else:
+            current_freqs = self.freqs
+            if current_freqs.device != device:
+                current_freqs = current_freqs.to(device)
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]  ## x is list of [B L D] or [B C F H W]
+        if self.is_audio_type:
+            # [B, 1]
+            grid_sizes = torch.stack([torch.tensor(u.shape[1:2], dtype=torch.long) for u in x])
+        else:
+            # [B, 3]
+            grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+            x = [u.flatten(2).transpose(1, 2) for u in x]  # [B C F H W] -> [B (F H W) C] -> [B L C]
+
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len, f"Sequence length {seq_lens.max()} exceeds maximum {seq_len}."
+        x = torch.cat(
+            [torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]
+        )  # single [B, L, C]
+
+        # time embeddings
+        if t.dim() == 1:
+            t = t.unsqueeze(1).expand(t.size(0), seq_len)
+        with amp.autocast("cuda", dtype=torch.bfloat16):
+            bt = t.size(0)
+            t = t.flatten()
+            e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, seq_len)).bfloat16())
+            e0 = self.time_projection(e).unflatten(2, (6, self.dim))  # [1, 26784, 6, 3072] - B, seq_len, 6, dim
+            assert e.dtype == torch.bfloat16 and e0.dtype == torch.bfloat16
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])
+        )
+
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=current_freqs,
+            context=context,
+            context_lens=context_lens,
+            ref_lengths=ref_lengths,
+            freqs_scaling=freqs_scaling,
+        )
+
+        return x, e, kwargs
+
+    def post_transformer_block_out(self, x, grid_sizes, e):
+        # head
+        x = self.head(x, e)
+        # unpatchify
+        if self.is_audio_type:
+            ## grid_sizes is [B 1] where 1 is L,
+            # converting grid_sizes from [B 1] -> [B]
+            grid_sizes = [gs[0] for gs in grid_sizes]
+            assert len(x) == len(grid_sizes)
+            x = [u[:gs] for u, gs in zip(x, grid_sizes)]
+        else:
+            ## grid_sizes is [B 3] where 3 is F H w
+            x = self.unpatchify(x, grid_sizes)
+
+        return [u.bfloat16() for u in x]
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError("WanModel model does not support forward pass.")
+
+    def unpatchify(self, x, grid_sizes) -> list[torch.Tensor]:
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            # v is [F H w] F * H * 80, 100, it was right padded by 20.
+            u = u[: math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum("fhwpqrc->cfphqwr", u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        # out is list of [C F H W]
+        return out

vllm_omni/diffusion/models/interface.py

@@ -13,6 +13,11 @@ class SupportImageInput(Protocol):
     color_format: ClassVar[str] = "RGB"  # Default color format
 
 
+@runtime_checkable
+class SupportAudioInput(Protocol):
+    support_audio_input: ClassVar[bool] = True
+
+
 @runtime_checkable
 class SupportAudioOutput(Protocol):
     support_audio_output: ClassVar[bool] = True

vllm_omni/diffusion/registry.py

@@ -137,6 +137,11 @@ _DIFFUSION_MODELS = {
         "pipeline_flux2",
         "Flux2Pipeline",
     ),
+    "DreamIDOmniPipeline": (
+        "dreamid_omni",
+        "pipeline_dreamid_omni",
+        "DreamIDOmniPipeline",
+    ),
 }
 
 
@@ -196,9 +201,9 @@ def initialize_model(
             od_config.vae_use_tiling = True
 
         # Configure VAE memory optimization settings from config
-        if hasattr(model.vae, "use_slicing"):
+        if hasattr(model, "vae") and hasattr(model.vae, "use_slicing"):
             model.vae.use_slicing = od_config.vae_use_slicing
-        if hasattr(model.vae, "use_tiling"):
+        if hasattr(model, "vae") and hasattr(model.vae, "use_tiling"):
             model.vae.use_tiling = od_config.vae_use_tiling
 
         if is_distributed_vae:

vllm_omni/entrypoints/omni_diffusion.py

@@ -61,18 +61,23 @@ class OmniDiffusion:
                 "model_index.json",
                 od_config.model,
             )
-            if config_dict is not None:
-                if od_config.model_class_name is None:
-                    od_config.model_class_name = config_dict.get("_class_name", None)
-                od_config.update_multimodal_support()
 
+            if config_dict is None:
+                raise FileNotFoundError("model_index.json not found")
+
+            if od_config.model_class_name is None:
+                od_config.model_class_name = config_dict.get("_class_name", None)
+            od_config.update_multimodal_support()
+
+            if od_config.model_class_name == "DreamIDOmniPipeline":
+                od_config.model_config = config_dict
+            else:
                 tf_config_dict = get_hf_file_to_dict(
                     "transformer/config.json",
                     od_config.model,
                 )
                 od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
-            else:
-                raise FileNotFoundError("model_index.json not found")
+
         except (AttributeError, OSError, ValueError, FileNotFoundError):
             cfg = get_hf_file_to_dict("config.json", od_config.model)
             if cfg is None: