From aeb348d8f8a555f0f9e91bee49f48778744a7bc6 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 26 Aug 2026 09:31:12 +0800 Subject: [PATCH] [CI]Remove DreamID-Omni and MagiHuman support (#6357) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] Co-authored-by: Cursor --- .claude/skills/add-diffusion-model/SKILL.md | 8 +- .../references/custom-model-patterns.md | 27 +- docs/.nav.yml | 1 - docs/design/feature/cfg_parallel.md | 2 +- docs/mkdocs/hooks/generate_examples.py | 1 - docs/models/supported_models.md | 2 - docs/user_guide/diffusion_features.md | 4 +- .../offline_inference/x_to_video_audio.md | 115 - .../x_to_video_audio/download_dreamid_omni.py | 108 - .../x_to_video_audio/x_to_video_audio.md | 163 -- .../x_to_video_audio/x_to_video_audio.py | 279 -- tests/diffusion/cache/test_cache_backends.py | 9 +- tests/diffusion/cache/test_cache_dit.py | 10 - .../dreamid_omni/test_dreamid_omni_hsdp.py | 27 - tests/diffusion/test_diffusion_ipc.py | 2 +- .../test_dreamid_omni_cache_dit.py | 290 --- .../test_magi_human_expansion.py | 162 -- tests/model_extras/test_model_extras.py | 17 - tests/model_tests/diffusion/test_alignment.py | 2 - .../diffusion/distributed/cfg_parallel.py | 4 +- vllm_omni/diffusion/ipc.py | 2 +- .../diffusion/models/dreamid_omni/__init__.py | 14 - .../diffusion/models/dreamid_omni/fusion.py | 460 ---- .../dreamid_omni/pipeline_dreamid_omni.py | 608 ----- .../diffusion/models/dreamid_omni/wan2_2.py | 694 ----- .../diffusion/models/magi_human/__init__.py | 1 - .../models/magi_human/magi_human_dit.py | 1624 ------------ .../models/magi_human/pipeline_magi_human.py | 2306 ----------------- vllm_omni/diffusion/registry.py | 13 - vllm_omni/model_extras/magi_human.py | 20 - vllm_omni/model_extras/registry.py | 8 - 31 files changed, 26 insertions(+), 6957 deletions(-) delete mode 100644 docs/user_guide/examples/offline_inference/x_to_video_audio.md delete mode 100644 examples/offline_inference/x_to_video_audio/download_dreamid_omni.py delete mode 100644 examples/offline_inference/x_to_video_audio/x_to_video_audio.md delete mode 100644 examples/offline_inference/x_to_video_audio/x_to_video_audio.py delete mode 100644 tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py delete mode 100644 tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py delete mode 100644 tests/e2e/offline_inference/test_magi_human_expansion.py delete mode 100644 vllm_omni/diffusion/models/dreamid_omni/__init__.py delete mode 100644 vllm_omni/diffusion/models/dreamid_omni/fusion.py delete mode 100644 vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py delete mode 100644 vllm_omni/diffusion/models/dreamid_omni/wan2_2.py delete mode 100644 vllm_omni/diffusion/models/magi_human/__init__.py delete mode 100644 vllm_omni/diffusion/models/magi_human/magi_human_dit.py delete mode 100644 vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py delete mode 100644 vllm_omni/model_extras/magi_human.py diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index 111f8e9bf..c31262d3a 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -105,7 +105,7 @@ For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py` ## Path B: Custom/Private Repo Model -For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns. +For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns. ### B1. Understand the reference repo @@ -126,9 +126,9 @@ This is the key design decision for custom models. Follow these placement rules: | **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models//pipeline_.py` | Always required | | **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models//_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` | | **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models//` 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//download_.py` | `download_dreamid_omni.py` | +| **Download/setup script** | `examples/offline_inference//download_.py` | `download_.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//` 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__`: diff --git a/.claude/skills/add-diffusion-model/references/custom-model-patterns.md b/.claude/skills/add-diffusion-model/references/custom-model-patterns.md index 2434e0b5d..4e25f9762 100644 --- a/.claude/skills/add-diffusion-model/references/custom-model-patterns.md +++ b/.claude/skills/add-diffusion-model/references/custom-model-patterns.md @@ -16,20 +16,19 @@ vllm_omni/diffusion/models/wan2_2/ The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`. -### Custom model with external deps (e.g., DreamID-Omni) +### Custom model with external deps ``` -vllm_omni/diffusion/models/dreamid_omni/ +vllm_omni/diffusion/models// ├── __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_.py # Pipeline: loads ALL weights in __init__ via custom helpers +└── .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// +└── download_.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: diff --git a/docs/.nav.yml b/docs/.nav.yml index ded87fe9c..7b4052e9c 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -24,7 +24,6 @@ nav: - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md - X-To-Text: user_guide/examples/offline_inference/x_to_text.md - - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md - Online Serving: - Image-To-Image: user_guide/examples/online_serving/image_to_image.md - Image-To-Video: user_guide/examples/online_serving/image_to_video.md diff --git a/docs/design/feature/cfg_parallel.md b/docs/design/feature/cfg_parallel.md index 053360494..db8f738c1 100644 --- a/docs/design/feature/cfg_parallel.md +++ b/docs/design/feature/cfg_parallel.md @@ -63,7 +63,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. ### N-Branch CFG (3+ branches) -Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches. +Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches. `predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`): diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py index 8ddb63561..4b72f9a95 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/hooks/generate_examples.py @@ -37,7 +37,6 @@ GENERAL_EXAMPLE_SLUGS = frozenset( "text_to_speech", "text_to_video", "x_to_text", - "x_to_video_audio", } ) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b9724e8eb..83b81d86e 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -58,7 +58,6 @@ th { | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — | -| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — | | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — | | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | @@ -94,7 +93,6 @@ th { | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — | | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — | | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — | -| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — | | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — | | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 73862febe..944b15ff2 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -129,7 +129,6 @@ The following tables show which models support each feature: | **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | @@ -147,7 +146,7 @@ The following tables show which models support each feature: > Notes: > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT. -> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary. +> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary. > 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload. > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`. > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). @@ -165,7 +164,6 @@ The following tables show which models support each feature: | **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | -| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ | diff --git a/docs/user_guide/examples/offline_inference/x_to_video_audio.md b/docs/user_guide/examples/offline_inference/x_to_video_audio.md deleted file mode 100644 index d5cd0ea46..000000000 --- a/docs/user_guide/examples/offline_inference/x_to_video_audio.md +++ /dev/null @@ -1,115 +0,0 @@ -# X-To-Video-Audio - -Source . - - -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 ": In the frame, a woman with black long hair is identified as .\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**: 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**: 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 maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \ - --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" - `````` diff --git a/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py b/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py deleted file mode 100644 index 2f66d5f77..000000000 --- a/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py +++ /dev/null @@ -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) diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md deleted file mode 100644 index 8305c785f..000000000 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ /dev/null @@ -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 ": In the frame, a woman with black long hair is identified as .\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**: 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**: 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 maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \ - --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... " \ - --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.) diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py deleted file mode 100644 index d7571a66c..000000000 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ /dev/null @@ -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() diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index fb912b08e..85bee131b 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/tests/diffusion/cache/test_cache_backends.py @@ -365,13 +365,10 @@ class TestCacheDiTBackend: @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") - def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter): - """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( diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 341d03d3a..165a98cc6 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -6,7 +6,6 @@ Model specific tests for CacheDiT enablement. """ import ast -import sys from pathlib import Path from unittest.mock import Mock, patch @@ -25,18 +24,9 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import L from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel from vllm_omni.platforms import current_omni_platform -# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect -# the class even though the dependency may not be set up correctly; this is ok for these -# tests because we just inspect it and never initialize the model. -for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"): - sys.modules.setdefault(mod, Mock()) -# isort: split -from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel # noqa: E402 - pytestmark = [pytest.mark.core_model, pytest.mark.cpu] SEPARATE_CFG_TRANSFORMERS = [ - DreamIdOmniModel, HeliosTransformer3DModel, LongCatImageTransformer2DModel, Cosmos3VFMTransformer, diff --git a/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py b/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py deleted file mode 100644 index c7e977ecf..000000000 --- a/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py +++ /dev/null @@ -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"] diff --git a/tests/diffusion/test_diffusion_ipc.py b/tests/diffusion/test_diffusion_ipc.py index f24767e18..aec062057 100644 --- a/tests/diffusion/test_diffusion_ipc.py +++ b/tests/diffusion/test_diffusion_ipc.py @@ -181,7 +181,7 @@ def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None: def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None: - # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output. + # LTX2 returns (video, audio) tuples as DiffusionOutput.output. video = torch.arange(300_000, dtype=torch.float32) audio = torch.arange(300_000, dtype=torch.float32) * 3 output = DiffusionOutput(output=(video, audio)) diff --git a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py b/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py deleted file mode 100644 index e22aa1174..000000000 --- a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py +++ /dev/null @@ -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 "" - pytest.skip( - "DreamID-Omni official example assets not found. " - f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. " - f"Searched: {searched}" - ) - - -def _load_prompt_from_json(path: Path) -> str: - raw_text = path.read_text(encoding="utf-8").strip() - parsed = json.loads(raw_text) - if isinstance(parsed, str): - prompt = parsed - elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str): - prompt = parsed["content"] - else: - raise AssertionError(f"Unexpected DreamID caption format in {path}") - - prompt = re.sub( - r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", - "", - prompt, - flags=re.DOTALL, - ).strip() - prompt = re.sub( - r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", - "", - prompt, - flags=re.DOTALL, - ).strip() - prompt = re.sub(r"\[[A-Z_]+\]", "", prompt) - prompt = re.sub(r"\n\s*\n", "\n", prompt).strip() - return prompt - - -def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]: - images: list[Image.Image] = [] - for path in example_case.image_paths: - with Image.open(path) as image: - images.append(image.convert("RGB")) - return images - - -def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]: - audios: list[np.ndarray] = [] - for path in example_case.audio_paths: - audio_array, sample_rate = sf.read(str(path), dtype="float32") - audio_np = np.asarray(audio_array, dtype=np.float32) - if audio_np.ndim == 2: - audio_np = np.mean(audio_np, axis=1) - audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)] - audios.append(audio_np.reshape(-1)) - return audios - - -def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams: - return OmniDiffusionSamplingParams( - height=704, - width=1280, - num_inference_steps=num_inference_steps, - seed=42, - extra_args={ - "solver_name": "unipc", - "shift": 5.0, - }, - ) - - -def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]: - return { - "prompt": example_case.prompt, - "video_negative_prompt": "jitter, bad hands, blur, distortion", - "audio_negative_prompt": "robotic, muffled, echo, distorted", - "multi_modal_data": { - "image": _load_example_images(example_case), - "audio": _load_example_audios(example_case), - }, - } - - -def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult: - """Extract video and audio from raw Omni outputs""" - assert outputs, "DreamID-Omni returned no outputs" - - result = outputs[0] - assert hasattr(result, "images") and result.images, "not return video frames" - - video = result.images[0] - assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}" - assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}" - - multimodal_output = getattr(result, "multimodal_output", None) or {} - assert "audio" in multimodal_output, "not return audio" - audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1) - assert audio.size > 0, "returned empty audio" - - audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000)) - fps = multimodal_output.get("fps", 24) - duration_s = float(audio.size) / float(audio_sample_rate) - - generation = DreamIDGenerationResult( - video=video.astype(np.float32, copy=False), - audio=audio, - fps=fps, - duration_s=duration_s, - elapsed_s=elapsed_s, - ) - return generation - - -def _run_generation( - *, - model_path: str, - example_case: DreamIDExampleCase, - use_cache_dit: bool, -) -> DreamIDGenerationResult: - runner_kwargs = { - "model_type": "dreamid-omni", - "parallel_config": DiffusionParallelConfig(cfg_parallel_size=1), - } - if use_cache_dit: - runner_kwargs["cache_backend"] = "cache_dit" - runner_kwargs["cache_config"] = CACHE_CONFIG - - with OmniRunner(model_path, **runner_kwargs) as runner: - warmup_payload = _request_payload(example_case) - _ = runner.omni.generate(warmup_payload, _sampling_params(2)) - - elapsed_samples: list[float] = [] - extracted_result: DreamIDGenerationResult | None = None - - request_payload = _request_payload(example_case) - start = time.perf_counter() - outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS)) - elapsed_s = time.perf_counter() - start - elapsed_samples.append(elapsed_s) - if extracted_result is None: - extracted_result = _extract_generation(outputs, elapsed_s) - - assert extracted_result is not None - median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64))) - return extracted_result._replace(elapsed_s=median_elapsed_s) - - -def _video_mae(a: np.ndarray, b: np.ndarray) -> float: - assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}" - return float(np.mean(np.abs(a - b))) - - -def _audio_rms(audio: np.ndarray) -> float: - audio = np.asarray(audio, dtype=np.float32).reshape(-1) - return float(np.sqrt(np.mean(audio**2))) - - -def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float: - rms_a = _audio_rms(a.audio) - rms_b = _audio_rms(b.audio) - return rms_b / max(rms_a, 1e-6) - - -@hardware_test(res={"cuda": "H100"}) -def test_dreamid_omni_cache_dit_e2e() -> None: - model_path = _require_model_path() - example_case = _load_example_case() - - baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False) - cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True) - - video_mae = _video_mae(baseline.video, cached.video) - audio_rms_ratio = _audio_rms_ratio(baseline, cached) - speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf") - - assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}" - assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}" - assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x" diff --git a/tests/e2e/offline_inference/test_magi_human_expansion.py b/tests/e2e/offline_inference/test_magi_human_expansion.py deleted file mode 100644 index e6b3c6a42..000000000 --- a/tests/e2e/offline_inference/test_magi_human_expansion.py +++ /dev/null @@ -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" - ": " - "\"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" - "" - ) - - 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) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 5f48b6eb4..e420a2076 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -177,23 +177,6 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False -@pytest.mark.core_model -@pytest.mark.cpu -def test_magi_human_extra_registry_declares_request_and_response_params() -> None: - assert get_extra_body_params("MagiHumanPipeline") == frozenset( - { - "seconds", - "audio_path", - "image_path", - "sr_height", - "sr_width", - "sr_num_inference_steps", - } - ) - assert get_extra_output_params("MagiHumanPipeline") == frozenset() - assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False - - @pytest.mark.core_model @pytest.mark.cpu def test_ltx_extra_registry_declares_official_guidance_params() -> None: diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 9d4b5a1ad..f9a5072f9 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -57,7 +57,6 @@ EXCLUDED_MODELS = [ "OmniGen2Pipeline", "HeliosPipeline", "HeliosPyramidPipeline", - "DreamIDOmniPipeline", "SenseNovaU1Pipeline", "HunyuanVideo15Pipeline", "HunyuanVideo15ImageToVideoPipeline", @@ -71,7 +70,6 @@ EXCLUDED_MODELS = [ # Requires camera/action assets and persistent AR state; covered by its # dedicated contract tests and real-checkpoint E2E smoke instead. "LingBotWorldCausalDMDPipeline", - "MagiHumanPipeline", "OmniVoicePipeline", "OmniVoice", "Cosmos3OmniDiffusersPipeline", diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index 55130e0cf..4c20f1587 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/vllm_omni/diffusion/distributed/cfg_parallel.py @@ -235,8 +235,8 @@ class CFGParallelMixin(metaclass=ABCMeta): Predict noise with N-branch CFG dispatch across M GPUs. This is the multi-branch counterpart of predict_noise_maybe_with_cfg(). - Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel, - 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: diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index 251e33502..6eb24c540 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -158,7 +158,7 @@ def _pack_value_if_large( Walks the container shapes pipelines return as ``DiffusionOutput.output``: bare tensors, dicts (e.g. Cosmos3 ``{"image"/"video": ...}``), and - tuples/lists (e.g. LTX2 and DreamID ``(video, audio)``). Other values pass + tuples/lists (e.g. LTX2 ``(video, audio)``). Other values pass through unchanged. ``_unpack_if_shm_handle`` must mirror these shapes — keep the two in sync. """ diff --git a/vllm_omni/diffusion/models/dreamid_omni/__init__.py b/vllm_omni/diffusion/models/dreamid_omni/__init__.py deleted file mode 100644 index 8c7f05bf8..000000000 --- a/vllm_omni/diffusion/models/dreamid_omni/__init__.py +++ /dev/null @@ -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}") diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py deleted file mode 100644 index 8bcbbbb64..000000000 --- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py +++ /dev/null @@ -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() diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py deleted file mode 100644 index 5555e3708..000000000 --- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py +++ /dev/null @@ -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)) diff --git a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py b/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py deleted file mode 100644 index c7e464cc8..000000000 --- a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py +++ /dev/null @@ -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 diff --git a/vllm_omni/diffusion/models/magi_human/__init__.py b/vllm_omni/diffusion/models/magi_human/__init__.py deleted file mode 100644 index 988131360..000000000 --- a/vllm_omni/diffusion/models/magi_human/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 diff --git a/vllm_omni/diffusion/models/magi_human/magi_human_dit.py b/vllm_omni/diffusion/models/magi_human/magi_human_dit.py deleted file mode 100644 index 491b1b3c4..000000000 --- a/vllm_omni/diffusion/models/magi_human/magi_human_dit.py +++ /dev/null @@ -1,1624 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 SandAI. All Rights Reserved. -# Ported from daVinci-MagiHuman inference/model/dit/dit_module.py -# Adaptations: removed Ulysses context-parallelism, inlined Modality/VarlenHandler. - -from __future__ import annotations - -import importlib -from collections.abc import Callable -from dataclasses import dataclass, field -from enum import Enum, IntEnum -from typing import TYPE_CHECKING, Any, Literal - -import torch -import torch.nn as nn -from einops import rearrange, repeat -from torch.nn import Parameter -from torch.nn import functional as F -from vllm.distributed import ( - get_tensor_model_parallel_world_size, -) -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.vllm_flash_attn import flash_attn_varlen_func as _vllm_fa_varlen - -try: - from magi_compiler.api import magi_register_custom_op - from magi_compiler.config import CompileConfig -except Exception: - - class CompileConfig: # type: ignore[no-redef] - pass - - def magi_register_custom_op(*args, **kwargs): # type: ignore[no-redef] - def decorator(func): - return func - - return decorator - - -def magi_compile(*args, **kwargs): - """No-op stub — vllm-omni handles execution; magi compilation is skipped.""" - - def decorator(cls_or_fn): - return cls_or_fn - - return decorator - - -# --------------------------------------------------------------------------- -# Inlined from inference/common/sequence_schema.py -# --------------------------------------------------------------------------- -class Modality(IntEnum): - VIDEO = 0 - AUDIO = 1 - TEXT = 2 - - -@dataclass -class VarlenHandler: - cu_seqlens_q: torch.Tensor - cu_seqlens_k: torch.Tensor - max_seqlen_q: int - max_seqlen_k: int - - -def _is_hopper_arch() -> bool: - if not torch.cuda.is_available(): - return False - return torch.cuda.get_device_capability()[0] == 9 - - -# --------------------------------------------------------------------------- -# FFA handler for local / flex attention -# --------------------------------------------------------------------------- -@dataclass -class FFAHandler: - q_ranges: torch.Tensor - k_ranges: torch.Tensor - max_seqlen_q: int - max_seqlen_k: int - attn_type_map: torch.Tensor - softmax_scale: float - - -# --------------------------------------------------------------------------- -# Activation helpers -# --------------------------------------------------------------------------- -class MLPActivationType(Enum): - SWIGLU7 = "swiglu7" - GELU7 = "gelu7" - - -def swiglu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: torch.dtype | None = None): - out_dtype = x.dtype if out_dtype is None else out_dtype - x = x.to(torch.float32) - x_glu, x_linear = x[..., ::2], x[..., 1::2] - x_glu = x_glu.clamp(min=None, max=limit) - x_linear = x_linear.clamp(min=-limit, max=limit) - out_glu = x_glu * torch.sigmoid(alpha * x_glu) - return (out_glu * (x_linear + 1)).to(out_dtype) - - -def gelu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: torch.dtype | None = None): - out_dtype = x.dtype if out_dtype is None else out_dtype - x = x.to(torch.float32) - x_glu = x.clamp(min=None, max=limit) - out_glu = x_glu * torch.sigmoid(alpha * x_glu) - return out_glu.to(out_dtype) - - -def create_activation_func(activation_type: MLPActivationType) -> Callable: - match activation_type: - case MLPActivationType.SWIGLU7: - return swiglu7 - case MLPActivationType.GELU7: - return gelu7 - case _: - raise ValueError(f"Unknown activation type: {activation_type}") - - -# --------------------------------------------------------------------------- -# Modality dispatcher (permutation helper) -# --------------------------------------------------------------------------- -class ModalityDispatcher: - permuted_modality_mapping: torch.Tensor - group_size: torch.Tensor - group_size_cpu: list[int] - num_modalities: int - - def __init__(self, modality_mapping: torch.Tensor, num_modalities: int): - self.modality_mapping = modality_mapping - self.num_modalities = num_modalities - self.permuted_modality_mapping = self._precompute_permute_mapping(modality_mapping) - self.group_size = torch.bincount(self.permuted_modality_mapping, minlength=num_modalities).to(torch.int32) - self.group_size_cpu: list[int] = [int(x) for x in self.group_size.to("cpu").tolist()] - - def _precompute_permute_mapping(self, modality_mapping): - self.permute_mapping = torch.argsort(modality_mapping) - self.inv_permute_mapping = torch.argsort(self.permute_mapping) - return modality_mapping[self.permute_mapping] - - def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]: - return list(torch.split(x, self.group_size_cpu, dim=0)) - - def undispatch(self, *processed_groups: list[torch.Tensor]) -> torch.Tensor: - return torch.cat(processed_groups, dim=0) - - @staticmethod - def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor: - return x[permute_mapping] - - @staticmethod - def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor: - return x[inv_permute_mapping] - - -# --------------------------------------------------------------------------- -# Positional / rotary embedding helpers -# --------------------------------------------------------------------------- -def freq_bands( - num_bands: int, temperature: float = 10000.0, step: int = 2, device: torch.device | None = None -) -> torch.Tensor: - exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands - return 1.0 / (temperature**exp) - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], dim=-1) - - -# --------------------------------------------------------------------------- -# Fourier positional embedding -# --------------------------------------------------------------------------- -class ElementWiseFourierEmbed(nn.Module): - def __init__( - self, - dim: int, - max_res: int = 224, - temperature: float = 10000.0, - in_pixels: bool = True, - linear_bands: bool = False, - learnable: bool = False, - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.float32, - ): - super().__init__() - self.dim = dim - self.in_pixels = in_pixels - self.learnable = learnable - self.temperature = temperature - self.max_res = max_res - self.linear_bands = linear_bands - self.device = device - self.dtype = dtype - bands = self.get_default_bands() - self.bands = nn.Parameter(bands, requires_grad=self.learnable) - - def forward(self, coords: torch.Tensor) -> torch.Tensor: - coords_xyz = coords[:, :3] - sizes = coords[:, 3:6] - refs = coords[:, 6:9] - - scales = (refs - 1) / (sizes - 1) - scales[(refs == 1) & (sizes == 1)] = 1 - assert not scales.isnan().any(), "scales has nan" - assert not scales.isinf().any(), "scales has inf" - - centers = (sizes - 1) / 2 - centers[:, 0] = 0 - coords_xyz = coords_xyz - centers - - bands = self.bands.to(coords.device, coords.dtype) - proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * bands - sin_proj = proj.sin() - cos_proj = proj.cos() - return torch.cat((sin_proj, cos_proj), dim=1).flatten(1) - - def reset_parameters(self): - self.bands.copy_(self.get_default_bands()) - - def get_default_bands(self): - if self.in_pixels: - raise NotImplementedError("in_pixels are not implemented yet") - return freq_bands(self.dim // 8, temperature=self.temperature, step=1, device=self.device).to(self.dtype) - - -# --------------------------------------------------------------------------- -# Multi-modality RMSNorm -# --------------------------------------------------------------------------- -class MultiModalityRMSNorm(nn.Module): - __constants__ = ["dim", "eps", "num_modality"] - - def __init__(self, dim: int, eps: float = 1e-6, device: torch.device | None = None, num_modality: int = 1): - super().__init__() - self.dim = dim - self.eps = eps - self.num_modality = num_modality - self.weight = nn.Parameter(torch.zeros(dim * num_modality, device=device, dtype=torch.float32)) - if num_modality > 1: - self.forward = self.forward_multi_experts - else: - self.forward = self.forward_single_expert - self.reset_parameters() - - def reset_parameters(self): - nn.init.zeros_(self.weight) - - def rms(self, x: torch.Tensor) -> torch.Tensor: - t = x.float() - return t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) - - def forward_multi_experts(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor: - original_dtype = x.dtype - t = self.rms(x) - weight_chunked = self.weight.chunk(self.num_modality, dim=0) - t_list = modality_dispatcher.dispatch(t) - for i in range(self.num_modality): - t_list[i] = t_list[i] * (weight_chunked[i] + 1) - t = modality_dispatcher.undispatch(*t_list) - return t.to(original_dtype) - - def forward_single_expert( - self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher | None = None - ) -> torch.Tensor: - t, original_dtype = x.float(), x.dtype - t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) - return (t * (self.weight + 1)).to(original_dtype) - - -# --------------------------------------------------------------------------- -# Linear layers with bf16 compute and MoE dispatch -# --------------------------------------------------------------------------- -class _BF16ComputeLinear(torch.autograd.Function): - @staticmethod - def forward( - ctx, - input: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None, - output_dtype: torch.dtype | None, - compute_dtype: torch.dtype = torch.bfloat16, - ): - input_cast = input.to(compute_dtype) - weight_cast = weight.to(compute_dtype) - output = torch.matmul(input_cast, weight_cast.t()) - if bias is not None: - output = output + bias.to(compute_dtype) - return output.to(output_dtype) - - -class BaseLinear(nn.Module): - __constants__ = ["in_features", "out_features", "num_layers", "num_experts"] - - def __init__( - self, in_features, out_features, num_layers_for_initialization, num_experts, bias=True, device=None, dtype=None - ): - super().__init__() - factory_kwargs = {"device": device, "dtype": torch.bfloat16} - self.in_features = in_features - self.out_features = out_features - self.num_layers_for_initialization = num_layers_for_initialization - self.num_experts = num_experts - self.use_bias = bias - self.weight = Parameter(torch.empty((out_features * num_experts, in_features), **factory_kwargs)) - if bias: - self.bias = Parameter(torch.empty(out_features * num_experts, **factory_kwargs)) - else: - self.register_parameter("bias", None) - - def forward( - self, - input: torch.Tensor, - output_dtype: torch.dtype | None = None, - modality_dispatcher: ModalityDispatcher | None = None, - ) -> torch.Tensor: - output_dtype = input.dtype if output_dtype is None else output_dtype - return _BF16ComputeLinear.apply(input, self.weight, self.bias, output_dtype, torch.bfloat16) - - -class NativeMoELinear(BaseLinear): - def forward( - self, - input: torch.Tensor, - output_dtype: torch.dtype | None = None, - modality_dispatcher: ModalityDispatcher | None = None, - ) -> torch.Tensor: - output_dtype = input.dtype if output_dtype is None else output_dtype - input_list = modality_dispatcher.dispatch(input) # type: ignore - weight_chunked = self.weight.chunk(self.num_experts, dim=0) - if self.bias is not None: - bias_chunked = self.bias.chunk(self.num_experts, dim=0) - for i in range(self.num_experts): - input_list[i] = _BF16ComputeLinear.apply( - input_list[i], - weight_chunked[i], - bias_chunked[i] if self.bias is not None else None, - output_dtype, - torch.bfloat16, - ) - return modality_dispatcher.undispatch(*input_list) # type: ignore - - -def create_linear( - in_features, out_features, num_layers=1, num_experts=1, bias=True, device=None, dtype=None -) -> BaseLinear | NativeMoELinear: - if num_experts == 1: - return BaseLinear(in_features, out_features, num_layers, num_experts, bias, device, dtype) - else: - return NativeMoELinear(in_features, out_features, num_layers, num_experts, bias, device, dtype) - - -# --------------------------------------------------------------------------- -# MoE TP parallel linear wrappers: per-expert vLLM parallel layers -# --------------------------------------------------------------------------- -class MoEQKVParallelLinear(nn.Module): - """Per-expert QKVParallelLinear with modality dispatch. - - Wraps ``num_experts`` independent QKVParallelLinear instances. - Forward: dispatch tokens by modality → per-expert QKV matmul (TP-sharded) - → undispatch. - """ - - def __init__( - self, - hidden_size: int, - head_size: int, - total_num_heads: int, - total_num_kv_heads: int, - num_experts: int, - bias: bool = False, - ): - super().__init__() - self.num_experts = num_experts - self.experts = nn.ModuleList( - [ - QKVParallelLinear( - hidden_size=hidden_size, - head_size=head_size, - total_num_heads=total_num_heads, - total_num_kv_heads=total_num_kv_heads, - bias=bias, - return_bias=False, - ) - for _ in range(num_experts) - ] - ) - # Expose per-rank head info from the first expert (all are identical). - self.num_heads = self.experts[0].num_heads - self.num_kv_heads = self.experts[0].num_kv_heads - self.head_size = head_size - - def forward( - self, - x: torch.Tensor, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - x_list = modality_dispatcher.dispatch(x) - out_list: list[torch.Tensor] = [] - for i in range(self.num_experts): - out = self.experts[i](x_list[i]) - out_list.append(out) - return modality_dispatcher.undispatch(*out_list) - - -class MoEColumnParallelLinear(nn.Module): - """Per-expert ColumnParallelLinear with modality dispatch. - - Forward: dispatch → per-expert column-parallel matmul → undispatch. - Output stays TP-local (no gather). - """ - - def __init__( - self, - input_size: int, - output_size: int, - num_experts: int, - bias: bool = False, - ): - super().__init__() - self.num_experts = num_experts - self.experts = nn.ModuleList( - [ - ColumnParallelLinear( - input_size=input_size, - output_size=output_size, - bias=bias, - gather_output=False, - return_bias=False, - ) - for _ in range(num_experts) - ] - ) - - def forward( - self, - x: torch.Tensor, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - x_list = modality_dispatcher.dispatch(x) - out_list: list[torch.Tensor] = [] - for i in range(self.num_experts): - out = self.experts[i](x_list[i]) - out_list.append(out) - return modality_dispatcher.undispatch(*out_list) - - -class MoERowParallelLinear(nn.Module): - """Per-expert RowParallelLinear with modality dispatch. - - Forward: dispatch → per-expert row-parallel matmul (includes all-reduce) - → undispatch. - """ - - def __init__( - self, - input_size: int, - output_size: int, - num_experts: int, - bias: bool = False, - ): - super().__init__() - self.num_experts = num_experts - self.experts = nn.ModuleList( - [ - RowParallelLinear( - input_size=input_size, - output_size=output_size, - bias=bias, - input_is_parallel=True, - return_bias=False, - ) - for _ in range(num_experts) - ] - ) - - def forward( - self, - x: torch.Tensor, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - x_list = modality_dispatcher.dispatch(x) - out_list: list[torch.Tensor] = [] - for i in range(self.num_experts): - out = self.experts[i](x_list[i]) - out_list.append(out) - return modality_dispatcher.undispatch(*out_list) - - -def validate_magi_human_tp_constraints( - *, - hidden_size: int, - num_heads_q: int, - num_heads_kv: int, - tensor_parallel_size: int, -) -> None: - """Validate MagiHuman TP divisibility constraints. - - Both shared layers (num_modality == 1) and MoE layers (num_modality == 3) - support TP via vLLM's parallel linear layers (QKVParallelLinear / - ColumnParallelLinear / RowParallelLinear). MoE layers use per-expert - parallel layers with modality dispatch. - - Supported tp_sizes given default config (hidden=5120, heads_q=40, kv=8): 1, 2, 4. - """ - tp = tensor_parallel_size - if tp <= 1: - return - errors: list[str] = [] - if num_heads_q % tp != 0: - errors.append(f"num_heads_q ({num_heads_q}) must be divisible by tensor_parallel_size ({tp})") - if num_heads_kv % tp != 0: - errors.append(f"num_heads_kv ({num_heads_kv}) must be divisible by tensor_parallel_size ({tp})") - # SWIGLU layers use intermediate = int(hidden * 8/3) // 4 * 4 - intermediate_swiglu = int(hidden_size * 4 * 2 / 3) // 4 * 4 - if intermediate_swiglu % tp != 0: - errors.append( - f"swiglu intermediate_size ({intermediate_swiglu}) must be divisible by " - f"tensor_parallel_size ({tp}). Supported tp values: 1, 2, 4" - ) - # GELU7 MoE layers use intermediate = hidden * 4 - intermediate_gelu = hidden_size * 4 - if intermediate_gelu % tp != 0: - errors.append(f"gelu intermediate_size ({intermediate_gelu}) must be divisible by tensor_parallel_size ({tp})") - if errors: - raise ValueError("MagiHuman TP constraint violations:\n" + "\n".join(f" - {e}" for e in errors)) - - -# --------------------------------------------------------------------------- -# Flash attention (no context-parallelism) — uses vllm's flash attention -# --------------------------------------------------------------------------- - -HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None - - -def _fa_varlen_simple( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, -) -> torch.Tensor: - had_batch = query.ndim == 4 - if had_batch: - query = query.squeeze(0) - key = key.squeeze(0) - value = value.squeeze(0) - seq_len = query.shape[0] - cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=query.device) - out = _vllm_fa_varlen( - q=query, - k=key, - v=value, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=seq_len, - max_seqlen_k=seq_len, - ) - if had_batch: - out = out.unsqueeze(0) - return out - - -@magi_register_custom_op(name="infra::flash_attn_func", is_subgraph_boundary=True) -def flash_attn_func(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: - return _fa_varlen_simple(query, key, value) - - -def _split_q_range_with_no_overlap( - q_ranges: torch.Tensor, k_ranges: torch.Tensor -) -> tuple[list[list[int]], list[list[list[int]]]]: - range_boundary = torch.unique(q_ranges, sorted=True).tolist() - candidates = [[start, end, []] for start, end in zip(range_boundary[:-1], range_boundary[1:])] - q_ranges = q_ranges.tolist() - k_ranges = k_ranges.tolist() - for q_range, k_range in zip(q_ranges, k_ranges): - q_start, q_end = q_range - for q_range_cand in candidates: - if q_start <= q_range_cand[0] and q_range_cand[1] <= q_end: - q_range_cand[2].append(k_range) - q_ranges_out = [] - k_ranges_out = [] - for q_range_cand in candidates: - if len(q_range_cand[2]) > 0: - q_ranges_out.append(q_range_cand[0:2]) - k_ranges_out.append(q_range_cand[2]) - return q_ranges_out, k_ranges_out - - -def _flash_attn_with_correction( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - q_ranges: list[list[int]], - k_range_list: list[list[list[int]]], -): - output = torch.zeros_like(query) - output_lse = torch.zeros((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device) - - for q_range, k_ranges in zip(q_ranges, k_range_list): - q_start, q_end = q_range - q_chunk = query[q_start:q_end] - q_len = q_chunk.shape[0] - - # Concatenate all k_ranges into a single key/value block, then run one - # flash-attention call. This avoids the need to merge per-chunk LSEs. - k_parts = [key[ks:ke] for ks, ke in k_ranges] - v_parts = [value[ks:ke] for ks, ke in k_ranges] - k_combined = torch.cat(k_parts, dim=0) if len(k_parts) > 1 else k_parts[0] - v_combined = torch.cat(v_parts, dim=0) if len(v_parts) > 1 else v_parts[0] - k_len = k_combined.shape[0] - - cu_q = torch.tensor([0, q_len], dtype=torch.int32, device=query.device) - cu_k = torch.tensor([0, k_len], dtype=torch.int32, device=query.device) - qo_out = _vllm_fa_varlen( - q=q_chunk, - k=k_combined, - v=v_combined, - cu_seqlens_q=cu_q, - cu_seqlens_k=cu_k, - max_seqlen_q=q_len, - max_seqlen_k=k_len, - ) - output[q_start:q_end] = qo_out - return output, output_lse - - -def _flex_flash_attn_func_infer_output_meta( - query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - output = torch.empty_like(query) - output_lse = torch.empty((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device) - return output, output_lse - - -@magi_register_custom_op( - name="infra::flex_flash_attn_func", - mutates_args=(), - infer_output_meta_fn=_flex_flash_attn_func_infer_output_meta, - is_subgraph_boundary=True, -) -def flex_flash_attn_func( - query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - if HAS_MAGI_ATTENTION and _is_hopper_arch(): - from magi_attention.api import flex_flash_attn_func as magi_flex_flash_attn_func - - return magi_flex_flash_attn_func(query, key, value, q_ranges, k_ranges) - else: - q_ranges_split, k_range_list = _split_q_range_with_no_overlap(q_ranges, k_ranges) - return _flash_attn_with_correction(query, key, value, q_ranges_split, k_range_list) - - -def flash_attn_no_cp(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16) - return flash_attn_func(q, k, v).squeeze(0) - - -def flex_flash_attn_no_cp( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - q_ranges: torch.Tensor, - k_ranges: torch.Tensor, -) -> torch.Tensor: - q, k, v = q.to(torch.bfloat16).squeeze(0), k.to(torch.bfloat16).squeeze(0), v.to(torch.bfloat16).squeeze(0) - out, _ = flex_flash_attn_func(q, k, v, q_ranges=q_ranges, k_ranges=k_ranges) - return out - - -# --------------------------------------------------------------------------- -# Attention module (no context-parallelism) -# --------------------------------------------------------------------------- -@dataclass -class AttentionConfig: - hidden_size: int - num_heads_q: int - num_heads_kv: int - head_dim: int - params_dtype: torch.dtype - checkpoint_qk_layernorm_rope: bool - num_modality: int - num_layers: int - use_local_attn: bool = False - enable_attn_gating: bool = False - - -class Attention(torch.nn.Module): - config: AttentionConfig - - def __init__(self, config: AttentionConfig): - super().__init__() - self.config = config - self.pre_norm = MultiModalityRMSNorm(config.hidden_size, eps=1e-6, num_modality=config.num_modality) - self.gating_size = config.num_heads_q if config.enable_attn_gating else 0 - - # Both shared blocks (num_modality == 1) and MoE blocks (num_modality > 1) - # use vLLM's parallel linear layers for TP support. - # MoE blocks wrap per-expert parallel layers with modality dispatch. - if config.num_modality == 1: - # QKVParallelLinear handles GQA head-sharding for any tp_size. - # The combined checkpoint weight [Q, K, V, G] is split during - # load_weights: Q+K+V → linear_qkv, G → linear_gating. - self.linear_qkv = QKVParallelLinear( - hidden_size=config.hidden_size, - head_size=config.head_dim, - total_num_heads=config.num_heads_q, - total_num_kv_heads=config.num_heads_kv, - bias=False, - return_bias=False, - ) - self.linear_proj = RowParallelLinear( - input_size=config.num_heads_q * config.head_dim, - output_size=config.hidden_size, - bias=False, - input_is_parallel=True, - return_bias=False, - ) - if config.enable_attn_gating: - self.linear_gating = ColumnParallelLinear( - input_size=config.hidden_size, - output_size=config.num_heads_q, - bias=False, - gather_output=False, - return_bias=False, - ) - else: - self.linear_gating = None - else: - # MoE blocks: per-expert TP-sharded parallel layers. - self.linear_qkv = MoEQKVParallelLinear( - hidden_size=config.hidden_size, - head_size=config.head_dim, - total_num_heads=config.num_heads_q, - total_num_kv_heads=config.num_heads_kv, - num_experts=config.num_modality, - bias=False, - ) - self.linear_proj = MoERowParallelLinear( - input_size=config.num_heads_q * config.head_dim, - output_size=config.hidden_size, - num_experts=config.num_modality, - bias=False, - ) - if config.enable_attn_gating: - self.linear_gating = MoEColumnParallelLinear( - input_size=config.hidden_size, - output_size=config.num_heads_q, - num_experts=config.num_modality, - bias=False, - ) - else: - self.linear_gating = None - - self.q_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality) - self.k_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality) - - # q_size / kv_size reflect the per-rank head count when tp > 1. - # Both shared and MoE QKV layers expose .num_heads / .num_kv_heads. - if config.num_modality == 1: - self.q_size = self.linear_qkv.num_heads * config.head_dim - self.kv_size = self.linear_qkv.num_kv_heads * config.head_dim - self._local_heads_q = self.linear_qkv.num_heads - self._local_heads_kv = self.linear_qkv.num_kv_heads - else: - self.q_size = self.linear_qkv.num_heads * config.head_dim - self.kv_size = self.linear_qkv.num_kv_heads * config.head_dim - self._local_heads_q = self.linear_qkv.num_heads - self._local_heads_kv = self.linear_qkv.num_kv_heads - - def forward( - self, - hidden_states: torch.Tensor, - rope: torch.Tensor, - permute_mapping: torch.Tensor, - inv_permute_mapping: torch.Tensor, - varlen_handler: VarlenHandler, - local_attn_handler: FFAHandler | None, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - hidden_states = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) - - if self.config.num_modality == 1: - # vLLM parallel layers with return_bias=False return a single tensor. - qkv = self.linear_qkv(hidden_states).to(torch.float32) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - if self.linear_gating is not None: - g = self.linear_gating(hidden_states).to(torch.float32) - else: - g = hidden_states.new_empty(hidden_states.shape[0], 0) - else: - # MoE TP path: per-expert QKV parallel layers. - qkv = self.linear_qkv(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - if self.linear_gating is not None: - g = self.linear_gating(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32) - else: - g = hidden_states.new_empty(hidden_states.shape[0], 0) - - q = q.view(-1, self._local_heads_q, self.config.head_dim) - k = k.view(-1, self._local_heads_kv, self.config.head_dim) - v = v.view(-1, self._local_heads_kv, self.config.head_dim) - g = g.view(k.shape[0], self._local_heads_q, -1) - - q = self.q_norm(q, modality_dispatcher=modality_dispatcher) - k = self.k_norm(k, modality_dispatcher=modality_dispatcher) - - q = ModalityDispatcher.inv_permute(q, inv_permute_mapping).unsqueeze(0) - k = ModalityDispatcher.inv_permute(k, inv_permute_mapping).unsqueeze(0) - v = ModalityDispatcher.inv_permute(v, inv_permute_mapping).unsqueeze(0) - - sin_emb, cos_emb = rope.tensor_split(2, -1) - q = apply_rotary_emb_torch(q, cos_emb, sin_emb) - k = apply_rotary_emb_torch(k, cos_emb, sin_emb) - - if self.config.use_local_attn and local_attn_handler is not None: - self_attn_out = flex_flash_attn_no_cp(q, k, v, local_attn_handler.q_ranges, local_attn_handler.k_ranges) - else: - self_attn_out = flash_attn_no_cp(q, k, v) - self_attn_out = ModalityDispatcher.permute(self_attn_out, permute_mapping) - - if self.config.enable_attn_gating: - self_attn_out = self_attn_out * torch.sigmoid(g) - - self_attn_out = self_attn_out.view(-1, self._local_heads_q * self.config.head_dim).to(torch.bfloat16) - if self.config.num_modality == 1: - return self.linear_proj(self_attn_out) - return self.linear_proj(self_attn_out, modality_dispatcher=modality_dispatcher) - - -# --------------------------------------------------------------------------- -# MLP module -# --------------------------------------------------------------------------- -@dataclass -class MLPConfig: - hidden_size: int - intermediate_size: int - activation_type: MLPActivationType - params_dtype: torch.dtype - num_modality: int = 1 - num_layers: int = 1 - gated_act: bool = False - - -class MLP(torch.nn.Module): - config: MLPConfig - - def __init__(self, config: MLPConfig): - super().__init__() - num_experts = config.num_modality - self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=config.num_modality) - intermediate_size_up = config.intermediate_size * 2 if config.gated_act else config.intermediate_size - - # Both shared blocks (num_experts == 1) and MoE blocks (num_experts > 1) - # use vLLM's parallel linear layers for TP support. - if num_experts == 1: - # ColumnParallelLinear shards the output dim uniformly. For - # SWIGLU7 the interleaved [up0, gate0, up1, gate1, ...] format - # is preserved within each rank's contiguous slice, so swiglu7 - # (which uses x[..., ::2] / x[..., 1::2]) still works correctly. - self.up_gate_proj = ColumnParallelLinear( - input_size=config.hidden_size, - output_size=intermediate_size_up, - bias=False, - gather_output=False, - return_bias=False, - ) - self.down_proj = RowParallelLinear( - input_size=config.intermediate_size, - output_size=config.hidden_size, - bias=False, - input_is_parallel=True, - return_bias=False, - ) - else: - # MoE blocks: per-expert TP-sharded parallel layers. - self.up_gate_proj = MoEColumnParallelLinear( - input_size=config.hidden_size, - output_size=intermediate_size_up, - num_experts=num_experts, - bias=False, - ) - self.down_proj = MoERowParallelLinear( - input_size=config.intermediate_size, - output_size=config.hidden_size, - num_experts=num_experts, - bias=False, - ) - self.activation_func = create_activation_func(config.activation_type) - - def forward(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor: - x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) - if isinstance(self.up_gate_proj, ColumnParallelLinear): - x = self.up_gate_proj(x).to(torch.float32) - x = self.activation_func(x).to(torch.bfloat16) - return self.down_proj(x).to(torch.float32) - # MoE TP path: per-expert column/row parallel layers. - x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32) - x = self.activation_func(x).to(torch.bfloat16) - x = self.down_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32) - return x - - -# --------------------------------------------------------------------------- -# Adapter (per-modality embedders + RoPE) -# --------------------------------------------------------------------------- -@dataclass -class AdapterConfig: - hidden_size: int - num_attention_heads: int - text_in_channels: int - video_in_channels: int - audio_in_channels: int - params_dtype: torch.dtype - - -class Adapter(torch.nn.Module): - config: AdapterConfig - - def __init__(self, config: AdapterConfig): - super().__init__() - self.config = config - self.video_embedder = nn.Linear(config.video_in_channels, config.hidden_size, bias=True, dtype=torch.float32) - self.text_embedder = nn.Linear(config.text_in_channels, config.hidden_size, bias=True, dtype=torch.float32) - self.audio_embedder = nn.Linear(config.audio_in_channels, config.hidden_size, bias=True, dtype=torch.float32) - self.rope = ElementWiseFourierEmbed( - config.hidden_size // config.num_attention_heads, in_pixels=False, learnable=False - ) - - def forward(self, x, coords_mapping, video_mask, audio_mask, text_mask): - rope = self.rope(coords_mapping) - - text_input = x[text_mask, : self.config.text_in_channels] - audio_input = x[audio_mask, : self.config.audio_in_channels] - video_input = x[video_mask, : self.config.video_in_channels] - - text_out = self.text_embedder(text_input) - audio_out = self.audio_embedder(audio_input) - video_out = self.video_embedder(video_input) - - output_x = torch.zeros(x.shape[0], self.config.hidden_size, device=x.device, dtype=x.dtype) - output_x[text_mask] = text_out - output_x[audio_mask] = audio_out - output_x[video_mask] = video_out - return output_x, rope - - -# --------------------------------------------------------------------------- -# Transformer layer (no CP) -# --------------------------------------------------------------------------- -class TransFormerLayer(torch.nn.Module): - def __init__(self, config: Any, layer_idx: int): - super().__init__() - num_modality = 3 if layer_idx in config.mm_layers else 1 - use_local_attn = layer_idx in config.local_attn_layers - self.post_norm = layer_idx in config.post_norm_layers - attention_config = AttentionConfig( - hidden_size=config.hidden_size, - num_heads_q=config.num_heads_q, - num_heads_kv=config.num_heads_kv, - head_dim=config.head_dim, - params_dtype=config.params_dtype, - checkpoint_qk_layernorm_rope=config.checkpoint_qk_layernorm_rope, - num_modality=num_modality, - num_layers=config.num_layers, - use_local_attn=use_local_attn, - enable_attn_gating=config.enable_attn_gating, - ) - self.attention: Attention = Attention(attention_config) - - activation_type = MLPActivationType.GELU7 if layer_idx in config.gelu7_layers else MLPActivationType.SWIGLU7 - if activation_type == MLPActivationType.SWIGLU7: - gated_act = True - intermediate_size = int(config.hidden_size * 4 * 2 / 3) // 4 * 4 - else: - gated_act = False - intermediate_size = config.hidden_size * 4 - mlp_config = MLPConfig( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - activation_type=activation_type, - params_dtype=config.params_dtype, - num_modality=num_modality, - num_layers=config.num_layers, - gated_act=gated_act, - ) - self.mlp: MLP = MLP(mlp_config) - if self.post_norm: - self.attn_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality) - self.mlp_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality) - - def forward( - self, - hidden_states: torch.Tensor, - rope: torch.Tensor, - permute_mapping: torch.Tensor, - inv_permute_mapping: torch.Tensor, - varlen_handler: VarlenHandler, - local_attn_handler: FFAHandler | None, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - attn_out = self.attention( - hidden_states, - rope, - permute_mapping, - inv_permute_mapping, - varlen_handler, - local_attn_handler, - modality_dispatcher, - ) - if self.post_norm: - attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher) - hidden_states = hidden_states + attn_out - - mlp_out = self.mlp(hidden_states, modality_dispatcher) - if self.post_norm: - mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher) - hidden_states = hidden_states + mlp_out - return hidden_states - - -# --------------------------------------------------------------------------- -# TransformerBlock with magi_compile -# --------------------------------------------------------------------------- -is_base_model = True - - -def config_patch(compile_config: CompileConfig) -> CompileConfig: - global is_base_model - if is_base_model: - is_base_model = False - else: - compile_config.offload_config.gpu_resident_weight_ratio = 0.0 - return compile_config - - -@magi_compile( - config_patch=config_patch, dynamic_arg_dims={"x": 0, "rope": 0, "permute_mapping": 0, "inv_permute_mapping": 0} -) -class TransformerBlock(torch.nn.Module): - def __init__(self, model_config: Any): - super().__init__() - self.layers: list[TransFormerLayer] = nn.ModuleList() - for layer_idx in range(model_config.num_layers): - self.layers.append(TransFormerLayer(model_config, layer_idx)) - - def forward( - self, - x: torch.Tensor, - rope: torch.Tensor, - permute_mapping: torch.Tensor, - inv_permute_mapping: torch.Tensor, - varlen_handler: VarlenHandler, - local_attn_handler: FFAHandler | None, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - for layer in self.layers: - x = layer( - x, rope, permute_mapping, inv_permute_mapping, varlen_handler, local_attn_handler, modality_dispatcher - ) - return x - - -# --------------------------------------------------------------------------- -# Internal config for TransformerBlock / DiTModel construction -# --------------------------------------------------------------------------- -@dataclass -class TransformerConfig: - hidden_size: int - video_in_channels: int - audio_in_channels: int - text_in_channels: int - params_dtype: torch.dtype - post_process_dtype: torch.dtype - - -# --------------------------------------------------------------------------- -# DiTModel (no context-parallelism) -# --------------------------------------------------------------------------- -class DiTModel(torch.nn.Module): - config: TransformerConfig - _layerwise_offload_blocks_attr = "blocks" - - @property - def blocks(self) -> nn.ModuleList: - return self.block.layers - - def __init__(self, model_config: Any): - super().__init__() - validate_magi_human_tp_constraints( - hidden_size=model_config.hidden_size, - num_heads_q=model_config.hidden_size // model_config.head_dim, - num_heads_kv=model_config.num_query_groups, - tensor_parallel_size=get_tensor_model_parallel_world_size(), - ) - self.config = TransformerConfig( - hidden_size=model_config.hidden_size, - video_in_channels=model_config.video_in_channels, - audio_in_channels=model_config.audio_in_channels, - text_in_channels=model_config.text_in_channels, - params_dtype=model_config.params_dtype, - post_process_dtype=torch.float32, - ) - adapter_config = AdapterConfig( - hidden_size=model_config.hidden_size, - num_attention_heads=model_config.num_heads_q, - text_in_channels=model_config.text_in_channels, - video_in_channels=model_config.video_in_channels, - audio_in_channels=model_config.audio_in_channels, - params_dtype=torch.float32, - ) - self.adapter: Adapter = Adapter(adapter_config) - self.block: TransformerBlock = TransformerBlock(model_config=model_config) - self.final_norm_video = MultiModalityRMSNorm(self.config.hidden_size) - self.final_norm_audio = MultiModalityRMSNorm(self.config.hidden_size) - self.final_linear_video = nn.Linear( - self.config.hidden_size, self.config.video_in_channels, bias=False, dtype=torch.float32 - ) - self.final_linear_audio = nn.Linear( - self.config.hidden_size, self.config.audio_in_channels, bias=False, dtype=torch.float32 - ) - - def forward( - self, - x: torch.Tensor, - coords_mapping: torch.Tensor, - modality_mapping: torch.Tensor, - varlen_handler: VarlenHandler, - local_attn_handler: FFAHandler | None, - ): - modality_dispatcher = ModalityDispatcher(modality_mapping, 3) - permute_mapping = modality_dispatcher.permute_mapping - inv_permute_mapping = modality_dispatcher.inv_permute_mapping - video_mask = modality_mapping == Modality.VIDEO - audio_mask = modality_mapping == Modality.AUDIO - text_mask = modality_mapping == Modality.TEXT - - x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask) - - x = x.to(self.config.params_dtype) - x = ModalityDispatcher.permute(x, permute_mapping) - - x = self.block( - x, - rope, - permute_mapping=permute_mapping, - inv_permute_mapping=inv_permute_mapping, - varlen_handler=varlen_handler, - local_attn_handler=local_attn_handler, - modality_dispatcher=modality_dispatcher, - ) - - x = ModalityDispatcher.inv_permute(x, inv_permute_mapping) - - x_video = x[video_mask].to(self.final_norm_video.weight.dtype) - x_video = self.final_norm_video(x_video) - x_video = self.final_linear_video(x_video) - - x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype) - x_audio = self.final_norm_audio(x_audio) - x_audio = self.final_linear_audio(x_audio) - - x_out = torch.zeros( - x.shape[0], - max(self.config.video_in_channels, self.config.audio_in_channels), - device=x.device, - dtype=x.dtype, - ) - x_out[video_mask, : self.config.video_in_channels] = x_video - x_out[audio_mask, : self.config.audio_in_channels] = x_audio - - return x_out - - -# --------------------------------------------------------------------------- -# Public config dataclass for building DiTModel from JSON -# --------------------------------------------------------------------------- -@dataclass -class MagiHumanDiTConfig: - num_layers: int = 40 - hidden_size: int = 5120 - head_dim: int = 128 - num_query_groups: int = 8 - video_in_channels: int = 48 * 4 - audio_in_channels: int = 64 - text_in_channels: int = 3584 - checkpoint_qk_layernorm_rope: bool = False - params_dtype: torch.dtype = torch.float32 - mm_layers: list = field(default_factory=lambda: [0, 1, 2, 3, 36, 37, 38, 39]) - local_attn_layers: list = field(default_factory=list) - enable_attn_gating: bool = True - gelu7_layers: list = field(default_factory=lambda: [0, 1, 2, 3]) - post_norm_layers: list = field(default_factory=list) - - def __post_init__(self): - self.num_heads_q = self.hidden_size // self.head_dim - self.num_heads_kv = self.num_query_groups - - -if TYPE_CHECKING: - from .pipeline_magi_human import EvalInput - - -# =========================================================================== -# Data proxy (ported from daVinci-MagiHuman inference/pipeline/data_proxy.py) -# =========================================================================== -def _unfold_3d( - x: torch.Tensor, - kernel_size: tuple[int, int, int], - stride: tuple[int, int, int], -) -> torch.Tensor: - """Pure-PyTorch 3D unfold matching UnfoldAnd behavior. - - After N unfold ops the shape is (batch, C, oD, oH, oW, kD, kH, kW). - UnfoldAnd permutes kernel dims next to channel before reshape so that the - col_dim axis is ordered as (C, kD, kH, kW) -- matching F.unfold semantics. - Without this permute, .view() interleaves spatial and kernel positions. - - Args: - x: (N, C, D, H, W) - kernel_size: (kD, kH, kW) - stride: (sD, sH, sW) - Returns: - (N, C*kD*kH*kW, L) where L = product of output spatial dims. - """ - ndim = len(kernel_size) - for d in range(ndim): - x = x.unfold(d + 2, kernel_size[d], stride[d]) - # x: (N, C, oD, oH, oW, kD, kH, kW) - # Permute to (N, C, kD, kH, kW, oD, oH, oW) so that view groups correctly - perm = [0, 1] + list(range(ndim + 2, 2 * ndim + 2)) + list(range(2, ndim + 2)) - x = x.permute(*perm).contiguous() - - batch_size = x.shape[0] - col_dim = 1 - for i in range(1, ndim + 2): - col_dim *= x.shape[i] - spatial = 1 - for i in range(ndim + 2, 2 * ndim + 2): - spatial *= x.shape[i] - return x.view(batch_size, col_dim, spatial) - - -def calc_local_qk_range( - num_video_tokens, - num_audio_and_txt_tokens, - num_frames, - frame_receptive_field, -): - token_per_frame = num_video_tokens // num_frames - total_tokens = num_video_tokens + num_audio_and_txt_tokens - - q_range_list = [] - k_range_list = [] - for i in range(num_frames): - q_range_list.append(torch.tensor([i * token_per_frame, (i + 1) * token_per_frame])) - k_range_list.append( - torch.tensor( - [ - (i - frame_receptive_field) * token_per_frame, - (i + frame_receptive_field + 1) * token_per_frame, - ] - ) - ) - local_q_range = torch.stack(q_range_list, dim=0) - local_k_range = torch.stack(k_range_list, dim=0) - - local_k_range[local_k_range < 0] = 0 - local_k_range[local_k_range > num_video_tokens] = num_video_tokens - - video_q_range = torch.tensor([[0, num_video_tokens]]) - video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]]) - - at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]]) - at_k_ranges = torch.tensor([[0, total_tokens]]) - - q_ranges = ( - torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) - ) - k_ranges = ( - torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) - ) - return q_ranges, k_ranges - - -def calc_local_attn_ffa_handler( - num_video_tokens, - num_audio_and_txt_tokens, - num_frames, - frame_receptive_field, -): - q_ranges, k_ranges = calc_local_qk_range( - num_video_tokens, - num_audio_and_txt_tokens, - num_frames, - frame_receptive_field, - ) - total = num_video_tokens + num_audio_and_txt_tokens - return FFAHandler( - q_ranges=q_ranges, - k_ranges=k_ranges, - max_seqlen_q=total, - max_seqlen_k=total, - attn_type_map=torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32), - softmax_scale=None, - ) - - -def get_coords( - shape: list[int], - ref_feat_shape: list[int], - offset_thw: list[int] | None = None, - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.float32, -): - if offset_thw is None: - offset_thw = [0, 0, 0] - ori_t, ori_h, ori_w = shape - ref_t, ref_h, ref_w = ref_feat_shape - - offset_t, offset_h, offset_w = offset_thw - time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t - height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h - width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w - - time_grid, height_grid, width_grid = torch.meshgrid( - time_rng, - height_rng, - width_rng, - indexing="ij", - ) - coords_flat = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3) - - meta = torch.tensor( - [ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], - device=device, - dtype=dtype, - ) - meta_expanded = meta.expand(coords_flat.size(0), -1) - return torch.cat([coords_flat, meta_expanded], dim=-1) - - -@dataclass -class SingleData: - video_x_t: torch.Tensor - audio_x_t: torch.Tensor - audio_feat_len: int - txt_feat: torch.Tensor - txt_feat_len: int - t: int - h: int - w: int - patch_size: int - t_patch_size: int - spatial_rope_interpolation: Literal["inter", "extra"] - ref_audio_offset: int - text_offset: int - coords_style: Literal["v1", "v2"] = "v1" - - def __post_init__(self): - self.video_token_num = self.video_x_t.shape[0] - self.audio_x_t = self.audio_x_t[: self.audio_feat_len] - self.txt_feat = self.txt_feat[: self.txt_feat_len] - self.video_channel = self.video_x_t.shape[-1] - self.audio_channel = self.audio_x_t.shape[-1] - self.txt_channel = self.txt_feat.shape[-1] - - @property - def device(self): - return self.video_x_t.device - - @property - def default_dtype(self): - return self.video_x_t.dtype - - @property - def total_token_num(self): - return self.video_token_num + self.audio_feat_len + self.txt_feat_len - - @property - def token_sequence(self): - tensors = [self.video_x_t, self.audio_x_t, self.txt_feat] - max_channel = max(t.shape[-1] for t in tensors) - padded = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors] - return torch.cat(padded, dim=0) - - @property - def modality_mapping(self): - v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device) - a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device) - t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device) - return torch.cat([v_map, a_map, t_map], dim=0) - - def default_coords(self, shape, ref_feat_shape, offset_thw=None): - if offset_thw is None: - offset_thw = [0, 0, 0] - return get_coords( - shape=shape, - ref_feat_shape=ref_feat_shape, - offset_thw=offset_thw, - device=self.device, - dtype=self.default_dtype, - ) - - @property - def coords_mapping(self): - if self.spatial_rope_interpolation == "inter": - video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32) - else: - video_ref_feat_shape = ( - self.t // self.t_patch_size, - self.h // self.patch_size, - self.w // self.patch_size, - ) - - video_coords = self.default_coords( - shape=( - self.t // self.t_patch_size, - self.h // self.patch_size, - self.w // self.patch_size, - ), - ref_feat_shape=video_ref_feat_shape, - ) - - if self.coords_style == "v1": - audio_coords = self.default_coords( - shape=(self.audio_feat_len, 1, 1), - ref_feat_shape=(self.t // self.t_patch_size, 1, 1), - ) - text_coords = self.default_coords( - shape=(self.txt_feat_len, 1, 1), - ref_feat_shape=(2, 1, 1), - offset_thw=[self.text_offset, 0, 0], - ) - elif self.coords_style == "v2": - magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1 - audio_coords = self.default_coords( - shape=(self.audio_feat_len, 1, 1), - ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1), - ) - text_coords = self.default_coords( - shape=(self.txt_feat_len, 1, 1), - ref_feat_shape=(1, 1, 1), - offset_thw=[-self.txt_feat_len, 0, 0], - ) - else: - raise ValueError(f"Unknown coords_style: {self.coords_style}") - - return torch.cat([video_coords, audio_coords, text_coords], dim=0) - - def depack_token_sequence(self, token_sequence): - video_x_t = token_sequence[: self.video_token_num, : self.video_channel] - video_x_t = rearrange( - video_x_t, - "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)", - H=self.h // self.patch_size, - W=self.w // self.patch_size, - pT=self.t_patch_size, - pH=self.patch_size, - pW=self.patch_size, - ).contiguous() - audio_x_t = token_sequence[ - self.video_token_num : self.video_token_num + self.audio_feat_len, - : self.audio_channel, - ] - return video_x_t, audio_x_t - - -@dataclass -class SimplePackedData: - items: list[SingleData] - - @property - def token_sequence(self): - return torch.cat([item.token_sequence for item in self.items], dim=0) - - @property - def modality_mapping(self): - return torch.cat([item.modality_mapping for item in self.items], dim=0) - - @property - def coords_mapping(self): - return torch.cat([item.coords_mapping for item in self.items], dim=0) - - @property - def total_token_num(self): - return sum(item.total_token_num for item in self.items) - - def __getitem__(self, index): - return self.items[index] - - @property - def cu_seqlen(self): - cu = torch.cumsum( - torch.tensor([item.total_token_num for item in self.items]), - dim=0, - ) - return F.pad(cu, (1, 0)) - - @property - def max_seqlen(self): - return torch.tensor(max(item.total_token_num for item in self.items)) - - def depack_token_sequence(self, token_sequence): - video_list, audio_list = [], [] - parts = torch.split( - token_sequence, - [item.total_token_num for item in self.items], - dim=0, - ) - for item, part in zip(self.items, parts): - v, a = item.depack_token_sequence(part) - video_list.append(v) - audio_list.append(a) - return torch.stack(video_list, dim=0), torch.stack(audio_list, dim=0) - - -class MagiDataProxy: - def __init__( - self, - patch_size: int = 2, - t_patch_size: int = 1, - frame_receptive_field: int = 11, - spatial_rope_interpolation: str = "extra", - ref_audio_offset: int = 1000, - text_offset: int = 0, - coords_style: str = "v2", - ): - self.patch_size = patch_size - self.t_patch_size = t_patch_size - self.frame_receptive_field = frame_receptive_field - self.spatial_rope_interpolation = spatial_rope_interpolation - self.ref_audio_offset = ref_audio_offset - self.text_offset = text_offset - self.coords_style = coords_style - self._kernel = (t_patch_size, patch_size, patch_size) - self._stride = (t_patch_size, patch_size, patch_size) - self._saved_data: dict[str, Any] = {} - - def saved_for_output(self, **kwargs): - self._saved_data.update(kwargs) - - def get_saved_data(self, key: str): - return self._saved_data[key] - - def img2tokens(self, x_t: torch.Tensor): - x_t_unfolded = _unfold_3d(x_t, self._kernel, self._stride) - return rearrange( - x_t_unfolded, - "N col_dim num_tokens -> N num_tokens col_dim", - ).contiguous() - - def process_input(self, transported_data: EvalInput): - batch_size, _, t, h, w = transported_data.x_t.shape - x_t = self.img2tokens(transported_data.x_t) - audio_x_t = transported_data.audio_x_t.contiguous() - text_in = transported_data.txt_feat.contiguous() - - simple_packed_data = SimplePackedData(items=[]) - for i in range(batch_size): - single_data = SingleData( - video_x_t=x_t[i], - audio_x_t=audio_x_t[i], - audio_feat_len=transported_data.audio_feat_len[i], - txt_feat=text_in[i], - txt_feat_len=transported_data.txt_feat_len[i], - t=t, - h=h, - w=w, - patch_size=self.patch_size, - t_patch_size=self.t_patch_size, - spatial_rope_interpolation=self.spatial_rope_interpolation, - ref_audio_offset=self.ref_audio_offset, - text_offset=self.text_offset, - coords_style=self.coords_style, - ) - simple_packed_data.items.append(single_data) - - if self.frame_receptive_field != -1: - assert batch_size == 1, "local attention only supports batch size 1" - local_attn_handler = calc_local_attn_ffa_handler( - num_video_tokens=simple_packed_data[0].video_token_num, - num_audio_and_txt_tokens=(simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len), - num_frames=t, - frame_receptive_field=self.frame_receptive_field, - ) - if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor): - local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item() - if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor): - local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item() - else: - local_attn_handler = None - - varlen_handler = VarlenHandler( - cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), - cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), - max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(), - max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(), - ) - - self.saved_for_output(simple_packed_data=simple_packed_data) - - x = simple_packed_data.token_sequence - coords_mapping = simple_packed_data.coords_mapping - modality_mapping = simple_packed_data.modality_mapping - return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler) - - def process_output(self, x: torch.Tensor): - simple_packed_data: SimplePackedData = self.get_saved_data("simple_packed_data") - return simple_packed_data.depack_token_sequence(x) diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py deleted file mode 100644 index 513d93335..000000000 --- a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py +++ /dev/null @@ -1,2306 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 SandAI. All Rights Reserved. -# Ported from daVinci-MagiHuman inference/pipeline/video_generate.py -# Adapted for vllm-omni: single-GPU, diffusers VAE, configurable dit_subfolder. - -from __future__ import annotations - -import json -import logging -import math -import os -from collections.abc import Iterable -from dataclasses import dataclass -from pathlib import Path -from typing import Any, ClassVar, Literal - -import numpy as np -import torch -import torch.nn as nn -import whisper -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.schedulers.scheduling_utils import ( - KarrasDiffusionSchedulers, - SchedulerMixin, - SchedulerOutput, -) -from diffusers.utils import deprecate, load_image -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from einops import rearrange -from PIL import Image -from safetensors.torch import load_file -from torch.nn import functional as F -from torch.nn.utils import weight_norm -from transformers import AutoTokenizer -from transformers.models.t5gemma import T5GemmaEncoderModel -from vllm.distributed import ( - 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.autoencoders.autoencoder_kl_wan import ( - DistributedAutoencoderKLWan, -) -from vllm_omni.diffusion.model_loader.diffusers_loader import ( - DiffusersPipelineLoader, -) -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.models.t5_encoder.t5_gemma_encoder import T5GemmaEncoderModelTP -from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( - DiffusionPipelineProfilerMixin, -) -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - -from .magi_human_dit import ( - DiTModel, - FFAHandler, - MagiHumanDiTConfig, - Modality, - VarlenHandler, -) - -logger = logging.getLogger(__name__) - - -# =========================================================================== -# Scheduler (ported from daVinci-MagiHuman inference/pipeline/scheduler_unipc.py) -# =========================================================================== -class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - solver_order: int = 2, - prediction_type: str = "flow_prediction", - shift: float = 1.0, - use_dynamic_shifting=False, - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - predict_x0: bool = True, - solver_type: str = "bh2", - lower_order_final: bool = True, - disable_corrector: list[int] = [], - solver_p: SchedulerMixin = None, - timestep_spacing: str = "linspace", - steps_offset: int = 0, - final_sigmas_type: str | None = "zero", - ): - if solver_type not in ["bh1", "bh2"]: - if solver_type in ["midpoint", "heun", "logrho"]: - self.register_to_config(solver_type="bh2") - else: - raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") - - self.predict_x0 = predict_x0 - self.num_inference_steps = None - alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy() - sigmas = 1.0 - alphas - sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) - - if not use_dynamic_shifting: - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) - - self.sigmas = sigmas - self.timesteps = sigmas * num_train_timesteps - - self.model_outputs = [None] * solver_order - self.timestep_list = [None] * solver_order - self.lower_order_nums = 0 - self.disable_corrector = disable_corrector - self.solver_p = solver_p - self.last_sample = None - self._step_index: int | None = None - self._begin_index: int | None = None - - self.sigmas = self.sigmas.to("cpu") - self.sigma_min = self.sigmas[-1].item() - self.sigma_max = self.sigmas[0].item() - - @property - def step_index(self): - return self._step_index - - @property - def begin_index(self): - return self._begin_index - - def set_begin_index(self, begin_index: int = 0): - self._begin_index = begin_index - - def set_timesteps( - self, - num_inference_steps: int | None = None, - device: str | torch.device = None, - sigmas: list[float] | None = None, - mu: float | None | None = None, - shift: float | None | None = None, - ): - if self.config.use_dynamic_shifting and mu is None: - raise ValueError(" you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`") - - if sigmas is None: - sigmas = np.linspace(self.sigma_max, self.sigma_min, num_inference_steps + 1).copy()[:-1] - - if self.config.use_dynamic_shifting: - sigmas = self.time_shift(mu, 1.0, sigmas) - else: - if shift is None: - shift = self.config.shift - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) - - if self.config.final_sigmas_type == "sigma_min": - sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - - timesteps = sigmas * self.config.num_train_timesteps - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas) - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) - - self.num_inference_steps = len(timesteps) - - self.model_outputs = [None] * self.config.solver_order - self.lower_order_nums = 0 - self.last_sample = None - if self.solver_p: - self.solver_p.set_timesteps(self.num_inference_steps, device=device) - - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") - - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() - - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - abs_sample = sample.abs() - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp(s, min=1, max=self.config.sample_max_value) - s = s.unsqueeze(1) - sample = torch.clamp(sample, -s, s) / s - sample = sample.reshape(batch_size, channels, *remaining_dims) - return sample.to(dtype) - - def _sigma_to_t(self, sigma): - return sigma * self.config.num_train_timesteps - - def _sigma_to_alpha_sigma_t(self, sigma): - return 1 - sigma, sigma - - def time_shift(self, mu: float, sigma: float, t: torch.Tensor): - return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) - - def convert_model_output( - self, model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs - ) -> torch.Tensor: - timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) - if sample is None: - if len(args) > 1: - sample = args[1] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep is not None: - deprecate( - "timesteps", - "1.0.0", - "Passing `timesteps` is deprecated and has no effect as model output " - "conversion is now handled via an internal counter `self.step_index`", - ) - - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - - if self.predict_x0: - if self.config.prediction_type == "flow_prediction": - sigma_t = self.sigmas[self.step_index] - x0_pred = sample - sigma_t * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," - " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." - ) - if self.config.thresholding: - x0_pred = self._threshold_sample(x0_pred) - return x0_pred - else: - if self.config.prediction_type == "flow_prediction": - sigma_t = self.sigmas[self.step_index] - epsilon = sample - (1 - sigma_t) * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," - " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." - ) - if self.config.thresholding: - sigma_t = self.sigmas[self.step_index] - x0_pred = sample - sigma_t * model_output - x0_pred = self._threshold_sample(x0_pred) - epsilon = model_output + x0_pred - return epsilon - - def multistep_uni_p_bh_update( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor | None = None, - order: int | None = None, - **kwargs, - ) -> torch.Tensor: - prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 1: - sample = args[1] - else: - raise ValueError(" missing `sample` as a required keyword argument") - if order is None: - if len(args) > 2: - order = args[2] - else: - raise ValueError(" missing `order` as a required keyword argument") - if prev_timestep is not None: - deprecate("prev_timestep", "1.0.0", "Passing `prev_timestep` is deprecated and has no effect.") - - model_output_list = self.model_outputs - s0 = self.timestep_list[-1] - m0 = model_output_list[-1] - x = sample - - if self.solver_p: - return self.solver_p.step(model_output, s0, x).prev_sample - - sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - h = lambda_t - lambda_s0 - device = sample.device - - rks = [] - D1s: list[Any] | None = [] - for i in range(1, order): - si = self.step_index - i - mi = model_output_list[-(i + 1)] - alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(alpha_si) - torch.log(sigma_si) - rk = (lambda_si - lambda_s0) / h - rks.append(rk) - D1s.append((mi - m0) / rk) - - rks.append(1.0) - rks = torch.tensor(rks, device=device) - - R = [] - b = [] - hh = -h if self.predict_x0 else h - h_phi_1 = torch.expm1(hh) - h_phi_k = h_phi_1 / hh - 1 - factorial_i = 1 - - if self.config.solver_type == "bh1": - B_h = hh - elif self.config.solver_type == "bh2": - B_h = torch.expm1(hh) - else: - raise NotImplementedError() - - for i in range(1, order + 1): - R.append(torch.pow(rks, i - 1)) - b.append(h_phi_k * factorial_i / B_h) - factorial_i *= i + 1 - h_phi_k = h_phi_k / hh - 1 / factorial_i - - R = torch.stack(R) - b = torch.tensor(b, device=device) - - if len(D1s) > 0: - D1s = torch.stack(D1s, dim=1) - if order == 2: - rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) - else: - rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) - else: - D1s = None - - if self.predict_x0: - x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 - pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 - x_t = x_t_ - alpha_t * B_h * pred_res - else: - x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 - pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 - x_t = x_t_ - sigma_t * B_h * pred_res - - return x_t.to(x.dtype) - - def multistep_uni_c_bh_update( - self, - this_model_output: torch.Tensor, - *args, - last_sample: torch.Tensor = None, - this_sample: torch.Tensor = None, - order: int | None = None, - **kwargs, - ) -> torch.Tensor: - this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None) - if last_sample is None: - if len(args) > 1: - last_sample = args[1] - else: - raise ValueError(" missing`last_sample` as a required keyword argument") - if this_sample is None: - if len(args) > 2: - this_sample = args[2] - else: - raise ValueError(" missing`this_sample` as a required keyword argument") - if order is None: - if len(args) > 3: - order = args[3] - else: - raise ValueError(" missing`order` as a required keyword argument") - if this_timestep is not None: - deprecate("this_timestep", "1.0.0", "Passing `this_timestep` is deprecated and has no effect.") - - model_output_list = self.model_outputs - m0 = model_output_list[-1] - x = last_sample - x_t = this_sample - model_t = this_model_output - - sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - h = lambda_t - lambda_s0 - device = this_sample.device - - rks = [] - D1s: list[Any] | None = [] - for i in range(1, order): - si = self.step_index - (i + 1) - mi = model_output_list[-(i + 1)] - alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(alpha_si) - torch.log(sigma_si) - rk = (lambda_si - lambda_s0) / h - rks.append(rk) - D1s.append((mi - m0) / rk) - - rks.append(1.0) - rks = torch.tensor(rks, device=device) - - R = [] - b = [] - hh = -h if self.predict_x0 else h - h_phi_1 = torch.expm1(hh) - h_phi_k = h_phi_1 / hh - 1 - factorial_i = 1 - - if self.config.solver_type == "bh1": - B_h = hh - elif self.config.solver_type == "bh2": - B_h = torch.expm1(hh) - else: - raise NotImplementedError() - - for i in range(1, order + 1): - R.append(torch.pow(rks, i - 1)) - b.append(h_phi_k * factorial_i / B_h) - factorial_i *= i + 1 - h_phi_k = h_phi_k / hh - 1 / factorial_i - - R = torch.stack(R) - b = torch.tensor(b, device=device) - - if len(D1s) > 0: - D1s = torch.stack(D1s, dim=1) - else: - D1s = None - - if order == 1: - rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) - else: - rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) - - if self.predict_x0: - x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 - corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0 - D1_t = model_t - m0 - x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) - else: - x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 - corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0 - D1_t = model_t - m0 - x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) - return x_t.to(x.dtype) - - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - indices = (schedule_timesteps == timestep).nonzero() - pos = 1 if len(indices) > 1 else 0 - return indices[pos].item() - - def _init_step_index(self, timestep): - if self.begin_index is None: - if isinstance(timestep, torch.Tensor): - timestep = timestep.to(self.timesteps.device) - self._step_index = self.index_for_timestep(timestep) - else: - self._step_index = self._begin_index - - def step( - self, - model_output: torch.Tensor, - timestep: int | torch.Tensor, - sample: torch.Tensor, - return_dict: bool = True, - generator=None, - ) -> SchedulerOutput | tuple: - if self.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - use_corrector = ( - self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None - ) - - model_output_convert = self.convert_model_output(model_output, sample=sample) - if use_corrector: - sample = self.multistep_uni_c_bh_update( - this_model_output=model_output_convert, - last_sample=self.last_sample, - this_sample=sample, - order=self.this_order, - ) - - for i in range(self.config.solver_order - 1): - self.model_outputs[i] = self.model_outputs[i + 1] - self.timestep_list[i] = self.timestep_list[i + 1] - - self.model_outputs[-1] = model_output_convert - self.timestep_list[-1] = timestep - - if self.config.lower_order_final: - this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index) - else: - this_order = self.config.solver_order - - self.this_order = min(this_order, self.lower_order_nums + 1) - assert self.this_order > 0 - - self.last_sample = sample - prev_sample = self.multistep_uni_p_bh_update(model_output=model_output, sample=sample, order=self.this_order) - - if self.lower_order_nums < self.config.solver_order: - self.lower_order_nums += 1 - - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - return SchedulerOutput(prev_sample=prev_sample) - - def step_ddim( - self, - velocity: torch.FloatTensor, - t: int, - curr_state: torch.FloatTensor, - prev_state: torch.FloatTensor | None = None, - generator: torch.Generator | None = None, - ): - device = curr_state.device - curr_t = self.sigmas[t] - prev_t = self.sigmas[t + 1] - variance_noise = randn_tensor(curr_state.shape, generator=generator, device=device, dtype=curr_state.dtype) - cur_clean_ = curr_state - curr_t * velocity - return prev_t * variance_noise + (1 - prev_t) * cur_clean_ - - def step_sde( - self, - velocity: torch.FloatTensor, - t: int, - curr_state: torch.FloatTensor, - noise_theta: float = 1.0, - prev_state: torch.FloatTensor | None = None, - generator: torch.Generator | None = None, - ): - device = curr_state.device - curr_t = self.sigmas[t] - prev_t = self.sigmas[t + 1] - cos = torch.cos(torch.tensor(noise_theta) * torch.pi / 2).to(device) - sin = torch.sin(torch.tensor(noise_theta) * torch.pi / 2).to(device) - prev_sample_mean = (1 - prev_t + prev_t * cos) * (curr_state - curr_t * velocity) + prev_t * cos * velocity - std_dev_t = prev_t * sin - std_dev_t = torch.ones((1, 1)).to(curr_state) * std_dev_t - if prev_state is None: - variance_noise = randn_tensor(curr_state.shape, generator=generator, device=device, dtype=curr_state.dtype) - prev_state = prev_sample_mean + std_dev_t * variance_noise - else: - prev_state = prev_sample_mean + (prev_state - prev_sample_mean.detach()) - return prev_state - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - return sample - - def add_noise( - self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor - ) -> torch.Tensor: - sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) - if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - step_indices = [self.step_index] * timesteps.shape[0] - else: - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - return alpha_t * original_samples + sigma_t * noise - - def __len__(self): - return self.config.num_train_timesteps - - -# =========================================================================== -# Audio VAE (ported from daVinci-MagiHuman inference/model/sa_audio/) -# =========================================================================== -def _snake_beta(x, alpha, beta): - return x + (1.0 / (beta + 1e-9)) * torch.pow(torch.sin(x * alpha), 2) - - -class _SnakeBeta(nn.Module): - def __init__(self, in_features: int, alpha: float = 1.0, alpha_trainable: bool = True, alpha_logscale: bool = True): - super().__init__() - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: - self.alpha = nn.Parameter(torch.zeros(in_features) * alpha) - self.beta = nn.Parameter(torch.zeros(in_features) * alpha) - else: - self.alpha = nn.Parameter(torch.ones(in_features) * alpha) - self.beta = nn.Parameter(torch.ones(in_features) * alpha) - self.alpha.requires_grad = alpha_trainable - self.beta.requires_grad = alpha_trainable - - def forward(self, x): - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) - beta = self.beta.unsqueeze(0).unsqueeze(-1) - if self.alpha_logscale: - alpha = torch.exp(alpha) - beta = torch.exp(beta) - return _snake_beta(x, alpha, beta) - - -def _vae_sample(mean, scale): - stdev = F.softplus(scale) + 1e-4 - var = stdev * stdev - logvar = torch.log(var) - latents = torch.randn_like(mean) * stdev + mean - kl = (mean * mean + var - logvar - 1).sum(1).mean() - return latents, kl - - -class _VAEBottleneck(nn.Module): - def encode(self, x, return_info=False, **kwargs): - info = {} - mean, scale = x.chunk(2, dim=1) - x, kl = _vae_sample(mean, scale) - info["kl"] = kl - return (x, info) if return_info else x - - def decode(self, x): - return x - - -def _WNConv1d(*args, **kwargs): - return weight_norm(nn.Conv1d(*args, **kwargs)) - - -def _WNConvTranspose1d(*args, **kwargs): - return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) - - -def _checkpoint(function, *args, **kwargs): - kwargs.setdefault("use_reentrant", False) - return torch.utils.checkpoint.checkpoint(function, *args, **kwargs) - - -def _get_activation(activation: Literal["elu", "snake", "none"], antialias: bool = False, channels=None) -> nn.Module: - if antialias: - raise NotImplementedError("antialias activation not supported") - if activation == "elu": - return nn.ELU() - if activation == "snake": - return _SnakeBeta(channels) - if activation == "none": - return nn.Identity() - raise ValueError(f"Unknown activation {activation}") - - -class _ResidualUnit(nn.Module): - def __init__(self, in_channels, out_channels, dilation, use_snake=False, antialias_activation=False): - super().__init__() - padding = (dilation * (7 - 1)) // 2 - self.layers = nn.Sequential( - _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels), - _WNConv1d(in_channels, out_channels, kernel_size=7, dilation=dilation, padding=padding), - _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels), - _WNConv1d(out_channels, out_channels, kernel_size=1), - ) - - def forward(self, x): - return (_checkpoint(self.layers, x) if self.training else self.layers(x)) + x - - -class _EncoderBlock(nn.Module): - def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False): - super().__init__() - self.layers = nn.Sequential( - _ResidualUnit(in_channels, in_channels, 1, use_snake=use_snake), - _ResidualUnit(in_channels, in_channels, 3, use_snake=use_snake), - _ResidualUnit(in_channels, in_channels, 9, use_snake=use_snake), - _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels), - _WNConv1d(in_channels, out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)), - ) - - def forward(self, x): - return self.layers(x) - - -class _DecoderBlock(nn.Module): - def __init__( - self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False, use_nearest_upsample=False - ): - super().__init__() - if use_nearest_upsample: - upsample_layer = nn.Sequential( - nn.Upsample(scale_factor=stride, mode="nearest"), - _WNConv1d(in_channels, out_channels, kernel_size=2 * stride, stride=1, bias=False, padding="same"), - ) - else: - upsample_layer = _WNConvTranspose1d( - in_channels, out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2) - ) - self.layers = nn.Sequential( - _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels), - upsample_layer, - _ResidualUnit(out_channels, out_channels, 1, use_snake=use_snake), - _ResidualUnit(out_channels, out_channels, 3, use_snake=use_snake), - _ResidualUnit(out_channels, out_channels, 9, use_snake=use_snake), - ) - - def forward(self, x): - return self.layers(x) - - -class _OobleckEncoder(nn.Module): - def __init__( - self, - in_channels=2, - channels=128, - latent_dim=32, - c_mults=[1, 2, 4, 8], - strides=[2, 4, 8, 8], - use_snake=False, - antialias_activation=False, - ): - super().__init__() - c_mults = [1] + c_mults - depth = len(c_mults) - layers = [_WNConv1d(in_channels, c_mults[0] * channels, kernel_size=7, padding=3)] - for i in range(depth - 1): - layers.append( - _EncoderBlock(c_mults[i] * channels, c_mults[i + 1] * channels, strides[i], use_snake=use_snake) - ) - layers.extend( - [ - _get_activation( - "snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[-1] * channels - ), - _WNConv1d(c_mults[-1] * channels, latent_dim, kernel_size=3, padding=1), - ] - ) - self.layers = nn.Sequential(*layers) - - def forward(self, x): - return self.layers(x) - - -class _OobleckDecoder(nn.Module): - def __init__( - self, - out_channels=2, - channels=128, - latent_dim=32, - c_mults=[1, 2, 4, 8], - strides=[2, 4, 8, 8], - use_snake=False, - antialias_activation=False, - use_nearest_upsample=False, - final_tanh=True, - ): - super().__init__() - c_mults = [1] + c_mults - depth = len(c_mults) - layers = [_WNConv1d(latent_dim, c_mults[-1] * channels, kernel_size=7, padding=3)] - for i in range(depth - 1, 0, -1): - layers.append( - _DecoderBlock( - c_mults[i] * channels, - c_mults[i - 1] * channels, - strides[i - 1], - use_snake=use_snake, - antialias_activation=antialias_activation, - use_nearest_upsample=use_nearest_upsample, - ) - ) - layers.extend( - [ - _get_activation( - "snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[0] * channels - ), - _WNConv1d(c_mults[0] * channels, out_channels, kernel_size=7, padding=3, bias=False), - nn.Tanh() if final_tanh else nn.Identity(), - ] - ) - self.layers = nn.Sequential(*layers) - - def forward(self, x): - return self.layers(x) - - -class _AudioAutoencoder(nn.Module): - def __init__( - self, - encoder, - decoder, - latent_dim, - downsampling_ratio, - sample_rate, - io_channels=2, - bottleneck=None, - in_channels=None, - out_channels=None, - soft_clip=False, - ): - super().__init__() - self.downsampling_ratio = downsampling_ratio - self.sample_rate = sample_rate - self.latent_dim = latent_dim - self.io_channels = io_channels - self.in_channels = in_channels if in_channels is not None else io_channels - self.out_channels = out_channels if out_channels is not None else io_channels - self.bottleneck = bottleneck - self.encoder = encoder - self.decoder = decoder - self.soft_clip = soft_clip - - def encode(self, audio, skip_bottleneck=False, return_info=False, **kwargs): - info = {} - latents = self.encoder(audio) - info["pre_bottleneck_latents"] = latents - if self.bottleneck is not None and not skip_bottleneck: - latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs) - info.update(bottleneck_info) - return (latents, info) if return_info else latents - - def decode(self, latents, skip_bottleneck=False, **kwargs): - if self.bottleneck is not None and not skip_bottleneck: - latents = self.bottleneck.decode(latents) - decoded = self.decoder(latents, **kwargs) - if self.soft_clip: - decoded = torch.tanh(decoded) - return decoded - - -def _create_encoder_from_config(cfg: dict[str, Any]): - assert cfg.get("type") == "oobleck", f"Only 'oobleck' encoder supported, got: {cfg.get('type')}" - enc = _OobleckEncoder(**cfg["config"]) - if not cfg.get("requires_grad", True): - for p in enc.parameters(): - p.requires_grad = False - return enc - - -def _create_decoder_from_config(cfg: dict[str, Any]): - assert cfg.get("type") == "oobleck", f"Only 'oobleck' decoder supported, got: {cfg.get('type')}" - dec = _OobleckDecoder(**cfg["config"]) - if not cfg.get("requires_grad", True): - for p in dec.parameters(): - p.requires_grad = False - return dec - - -def _create_bottleneck_from_config(cfg: dict[str, Any]): - assert cfg.get("type") == "vae", f"Only 'vae' bottleneck supported, got: {cfg.get('type')}" - bn = _VAEBottleneck() - if not cfg.get("requires_grad", True): - for p in bn.parameters(): - p.requires_grad = False - return bn - - -def _create_autoencoder_from_config(config: dict[str, Any]): - ae_config = config["model"] - if ae_config.get("pretransform") is not None: - raise NotImplementedError("Nested pretransform not supported") - encoder = _create_encoder_from_config(ae_config["encoder"]) - decoder = _create_decoder_from_config(ae_config["decoder"]) - bottleneck_cfg = ae_config.get("bottleneck") - bottleneck = _create_bottleneck_from_config(bottleneck_cfg) if bottleneck_cfg else None - return _AudioAutoencoder( - encoder=encoder, - decoder=decoder, - latent_dim=ae_config["latent_dim"], - downsampling_ratio=ae_config["downsampling_ratio"], - sample_rate=config["sample_rate"], - io_channels=ae_config["io_channels"], - bottleneck=bottleneck, - in_channels=ae_config.get("in_channels"), - out_channels=ae_config.get("out_channels"), - soft_clip=ae_config["decoder"].get("soft_clip", False), - ) - - -class SAAudioFeatureExtractor: - def __init__(self, device, model_path): - self.device = device - self.vae_model, self.sample_rate = self._load_vae(model_path) - self.resampler = None - - def _load_vae(self, model_path): - if not (isinstance(model_path, str) and Path(model_path).is_dir()): - raise ValueError("model_path must be a local directory") - - model_config_path = os.path.join(model_path, "model_config.json") - with open(model_config_path) as f: - full_config = json.load(f) - - vae_config = full_config["model"]["pretransform"]["config"] - sample_rate = full_config["sample_rate"] - - autoencoder_config = { - "model_type": "autoencoder", - "sample_rate": sample_rate, - "model": vae_config, - } - vae_model = _create_autoencoder_from_config(autoencoder_config) - - weights_path = Path(model_path) / "model.safetensors" - if not weights_path.exists(): - raise FileNotFoundError(f"Weight file does not exist: {weights_path}") - - full_state_dict = load_file(weights_path, device=str(self.device)) - vae_state_dict = {} - for key, value in full_state_dict.items(): - if key.startswith("pretransform.model."): - vae_state_dict[key[len("pretransform.model.") :]] = value - - model_keys = set(vae_model.state_dict().keys()) - vae_keys = set(vae_state_dict.keys()) - missing = model_keys - vae_keys - extra = vae_keys - model_keys - if missing: - logger.warning("Audio VAE missing keys (%d): %s", len(missing), list(missing)[:5]) - if extra: - logger.warning("Audio VAE unexpected keys (%d): %s", len(extra), list(extra)[:5]) - - vae_model.load_state_dict(vae_state_dict) - vae_model.to(self.device) - return vae_model, sample_rate - - def decode(self, latents): - with torch.no_grad(): - return self.vae_model.decode(latents) - - def encode(self, waveform): - with torch.no_grad(): - return self.vae_model.encode(waveform) - - -# =========================================================================== -# Audio utilities (ported from daVinci-MagiHuman inference/pipeline/video_process.py) -# =========================================================================== -_SAMPLE_RATE = 51200 -_AUDIO_CHUNK_DURATION = 29 -_OVERLAP_RATIO = 0.5 - - -def _merge_overlapping_vae_features(audio_feats: list[torch.Tensor], overlap_ratio: float = 0.5) -> torch.Tensor | None: - if not audio_feats: - return None - if len(audio_feats) == 1: - return audio_feats[0] - - batch_size, total_frames, feature_dim = audio_feats[0].shape - overlap_frames = int(total_frames * overlap_ratio) - step_frames = total_frames - overlap_frames - final_length = (len(audio_feats) - 1) * step_frames + total_frames - output_feat = torch.zeros( - batch_size, final_length, feature_dim, device=audio_feats[0].device, dtype=audio_feats[0].dtype - ) - - for block_idx, current_feat in enumerate(audio_feats): - output_start = block_idx * step_frames - if block_idx == 0: - output_feat[:, output_start : output_start + total_frames, :] = current_feat - continue - - non_overlap_start = output_start + overlap_frames - non_overlap_end = output_start + total_frames - output_feat[:, non_overlap_start:non_overlap_end, :] = current_feat[:, overlap_frames:, :] - - for frame_idx in range(overlap_frames): - output_pos = output_start + frame_idx - prev_weight = (overlap_frames - frame_idx) / overlap_frames - curr_weight = frame_idx / overlap_frames - output_feat[:, output_pos, :] = ( - prev_weight * output_feat[:, output_pos, :] + curr_weight * current_feat[:, frame_idx, :] - ) - return output_feat - - -def load_audio_and_encode(audio_vae, audio_path: str, seconds: int | None = None) -> torch.Tensor: - """Load audio from file and encode to latent space using the Stable Audio VAE.""" - audio_full = whisper.load_audio(audio_path, sr=_SAMPLE_RATE) - if seconds is not None: - audio_full = audio_full[: min(int(seconds * _SAMPLE_RATE), audio_full.shape[0])] - total_samples = audio_full.shape[0] - - window_size = int(_AUDIO_CHUNK_DURATION * _SAMPLE_RATE) - step_size = int(window_size * (1 - _OVERLAP_RATIO)) - if total_samples <= window_size: - audio = torch.from_numpy(audio_full).cuda() - audio = audio.unsqueeze(0).expand(2, -1) - return audio_vae.vae_model.encode(audio) - - encoded_chunks = [] - latent_to_audio_ratio = None - for offset_start in range(0, total_samples, step_size): - offset_end = min(offset_start + window_size, total_samples) - chunk = whisper.pad_or_trim(audio_full[offset_start:offset_end], length=window_size) - chunk_tensor = torch.from_numpy(chunk).cuda().unsqueeze(0).expand(2, -1) - encoded_chunk = audio_vae.vae_model.encode(chunk_tensor) - - if latent_to_audio_ratio is None: - latent_to_audio_ratio = encoded_chunk.shape[-1] / window_size - - encoded_chunks.append(encoded_chunk.permute(0, 2, 1)) - if offset_end >= total_samples: - break - - final_feat = _merge_overlapping_vae_features(encoded_chunks, overlap_ratio=_OVERLAP_RATIO).permute(0, 2, 1) - final_target_len = math.ceil(total_samples * latent_to_audio_ratio) - return final_feat[:, :, :final_target_len] - - -# =========================================================================== -# Data proxy (ported from daVinci-MagiHuman inference/pipeline/data_proxy.py) -# =========================================================================== -def _unfold_3d(x: torch.Tensor, kernel_size: tuple[int, int, int], stride: tuple[int, int, int]) -> torch.Tensor: - """Pure-PyTorch 3D unfold matching UnfoldAnd behavior. - - After N unfold ops the shape is (batch, C, oD, oH, oW, kD, kH, kW). - UnfoldAnd permutes kernel dims next to channel before reshape so that the - col_dim axis is ordered as (C, kD, kH, kW) -- matching F.unfold semantics. - Without this permute, .view() interleaves spatial and kernel positions. - - Args: - x: (N, C, D, H, W) - kernel_size: (kD, kH, kW) - stride: (sD, sH, sW) - Returns: - (N, C*kD*kH*kW, L) where L = product of output spatial dims. - """ - ndim = len(kernel_size) - for d in range(ndim): - x = x.unfold(d + 2, kernel_size[d], stride[d]) - perm = [0, 1] + list(range(ndim + 2, 2 * ndim + 2)) + list(range(2, ndim + 2)) - x = x.permute(*perm).contiguous() - - batch_size = x.shape[0] - col_dim = 1 - for i in range(1, ndim + 2): - col_dim *= x.shape[i] - spatial = 1 - for i in range(ndim + 2, 2 * ndim + 2): - spatial *= x.shape[i] - return x.view(batch_size, col_dim, spatial) - - -def _calc_local_qk_range(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field): - token_per_frame = num_video_tokens // num_frames - total_tokens = num_video_tokens + num_audio_and_txt_tokens - - q_range_list = [] - k_range_list = [] - for i in range(num_frames): - q_range_list.append(torch.tensor([i * token_per_frame, (i + 1) * token_per_frame])) - k_range_list.append( - torch.tensor( - [ - (i - frame_receptive_field) * token_per_frame, - (i + frame_receptive_field + 1) * token_per_frame, - ] - ) - ) - local_q_range = torch.stack(q_range_list, dim=0) - local_k_range = torch.stack(k_range_list, dim=0) - - local_k_range[local_k_range < 0] = 0 - local_k_range[local_k_range > num_video_tokens] = num_video_tokens - - video_q_range = torch.tensor([[0, num_video_tokens]]) - video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]]) - - at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]]) - at_k_ranges = torch.tensor([[0, total_tokens]]) - - q_ranges = ( - torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) - ) - k_ranges = ( - torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) - ) - return q_ranges, k_ranges - - -def _calc_local_attn_ffa_handler(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field): - q_ranges, k_ranges = _calc_local_qk_range( - num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field - ) - total = num_video_tokens + num_audio_and_txt_tokens - return FFAHandler( - q_ranges=q_ranges, - k_ranges=k_ranges, - max_seqlen_q=total, - max_seqlen_k=total, - attn_type_map=torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32), - softmax_scale=None, - ) - - -def _get_coords( - shape: list[int], - ref_feat_shape: list[int], - offset_thw: list[int] | None = None, - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.float32, -): - if offset_thw is None: - offset_thw = [0, 0, 0] - ori_t, ori_h, ori_w = shape - ref_t, ref_h, ref_w = ref_feat_shape - - offset_t, offset_h, offset_w = offset_thw - time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t - height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h - width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w - - time_grid, height_grid, width_grid = torch.meshgrid(time_rng, height_rng, width_rng, indexing="ij") - coords_flat = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3) - - meta = torch.tensor([ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], device=device, dtype=dtype) - meta_expanded = meta.expand(coords_flat.size(0), -1) - return torch.cat([coords_flat, meta_expanded], dim=-1) - - -@dataclass -class _SingleData: - video_x_t: torch.Tensor - audio_x_t: torch.Tensor - audio_feat_len: int - txt_feat: torch.Tensor - txt_feat_len: int - t: int - h: int - w: int - patch_size: int - t_patch_size: int - spatial_rope_interpolation: Literal["inter", "extra"] - ref_audio_offset: int - text_offset: int - coords_style: Literal["v1", "v2"] = "v1" - - def __post_init__(self): - self.video_token_num = self.video_x_t.shape[0] - self.audio_x_t = self.audio_x_t[: self.audio_feat_len] - self.txt_feat = self.txt_feat[: self.txt_feat_len] - self.video_channel = self.video_x_t.shape[-1] - self.audio_channel = self.audio_x_t.shape[-1] - self.txt_channel = self.txt_feat.shape[-1] - - @property - def device(self): - return self.video_x_t.device - - @property - def default_dtype(self): - return self.video_x_t.dtype - - @property - def total_token_num(self): - return self.video_token_num + self.audio_feat_len + self.txt_feat_len - - @property - def token_sequence(self): - tensors = [self.video_x_t, self.audio_x_t, self.txt_feat] - max_channel = max(t.shape[-1] for t in tensors) - padded = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors] - return torch.cat(padded, dim=0) - - @property - def modality_mapping(self): - v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device) - a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device) - t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device) - return torch.cat([v_map, a_map, t_map], dim=0) - - def _default_coords(self, shape, ref_feat_shape, offset_thw=None): - if offset_thw is None: - offset_thw = [0, 0, 0] - return _get_coords( - shape=shape, - ref_feat_shape=ref_feat_shape, - offset_thw=offset_thw, - device=self.device, - dtype=self.default_dtype, - ) - - @property - def coords_mapping(self): - if self.spatial_rope_interpolation == "inter": - video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32) - else: - video_ref_feat_shape = (self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size) - - video_coords = self._default_coords( - shape=(self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size), - ref_feat_shape=video_ref_feat_shape, - ) - - if self.coords_style == "v1": - audio_coords = self._default_coords( - shape=(self.audio_feat_len, 1, 1), - ref_feat_shape=(self.t // self.t_patch_size, 1, 1), - ) - text_coords = self._default_coords( - shape=(self.txt_feat_len, 1, 1), - ref_feat_shape=(2, 1, 1), - offset_thw=[self.text_offset, 0, 0], - ) - elif self.coords_style == "v2": - magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1 - audio_coords = self._default_coords( - shape=(self.audio_feat_len, 1, 1), - ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1), - ) - text_coords = self._default_coords( - shape=(self.txt_feat_len, 1, 1), - ref_feat_shape=(1, 1, 1), - offset_thw=[-self.txt_feat_len, 0, 0], - ) - else: - raise ValueError(f"Unknown coords_style: {self.coords_style}") - - return torch.cat([video_coords, audio_coords, text_coords], dim=0) - - def depack_token_sequence(self, token_sequence): - video_x_t = token_sequence[: self.video_token_num, : self.video_channel] - video_x_t = rearrange( - video_x_t, - "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)", - H=self.h // self.patch_size, - W=self.w // self.patch_size, - pT=self.t_patch_size, - pH=self.patch_size, - pW=self.patch_size, - ).contiguous() - audio_x_t = token_sequence[ - self.video_token_num : self.video_token_num + self.audio_feat_len, : self.audio_channel - ] - return video_x_t, audio_x_t - - -@dataclass -class _SimplePackedData: - items: list[_SingleData] - - @property - def token_sequence(self): - return torch.cat([item.token_sequence for item in self.items], dim=0) - - @property - def modality_mapping(self): - return torch.cat([item.modality_mapping for item in self.items], dim=0) - - @property - def coords_mapping(self): - return torch.cat([item.coords_mapping for item in self.items], dim=0) - - @property - def total_token_num(self): - return sum(item.total_token_num for item in self.items) - - def __getitem__(self, index): - return self.items[index] - - @property - def cu_seqlen(self): - cu = torch.cumsum(torch.tensor([item.total_token_num for item in self.items]), dim=0) - return F.pad(cu, (1, 0)) - - @property - def max_seqlen(self): - return torch.tensor(max(item.total_token_num for item in self.items)) - - def depack_token_sequence(self, token_sequence): - video_list, audio_list = [], [] - parts = torch.split(token_sequence, [item.total_token_num for item in self.items], dim=0) - for item, part in zip(self.items, parts): - v, a = item.depack_token_sequence(part) - video_list.append(v) - audio_list.append(a) - return torch.stack(video_list, dim=0), torch.stack(audio_list, dim=0) - - -class MagiDataProxy: - def __init__( - self, - patch_size: int = 2, - t_patch_size: int = 1, - frame_receptive_field: int = 11, - spatial_rope_interpolation: str = "extra", - ref_audio_offset: int = 1000, - text_offset: int = 0, - coords_style: str = "v2", - ): - self.patch_size = patch_size - self.t_patch_size = t_patch_size - self.frame_receptive_field = frame_receptive_field - self.spatial_rope_interpolation = spatial_rope_interpolation - self.ref_audio_offset = ref_audio_offset - self.text_offset = text_offset - self.coords_style = coords_style - self._kernel = (t_patch_size, patch_size, patch_size) - self._stride = (t_patch_size, patch_size, patch_size) - self._saved_data: dict[str, Any] = {} - - def saved_for_output(self, **kwargs): - self._saved_data.update(kwargs) - - def get_saved_data(self, key: str): - return self._saved_data[key] - - def img2tokens(self, x_t: torch.Tensor): - x_t_unfolded = _unfold_3d(x_t, self._kernel, self._stride) - return rearrange(x_t_unfolded, "N col_dim num_tokens -> N num_tokens col_dim").contiguous() - - def process_input(self, transported_data: EvalInput): - batch_size, _, t, h, w = transported_data.x_t.shape - x_t = self.img2tokens(transported_data.x_t) - audio_x_t = transported_data.audio_x_t.contiguous() - text_in = transported_data.txt_feat.contiguous() - - simple_packed_data = _SimplePackedData(items=[]) - for i in range(batch_size): - single_data = _SingleData( - video_x_t=x_t[i], - audio_x_t=audio_x_t[i], - audio_feat_len=transported_data.audio_feat_len[i], - txt_feat=text_in[i], - txt_feat_len=transported_data.txt_feat_len[i], - t=t, - h=h, - w=w, - patch_size=self.patch_size, - t_patch_size=self.t_patch_size, - spatial_rope_interpolation=self.spatial_rope_interpolation, - ref_audio_offset=self.ref_audio_offset, - text_offset=self.text_offset, - coords_style=self.coords_style, - ) - simple_packed_data.items.append(single_data) - - if self.frame_receptive_field != -1: - assert batch_size == 1, "local attention only supports batch size 1" - local_attn_handler = _calc_local_attn_ffa_handler( - num_video_tokens=simple_packed_data[0].video_token_num, - num_audio_and_txt_tokens=simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len, - num_frames=t, - frame_receptive_field=self.frame_receptive_field, - ) - if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor): - local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item() - if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor): - local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item() - else: - local_attn_handler = None - - varlen_handler = VarlenHandler( - cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), - cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), - max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(), - max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(), - ) - - self.saved_for_output(simple_packed_data=simple_packed_data) - - x = simple_packed_data.token_sequence - coords_mapping = simple_packed_data.coords_mapping - modality_mapping = simple_packed_data.modality_mapping - return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler) - - def process_output(self, x: torch.Tensor): - simple_packed_data: _SimplePackedData = self.get_saved_data("simple_packed_data") - return simple_packed_data.depack_token_sequence(x) - - -# =========================================================================== -# Pipeline helpers -# =========================================================================== -@dataclass -class EvalInput: - x_t: torch.Tensor - audio_x_t: torch.Tensor - audio_feat_len: torch.Tensor | list[int] - txt_feat: torch.Tensor - txt_feat_len: torch.Tensor | list[int] - - -class _T5GemmaEncoder: - def __init__(self, model_path: str, device: str, weight_dtype: torch.dtype, subfolder: str | None = None): - from vllm.distributed import get_tensor_model_parallel_world_size - - self.device = device - hf_kwargs: dict[str, Any] = {} - if subfolder is not None: - hf_kwargs["subfolder"] = subfolder - self.tokenizer = AutoTokenizer.from_pretrained(model_path, **hf_kwargs) - - tp_size = get_tensor_model_parallel_world_size() - if tp_size > 1: - from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig - - config = T5GemmaConfig.from_pretrained(model_path, **hf_kwargs) - # The config we need is the encoder config - config_encoder = config.encoder - # Propagate some outer config values - config_encoder.vocab_size = config.vocab_size - config_encoder.rms_norm_eps = getattr(config, "rms_norm_eps", config_encoder.rms_norm_eps) - self.model = T5GemmaEncoderModelTP(config_encoder).to(device).to(weight_dtype) - self.is_tp = True - else: - self.model = T5GemmaEncoderModel.from_pretrained( - model_path, is_encoder_decoder=False, dtype=weight_dtype, **hf_kwargs - ).to(device) - self.is_tp = False - - @torch.inference_mode() - def encode(self, prompt: str) -> torch.Tensor: - inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device) - outputs = self.model(**inputs) - - if self.is_tp: - # T5GemmaEncoderModelTP just returns the hidden states tensor - return outputs.half() - else: - # HF model returns BaseModelOutput - return outputs["last_hidden_state"].half() - - -def _pad_or_trim(tensor: torch.Tensor, target_size: int, dim: int, pad_value: float = 0.0) -> tuple[torch.Tensor, int]: - current_size = tensor.size(dim) - if current_size < target_size: - padding_amount = target_size - current_size - padding_tuple = [0] * (2 * tensor.dim()) - padding_dim_index = tensor.dim() - 1 - dim - padding_tuple[2 * padding_dim_index + 1] = padding_amount - return F.pad(tensor, tuple(padding_tuple), "constant", pad_value), current_size - slicing = [slice(None)] * tensor.dim() - slicing[dim] = slice(0, target_size) - return tensor[tuple(slicing)], target_size - - -def _get_padded_t5_gemma_embedding( - prompt: str, - encoder: _T5GemmaEncoder, - target_length: int, -) -> tuple[torch.Tensor, int]: - txt_feat = encoder.encode(prompt) - txt_feat, original_len = _pad_or_trim(txt_feat, target_size=target_length, dim=1) - return txt_feat.to(torch.float32), original_len - - -def _resizecrop(img: Image.Image, target_height: int, target_width: int) -> Image.Image: - """Centre-crop resize keeping aspect ratio then letterbox to target.""" - pil_image = img.convert("RGB") - original_width, original_height = pil_image.size - scale_x = target_width / original_width - scale_y = target_height / original_height - scale = max(scale_x, scale_y) - new_width = int(round(original_width * scale)) - new_height = int(round(original_height * scale)) - resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS) - left = (new_width - target_width) // 2 - top = (new_height - target_height) // 2 - return resized_image.crop((left, top, left + target_width, top + target_height)) - - -class ZeroSNRDDPMDiscretization: - """ZeroSNR DDPM sigma schedule, ported from daVinci-MagiHuman. - Used to compute sigma values for SR noise injection. - """ - - def __init__( - self, - linear_start: float = 0.00085, - linear_end: float = 0.0120, - num_timesteps: int = 1000, - shift_scale: float = 1.0, - keep_start: bool = False, - post_shift: bool = False, - ): - from functools import partial - - if keep_start and not post_shift: - linear_start = linear_start / (shift_scale + (1 - shift_scale) * linear_start) - self.num_timesteps = num_timesteps - betas = torch.linspace(linear_start**0.5, linear_end**0.5, num_timesteps, dtype=torch.float64) ** 2 - alphas = 1.0 - betas.cpu().numpy() - self.alphas_cumprod = np.cumprod(alphas, axis=0) - self.to_torch = partial(torch.tensor, dtype=torch.float32) - if not post_shift: - self.alphas_cumprod = self.alphas_cumprod / (shift_scale + (1 - shift_scale) * self.alphas_cumprod) - self.post_shift = post_shift - self.shift_scale = shift_scale - - def __call__( - self, - n: int, - do_append_zero: bool = True, - device: str = "cpu", - flip: bool = False, - return_idx: bool = False, - ): - from functools import partial - - if n < self.num_timesteps: - timesteps = np.linspace(self.num_timesteps - 1, 0, n, endpoint=False).astype(int)[::-1] - alphas_cumprod = self.alphas_cumprod[timesteps] - elif n == self.num_timesteps: - alphas_cumprod = self.alphas_cumprod - else: - raise ValueError(f"n={n} > num_timesteps={self.num_timesteps}") - - to_torch = partial(torch.tensor, dtype=torch.float32, device=device) - alphas_cumprod = to_torch(alphas_cumprod) - alphas_cumprod_sqrt = alphas_cumprod.sqrt() - alphas_cumprod_sqrt_0 = alphas_cumprod_sqrt[0].clone() - alphas_cumprod_sqrt_T = alphas_cumprod_sqrt[-1].clone() - alphas_cumprod_sqrt -= alphas_cumprod_sqrt_T - alphas_cumprod_sqrt *= alphas_cumprod_sqrt_0 / (alphas_cumprod_sqrt_0 - alphas_cumprod_sqrt_T) - - if self.post_shift: - alphas_cumprod_sqrt = ( - alphas_cumprod_sqrt**2 / (self.shift_scale + (1 - self.shift_scale) * alphas_cumprod_sqrt**2) - ) ** 0.5 - - sigmas = torch.flip(alphas_cumprod_sqrt, (0,)) - sigmas = torch.cat([sigmas, sigmas.new_zeros([1])]) if do_append_zero else sigmas - if return_idx: - return sigmas if not flip else torch.flip(sigmas, (0,)), timesteps - return sigmas if not flip else torch.flip(sigmas, (0,)) - - -def _schedule_latent_step( - *, - video_scheduler: FlowUniPCMultistepScheduler, - audio_scheduler: FlowUniPCMultistepScheduler, - latent_video: torch.Tensor, - latent_audio: torch.Tensor, - t, - idx: int, - steps, - v_cfg_video: torch.Tensor, - v_cfg_audio: torch.Tensor, - is_a2v: bool, - cfg_number: int, - using_sde_flag: bool, - use_sr_model: bool = False, -): - # Fast DDIM path for cfg_number==1, only used during the BR stage - if cfg_number == 1 and not use_sr_model: - latent_video = video_scheduler.step_ddim(v_cfg_video, idx, latent_video) - latent_audio = audio_scheduler.step_ddim(v_cfg_audio, idx, latent_audio) - return latent_video, latent_audio - - if using_sde_flag: - if use_sr_model: - # SR stage with SDE: only update video, keep audio unchanged - latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0] - return latent_video, latent_audio - if idx < int(len(steps) * (3 / 4)): - noise_theta = 1.0 if (idx + 1) % 2 == 0 else 0.0 - else: - noise_theta = 1.0 if idx % 3 == 0 else 0.0 - latent_video = video_scheduler.step_sde(v_cfg_video, idx, latent_video, noise_theta=noise_theta) - if not is_a2v: - latent_audio = audio_scheduler.step_sde(v_cfg_audio, idx, latent_audio, noise_theta=noise_theta) - return latent_video, latent_audio - - latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0] - # Do not update audio latent during the SR stage - if not is_a2v and not use_sr_model: - latent_audio = audio_scheduler.step(v_cfg_audio, t, latent_audio, return_dict=False)[0] - return latent_video, latent_audio - - -_NEGATIVE_PROMPT = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, " - "overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, " - "poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, " - "still picture, messy background, three legs, many people in the background, walking backwards" - ", low quality, worst quality, poor quality, noise, background noise, hiss, hum, buzz, crackle, static, " - "compression artifacts, MP3 artifacts, digital clipping, distortion, muffled, muddy, unclear, echo, " - "reverb, room echo, over-reverberated, hollow sound, distant, washed out, harsh, shrill, piercing, " - "grating, tinny, thin sound, boomy, bass-heavy, flat EQ, over-compressed, abrupt cut, jarring transition, " - "sudden silence, looping artifact, music, instrumental, sirens, alarms, crowd noise, unrelated sound " - "effects, chaotic, disorganized, messy, cheap sound" - ", emotionless, flat delivery, deadpan, lifeless, apathetic, robotic, mechanical, monotone, flat " - "intonation, undynamic, boring, reading from a script, AI voice, synthetic, text-to-speech, TTS, " - "insincere, fake emotion, exaggerated, overly dramatic, melodramatic, cheesy, cringey, hesitant, " - "unconfident, tired, weak voice, stuttering, stammering, mumbling, slurred speech, mispronounced, " - "bad articulation, lisp, vocal fry, creaky voice, mouth clicks, lip smacks, wet mouth sounds, heavy " - "breathing, audible inhales, plosives, p-pops, coughing, clearing throat, sneezing, speaking too fast, " - "rushed, speaking too slow, dragged out, unnatural pauses, awkward silence, choppy, disjointed, multiple " - "speakers, two voices, background talking, out of tune, off-key, autotune artifacts" -) - - -# =========================================================================== -# Pre/post process funcs (registered in registry) -# =========================================================================== -def get_magi_human_pre_process_func(*args, **kwargs): - def pre_process(request: OmniDiffusionRequest): - return request - - return pre_process - - -def get_magi_human_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": 44100, - "fps": 25, - } - return output - - return post_process - - -# =========================================================================== -# HF Hub / local path helpers -# =========================================================================== - - -def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict: - """Load a JSON config file from a local path or HuggingFace Hub repo.""" - if local_files_only: - path = os.path.join(model_path, *filename.split("/")) - with open(path) as f: - return json.load(f) - else: - from huggingface_hub import hf_hub_download - - cached = hf_hub_download(repo_id=model_path, filename=filename) - with open(cached) as f: - return json.load(f) - - -def _resolve_subdir( - model_path: str, - subfolder: str, - local_files_only: bool = True, - required_files: list[str] | None = None, -) -> str: - """Resolve a model subfolder to a local directory path. - - For HF Hub repos, downloads all ``required_files`` (default: ``["config.json"]``) - into the HF cache and returns the parent directory. - """ - if local_files_only: - return os.path.join(model_path, subfolder) - from huggingface_hub import hf_hub_download - - files = required_files or ["config.json"] - last_cached: str | None = None - for fname in files: - last_cached = hf_hub_download(repo_id=model_path, filename=f"{subfolder}/{fname}") - return os.path.dirname(last_cached) - - -# =========================================================================== -# Main Pipeline -# =========================================================================== -class MagiHumanPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin): - _dit_modules: ClassVar[list[str]] = ["dit", "sr_dit"] - _encoder_modules: ClassVar[list[str]] = ["text_encoder"] - _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] - - def __init__(self, od_config: OmniDiffusionConfig, **kwargs): - super().__init__() - model_path = od_config.model - local_files_only = os.path.exists(model_path) - device = f"cuda:{torch.accelerator.current_device_index()}" - self.device_str = device - self.dtype = od_config.dtype or torch.bfloat16 - - model_index = _load_json(model_path, "model_index.json", local_files_only) - eval_cfg = model_index - dp_cfg = model_index.get("data_proxy", {}) - - dit_subfolder = "transformer" - - dit_json = _load_json(model_path, f"{dit_subfolder}/config.json", local_files_only) - dit_model_config = MagiHumanDiTConfig(**dit_json) - - self.dit = DiTModel(dit_model_config) - self.dit.eval() - - self.vae = DistributedAutoencoderKLWan.from_pretrained(model_path, subfolder="vae") - self.vae.to(device) - self.vae.eval() - vae_cfg = _load_json(model_path, "vae/config.json", local_files_only) - self.vae_latent_mean = torch.tensor(vae_cfg["latents_mean"], dtype=torch.float32) - self.vae_latent_std = torch.tensor(vae_cfg["latents_std"], dtype=torch.float32) - - self.audio_vae = SAAudioFeatureExtractor( - device=device, - model_path=_resolve_subdir( - model_path, - "audio_vae", - local_files_only, - required_files=["config.json", "model_config.json", "model.safetensors"], - ), - ) - - logger.info("Loading T5Gemma text encoder from %s (subfolder=text_encoder)", model_path) - if local_files_only: - txt_enc_path = os.path.join(model_path, "text_encoder") - txt_enc_subfolder = None - else: - txt_enc_path = model_path - txt_enc_subfolder = "text_encoder" - self.text_encoder = _T5GemmaEncoder( - model_path=txt_enc_path, - device=device, - weight_dtype=self.dtype, - subfolder=txt_enc_subfolder, - ) - - self.data_proxy = MagiDataProxy( - patch_size=dp_cfg.get("patch_size", 2), - t_patch_size=dp_cfg.get("t_patch_size", 1), - frame_receptive_field=dp_cfg.get("frame_receptive_field", 11), - spatial_rope_interpolation=dp_cfg.get("spatial_rope_interpolation", "extra"), - ref_audio_offset=dp_cfg.get("ref_audio_offset", 1000), - text_offset=dp_cfg.get("text_offset", 0), - coords_style=dp_cfg.get("coords_style", "v2"), - ) - # SR DataProxy forces v1 coordinate style (consistent with the original) - self.sr_data_proxy = MagiDataProxy( - patch_size=dp_cfg.get("patch_size", 2), - t_patch_size=dp_cfg.get("t_patch_size", 1), - frame_receptive_field=dp_cfg.get("frame_receptive_field", 11), - spatial_rope_interpolation=dp_cfg.get("spatial_rope_interpolation", "extra"), - ref_audio_offset=dp_cfg.get("ref_audio_offset", 1000), - text_offset=dp_cfg.get("text_offset", 0), - coords_style="v1", - ) - - self.fps = eval_cfg.get("fps", 25) - self.num_inference_steps_default = eval_cfg.get("num_inference_steps", 32) - self.video_txt_guidance_scale = eval_cfg.get("video_txt_guidance_scale", 5.0) - self.audio_txt_guidance_scale = eval_cfg.get("audio_txt_guidance_scale", 5.0) - self.shift = eval_cfg.get("shift", 5.0) - self.cfg_number = eval_cfg.get("cfg_number", 2) - self.use_cfg_trick = eval_cfg.get("use_cfg_trick", True) - self.cfg_trick_start_frame = eval_cfg.get("cfg_trick_start_frame", 13) - self.cfg_trick_value = eval_cfg.get("cfg_trick_value", 2.0) - self.using_sde_flag = eval_cfg.get("using_sde_flag", False) - self.t5_gemma_target_length = eval_cfg.get("t5_gemma_target_length", 640) - self.vae_stride = eval_cfg.get("vae_stride", [4, 16, 16]) - self.z_dim = eval_cfg.get("z_dim", 48) - self.patch_size = eval_cfg.get("patch_size", [1, 2, 2]) - # SR-specific hyperparameters - self.sr_num_inference_steps_default = eval_cfg.get("sr_num_inference_steps", 5) - self.sr_cfg_number = eval_cfg.get("sr_cfg_number", 2) - self.sr_video_txt_guidance_scale = eval_cfg.get("sr_video_txt_guidance_scale", 3.5) - self.noise_value = eval_cfg.get("noise_value", 220) - self.sr_audio_noise_scale = eval_cfg.get("sr_audio_noise_scale", 0.7) - # ZeroSNR sigma schedule for SR noise injection (flip=True, high to low) - self.zerosnr_sigmas = ZeroSNRDDPMDiscretization()(1000, do_append_zero=False, flip=True) - - self.context_null, self.original_context_null_len = _get_padded_t5_gemma_embedding( - _NEGATIVE_PROMPT, - self.text_encoder, - self.t5_gemma_target_length, - ) - self.video_processor = VideoProcessor(vae_scale_factor=16) - - # SR DiT model (loaded from the sr/ subdirectory) - sr_dit_subfolder = "sr" - sr_dit_json = _load_json(model_path, f"{sr_dit_subfolder}/config.json", local_files_only) - sr_dit_model_config = MagiHumanDiTConfig(**sr_dit_json) - self.sr_dit = DiTModel(sr_dit_model_config) - self.sr_dit.eval() - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=model_path, - subfolder=dit_subfolder, - revision=None, - prefix="dit.", - fall_back_to_pt=True, - ), - DiffusersPipelineLoader.ComponentSource( - model_or_path=model_path, - subfolder=sr_dit_subfolder, - revision=None, - prefix="sr_dit.", - fall_back_to_pt=True, - ), - ] - if getattr(self.text_encoder, "is_tp", False): - self.weights_sources.append( - DiffusersPipelineLoader.ComponentSource( - model_or_path=model_path, - subfolder="text_encoder", - revision=None, - prefix="text_encoder.", - fall_back_to_pt=True, - ), - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Weight loading for MagiHuman DiT with TP support. - # - # The checkpoint stores weights with these naming patterns: - # - attention.linear_qkv.weight: fused [Q, K, V, G] for shared layers, - # or stacked per-expert [expert0_Q|K|V|G, expert1_..., expert2_...] for MoE. - # - attention.linear_proj.weight: single for shared, stacked per-expert for MoE. - # - mlp.up_gate_proj.weight / mlp.down_proj.weight: similarly stacked for MoE. - # - # The model now uses per-expert vLLM parallel layers for MoE blocks: - # attention.linear_qkv.experts.{i}.weight (QKVParallelLinear per expert) - # attention.linear_gating.experts.{i}.weight (ColumnParallelLinear per expert) - # attention.linear_proj.experts.{i}.weight (RowParallelLinear per expert) - # mlp.up_gate_proj.experts.{i}.weight (ColumnParallelLinear per expert) - # mlp.down_proj.experts.{i}.weight (RowParallelLinear per expert) - # - # Shared layers keep the same naming (no .experts.). - params_dict = dict(self.named_parameters()) - modules_dict = dict(self.named_modules()) - loaded_params: set[str] = set() - - for name, loaded_weight in weights: - # ── Text Encoder weights ── - if name.startswith("text_encoder."): - if getattr(self.text_encoder, "is_tp", False): - # Strip "text_encoder." prefix for the T5Gemma TP model - # The T5GemmaEncoderModelTP load_weights handles the "encoder." prefix itself - sub_name = name[len("text_encoder.") :] - loaded_params.update( - f"text_encoder.{k}" for k in self.text_encoder.model.load_weights([(sub_name, loaded_weight)]) - ) - else: - loaded_params.add(name) - continue - - # ── Shared attention QKV + Gating split ── - # Checkpoint: attention.linear_qkv.weight = [Q, K, V, G] fused. - # Model: attention.linear_qkv.weight (QKVParallelLinear) + attention.linear_gating.weight. - if "attention.linear_qkv.weight" in name: - gating_name = name.replace("attention.linear_qkv.weight", "attention.linear_gating.weight") - # Check if this is a shared layer (direct param exists, no .experts.) - if name in params_dict and gating_name in params_dict: - qkv_param = params_dict[name] - gating_param = params_dict[gating_name] - - mod_path = name[: -len(".weight")] - qkv_mod = modules_dict.get(mod_path) - if qkv_mod is not None and hasattr(qkv_mod, "total_num_heads"): - total_heads_q = qkv_mod.total_num_heads - total_heads_kv = qkv_mod.total_num_kv_heads - head_dim = qkv_mod.head_size - else: - head_dim = 128 - tp_size = get_tensor_model_parallel_world_size() - total_heads_q = gating_param.data.shape[0] * tp_size - total_heads_kv = (loaded_weight.shape[0] - total_heads_q * head_dim - total_heads_q) // ( - 2 * head_dim - ) - - q_size = total_heads_q * head_dim - kv_size = total_heads_kv * head_dim - - q_w = loaded_weight[:q_size] - k_w = loaded_weight[q_size : q_size + kv_size] - v_w = loaded_weight[q_size + kv_size : q_size + 2 * kv_size] - g_w = loaded_weight[q_size + 2 * kv_size :] - - qkv_loader = getattr(qkv_param, "weight_loader", default_weight_loader) - qkv_loader(qkv_param, q_w, "q") - qkv_loader(qkv_param, k_w, "k") - qkv_loader(qkv_param, v_w, "v") - - gating_loader = getattr(gating_param, "weight_loader", default_weight_loader) - gating_loader(gating_param, g_w) - - loaded_params.add(name) - loaded_params.add(gating_name) - continue - - # ── MoE attention QKV + Gating split ── - # Checkpoint: attention.linear_qkv.weight = stacked [expert0_QKVG, expert1_QKVG, ...]. - # Model: attention.linear_qkv.experts.{i}.weight (QKVParallelLinear per expert) - # + attention.linear_gating.experts.{i}.weight (ColumnParallelLinear per expert). - expert0_name = name.replace("attention.linear_qkv.weight", "attention.linear_qkv.experts.0.weight") - if expert0_name in params_dict: - # Determine num_experts by checking which expert indices exist. - moe_qkv_mod_path = name[: -len(".weight")] - moe_qkv_mod = modules_dict.get(moe_qkv_mod_path) - num_experts = moe_qkv_mod.num_experts if moe_qkv_mod is not None else 3 - - # Get head info from the first expert's QKVParallelLinear. - expert0_mod_path = name.replace("attention.linear_qkv.weight", "attention.linear_qkv.experts.0") - expert0_mod = modules_dict.get(expert0_mod_path) - if expert0_mod is not None and hasattr(expert0_mod, "total_num_heads"): - total_heads_q = expert0_mod.total_num_heads - total_heads_kv = expert0_mod.total_num_kv_heads - head_dim = expert0_mod.head_size - else: - head_dim = 128 - # Infer from checkpoint weight shape. - # We'll get exact sizes from model config below. - total_heads_q = 40 # fallback for default config - total_heads_kv = 8 - - q_size = total_heads_q * head_dim - kv_size = total_heads_kv * head_dim - # Check if gating is present. - gating_expert0_name = name.replace( - "attention.linear_qkv.weight", "attention.linear_gating.experts.0.weight" - ) - has_gating = gating_expert0_name in params_dict - - # Split stacked checkpoint weight into per-expert chunks. - expert_weights = loaded_weight.chunk(num_experts, dim=0) - - for i in range(num_experts): - expert_w = expert_weights[i] - # Each expert chunk: [Q, K, V, G (optional)]. - q_w = expert_w[:q_size] - k_w = expert_w[q_size : q_size + kv_size] - v_w = expert_w[q_size + kv_size : q_size + 2 * kv_size] - - expert_param_name = name.replace( - "attention.linear_qkv.weight", - f"attention.linear_qkv.experts.{i}.weight", - ) - expert_param = params_dict[expert_param_name] - expert_loader = getattr(expert_param, "weight_loader", default_weight_loader) - expert_loader(expert_param, q_w, "q") - expert_loader(expert_param, k_w, "k") - expert_loader(expert_param, v_w, "v") - loaded_params.add(expert_param_name) - - if has_gating: - g_w = expert_w[q_size + 2 * kv_size :] - gating_param_name = name.replace( - "attention.linear_qkv.weight", - f"attention.linear_gating.experts.{i}.weight", - ) - gating_param = params_dict[gating_param_name] - gating_loader = getattr(gating_param, "weight_loader", default_weight_loader) - gating_loader(gating_param, g_w) - loaded_params.add(gating_param_name) - continue - - # ── MoE stacked weight splitting for proj / MLP layers ── - # Checkpoint: x.y.weight (stacked [expert0, expert1, ...]). - # Model: x.y.experts.{i}.weight. - if name not in params_dict: - # Check if this is a stacked MoE weight by looking for .experts.0. - base, _, suffix = name.rpartition(".") - expert0_name = f"{base}.experts.0.{suffix}" if base else None - if expert0_name and expert0_name in params_dict: - # Determine num_experts. - moe_mod = modules_dict.get(base) - num_experts = getattr(moe_mod, "num_experts", 3) if moe_mod is not None else 3 - - # Split stacked weight into per-expert chunks. - expert_weights = loaded_weight.chunk(num_experts, dim=0) - for i in range(num_experts): - expert_param_name = f"{base}.experts.{i}.{suffix}" - if expert_param_name not in params_dict: - continue - expert_param = params_dict[expert_param_name] - expert_loader = getattr(expert_param, "weight_loader", default_weight_loader) - expert_loader(expert_param, expert_weights[i]) - loaded_params.add(expert_param_name) - continue - # Truly unknown weight — skip. - continue - - # ── Standard weight loading (shared layers + non-MoE params) ── - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - if getattr(self.text_encoder, "is_tp", False): - self.context_null, self.original_context_null_len = _get_padded_t5_gemma_embedding( - _NEGATIVE_PROMPT, - self.text_encoder, - self.t5_gemma_target_length, - ) - - return loaded_params - - def _dit_forward(self, eval_input: EvalInput) -> tuple[torch.Tensor, torch.Tensor]: - packed = self.data_proxy.process_input(eval_input) - noise_pred = self.dit(*packed) - return self.data_proxy.process_output(noise_pred) - - def _sr_dit_forward(self, eval_input: EvalInput) -> tuple[torch.Tensor, torch.Tensor]: - """SR stage uses sr_data_proxy (coords_style=v1) and sr_dit model.""" - packed = self.sr_data_proxy.process_input(eval_input) - noise_pred = self.sr_dit(*packed) - return self.sr_data_proxy.process_output(noise_pred) - - @torch.inference_mode() - def _evaluate_with_latent( - self, - context: torch.Tensor, - original_context_len: int, - latent_image: torch.Tensor | None, - latent_video: torch.Tensor, - latent_audio: torch.Tensor, - num_inference_steps: int, - is_a2v: bool = False, - use_sr_model: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Select cfg_number and guidance_scale based on BR/SR stage - cfg_number = self.sr_cfg_number if use_sr_model else self.cfg_number - video_guidance = self.sr_video_txt_guidance_scale if use_sr_model else self.video_txt_guidance_scale - forward_fn = self._sr_dit_forward if use_sr_model else self._dit_forward - - video_scheduler = FlowUniPCMultistepScheduler() - audio_scheduler = FlowUniPCMultistepScheduler() - video_scheduler.set_timesteps(num_inference_steps, device=self.device_str, shift=self.shift) - audio_scheduler.set_timesteps(num_inference_steps, device=self.device_str, shift=self.shift) - timesteps = video_scheduler.timesteps - - latent_length = latent_video.shape[2] - cfg_trick_guidance = ( - torch.tensor(video_guidance, device=self.device_str).expand(1, 1, latent_length, 1, 1).clone() - ) - if self.use_cfg_trick: - cfg_trick_guidance[:, :, : self.cfg_trick_start_frame] = min(self.cfg_trick_value, video_guidance) - - with self.progress_bar(total=len(timesteps)) as pbar: - for idx, t in enumerate(timesteps): - if latent_image is not None: - latent_video[:, :, :1] = latent_image[:, :, :1] - - # Reduce guidance when t<=500 during BR stage (original behavior) - cur_video_guidance = video_guidance if (use_sr_model or t > 500) else 2.0 - - eval_input_cond = EvalInput( - x_t=latent_video, - audio_x_t=latent_audio, - audio_feat_len=[latent_audio.shape[1]], - txt_feat=context, - txt_feat_len=[original_context_len], - ) - - v_cond_video, v_cond_audio = forward_fn(eval_input_cond) - - if cfg_number == 1: - v_cfg_video = v_cond_video - v_cfg_audio = v_cond_audio - elif cfg_number == 2: - eval_input_uncond = EvalInput( - x_t=latent_video, - audio_x_t=latent_audio, - audio_feat_len=[latent_audio.shape[1]], - txt_feat=self.context_null, - txt_feat_len=[self.original_context_null_len], - ) - v_uncond_video, v_uncond_audio = forward_fn(eval_input_uncond) - v_cfg_video = v_uncond_video + cur_video_guidance * (v_cond_video - v_uncond_video) - v_cfg_audio = v_uncond_audio + self.audio_txt_guidance_scale * (v_cond_audio - v_uncond_audio) - else: - raise ValueError(f"Invalid cfg_number: {cfg_number}") - - latent_video, latent_audio = _schedule_latent_step( - video_scheduler=video_scheduler, - audio_scheduler=audio_scheduler, - latent_video=latent_video, - latent_audio=latent_audio, - t=t, - idx=idx, - steps=timesteps, - v_cfg_video=v_cfg_video, - v_cfg_audio=v_cfg_audio, - is_a2v=is_a2v, - cfg_number=cfg_number, - using_sde_flag=self.using_sde_flag, - use_sr_model=use_sr_model, - ) - - pbar.update() - - if latent_image is not None: - latent_video[:, :, :1] = latent_image[:, :, :1] - return latent_video, latent_audio - - def _encode_image(self, image: Image.Image, height: int, width: int) -> torch.Tensor: - image = load_image(image) - image = _resizecrop(image, height, width) - image = self.video_processor.preprocess(image, height=height, width=width) - image = image.to(device=self.device_str, dtype=self.dtype).unsqueeze(2) - vae_out = self.vae.encode(image) - if hasattr(vae_out, "latent_dist"): - return vae_out.latent_dist.mode().to(torch.float32) - return vae_out.to(torch.float32) - - def _decode_video(self, latent: torch.Tensor) -> list[np.ndarray]: - mean = self.vae_latent_mean.to(latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1) - std = self.vae_latent_std.to(latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1) - latent = latent * std + mean - - videos = self.vae.decode(latent.to(self.dtype)) - if hasattr(videos, "sample"): - videos = videos.sample - videos.mul_(0.5).add_(0.5).clamp_(0, 1) - videos = [v.float().cpu().permute(1, 2, 3, 0) * 255 for v in videos] - return [v.numpy().astype(np.uint8) for v in videos] - - def _decode_audio(self, latent_audio: torch.Tensor) -> np.ndarray: - latent_audio = latent_audio.squeeze(0).to(self.dtype) - audio_output = self.audio_vae.decode(latent_audio.T) - audio_np = audio_output.squeeze(0).T.float().cpu().numpy() - target_len = int(audio_np.shape[0] * 441 / 512) - from scipy.signal import resample - - return resample(audio_np, target_len) - - def encode_prompt( - self, - prompt: str, - target_length: int | None = None, - ) -> tuple[torch.Tensor, int]: - """Encode *prompt* with the T5-Gemma text encoder and pad to fixed length. - - This is the single text-encoder entrypoint so the runner-level - prompt-embedding cache (see - ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently - memoize results when the same prompt is submitted repeatedly. - - Returns: - ``(context, original_context_len)`` matching - :func:`_get_padded_t5_gemma_embedding`. - """ - return _get_padded_t5_gemma_embedding( - prompt, - self.text_encoder, - target_length if target_length is not None else self.t5_gemma_target_length, - ) - - @torch.inference_mode() - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | None = None, - height: int = 256, - width: int = 448, - num_inference_steps: int | None = None, - seconds: int = 10, - seed: int | None = None, - image_path: str | None = None, - audio_path: str | None = None, - **kwargs, - ) -> DiffusionOutput: - if len(req.prompts) >= 1: - p = req.prompts[0] - prompt = p if isinstance(p, str) else p.get("prompt", prompt) - if not isinstance(p, str): - image_path = p.get("image_path", image_path) - audio_path = p.get("audio_path", audio_path) - if prompt is None: - raise ValueError("prompt is required") - - height = req.sampling_params.height or height - width = req.sampling_params.width or width - seed = req.sampling_params.seed if req.sampling_params.seed is not None else seed - num_steps = req.sampling_params.num_inference_steps or num_inference_steps or self.num_inference_steps_default - sr_height: int | None = None - sr_width: int | None = None - sr_num_steps: int | None = None - if hasattr(req.sampling_params, "extra_args") and req.sampling_params.extra_args: - seconds = req.sampling_params.extra_args.get("seconds", seconds) - audio_path = req.sampling_params.extra_args.get("audio_path", audio_path) - image_path = req.sampling_params.extra_args.get("image_path", image_path) - sr_height = req.sampling_params.extra_args.get("sr_height", None) - sr_width = req.sampling_params.extra_args.get("sr_width", None) - sr_num_steps = req.sampling_params.extra_args.get("sr_num_inference_steps", None) - - device = self.device_str - - br_latent_height = height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1] - br_latent_width = width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2] - br_height = br_latent_height * self.vae_stride[1] - br_width = br_latent_width * self.vae_stride[2] - - if seed is not None: - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - if audio_path is not None: - latent_audio = load_audio_and_encode(self.audio_vae, audio_path, seconds) - latent_audio = latent_audio.permute(0, 2, 1) - num_frames = latent_audio.shape[1] - is_a2v = True - else: - num_frames = seconds * self.fps + 1 - latent_audio = torch.randn(1, num_frames, 64, dtype=torch.float32, device=device) - is_a2v = False - - latent_length = (num_frames - 1) // 4 + 1 - latent_video = torch.randn( - 1, - self.z_dim, - latent_length, - br_latent_height, - br_latent_width, - dtype=torch.float32, - device=device, - ) - - context, original_context_len = self.encode_prompt(prompt) - - if image_path is not None: - br_image = self._encode_image(load_image(image_path), br_height, br_width) - else: - br_image = None - - # ── BR stage ───────────────────────────────────────────────────────── - br_latent_video, br_latent_audio = self._evaluate_with_latent( - context, - original_context_len, - br_image, - latent_video.clone(), - latent_audio.clone(), - num_steps, - is_a2v, - use_sr_model=False, - ) - - # ── SR stage (optional, triggered when sr_height/sr_width are provided) ── - if sr_height is not None and sr_width is not None: - sr_latent_height = sr_height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1] - sr_latent_width = sr_width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2] - sr_height = sr_latent_height * self.vae_stride[1] - sr_width = sr_latent_width * self.vae_stride[2] - - # Image condition (at SR resolution) - if image_path is not None: - sr_image = self._encode_image(load_image(image_path), sr_height, sr_width) - else: - sr_image = None - - # Trilinear interpolation of BR latent to SR resolution - sr_latent_video = torch.nn.functional.interpolate( - br_latent_video, - size=(latent_length, sr_latent_height, sr_latent_width), - mode="trilinear", - align_corners=True, - ) - - # Noise injection: sigma-weighted blend (noise_value indexes the ZeroSNR sigma schedule) - if self.noise_value != 0: - noise = torch.randn_like(sr_latent_video) - sigma = self.zerosnr_sigmas.to(sr_latent_video.device)[self.noise_value] - sr_latent_video = sr_latent_video * sigma + noise * (1 - sigma**2) ** 0.5 - - # Audio: blend with noise (noised version used during SR inference; final audio keeps BR result) - sr_latent_audio = torch.randn_like(br_latent_audio) * self.sr_audio_noise_scale + br_latent_audio * ( - 1 - self.sr_audio_noise_scale - ) - - torch.accelerator.empty_cache() - sr_steps = sr_num_steps or self.sr_num_inference_steps_default - final_latent_video, _ = self._evaluate_with_latent( - context, - original_context_len, - sr_image, - sr_latent_video.clone(), - sr_latent_audio.clone(), - sr_steps, - is_a2v, - use_sr_model=True, - ) - # SR stage does not update audio; keep the BR result - final_latent_video = final_latent_video - final_latent_audio = br_latent_audio - else: - final_latent_video = br_latent_video - final_latent_audio = br_latent_audio - - torch.accelerator.empty_cache() - videos_np = self._decode_video(final_latent_video) - torch.accelerator.empty_cache() - audio_np = self._decode_audio(final_latent_audio) - - return DiffusionOutput(output=(videos_np, audio_np)) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index f5d6ffe07..8a707711a 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -266,11 +266,6 @@ _DIFFUSION_MODELS = { "pipeline_flux2", "Flux2Pipeline", ), - "DreamIDOmniPipeline": ( - "dreamid_omni", - "pipeline_dreamid_omni", - "DreamIDOmniPipeline", - ), "SenseNovaU1Pipeline": ( "sensenova_u1", "pipeline_sensenova_u1", @@ -301,11 +296,6 @@ _DIFFUSION_MODELS = { "pipeline_sana_video_i2v", "SanaImageToVideoPipeline", ), - "MagiHumanPipeline": ( - "magi_human", - "pipeline_magi_human", - "MagiHumanPipeline", - ), "OmniVoicePipeline": ( "omnivoice", "pipeline_omnivoice", @@ -598,9 +588,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "LingBotVideoPipeline": "get_lingbot_video_post_process_func", "SanaVideoPipeline": "get_sana_video_post_process_func", "SanaImageToVideoPipeline": "get_sana_video_i2v_post_process_func", - "MagiHumanPipeline": "get_magi_human_post_process_func", "OmniVoicePipeline": "get_omnivoice_post_process_func", - "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func", "Cosmos3OmniPipeline": "get_cosmos3_post_process_func", @@ -645,7 +633,6 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "LingBotVideoPipeline": "get_lingbot_video_pre_process_func", "SanaImageToVideoPipeline": "get_sana_video_i2v_pre_process_func", "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func", - "MagiHumanPipeline": "get_magi_human_pre_process_func", "SanaWmPipeline": "get_sana_wm_pre_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func", "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func", diff --git a/vllm_omni/model_extras/magi_human.py b/vllm_omni/model_extras/magi_human.py deleted file mode 100644 index 470a1455b..000000000 --- a/vllm_omni/model_extras/magi_human.py +++ /dev/null @@ -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() diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index b8eb6bbe8..e19530445 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -36,10 +36,6 @@ from vllm_omni.model_extras.ltx2 import ( ltx_preserves_reference_image_size, ltx_transformer_config_subfolder, ) -from vllm_omni.model_extras.magi_human import ( - MAGI_HUMAN_EXTRA_BODY_PARAMS, - MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, -) from vllm_omni.model_extras.ming_flash_omni import ( MING_FLASH_OMNI_EXTRA_BODY_PARAMS, MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS, @@ -182,10 +178,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS, "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS, }, - "MagiHumanPipeline": { - "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS, - "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, - }, "HeliosPipeline": { "extra_body_params": HELIOS_EXTRA_BODY_PARAMS, "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS,