[CI]Remove DreamID-Omni and MagiHuman support (#6357)

Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: [Your Name] <your.email@example.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

.claude/skills/add-diffusion-model/SKILL.md

@@ -105,7 +105,7 @@ For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py`
 
 ## Path B: Custom/Private Repo Model
 
-For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
+For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
 
 ### B1. Understand the reference repo
 
@@ -126,9 +126,9 @@ This is the key design decision for custom models. Follow these placement rules:
 | **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
 | **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
 | **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
-| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
+| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | original-repo package via git clone |
 | **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
-| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
+| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_<name>.py` |
 | **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
 
 ### B3. Handle external dependencies
@@ -157,7 +157,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
 
 Custom models have two common patterns for weight loading:
 
-**Pattern 1: Bypass standard loader** (DreamID-Omni style)
+**Pattern 1: Bypass standard loader** (eager custom init)
 
 When the original model has complex custom init functions that load weights in `__init__`:
 

.claude/skills/add-diffusion-model/references/custom-model-patterns.md

@@ -16,20 +16,19 @@ vllm_omni/diffusion/models/wan2_2/
 
 The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
 
-### Custom model with external deps (e.g., DreamID-Omni)
+### Custom model with external deps
 
 ```
-vllm_omni/diffusion/models/dreamid_omni/
+vllm_omni/diffusion/models/<name>/
 ├── __init__.py                    # Exports pipeline only
-├── pipeline_dreamid_omni.py       # Pipeline: loads ALL weights in __init__ via custom helpers
+├── pipeline_<name>.py             # Pipeline: loads ALL weights in __init__ via custom helpers
-├── fusion.py                      # Custom fusion architecture (video + audio cross-attention)
+└── <backbone>.py                  # Custom transformer / fusion architecture
-└── wan2_2.py                      # Re-implemented Wan backbone with split API
 
-examples/offline_inference/x_to_video_audio/
+examples/offline_inference/<task>/
-└── download_dreamid_omni.py       # Downloads weights from 3 HF repos + clones code repo
+└── download_<name>.py             # Downloads weights and/or clones code repo
 ```
 
-All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
+All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency imported with try/except.
 
 ### Custom model with ported code (e.g., BAGEL)
 
@@ -76,7 +75,7 @@ self.weights_sources = [
 ]
 ```
 
-### Pattern 3: Fully custom loading (DreamID-Omni)
+### Pattern 3: Fully custom loading
 
 ```
 init → load ALL weights eagerly via custom helpers → load_weights() = no-op
@@ -108,8 +107,8 @@ Standard diffusers `model_index.json`:
 Custom model `model_index.json` (minimal):
 ```json
 {
-    "_class_name": "DreamIDOmniPipeline",
+    "_class_name": "YourCustomPipeline",
-    "fusion": "DreamID-Omni/dreamid_omni.safetensors"
+    "fusion": "weights/model.safetensors"
 }
 ```
 
@@ -117,7 +116,7 @@ The only **required** field is `_class_name` — it must match a key in `_DIFFUS
 
 ## External Dependency Management
 
-### Git clone + .pth injection (DreamID-Omni pattern)
+### Git clone + .pth injection
 
 ```python
 def download_dependency():
@@ -163,7 +162,7 @@ The engine checks `isinstance(pipeline, SupportImageInput)` at startup to config
 
 Diffusers models use `config.json` in each subfolder. Custom models often use:
 
-**Module-level config dicts** (DreamID-Omni):
+**Module-level config dicts**:
 ```python
 VIDEO_CONFIG = {
     "patch_size": [1, 2, 2], "model_type": "ti2v",
@@ -181,7 +180,7 @@ vae_cfg = bagel_cfg.get("vae_config", {})
 
 ## Custom Architecture Patterns
 
-### Split forward API (DreamID-Omni)
+### Split forward API
 
 When a fusion model needs to interleave blocks from two backbones:
 

docs/.nav.yml

@@ -24,7 +24,6 @@ nav:
       - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
       - X-To-Text: user_guide/examples/offline_inference/x_to_text.md
-      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md

docs/design/feature/cfg_parallel.md

@@ -63,7 +63,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 
 ### N-Branch CFG (3+ branches)
 
-Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches.
+Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches.
 
 `predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`):
 

docs/mkdocs/hooks/generate_examples.py

@@ -37,7 +37,6 @@ GENERAL_EXAMPLE_SLUGS = frozenset(
         "text_to_speech",
         "text_to_video",
         "x_to_text",
-        "x_to_video_audio",
     }
 )
 

docs/models/supported_models.md

@@ -58,7 +58,6 @@ th {
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
-| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -94,7 +93,6 @@ th {
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
 | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
 | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
 | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
 | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |

docs/user_guide/diffusion_features.md

@@ -129,7 +129,6 @@ The following tables show which models support each feature:
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -147,7 +146,7 @@ The following tables show which models support each feature:
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
-> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
+> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
 > 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
 > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
@@ -165,7 +164,6 @@ The following tables show which models support each feature:
 | **SANA-Video-2B T2V I2V**   |     ❓     |     ❓      |           ❓           |       ❓        |         ❓         |         ❓         |   ❓    |             ❓             |          ❓           |       ❓        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
-| **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **LongCat-Video-Avatar-1.5** |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **MiniMax-H3**               | ✅ (FL2VA) |     ✅      |           ✅           |       ❌        |       ✅ (DiT/TE)  |         ❌         |   ✅    |             ✅             |       ✅ (tile)       |      ✅ (DiT)      |        ❌         |

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -1,115 +0,0 @@
-# X-To-Video-Audio
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/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
-python x_to_video_audio.py \
-  --model /path/to/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 out_dreamid_omni_twoip.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.
-For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
-
-```python
-python x_to_video_audio.py \
-  --model /path/to/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 \
-  --use-hsdp \
-  --hsdp-shard-size 2 \
-  --output out_dreamid_omni_twoip.mp4
-```
-
-
-You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
-
-For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
-
-```python
-# Example usage for oneip, ref media from the official repo DreamID-Omni
-python x_to_video_audio.py \
-  --model /path/to/dreamid_omni \
-  --prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
-  --image-path 9.png \
-  --audio-path 9.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 out_dreamid_omni_oneip.mp4
-```
-
-
-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).
-- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
-- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
-- `--hsdp-replicate-size`: number of HSDP replica groups.
-- `--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).
-
-## Example materials
-
-??? abstract "download_dreamid_omni.py"
-    ``````py
-    --8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py"
-    ``````
-??? abstract "x_to_video_audio.py"
-    ``````py
-    --8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py"
-    ``````

