[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>
This commit is contained in:
wangyuandCursor authored and GitHub committed
1 parent
40dcfdbb43
commit
aeb348d8f8
31 files changed
+26
-6957
No files matched your search
@@ -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__`:
|
||||
|
||||
|
||||
@@ -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
|
||||
├── fusion.py # Custom fusion architecture (video + audio cross-attention)
|
||||
└── wan2_2.py # Re-implemented Wan backbone with split API
|
||||
├── pipeline_<name>.py # Pipeline: loads ALL weights in __init__ via custom helpers
|
||||
└── <backbone>.py # Custom transformer / fusion architecture
|
||||
|
||||
examples/offline_inference/x_to_video_audio/
|
||||
└── download_dreamid_omni.py # Downloads weights from 3 HF repos + clones code repo
|
||||
examples/offline_inference/<task>/
|
||||
└── 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",
|
||||
"fusion": "DreamID-Omni/dreamid_omni.safetensors"
|
||||
"_class_name": "YourCustomPipeline",
|
||||
"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:
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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`):
|
||||
|
||||
|
||||
@@ -37,7 +37,6 @@ GENERAL_EXAMPLE_SLUGS = frozenset(
|
||||
"text_to_speech",
|
||||
"text_to_video",
|
||||
"x_to_text",
|
||||
"x_to_video_audio",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -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` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
@@ -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) | ❌ |
|
||||
|
||||
@@ -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"
|
||||
``````
|
||||
@@ -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)
|
||||
@@ -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.)
|
||||
@@ -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()
|
||||
+3
-6
@@ -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):
|
||||
"""Test DreamID uses pipeline.transformer for cache enable/refresh.
|
||||
|
||||
NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
|
||||
"""
|
||||
def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Generic cache path should pick up ``transformer.fused_blocks``."""
|
||||
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(
|
||||
|
||||
-10
@@ -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,
|
||||
|
||||
@@ -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"]
|
||||
@@ -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))
|
||||
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
DreamID). Existing 2-branch models should continue using
|
||||
Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel).
|
||||
Existing 2-branch models should continue using
|
||||
predict_noise_maybe_with_cfg().
|
||||
|
||||
Args:
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
|
||||
@@ -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}")
|
||||
@@ -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()
|
||||
@@ -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))
|
||||
@@ -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
|
||||
@@ -1 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user