examples/offline_inference/x_to_video_audio/download_dreamid_omni.py

@@ -1,108 +0,0 @@
-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",
-    }
-
-    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')}")
-
-    transformer_dir = os.path.join(output_dir, "transformer")
-    os.makedirs(transformer_dir, exist_ok=True)
-    with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
-        json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
-    print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.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

@@ -1,163 +0,0 @@
-# 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
-└── transformer/
-    └── config.json   # create by download_dreamid_omni.py
-```
-
-### Run the Inference
-```python
-python x_to_video_audio.py \
-  --model /path/to/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 out_dreamid_omni_twoip.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.
-For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
-
-```python
-python x_to_video_audio.py \
-  --model /path/to/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 \
-  --use-hsdp \
-  --hsdp-shard-size 2 \
-  --output out_dreamid_omni_twoip.mp4
-```
-
-
-You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
-
-For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
-
-```python
-# Example usage for oneip, ref media from the official repo DreamID-Omni
-python x_to_video_audio.py \
-  --model /path/to/dreamid_omni \
-  --prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
-  --image-path 9.png \
-  --audio-path 9.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 out_dreamid_omni_oneip.mp4
-```
-
-
-## MagiHuman (`--model-type magi-human`)
-
-MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
-T5Gemma text encoder. A detailed text prompt is the only required input; an optional
-image and/or audio file may be supplied for conditioning. Natively supports Tensor
-Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
-the MoE weights and the text encoder.
-
-> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
-> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
-
-Text-only generation:
-
-```bash
-python x_to_video_audio.py \
-  --model-type magi-human \
-  --model /path/to/daVinci-MagiHuman \
-  --prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \
-  --tensor-parallel-size 4 \
-  --height 256 --width 448 \
-  --num-inference-steps 8 \
-  --seed 52 \
-  --extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
-  --output output_magihuman.mp4
-```
-
-With optional image and audio conditioning:
-
-```bash
-python x_to_video_audio.py \
-  --model-type magi-human \
-  --model /path/to/daVinci-MagiHuman \
-  --prompt "A young woman..." \
-  --tensor-parallel-size 4 \
-  --height 256 --width 448 \
-  --num-inference-steps 8 \
-  --seed 52 \
-  --extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
-  --output output_magihuman.mp4
-```
-
-MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
-in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
-
-- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
-  because the audio length then determines the number of frames).
-- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
-  the audio drives the frame count and is encoded as a latent condition into the DiT
-  (`is_a2v` mode). Omit for pure text-to-video+audio generation.
-- `image_path`: path to an image file for visual conditioning. Applied at both the
-  base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
-- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
-  when these are omitted.
-- `sr_num_inference_steps`: denoising steps for the SR stage.
-
-## DreamID-Omni arguments
-
-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).
-- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
-- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
-- `--hsdp-replicate-size`: number of HSDP replica groups.
-- `--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).
-- `--cache-backend`: enable `cache_dit` for acceleration.
-- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
-  (VAEs, the T5 text encoder, norms and modulation stay bf16.)

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -1,279 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import json
-import re
-import time
-from typing import Any
-
-import numpy as np
-from PIL import Image
-from vllm.multimodal.media.audio import load_audio
-
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-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 X -> video+audio models.")
-    parser.add_argument("--model", required=True, help="Model ckpt root directory.")
-    parser.add_argument(
-        "--model-type",
-        default="dreamid-omni",
-        choices=["dreamid-omni", "magi-human"],
-        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(
-        "--extra-body",
-        type=str,
-        default=None,
-        help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
-        'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
-    )
-
-    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, 3, 4],
-        help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
-    )
-    parser.add_argument(
-        "--tensor-parallel-size",
-        type=int,
-        default=1,
-        help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
-    )
-    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(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "int8"],
-        help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
-    )
-    parser.add_argument(
-        "--enable-cpu-offload",
-        action="store_true",
-        default=False,
-        help="Enable CPU offloading for diffusion models.",
-    )
-    parser.add_argument(
-        "--enable-layerwise-offload",
-        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",
-        help="Enable HSDP for DreamID-Omni fused transformer blocks.",
-    )
-    parser.add_argument(
-        "--hsdp-shard-size",
-        type=int,
-        default=1,
-        help="Number of GPUs used for HSDP sharding when HSDP is enabled.",
-    )
-    parser.add_argument(
-        "--hsdp-replicate-size",
-        type=int,
-        default=1,
-        help="Number of HSDP replica groups. Default 1 means pure sharding.",
-    )
-    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 = load_audio(path, sr=16000)
-        audio_array = audio_array[int(sr * 1) : int(sr * 3)]
-        audio.append(audio_array)
-    return image, audio
-
-
-def _extract_peak_memory_mb(result: Any) -> float:
-    """Pull worker-reported peak VRAM (MiB) generation result.
-
-    Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
-    """
-    if isinstance(result, list):
-        result = result[0] if result else None
-    if result is None:
-        return 0.0
-    val = getattr(result, "peak_memory_mb", 0.0)
-    if not val:
-        inner = result
-        if isinstance(inner, list):
-            inner = inner[0] if inner else None
-        val = getattr(inner, "peak_memory_mb", 0.0)
-    try:
-        return float(val or 0.0)
-    except (TypeError, ValueError):
-        return 0.0
-
-
-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:
-        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()
-
-    if args.model_type == "magi-human":
-        prompt = text_prompt
-        extra_args = json.loads(args.extra_body) if args.extra_body else {}
-        sampling_params = OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_inference_steps=args.num_inference_steps,
-            seed=args.seed,
-            extra_args=extra_args,
-        )
-    else:
-        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,
-            },
-        )
-
-    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_kwargs = dict(
-        model=args.model,
-        cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
-        use_hsdp=args.use_hsdp,
-        hsdp_shard_size=args.hsdp_shard_size,
-        hsdp_replicate_size=args.hsdp_replicate_size,
-        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,
-    )
-    if args.quantization is not None:
-        omni_kwargs["quantization"] = args.quantization
-    omni = Omni(**omni_kwargs)
-    start = time.perf_counter()
-    outputs = omni.generate(prompt, sampling_params)
-    elapsed = time.perf_counter() - start
-
-    if not outputs:
-        raise RuntimeError("No output returned from the model.")
-    result = outputs[0]
-    if not result.images:
-        raise RuntimeError("No video frames found in model output.")
-    generated_video = result.images[0]
-    mm = result.multimodal_output or {}
-    generated_audio = mm.get("audio")
-
-    if args.model_type == "magi-human":
-        # MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
-        fps = float(mm.get("fps", 25))
-        sample_rate = int(mm.get("audio_sample_rate", 24000))
-        frames = generated_video
-    else:
-        fps = float(mm.get("fps", 24))
-        sample_rate = int(mm.get("audio_sample_rate", 16000))
-        # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
-        # mux_video_audio_bytes expects (F, H, W, C) uint8.
-        if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
-            raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
-        frames = generated_video.transpose(1, 2, 3, 0)
-        frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
-
-    audio_np = None
-    if generated_audio is not None:
-        audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
-
-    output_path = args.output
-    video_bytes = mux_video_audio_bytes(
-        frames,
-        audio_np,
-        fps=float(fps),
-        audio_sample_rate=sample_rate,
-    )
-    with open(output_path, "wb") as f:
-        f.write(video_bytes)
-    print(f"Saved generated video to {output_path}")
-    print(f"Total time: {elapsed:.2f}s")
-    peak_mb = _extract_peak_memory_mb(outputs)
-    if peak_mb:
-        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
-
-
-if __name__ == "__main__":
-    main()

tests/diffusion/cache/test_cache_backends.py

@@ -365,13 +365,10 @@ class TestCacheDiTBackend:
 
     @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
-    def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
+    def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
-        """Test DreamID uses pipeline.transformer for cache enable/refresh.
+        """Generic cache path should pick up ``transformer.fused_blocks``."""
-
-        NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
-        """
         mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+        mock_pipeline.__class__.__name__ = "FakeFusedBlocksPipeline"
         mock_pipeline.transformer = Mock()
         mock_pipeline.transformer.fused_blocks = Mock()
         mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(

tests/diffusion/cache/test_cache_dit.py

@@ -6,7 +6,6 @@ Model specific tests for CacheDiT enablement.
 """
 
 import ast
-import sys
 from pathlib import Path
 from unittest.mock import Mock, patch
 
@@ -25,18 +24,9 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import L
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
 from vllm_omni.platforms import current_omni_platform
 
-# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
-# the class even though the dependency may not be set up correctly; this is ok for these
-# tests because we just inspect it and never initialize the model.
-for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
-    sys.modules.setdefault(mod, Mock())
-# isort: split
-from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel  # noqa: E402
-
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 SEPARATE_CFG_TRANSFORMERS = [
-    DreamIdOmniModel,
     HeliosTransformer3DModel,
     LongCatImageTransformer2DModel,
     Cosmos3VFMTransformer,

tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py

@@ -1,27 +0,0 @@
-import pytest
-import torch.nn as nn
-
-pytest.importorskip("dreamid_omni.modules.model")
-
-from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks():
-    model = object.__new__(FusionModel)
-    nn.Module.__init__(model)
-    model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
-    model.attn = nn.Identity()
-
-    conditions = getattr(model, "_hsdp_shard_conditions", None)
-
-    assert conditions is not None
-    assert len(conditions) == 1
-
-    matched = []
-    for name, module in model.named_modules():
-        if any(cond(name, module) for cond in conditions):
-            matched.append(name)
-
-    assert matched == ["fused_blocks.0", "fused_blocks.1"]

tests/diffusion/test_diffusion_ipc.py

@@ -181,7 +181,7 @@ def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
 
 
 def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None:
-    # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output.
+    # LTX2 returns (video, audio) tuples as DiffusionOutput.output.
     video = torch.arange(300_000, dtype=torch.float32)
     audio = torch.arange(300_000, dtype=torch.float32) * 3
     output = DiffusionOutput(output=(video, audio))

tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py

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

tests/e2e/offline_inference/test_magi_human_expansion.py

@@ -1,162 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""
-
-import io
-
-import av
-import numpy as np
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-MODEL = "SII-GAIR/daVinci-MagiHuman-Base-1080p"
-
-# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
-_OMNI_RUNNER_PARAM = (
-    MODEL,
-    None,
-    {
-        "init_timeout": 1200,
-        "tensor_parallel_size": 2,
-        "vae_use_tiling": True,
-    },
-)
-
-_OMNI_RUNNER_PARAM_OFFLOAD = (
-    MODEL,
-    None,
-    {
-        "init_timeout": 1200,
-        "tensor_parallel_size": 2,
-        "enable_cpu_offload": True,
-        "vae_use_tiling": True,
-    },
-)
-
-pytestmark = [
-    pytest.mark.slow,
-    pytest.mark.diffusion,
-    pytest.mark.parametrize(
-        "omni_runner",
-        [_OMNI_RUNNER_PARAM, _OMNI_RUNNER_PARAM_OFFLOAD],
-        indirect=True,
-    ),
-]
-
-
-def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
-    """Validate that the MP4 contains meaningful video and audio tracks."""
-    container = av.open(io.BytesIO(video_bytes))
-
-    v_streams = [s for s in container.streams if s.type == "video"]
-    assert len(v_streams) >= 1, "No video stream found in MP4"
-
-    a_streams = [s for s in container.streams if s.type == "audio"]
-    assert len(a_streams) >= 1, "No audio stream found in MP4"
-
-    v_stream = v_streams[0]
-    assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
-    assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"
-
-    frame_count = 0
-    for frame in container.decode(video=0):
-        frame_count += 1
-        if frame_count >= min_frames:
-            break
-    assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"
-
-    container.close()
-
-
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_magi_human_e2e(omni_runner: OmniRunner):
-    """End-to-end test for MagiHuman generating video and audio."""
-    prompt = (
-        "A young woman with long, wavy golden blonde hair and bright blue eyes, "
-        "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
-        "stationary in front of a softly lit, blurred warm-toned interior. Her "
-        "overall disposition is warm, composed, and gently confident. The camera "
-        "holds a static medium close-up, framing her from the shoulders up, "
-        "with shallow depth of field keeping her face in sharp focus. Soft "
-        "directional key light falls from the upper left, casting a gentle "
-        "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
-        "the levator labii superiors relaxing as her lips part. She speaks in "
-        "clear, warm, unhurried American English: "
-        "\"The most beautiful things in life aren't things at all — "
-        "they're moments, feelings, and the people who make you feel truly alive.\" "
-        "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
-        "shapes each vowel with precision. A faint, genuine smile engages the "
-        "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
-        "in a soft, neutral arch throughout. She maintains steady, forward-facing "
-        "eye contact. Head position remains level; no torso displacement occurs.\n\n"
-        "Dialogue:\n"
-        "<Young blonde woman, American English>: "
-        "\"The most beautiful things in life aren't things at all — "
-        "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
-        "Background Sound:\n"
-        "<Soft, warm indoor ambience with a faint distant piano melody>"
-    )
-
-    sampling_params = OmniDiffusionSamplingParams(
-        height=256,
-        width=448,
-        num_inference_steps=8,
-        seed=52,
-        extra_args={
-            "seconds": 5,
-            "sr_height": 1080,
-            "sr_width": 1920,
-            "sr_num_inference_steps": 5,
-        },
-    )
-
-    omni = omni_runner.omni
-    outputs = list(
-        omni.generate(
-            prompts=[prompt],
-            sampling_params_list=[sampling_params],
-        )
-    )
-
-    assert len(outputs) > 0, "No outputs returned"
-    first = outputs[0]
-
-    assert hasattr(first, "images") and first.images, "No video frames in output"
-    video_frames = first.images[0]
-    assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
-    assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
-
-    mm = first.multimodal_output
-    assert mm, "multimodal_output is empty or missing"
-
-    audio_waveform = mm.get("audio")
-    assert audio_waveform is not None, "No audio waveform in multimodal_output"
-
-    audio_sample_rate = mm.get("audio_sample_rate")
-    assert audio_sample_rate is not None, (
-        "audio_sample_rate not found in multimodal_output; model post-process must propagate it"
-    )
-    assert isinstance(audio_sample_rate, (int, float)), (
-        f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
-    )
-    assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
-
-    fps = mm.get("fps")
-    assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
-    assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
-    assert int(fps) > 0, f"fps must be positive, got {fps}"
-
-    video_bytes = mux_video_audio_bytes(
-        video_frames,
-        audio_waveform,
-        fps=float(fps),
-        audio_sample_rate=int(audio_sample_rate),
-    )
-    assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
-    assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
-
-    _validate_mp4(video_bytes)

tests/model_extras/test_model_extras.py

@@ -177,23 +177,6 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
     assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_magi_human_extra_registry_declares_request_and_response_params() -> None:
-    assert get_extra_body_params("MagiHumanPipeline") == frozenset(
-        {
-            "seconds",
-            "audio_path",
-            "image_path",
-            "sr_height",
-            "sr_width",
-            "sr_num_inference_steps",
-        }
-    )
-    assert get_extra_output_params("MagiHumanPipeline") == frozenset()
-    assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_ltx_extra_registry_declares_official_guidance_params() -> None:

tests/model_tests/diffusion/test_alignment.py

@@ -57,7 +57,6 @@ EXCLUDED_MODELS = [
     "OmniGen2Pipeline",
     "HeliosPipeline",
     "HeliosPyramidPipeline",
-    "DreamIDOmniPipeline",
     "SenseNovaU1Pipeline",
     "HunyuanVideo15Pipeline",
     "HunyuanVideo15ImageToVideoPipeline",
@@ -71,7 +70,6 @@ EXCLUDED_MODELS = [
     # Requires camera/action assets and persistent AR state; covered by its
     # dedicated contract tests and real-checkpoint E2E smoke instead.
     "LingBotWorldCausalDMDPipeline",
-    "MagiHumanPipeline",
     "OmniVoicePipeline",
     "OmniVoice",
     "Cosmos3OmniDiffusersPipeline",

vllm_omni/diffusion/distributed/cfg_parallel.py

@@ -235,8 +235,8 @@ class CFGParallelMixin(metaclass=ABCMeta):
         Predict noise with N-branch CFG dispatch across M GPUs.
 
         This is the multi-branch counterpart of predict_noise_maybe_with_cfg().
-        Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel,
+        Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel).
-        DreamID). Existing 2-branch models should continue using
+        Existing 2-branch models should continue using
         predict_noise_maybe_with_cfg().
 
         Args:

vllm_omni/diffusion/ipc.py

@@ -158,7 +158,7 @@ def _pack_value_if_large(
 
     Walks the container shapes pipelines return as ``DiffusionOutput.output``:
     bare tensors, dicts (e.g. Cosmos3 ``{"image"/"video": ...}``), and
-    tuples/lists (e.g. LTX2 and DreamID ``(video, audio)``). Other values pass
+    tuples/lists (e.g. LTX2 ``(video, audio)``). Other values pass
     through unchanged. ``_unpack_if_shm_handle`` must mirror these shapes — keep
     the two in sync.
     """

vllm_omni/diffusion/models/dreamid_omni/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-__all__ = [
-    "DreamIDOmniPipeline",
-]
-
-
-def __getattr__(name: str):
-    if name == "DreamIDOmniPipeline":
-        from .pipeline_dreamid_omni import DreamIDOmniPipeline
-
-        return DreamIDOmniPipeline
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

vllm_omni/diffusion/models/dreamid_omni/fusion.py

@@ -1,460 +0,0 @@
-import re
-from typing import TYPE_CHECKING
-
-import torch
-import torch.nn as nn
-from cache_dit import ForwardPattern
-from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import ColumnParallelLinear
-
-from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
-
-if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
-try:
-    from dreamid_omni.modules.model import WanLayerNorm
-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 (
-    DistributedRMSNorm,
-    WanModel,
-    rope_apply,
-)
-
-logger = init_logger(__name__)
-
-
-class FusedBlock(nn.Module):
-    """Wrapper pairing a video block and audio block for layerwise offloading.
-
-    Registers both blocks as submodules so their parameters are visible to the offload hooks.
-    """
-
-    def __init__(
-        self,
-        vid_block: nn.Module,
-        audio_block: nn.Module,
-        device: torch.device,
-    ):
-        super().__init__()
-        self.vid_block = vid_block
-        self.audio_block = audio_block
-        self.device = device
-
-    def _cross_attention_forward(
-        self,
-        attn: Attention,
-        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"):
-            q, k, v, k_img, v_img = cross_attn_block.qkv_fn(src_seq, context)
-        else:
-            q, k, v = cross_attn_block.qkv_fn(src_seq, context)
-            k_img = v_img = None
-
-        x = attn(q, k, v)
-
-        if k_img is not None:
-            img_x = attn(q, k_img, v_img)
-            x = x + img_x
-
-        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 = attn(q, k_target, v_target)
-
-        x = x + target_x
-        x = x.flatten(2)
-        x = cross_attn_block.o(x)
-        return x
-
-    def _cross_attention_ffn_forward(
-        self,
-        attn: Attention,
-        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._cross_attention_forward(
-            attn,
-            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 forward(
-        self,
-        hidden_states,
-        encoder_hidden_states,
-        attn: Attention,
-        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,
-    ):
-        vid = hidden_states
-        audio = encoder_hidden_states
-        vid_block = self.vid_block
-        audio_block = self.audio_block
-
-        ## 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._cross_attention_ffn_forward(
-            attn,
-            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._cross_attention_ffn_forward(
-            attn,
-            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
-
-
-class FusionModel(nn.Module):
-    _cache_dit_adapter_config = CacheDiTAdapterConfig(
-        block_forward_patterns={
-            "fused_blocks": ForwardPattern.Pattern_0,
-        },
-        has_separate_cfg=True,
-    )
-
-    _layerwise_offload_blocks_attrs = ["fused_blocks"]
-
-    packed_modules_mapping = {
-        "to_qkv": ["q", "k", "v"],
-    }
-
-    @staticmethod
-    def _is_fused_block(name: str, module) -> bool:
-        return "fused_blocks" in name and name.split(".")[-1].isdigit()
-
-    _hsdp_shard_conditions = [_is_fused_block]
-
-    def __init__(
-        self,
-        video_config=None,
-        audio_config=None,
-        *,
-        quant_config: "QuantizationConfig | None" = None,
-    ):
-        super().__init__()
-        has_video = True
-        has_audio = True
-        self.device = get_local_device()
-        if video_config is not None:
-            self.video_model = WanModel(quant_config=quant_config, prefix="video_model", **video_config)
-        else:
-            has_video = False
-            self.video_model = None
-            logger.warning("No video model is provided!")
-
-        if audio_config is not None:
-            self.audio_model = WanModel(quant_config=quant_config, prefix="audio_model", **audio_config)
-        else:
-            has_audio = False
-            self.audio_model = None
-            logger.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(quant_config=quant_config)
-
-        tp_size = get_tensor_model_parallel_world_size()
-        self.full_num_heads = self.video_model.num_heads
-        self.num_heads = self.full_num_heads // tp_size
-        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,
-        )
-
-        if has_video and has_audio:
-            self.fused_blocks = nn.ModuleList(
-                [
-                    FusedBlock(
-                        self.video_model.blocks[i],
-                        self.audio_model.blocks[i],
-                        self.device,
-                    )
-                    for i in range(self.num_blocks)
-                ]
-            )
-
-    def load_state_dict(self, state_dict, strict=True, assign=False):
-        """Remap checkpoints where blocks are stored under
-        `video_model.blocks.N.*` / `audio_model.blocks.N.*` to the current
-        `fused_blocks.N.vid_block.*` / `fused_blocks.N.audio_block.*`.
-        """
-        needs_remap = any(re.match(r"^(video_model|audio_model)\.blocks\.\d+\.", k) for k in state_dict)
-        if needs_remap:
-            remapped = {}
-            for k, v in state_dict.items():
-                new_k = re.sub(r"^video_model\.blocks\.(\d+)\.", r"fused_blocks.\1.vid_block.", k)
-                new_k = re.sub(r"^audio_model\.blocks\.(\d+)\.", r"fused_blocks.\1.audio_block.", new_k)
-                remapped[new_k] = v
-            state_dict = remapped
-
-        self._detach_blocks_from_backbones()
-
-        return super().load_state_dict(state_dict, strict=strict, assign=assign)
-
-    def inject_cross_attention_kv_projections(
-        self,
-        quant_config: "QuantizationConfig | None" = None,
-    ):
-        def _make_kv_proj(dim: int, prefix: str) -> ColumnParallelLinear:
-            # K/V projections for the cross-modal (a2v / v2a) attention path.
-            # Inputs are full latents, so we quantize matched to self_attn.
-            return ColumnParallelLinear(
-                dim,
-                dim,
-                bias=True,
-                gather_output=False,
-                return_bias=False,
-                quant_config=quant_config,
-                prefix=prefix,
-            )
-
-        tp_size = get_tensor_model_parallel_world_size()
-        for i, vid_block in enumerate(self.video_model.blocks):
-            cross = vid_block.cross_attn
-            block_prefix = f"video_model.blocks.{i}.cross_attn"
-            cross.k_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.k_fusion")
-            cross.v_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.v_fusion")
-            cross.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
-            cross.norm_k_fusion = (
-                DistributedRMSNorm(vid_block.dim // tp_size, eps=1e-6) if vid_block.qk_norm else nn.Identity()
-            )
-
-        for i, audio_block in enumerate(self.audio_model.blocks):
-            cross = audio_block.cross_attn
-            block_prefix = f"audio_model.blocks.{i}.cross_attn"
-            cross.k_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.k_fusion")
-            cross.v_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.v_fusion")
-            cross.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
-            cross.norm_k_fusion = (
-                DistributedRMSNorm(audio_block.dim // tp_size, eps=1e-6) if audio_block.qk_norm else nn.Identity()
-            )
-
-    def _detach_blocks_from_backbones(self) -> None:
-        """Keep offloadable blocks owned only by a single place.
-
-        NOTE: This is a special workaround to support layerwise offloading.
-        The model registers the same Wan blocks under both the video/audio
-        backbones and `fused_blocks` which is a wrapper for unified blocks
-        walking through. However, layerwise offloading will only consider
-        `fused_blocks` as offloadable components and will materialize all
-        other modules onto device, including the same blocks owned by both
-        `fused_blocks` and `video_model` and `audio_model`.
-        """
-        video_blocks = list(self.video_model.blocks)
-        audio_blocks = list(self.audio_model.blocks)
-        self.video_model._modules.pop("blocks", None)
-        self.audio_model._modules.pop("blocks", None)
-        self.video_model.blocks = tuple(video_blocks)
-        self.audio_model.blocks = tuple(audio_blocks)
-
-    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 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 fused_block in self.fused_blocks:
-            vid, audio = fused_block(vid, audio, self.attn, **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

@@ -1,608 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import logging
-import math
-import os
-from collections.abc import Iterable
-from typing import ClassVar
-
-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.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput, SupportsComponentDiscovery
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-
-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,
-    )
-    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__)
-
-
-def get_dreamid_omni_post_process_func(*args, **kwargs):
-    def post_process(output):
-        if isinstance(output, tuple) and len(output) == 2:
-            video, audio = output
-            return {
-                "video": video,
-                "audio": audio,
-                "audio_sample_rate": 16000,
-                "fps": 24,
-            }
-        return output
-
-    return post_process
-
-
-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, SupportsComponentDiscovery
-):
-    """DreamID-Omni pipeline for vLLM-Omni."""
-
-    _dit_modules: ClassVar[list[str]] = ["model"]
-    _encoder_modules: ClassVar[list[str]] = ["text_model"]
-    _vae_modules: ClassVar[list[str]] = ["vae_model_video", "vae_model_audio"]
-
-    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 — weights are loaded later via load_weights()
-        quant_config = getattr(self.od_config, "quantization_config", None)
-        self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG, quant_config=quant_config)
-        self.transformer = self.model
-
-        fusion_path = self.od_config.tf_model_config.get("fusion", None)
-        assert fusion_path is not None, "fusion checkpoint path is None in transformer config"
-        fusion_subfolder = os.path.dirname(fusion_path) or None
-        fusion_filename = os.path.basename(fusion_path)
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=model,
-                subfolder=fusion_subfolder,
-                revision=None,
-                prefix="model.",
-                allow_patterns_overrides=[fusion_filename],
-            )
-        ]
-
-        # 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: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        pipeline_prefix = "model."
-
-        qkv_stacked_mapping = [
-            (".self_attn.to_qkv", ".self_attn.q", "q"),
-            (".self_attn.to_qkv", ".self_attn.k", "k"),
-            (".self_attn.to_qkv", ".self_attn.v", "v"),
-        ]
-
-        params_dict = dict(self.model.named_parameters())
-        loaded_local: set[str] = set()
-
-        tp_size = get_tensor_model_parallel_world_size()
-        tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
-
-        def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
-            """Slice a checkpoint tensor to this rank's TP shard when the
-            param is a TP-local shape but the checkpoint stores the full one.
-            """
-            if tp_size <= 1 or weight.shape == param.shape:
-                return weight
-            if weight.ndim == 1 and weight.numel() == param.numel() * tp_size:
-                return weight.chunk(tp_size, dim=0)[tp_rank]
-            if weight.ndim == 2:
-                if weight.shape[0] == param.shape[0] * tp_size:
-                    return weight.chunk(tp_size, dim=0)[tp_rank]
-                if weight.shape[1] == param.shape[1] * tp_size:
-                    return weight.chunk(tp_size, dim=1)[tp_rank]
-            return weight
-
-        for full_name, loaded_weight in weights:
-            if not full_name.startswith(pipeline_prefix):
-                continue
-            name = full_name[len(pipeline_prefix) :]
-
-            # Fused-QKV stacking: route checkpoint's separate
-            # - ...self_attn.{q,k,v}: weights into the corresponding
-            # - ...self_attn.to_qkv : via its sharded weight_loader.
-            stacked_hit = False
-            for fused_name, raw_name, shard_id in qkv_stacked_mapping:
-                if raw_name not in name:
-                    continue
-                stacked = name.replace(raw_name, fused_name)
-                if stacked not in params_dict:
-                    continue
-                param = params_dict[stacked]
-                param.weight_loader(param, loaded_weight, shard_id)
-                loaded_local.add(stacked)
-                stacked_hit = True
-                break
-            if stacked_hit:
-                continue
-
-            if name not in params_dict:
-                logger.warning(f"Skipping DreamID-Omni weight {full_name} -- not found in model parameters")
-                continue
-            param = params_dict[name]
-            weight_loader = getattr(param, "weight_loader", None)
-            if weight_loader is not None:
-                # Parallel Linears handle TP sharding internally
-                weight_loader(param, loaded_weight)
-            else:
-                default_weight_loader(param, _maybe_shard_weight(loaded_weight, param))
-            loaded_local.add(name)
-
-        # Detach must happen before `DiffusersPipelineLoader._process_weights_after_loading`
-        # and before any forward pass that triggers layerwise-offload hooks.
-        self.model._detach_blocks_from_backbones()
-
-        return {pipeline_prefix + k for k in loaded_local}
-
-    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 predict_noise(self, **kwargs):
-        pred_vid, pred_audio = self.model(**kwargs)
-        return (pred_vid[0], pred_audio[0])
-
-    def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
-        vid_pos, audio_pos = predictions[0]
-        vid_neg, audio_neg = predictions[1]
-        vid_ip_neg, _ = predictions[2]
-        _, refaudio_neg = predictions[3]
-
-        pred_video = (
-            vid_neg
-            + true_cfg_scale["video_cfg_scale"] * (vid_pos - vid_neg)
-            + true_cfg_scale["video_ref_cfg_scale"] * (vid_pos - vid_ip_neg)
-        )
-        pred_audio = (
-            audio_neg
-            + true_cfg_scale["audio_cfg_scale"] * (audio_pos - audio_neg)
-            + true_cfg_scale["audio_ref_cfg_scale"] * (audio_pos - refaudio_neg)
-        )
-        return (pred_video, pred_audio)
-
-    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],
-            }
-
-            branches_kwargs = [
-                {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **pos_args},
-                {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **neg_args},
-                {"vid": [model_input_video_neg], "audio": [model_input_audio], "t": timestep_input, **pos_args},
-                {"vid": [model_input_video], "audio": [model_input_audio_neg], "t": timestep_input, **pos_args},
-            ]
-
-            pred_video_guided, pred_audio_guided = self.predict_noise_with_multi_branch_cfg(
-                do_true_cfg=True,
-                true_cfg_scale={
-                    "video_cfg_scale": self.video_cfg_scale,
-                    "video_ref_cfg_scale": self.video_ref_cfg_scale,
-                    "audio_cfg_scale": self.audio_cfg_scale,
-                    "audio_ref_cfg_scale": self.audio_ref_cfg_scale,
-                },
-                branches_kwargs=branches_kwargs,
-            )
-            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 encode_prompt(
-        self,
-        prompt: str,
-        video_negative_prompt: str = "",
-        audio_negative_prompt: str = "",
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Encode the positive and negative prompts via ``self.text_model``.
-
-        This is the single text-encoder entrypoint so that the runner-level
-        prompt-embedding cache (see
-        ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently
-        memoize results when the same prompts are submitted repeatedly (e.g.
-        GRPO rollouts that sample the same prompt with different seeds).
-
-        Returns:
-            (audio_pos, video_pos, video_neg, audio_neg) embeddings cast to
-            ``self.target_dtype``. ``audio_pos`` and ``video_pos`` are the
-            same tensor, matching the original inline behavior.
-        """
-        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]
-        # Index 0 is the positive embedding, shared between video and audio branches.
-        return text_embeddings[0], text_embeddings[0], text_embeddings[1], text_embeddings[2]
-
-    def forward(
-        self,
-        request: DiffusionRequestBatch,
-        **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_audio_pos, text_embeddings_video_pos, text_embeddings_video_neg, text_embeddings_audio_neg = (
-            self.encode_prompt(
-                prompt=prompt,
-                video_negative_prompt=video_negative_prompt,
-                audio_negative_prompt=audio_negative_prompt,
-            )
-        )
-
-        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

@@ -1,694 +0,0 @@
-# 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.distributed import (
-    get_tensor_model_parallel_world_size,
-    tensor_model_parallel_all_reduce,
-)
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
-from vllm_omni.diffusion.attention.layer import Attention
-
-try:
-    from dreamid_omni.modules.model import (
-        ChannelLastConv1d,
-        ConvMLP,
-        Head,
-        MLPProj,
-        ModulationAdd,
-        WanLayerNorm,
-        rope_apply,  # diff from  wan2.2
-        rope_params,
-        sinusoidal_embedding_1d,
-    )
-except ImportError:
-    raise ImportError("Failed to import from dependency 'dreamid_omni'.")
-
-
-class DistributedRMSNorm(nn.Module):
-    """RMSNorm that computes global RMS across tensor parallel ranks.
-
-    Mirrors `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py::DistributedRMSNorm`
-    """
-
-    def __init__(self, hidden_size: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        tp_size = get_tensor_model_parallel_world_size()
-        x_dtype = x.dtype
-        x_float = x.float()
-        local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
-        local_count = x.shape[-1]
-
-        if tp_size > 1:
-            local_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
-            local_count = local_count * tp_size
-
-        rms = torch.sqrt(local_sum_sq / local_count + self.eps)
-        return (x_float / rms * self.weight.float()).to(x_dtype)
-
-
-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,
-        *,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        assert dim % num_heads == 0
-        super().__init__()
-        self.dim = dim
-        self.head_dim = dim // num_heads
-        self.window_size = window_size
-        self.qk_norm = qk_norm
-        self.eps = eps
-
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=dim,
-            head_size=self.head_dim,
-            total_num_heads=num_heads,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
-        )
-        self.num_heads = self.to_qkv.num_heads
-        self.num_kv_heads = self.to_qkv.num_kv_heads
-        self.tp_inner_dim = self.num_heads * self.head_dim
-        self.tp_kv_dim = self.num_kv_heads * self.head_dim
-
-        self.o = RowParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            input_is_parallel=True,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o" if prefix else "o",
-        )
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = DistributedRMSNorm(self.tp_kv_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_kv_heads,
-            softmax_scale=1.0 / (self.head_dim**0.5),
-            causal=False,
-            prefix=prefix,
-        )
-
-    # query, key, value function
-    def qkv_fn(self, x):
-        b, s = x.shape[:2]
-        qkv, _ = self.to_qkv(x)
-        q_size = self.num_heads * self.head_dim
-        kv_size = self.num_kv_heads * self.head_dim
-        q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
-
-        q = self.norm_q(q).view(b, s, self.num_heads, self.head_dim)
-        k = self.norm_k(k).view(b, s, self.num_kv_heads, self.head_dim)
-        v = v.view(b, s, self.num_kv_heads, self.head_dim)
-        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(nn.Module):
-    """Text -> latent cross-attention."""
-
-    def __init__(
-        self,
-        dim,
-        num_heads,
-        window_size=(-1, -1),
-        qk_norm=True,
-        eps=1e-6,
-        *,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        assert dim % num_heads == 0
-        self.dim = dim
-        self.head_dim = dim // num_heads
-        self.window_size = window_size
-        self.qk_norm = qk_norm
-        self.eps = eps
-
-        # Separate Q/K/V — Q's input differs from K/V's input (text context)
-        self.q = ColumnParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.q" if prefix else "q",
-        )
-        self.k = ColumnParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.k" if prefix else "k",
-        )
-        self.v = ColumnParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.v" if prefix else "v",
-        )
-
-        tp_size = get_tensor_model_parallel_world_size()
-        self.num_heads = num_heads // tp_size
-        self.num_kv_heads = self.num_heads
-        self.tp_inner_dim = self.num_heads * self.head_dim
-
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
-
-        self.o = RowParallelLinear(
-            dim,
-            dim,
-            bias=True,
-            input_is_parallel=True,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o" if prefix else "o",
-        )
-
-        # Kept for interface parity with WanSelfAttention;
-        # the active cross-attention path in DreamID-Omni uses the shared FusionModel.attn.
-        # Disable kv quant to prevent affecting quality.
-        self.attn = Attention(
-            num_heads=self.num_heads,
-            head_size=self.head_dim,
-            num_kv_heads=self.num_kv_heads,
-            softmax_scale=1.0 / (self.head_dim**0.5),
-            causal=False,
-            prefix=prefix,
-            disable_kv_quant=True,
-        )
-
-    def qkv_fn(self, x, context):
-        b = x.size(0)
-
-        q = self.q(x)
-        k = self.k(context)
-        v = self.v(context)
-        q = self.norm_q(q).view(b, -1, self.num_heads, self.head_dim)
-        k = self.norm_k(k).view(b, -1, self.num_kv_heads, self.head_dim)
-        v = v.view(b, -1, self.num_kv_heads, self.head_dim)
-
-        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,
-        *,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        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,
-            quant_config=quant_config,
-            prefix=f"{prefix}.self_attn" if prefix else "self_attn",
-        )
-        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,
-                quant_config=quant_config,
-                prefix=f"{prefix}.cross_attn" if prefix else "cross_attn",
-            )
-        self.norm2 = WanLayerNorm(dim, eps)
-        # FFN layout: Stay compatible with upstream DreamID-Omni checkpoints
-        self.ffn = nn.Sequential(
-            ColumnParallelLinear(
-                dim,
-                ffn_dim,
-                bias=True,
-                gather_output=False,
-                return_bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.ffn.0" if prefix else "ffn.0",
-            ),
-            nn.GELU(approximate="tanh"),
-            RowParallelLinear(
-                ffn_dim,
-                dim,
-                bias=True,
-                input_is_parallel=True,
-                return_bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.ffn.2" if prefix else "ffn.2",
-            ),
-        )
-
-        # 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,
-        *,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ):
-        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,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.blocks.{layer_idx}" if prefix else f"blocks.{layer_idx}",
-                )
-                for layer_idx 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/magi_human/__init__.py

@@ -1 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0

vllm_omni/diffusion/registry.py

@@ -266,11 +266,6 @@ _DIFFUSION_MODELS = {
         "pipeline_flux2",
         "Flux2Pipeline",
     ),
-    "DreamIDOmniPipeline": (
-        "dreamid_omni",
-        "pipeline_dreamid_omni",
-        "DreamIDOmniPipeline",
-    ),
     "SenseNovaU1Pipeline": (
         "sensenova_u1",
         "pipeline_sensenova_u1",
@@ -301,11 +296,6 @@ _DIFFUSION_MODELS = {
         "pipeline_sana_video_i2v",
         "SanaImageToVideoPipeline",
     ),
-    "MagiHumanPipeline": (
-        "magi_human",
-        "pipeline_magi_human",
-        "MagiHumanPipeline",
-    ),
     "OmniVoicePipeline": (
         "omnivoice",
         "pipeline_omnivoice",
@@ -598,9 +588,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LingBotVideoPipeline": "get_lingbot_video_post_process_func",
     "SanaVideoPipeline": "get_sana_video_post_process_func",
     "SanaImageToVideoPipeline": "get_sana_video_i2v_post_process_func",
-    "MagiHumanPipeline": "get_magi_human_post_process_func",
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
-    "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",
     "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_post_process_func",
@@ -645,7 +633,6 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "LingBotVideoPipeline": "get_lingbot_video_pre_process_func",
     "SanaImageToVideoPipeline": "get_sana_video_i2v_pre_process_func",
     "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func",
-    "MagiHumanPipeline": "get_magi_human_pre_process_func",
     "SanaWmPipeline": "get_sana_wm_pre_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",

vllm_omni/model_extras/magi_human.py

@@ -1,20 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-# MagiHuman is a text -> video+audio model. Its model-specific knobs are read
-# from sampling_params.extra_args in pipeline_magi_human.py; declaring them here
-# routes request `extra_body` fields into OmniDiffusionSamplingParams.extra_args
-# so the model can be driven through the standard task example.
-MAGI_HUMAN_EXTRA_BODY_PARAMS = frozenset(
-    {
-        "seconds",
-        "audio_path",
-        "image_path",
-        "sr_height",
-        "sr_width",
-        "sr_num_inference_steps",
-    }
-)
-MAGI_HUMAN_EXTRA_OUTPUT_PARAMS = frozenset()

vllm_omni/model_extras/registry.py

@@ -36,10 +36,6 @@ from vllm_omni.model_extras.ltx2 import (
     ltx_preserves_reference_image_size,
     ltx_transformer_config_subfolder,
 )
-from vllm_omni.model_extras.magi_human import (
-    MAGI_HUMAN_EXTRA_BODY_PARAMS,
-    MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
-)
 from vllm_omni.model_extras.ming_flash_omni import (
     MING_FLASH_OMNI_EXTRA_BODY_PARAMS,
     MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS,
@@ -182,10 +178,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
     },
-    "MagiHumanPipeline": {
-        "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
-        "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
-    },
     "HeliosPipeline": {
         "extra_body_params": HELIOS_EXTRA_BODY_PARAMS,
         "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS,