Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 jiangmengyu18andbetta18
|
88cff6c3eb |
No files matched your search
@@ -399,29 +399,6 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Higgs-Audio-V3 Test"
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_USE_DEEP_GEMM=0
|
||||
export VLLM_MOE_USE_DEEP_GEMM=0
|
||||
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
commands:
|
||||
- |
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
---
|
||||
name: precheck-pr
|
||||
description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
|
||||
---
|
||||
|
||||
# PR Pre-Check
|
||||
|
||||
Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only.
|
||||
|
||||
## Mode Selection
|
||||
|
||||
| Mode | When | Time |
|
||||
|------|------|------|
|
||||
| **Quick** | About to push, final sanity check | ~3 min |
|
||||
| **Full** | Ready for review, want maintainer-level scan | ~10 min |
|
||||
|
||||
Default to quick if unsure. Run full before marking a PR "ready for review."
|
||||
|
||||
## Workflow
|
||||
|
||||
### Step 1: Detect Base Branch
|
||||
|
||||
```bash
|
||||
BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \
|
||||
|| git merge-base HEAD main 2>/dev/null \
|
||||
|| echo origin/main)
|
||||
echo "diffing against ${BASE_SHA}"
|
||||
git diff --name-only ${BASE_SHA}...HEAD
|
||||
```
|
||||
|
||||
### Step 2: Validate PR Title
|
||||
|
||||
Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes:
|
||||
|
||||
| Prefix | Applies to |
|
||||
|--------|-----------|
|
||||
| `[Bugfix]` | Bug fixes |
|
||||
| `[CI/Build]` | Build or CI improvements |
|
||||
| `[Doc]` | Documentation changes |
|
||||
| `[Model]` | New/improved models (include model name) |
|
||||
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
|
||||
| `[Kernel]` | CUDA/kernel changes |
|
||||
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
|
||||
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
|
||||
| `[Misc]` | Other changes (use sparingly) |
|
||||
|
||||
✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title.
|
||||
⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add <ModelName> ...`).
|
||||
|
||||
### Step 3: Categorize the PR
|
||||
|
||||
| Diff contains | PR type |
|
||||
|---------------|---------|
|
||||
| New files under `vllm_omni/model_executor/models/<name>/` | **New Model** |
|
||||
| Changes to `vllm_omni/diffusion/` | **Diffusion Model** |
|
||||
| `[Bugfix]` prefix or single-file fix | **Bug Fix** |
|
||||
| Perf/benchmark/throughput claims in commit msg or diff | **Performance** |
|
||||
| Everything else | **General** |
|
||||
|
||||
If multiple rows apply (e.g., a diffusion model is also a new model), union the checklists.
|
||||
|
||||
### Step 4: Run Checklist
|
||||
|
||||
Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
|
||||
|
||||
### Step 5: Print Report
|
||||
|
||||
```
|
||||
Pre-check report for <branch>
|
||||
|
||||
Mode: quick | full
|
||||
Type: <new-model | diffusion-model | bug-fix | perf | general>
|
||||
|
||||
Dimension Result
|
||||
───────────────── ──────
|
||||
PR title format ✓
|
||||
PR desc integrity ✓
|
||||
Registry/config ✓
|
||||
Dead code ⚠ 2 warnings
|
||||
Accuracy ✓
|
||||
Benchmark ✗ missing software versions
|
||||
|
||||
Verdict: 1 blocking | 2 warnings | recommend fixing ✗ before PR
|
||||
```
|
||||
|
||||
**Severity:**
|
||||
|
||||
| Mark | Meaning |
|
||||
|------|---------|
|
||||
| ✗ | Blocking — fix before opening PR |
|
||||
| ⚠ | Warning — consider fixing |
|
||||
| ✓ | Pass |
|
||||
| — | Skipped (not applicable) |
|
||||
|
||||
### Stop Here
|
||||
|
||||
Do not post comments, open PRs, or modify files. The report is for the contributor's terminal only.
|
||||
@@ -1,135 +0,0 @@
|
||||
# Pre-Check Checklists
|
||||
|
||||
Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing.
|
||||
|
||||
---
|
||||
|
||||
## All PRs: Title Format (check first)
|
||||
|
||||
PR title must follow the project convention documented in `docs/contributing/README.md`:
|
||||
|
||||
| Prefix | Applies to |
|
||||
|--------|-----------|
|
||||
| `[Bugfix]` | Bug fixes |
|
||||
| `[CI/Build]` | Build or CI improvements |
|
||||
| `[Doc]` | Documentation changes |
|
||||
| `[Model]` | New/improved models (include model name) |
|
||||
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
|
||||
| `[Kernel]` | CUDA/kernel changes |
|
||||
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
|
||||
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
|
||||
| `[Misc]` | Other changes (use sparingly) |
|
||||
|
||||
**Examples:**
|
||||
```
|
||||
✓ [Model] Add MiniCPM-o 4.5 support
|
||||
✓ [Bugfix] Fix Ovis image text encoder dtype
|
||||
✓ [Core] Qwen3-Omni performance optimization
|
||||
✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend
|
||||
✗ Fix Ovis image text encoder dtype ← missing prefix
|
||||
✗ [bugfix] Fix Ovis image text encoder ← wrong case
|
||||
✗ [WIP] Add model ← WIP in title
|
||||
```
|
||||
|
||||
- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add <ModelName> ...`). ⚠ if missing the model name.
|
||||
|
||||
---
|
||||
|
||||
## Bug Fix PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body has 4 sections:** what broke (error message), repro steps, root cause, fix description
|
||||
- [ ] **Fix is minimal:** diff touches only files directly related to the bug. No unrelated refactors or style changes.
|
||||
- [ ] **Repro is runnable:** the PR body includes a copy-paste command that reproduces the bug
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
|
||||
- [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
|
||||
- [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
|
||||
- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
|
||||
- [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
|
||||
- [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
|
||||
|
||||
---
|
||||
|
||||
## Performance PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **Before/after numbers in PR body:** at minimum, one metric (RTF, RPS, latency, VRAM) with main vs PR values
|
||||
- [ ] **Hardware specified:** GPU model, count, VRAM. "L4" or "H20" is sufficient. "tested on GPU" is not.
|
||||
- [ ] **Concurrency scaling:** if the optimization claims better throughput, numbers at 2+ concurrency levels
|
||||
- [ ] **No unexplained regressions:** if any metric goes down at any concurrency level, explain why
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Software versions:** torch, CUDA, vllm, vllm-omni versions in PR body
|
||||
- [ ] **Warmup stated:** "1 warmup + 3 measured" or equivalent
|
||||
- [ ] **Stddev or range:** not just a single best run. "average of 2 rounds" is minimum.
|
||||
- [ ] **VRAM measured:** peak GPU memory before and after
|
||||
- [ ] **Benchmark script checked in:** under `tests/` or `examples/`, runnable by reviewer
|
||||
- [ ] **Exact command line in PR:** not pseudocode, not "run the benchmark"
|
||||
- [ ] **Pytest summary included:** `36 passed, 3 skipped, 17 warnings in 2897.66s` — proves the whole suite passes
|
||||
- [ ] **Config changes documented:** if `max_num_seqs`, `gpu_memory_utilization`, or deploy YAML changed, old and new values stated with reasoning
|
||||
|
||||
---
|
||||
|
||||
## New Model PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body file list matches `git diff --name-only`:** every file path in the description exists in the diff
|
||||
- [ ] **Registry entries match claimed count:** count `_OMNI_MODELS` or `_DIFFUSION_MODELS` entries added vs claimed in body
|
||||
- [ ] **Pipeline config `custom_process_input_func` paths resolve:** each referenced module.function exists
|
||||
- [ ] **`__init__.py` exports match registry:** each registered class appears in `__init__.py` and `__all__`
|
||||
- [ ] **Shell script modes match README:** `run_curl.sh` supports every mode the README documents
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Dead code scan (5 patterns):**
|
||||
1. Dead `forward()` — defined but never called through the pipeline path
|
||||
2. Dead factory/builder functions — instantiated nowhere
|
||||
3. Dead wrappers — thin wrappers with no callers
|
||||
4. Dead branch guards — `if version > X` where X is always satisfied
|
||||
5. Unused parameters — `__init__` args never accessed
|
||||
- [ ] **Copy-paste detection:** no string constant defined in 3+ files, no cross-module validation duplication, no near-identical shape coercion fns
|
||||
- [ ] **Import hygiene:** no re-export-only imports, no module-level side effects without docs, no `import os` in function body when already at top. Function-body imports are acceptable when the import is only reachable from a closure or factory return value (i.e., `import` would otherwise execute at module load even though the user may never trigger that path). Either lift to top-of-file or annotate with `# noqa: PLC0415`.
|
||||
- [ ] **Accuracy:** at least one test compares output against reference; audio = valid WAV at correct sample rate; determinism verified if seed claimed. ⚠ if missing and not explicitly deferred to a follow-up PR or RFC section; — if explicitly deferred (the deferral target SHOULD be linked in the PR body).
|
||||
- [ ] **Performance:** hardware, software versions, warmup, and metrics (RTF, VRAM, RPS, latency) stated. ⚠ if missing and not explicitly deferred; — if explicitly deferred.
|
||||
- [ ] **Benchmark settings:** model config (TP, PP, max_model_len, enforce_eager, quant), runtime config (batch, input/output spec), environment versions
|
||||
- [ ] **Benchmark script checked in + exact command line + pytest summary**
|
||||
|
||||
---
|
||||
|
||||
## Diffusion Model PRs
|
||||
|
||||
Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model checklist above, verify:
|
||||
|
||||
- [ ] **Transformer adapter follows contract:** inherits from `nn.Module`, implements `load_weights()`, uses `vllm_omni.diffusion.attention.layer.Attention`
|
||||
- [ ] **Pipeline has both offline and online paths:** `forward()` for batch generation, streaming path for online serving
|
||||
- [ ] **Cache-DiT integration:** at least one acceleration (TeaCache, FBCache, etc.) is wired up. ✗ for new **top-level** pipelines. For **subclass** pipelines that inherit `__call__` from a parent where Cache-DiT is already absent, downgrade to ⚠ and recommend filing a follow-up issue against the parent pipeline family.
|
||||
- [ ] **Parallelism config:** TP/SP/USP/CFG-Parallel options exposed in pipeline config
|
||||
- [ ] **Docs table updated:** model added to `docs/models/supported_models.md` with correct metadata
|
||||
- [ ] **E2E test exists:** at least one test exercises the full generate path
|
||||
|
||||
---
|
||||
|
||||
## General PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body matches diff:** description doesn't claim changes not in the diff
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
- [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **No new dead code in changed files:** scan for unreachable paths, unused imports
|
||||
- [ ] **Test coverage:** if core code changed (`engine/`, `stages/`, `connectors/`), tests are added or existing tests cover the path
|
||||
- [ ] **Import hygiene:** no redundant imports added
|
||||
- [ ] **No unrelated changes:** diff scope matches PR description
|
||||
@@ -25,12 +25,7 @@ include:
|
||||
- `add-tts-model`: covers integration of new TTS models and related serving
|
||||
workflows
|
||||
- `generate-release-note`: helps prepare release notes for repository changes
|
||||
- `precheck-pr`: self-check a branch before creating a PR — validates PR title
|
||||
format, catches dead code, verifies accuracy/perf claims, and confirms merge
|
||||
readiness
|
||||
- `review-pr`: provides a structured workflow for reviewing pull requests
|
||||
- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
|
||||
latest vllm-ascend NPUModelRunner
|
||||
|
||||
## Maintenance Guidelines
|
||||
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
<!-- markdownlint-disable -->
|
||||
**Before submitting:** run the precheck-pr skill with code agent for a self-check against project conventions — catches dead code, missing benchmarks, and title format issues.
|
||||
|
||||
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
|
||||
|
||||
## Purpose
|
||||
@@ -20,5 +18,4 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
|
||||
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
|
||||
</details>
|
||||
|
||||
**BEFORE SUBMITTING,** please read <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>
|
||||
(anything written below this line will be removed by GitHub Actions)
|
||||
**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
|
||||
@@ -16,7 +16,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
---
|
||||
|
||||
*Latest News* 🔥
|
||||
- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
|
||||
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
|
||||
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
|
||||
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
|
||||
|
||||
@@ -92,6 +92,7 @@ Core scheduling and caching components.
|
||||
Configuration classes.
|
||||
|
||||
- [vllm_omni.config.model.OmniModelConfig][]
|
||||
- [vllm_omni.config.stage_config.ModelPipeline][]
|
||||
- [vllm_omni.config.stage_config.StageConfig][]
|
||||
- [vllm_omni.config.stage_config.StageConfigFactory][]
|
||||
- [vllm_omni.config.stage_config.StageType][]
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 403 KiB After Width: | Height: | Size: 408 KiB |
@@ -116,20 +116,6 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
!!! note
|
||||
If the PR spans more than one category, please include all relevant prefixes.
|
||||
|
||||
### Pre-Check Before Submitting
|
||||
|
||||
Before submitting a PR, run the precheck-pr skill with code agent for a self-review against project conventions:
|
||||
|
||||
```
|
||||
/precheck-pr
|
||||
```
|
||||
|
||||
The skill offers two modes:
|
||||
- **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
|
||||
- **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene
|
||||
|
||||
The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the skill in `.claude/skills/precheck-pr/` for the full checklist.
|
||||
|
||||
### Local Test
|
||||
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
|
||||
|
||||
|
||||
@@ -305,8 +305,8 @@ class GPUWorker:
|
||||
)
|
||||
|
||||
# Load model
|
||||
model_loader = DiffusersPipelineLoader(load_config, self.od_config)
|
||||
self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}")
|
||||
model_loader = DiffusersPipelineLoader(load_config)
|
||||
self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}")
|
||||
|
||||
# Setup cache backend
|
||||
from vllm_omni.diffusion.cache.selector import get_cache_backend
|
||||
|
||||
@@ -45,7 +45,7 @@ If you do not need to modify source code of vLLM, you can directly install the s
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -70,8 +70,8 @@ git checkout v0.22.0
|
||||
```
|
||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||
```bash
|
||||
#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
|
||||
@@ -22,7 +22,7 @@ source .venv/bin/activate
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
|
||||
@@ -33,7 +33,7 @@ th {
|
||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
|
||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
@@ -50,8 +50,6 @@ th {
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
|
||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
|
||||
@@ -64,10 +62,6 @@ th {
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
|
||||
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
|
||||
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
|
||||
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
|
||||
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
|
||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
|
||||
@@ -75,13 +69,11 @@ th {
|
||||
| `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` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
|
||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
|
||||
|
||||
|
||||
@@ -116,7 +116,7 @@ Content-Type: application/json
|
||||
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
|
||||
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
|
||||
|
||||
**Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
|
||||
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
|
||||
|
||||
#### Voice Clone Parameters (Base task)
|
||||
|
||||
|
||||
@@ -67,9 +67,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
|
||||
| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
|
||||
| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
|
||||
| `width` | integer | model default | Output video width |
|
||||
| `height` | integer | model default | Output video height |
|
||||
| `num_frames` | integer | model default | Number of generated frames |
|
||||
@@ -81,8 +80,6 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
| `flow_shift` | number | null | Scheduler flow-shift value |
|
||||
| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
|
||||
| `seed` | integer | null | Random seed for reproducibility |
|
||||
| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation |
|
||||
| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration |
|
||||
| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
|
||||
| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
|
||||
| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
|
||||
@@ -126,41 +123,6 @@ curl -s http://localhost:8091/v1/videos \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
### Video-to-Video
|
||||
|
||||
For models that support video conditioning, upload the reference video with
|
||||
`input_reference`:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=continue this motion with consistent subjects and lighting" \
|
||||
-F "input_reference=@input.mp4;type=video/mp4" \
|
||||
-F "width=1280" \
|
||||
-F "height=720" \
|
||||
-F "num_frames=80" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload
|
||||
through `video_reference`. Do not send `video_reference` together with
|
||||
`input_reference` or `image_reference`.
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=continue this motion with consistent subjects and lighting" \
|
||||
-F 'video_reference={"video_url":"https://example.com/input.mp4"}' \
|
||||
-F "width=1280" \
|
||||
-F "height=720" \
|
||||
-F "num_frames=80" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
JSON references currently support `image_url`/`video_url`; `file_id` references
|
||||
are not implemented yet. Models may expose additional V2V controls through
|
||||
`extra_params`. For example, Cosmos3 supports
|
||||
`condition_frame_indexes_vision` and `condition_video_keep` to select which
|
||||
decoded reference frames are used as clean conditioning.
|
||||
|
||||
### Synchronous Generation
|
||||
|
||||
```bash
|
||||
@@ -184,10 +146,7 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
|
||||
## Model-Specific Examples
|
||||
|
||||
For complete text-to-video, image-to-video, and model-specific video-to-video
|
||||
walkthroughs, see:
|
||||
For complete text-to-video and image-to-video walkthroughs, see:
|
||||
|
||||
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
|
||||
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
|
||||
- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
|
||||
for model-specific video-to-video examples and conditioning controls
|
||||
@@ -148,7 +148,7 @@ The following tables show which models support each feature:
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
@@ -28,50 +28,6 @@ in deep DiT blocks.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
|
||||
### Faster FP8 GEMM on Blackwell (quack)
|
||||
|
||||
On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which
|
||||
applies the bias as a separate kernel after the GEMM. On the small GEMMs in video
|
||||
DiTs this bias add is a significant overhead. Installing the optional `quack` kernel
|
||||
lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering
|
||||
that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster).
|
||||
|
||||
```bash
|
||||
# CUDA 12.9
|
||||
pip install vllm-omni[quack]
|
||||
|
||||
# CUDA 13.x
|
||||
pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
|
||||
```
|
||||
|
||||
It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
|
||||
on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
|
||||
Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
|
||||
not installed, FP8 still works — it just keeps the unfused FlashInfer path.
|
||||
|
||||
#### Compile cache and warmup
|
||||
|
||||
quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer
|
||||
the first time across all autotuned configs). The compiled `.o` files are cached on
|
||||
disk and reused on later runs, so this is a one-time cost — **not** per request.
|
||||
|
||||
vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
|
||||
`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
|
||||
In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
|
||||
the image — so the first cold start does not recompile. The engine's startup dummy
|
||||
run already exercises the kernels, so with a warm cache the first real request is fast.
|
||||
|
||||
To pre-warm specific shapes (e.g. at image build time):
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8
|
||||
# (M, K, N) per linear; M = number of tokens for your resolution/frame count
|
||||
warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
```
|
||||
|
||||
> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install
|
||||
> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
{
|
||||
"prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
|
||||
"negative_prompt": "",
|
||||
"vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 10.0,
|
||||
"fps": 24,
|
||||
"condition_frame_indexes_vision": [0, 1],
|
||||
"condition_video_keep": "first"
|
||||
}
|
||||
@@ -1,204 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline higgs-audio v3 inference example.
|
||||
|
||||
Runs Stage 0 (Qwen3 talker) + Stage 1 (HiggsAudio codec) end-to-end
|
||||
through the vLLM-Omni engine without going through the HTTP server, and
|
||||
saves a 24 kHz mono WAV per prompt.
|
||||
|
||||
Example (plain TTS):
|
||||
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
|
||||
--texts "Hello world." \\
|
||||
--output-dir results/higgs_v3_wavs
|
||||
|
||||
Example (voice clone):
|
||||
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
|
||||
--texts "Hello world." \\
|
||||
--ref-audio path/to/reference.wav \\
|
||||
--ref-text "Transcript of the reference clip." \\
|
||||
--output-dir results/higgs_v3_wavs
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
||||
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
SAMPLE_RATE = 24_000
|
||||
DEFAULT_TEXTS = (
|
||||
"Hello world.",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
"Today is a beautiful day for a walk in the park.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="bosonai/higgs-audio-v3-tts-4b",
|
||||
help="Stage-0 talker model id or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--texts",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=list(DEFAULT_TEXTS),
|
||||
help="One or more plain-text prompts to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="results/higgs_v3_wavs",
|
||||
help="Directory to write per-prompt WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="Cap on Stage-0 codec frames per request.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of the reference audio. Optional but improves fidelity.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _slugify(text: str) -> str:
|
||||
slug = re.sub(r"\s+", "_", text.strip().lower())
|
||||
slug = re.sub(r"[^a-z0-9_]+", "", slug)
|
||||
return slug[:48] or "prompt"
|
||||
|
||||
|
||||
def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
|
||||
if isinstance(audio, list):
|
||||
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
||||
if not valid:
|
||||
raise ValueError("audio list is empty")
|
||||
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
|
||||
arr = pcm.numpy()
|
||||
if arr.dtype.kind == "f":
|
||||
arr = np.clip(arr, -1.0, 1.0)
|
||||
arr = (arr * 32767.0).astype(np.int16)
|
||||
else:
|
||||
arr = arr.astype(np.int16)
|
||||
return arr
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
|
||||
HiggsAudioV3TokenizerAdapter,
|
||||
apply_delay_pattern,
|
||||
encode_reference_audio,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
|
||||
|
||||
# Load and encode reference audio once (if voice cloning)
|
||||
ref_codes_delayed: torch.Tensor | None = None
|
||||
if args.ref_audio is not None:
|
||||
ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
|
||||
if ref_wav.ndim == 2:
|
||||
ref_wav = ref_wav.mean(axis=1)
|
||||
ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
|
||||
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
|
||||
print(f"Reference : {args.ref_audio}")
|
||||
print(f"Ref codes : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
|
||||
if args.ref_text:
|
||||
print(f"Ref text : {args.ref_text}")
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Prompts : {len(args.texts)}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")
|
||||
|
||||
total_elapsed = 0.0
|
||||
total_dur = 0.0
|
||||
for text in args.texts:
|
||||
if ref_codes_delayed is not None:
|
||||
prompt_ids = adapter.build_prompt(
|
||||
text,
|
||||
num_ref_tokens=int(ref_codes_delayed.shape[0]),
|
||||
reference_text=args.ref_text,
|
||||
)
|
||||
prompt = {
|
||||
"prompt_token_ids": prompt_ids,
|
||||
"additional_information": {
|
||||
"audio_input_ids": ref_codes_delayed.to(torch.long),
|
||||
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
|
||||
},
|
||||
}
|
||||
else:
|
||||
prompt_ids = adapter.build_prompt(text)
|
||||
prompt = {"prompt_token_ids": prompt_ids}
|
||||
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([prompt])
|
||||
elapsed = time.perf_counter() - t_start
|
||||
total_elapsed += elapsed
|
||||
|
||||
mm = outputs[0].outputs[0].multimodal_output
|
||||
pcm = _extract_pcm(mm)
|
||||
slug = _slugify(text)
|
||||
out_path = output_dir / f"{slug}.wav"
|
||||
sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
|
||||
dur = pcm.numel() / SAMPLE_RATE
|
||||
total_dur += dur
|
||||
print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}")
|
||||
|
||||
rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
|
||||
print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -51,18 +51,13 @@ def build_payload(
|
||||
response_format: str = "wav",
|
||||
stream: bool = False,
|
||||
) -> dict:
|
||||
"""Build the /v1/audio/speech request payload.
|
||||
|
||||
Fish Speech has no built-in speakers; omit ``voice`` so the model
|
||||
generates with a random timbre (controllable via text tags like
|
||||
``[cool]``, ``[warm]``, etc.). When ``ref_audio`` is supplied we
|
||||
enter voice-cloning mode instead.
|
||||
"""
|
||||
"""Build the /v1/audio/speech request payload."""
|
||||
if not text or not text.strip():
|
||||
raise gr.Error("Please enter text to synthesize.")
|
||||
|
||||
payload: dict = {
|
||||
"input": text.strip(),
|
||||
"voice": "default",
|
||||
"response_format": "pcm" if stream else response_format,
|
||||
"stream": stream,
|
||||
}
|
||||
@@ -198,8 +193,7 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
|
||||
|
||||
with gr.Accordion("Voice Cloning (optional)", open=False):
|
||||
gr.Markdown(
|
||||
"Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice. "
|
||||
"Leave empty for random timbre; use text tags like `[cool]`, `[warm]`, `[happy]` to steer the output."
|
||||
"Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice."
|
||||
)
|
||||
ref_audio = gr.Audio(
|
||||
label="Reference Audio",
|
||||
|
||||
@@ -39,6 +39,7 @@ def run_tts(args) -> None:
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"voice": "default",
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
|
||||
|
||||
@@ -1,40 +0,0 @@
|
||||
# Higgs-Audio V3 Online Serving
|
||||
|
||||
## Start the server
|
||||
|
||||
```bash
|
||||
# Default: GPU 0, port 8095
|
||||
./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
|
||||
|
||||
# Custom GPU / port
|
||||
PORT=8096 GPUS=0,1 ./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
|
||||
```
|
||||
|
||||
## Plain text TTS
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--output-dir /tmp/higgs_v3_batch \
|
||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
||||
```
|
||||
|
||||
## Voice clone
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--output-dir /tmp/higgs_v3_clone \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "transcript of the reference clip" \
|
||||
--prompts "Text to synthesize in the cloned voice."
|
||||
```
|
||||
|
||||
## curl example
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8095/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "higgs_audio_v3", "input": "Hello world."}' \
|
||||
--output hello.wav
|
||||
```
|
||||
@@ -1,115 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Batch client for the higgs-audio v3 online server.
|
||||
|
||||
Sends prompts to ``/v1/audio/speech`` and saves the returned WAV files.
|
||||
|
||||
Usage (plain text -> speech):
|
||||
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--output-dir /tmp/higgs_v3_batch \
|
||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
||||
|
||||
Usage (voice clone):
|
||||
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--output-dir /tmp/higgs_v3_clone \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "the transcript of the reference clip" \
|
||||
--prompts "Hello world."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
DEFAULT_PROMPTS = (
|
||||
"Hello world.",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
"Today is a beautiful day for a walk in the park.",
|
||||
"Innovation distinguishes between a leader and a follower.",
|
||||
)
|
||||
|
||||
|
||||
def _slug(text: str) -> str:
|
||||
import re
|
||||
|
||||
s = re.sub(r"\s+", "_", text.strip().lower())
|
||||
return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--base-url", default="http://localhost:8095")
|
||||
parser.add_argument("--model", default="higgs_audio_v3")
|
||||
parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
|
||||
parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_v3_batch"))
|
||||
parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=2048)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--timeout-s", type=float, default=120.0)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Reference clip for voice clone (WAV/FLAC/MP3 path). Pair with --ref-text.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of the reference clip. Optional but improves fidelity.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
ref_audio_data_url: str | None = None
|
||||
if args.ref_audio is not None:
|
||||
if not args.ref_audio.exists():
|
||||
print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
|
||||
return 2
|
||||
mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
|
||||
ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
|
||||
ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
|
||||
|
||||
try:
|
||||
import httpx
|
||||
except ImportError:
|
||||
print("this client needs `httpx`. Install with `pip install httpx`.", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
url = args.base_url.rstrip("/") + "/v1/audio/speech"
|
||||
failures = 0
|
||||
with httpx.Client(timeout=args.timeout_s) as client:
|
||||
for prompt in args.prompts:
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": prompt,
|
||||
"response_format": args.format,
|
||||
"max_new_tokens": args.max_new_tokens,
|
||||
"seed": args.seed,
|
||||
}
|
||||
if ref_audio_data_url is not None:
|
||||
payload["ref_audio"] = ref_audio_data_url
|
||||
if args.ref_text:
|
||||
payload["ref_text"] = args.ref_text
|
||||
resp = client.post(url, json=payload)
|
||||
if resp.status_code != 200:
|
||||
print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
|
||||
failures += 1
|
||||
continue
|
||||
suffix = ".wav" if args.format == "wav" else ".pcm"
|
||||
out = args.output_dir / f"{_slug(prompt)}{suffix}"
|
||||
out.write_bytes(resp.content)
|
||||
print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
|
||||
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,32 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for higgs-audio v3.
|
||||
#
|
||||
# Supports plain text TTS and voice cloning via /v1/audio/speech.
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh # default port 8095, GPU 0
|
||||
# PORT=8096 GPUS=0,1 ./run_server.sh
|
||||
# MODEL=/path/to/local/checkpoint ./run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
MODEL="${MODEL:-bosonai/higgs-audio-v3-tts-4b}"
|
||||
PORT="${PORT:-8095}"
|
||||
GPUS="${GPUS:-0}"
|
||||
GPU_UTIL="${GPU_UTIL:-0.6}"
|
||||
|
||||
echo "Starting higgs-audio v3 server"
|
||||
echo " MODEL=$MODEL"
|
||||
echo " PORT=$PORT"
|
||||
echo " CUDA_VISIBLE_DEVICES=$GPUS"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="$GPUS" \
|
||||
VLLM_USE_DEEP_GEMM=0 \
|
||||
VLLM_MOE_USE_DEEP_GEMM=0 \
|
||||
vllm-omni serve "$MODEL" \
|
||||
--deploy-config vllm_omni/deploy/higgs_multimodal_qwen3.yaml \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--gpu-memory-utilization "$GPU_UTIL" \
|
||||
--trust-remote-code \
|
||||
--omni
|
||||
@@ -1,247 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Verify a ModelOpt FP8 diffusers checkpoint exported by
|
||||
quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
|
||||
|
||||
Three checks:
|
||||
A. transformer/config.json has a sane quantization_config block.
|
||||
B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
|
||||
C. transformer disk size is materially smaller than a BF16 baseline.
|
||||
|
||||
Example:
|
||||
python examples/quantization/check_modelopt_fp8_export.py \\
|
||||
--output ./hv15-480p-modelopt-fp8
|
||||
|
||||
# Optional: compare disk size against a local or HF BF16 baseline.
|
||||
python examples/quantization/check_modelopt_fp8_export.py \\
|
||||
--output ./hv15-480p-modelopt-fp8 \\
|
||||
--baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
SUPPORTED_ALGOS = ("FP8",)
|
||||
|
||||
|
||||
def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
|
||||
"""Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
|
||||
cfg_path = transformer_dir / "config.json"
|
||||
if not cfg_path.exists():
|
||||
print(f"[FAIL] {cfg_path} missing.")
|
||||
return 1, None
|
||||
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
|
||||
qc = cfg.get("quantization_config")
|
||||
if not isinstance(qc, dict):
|
||||
print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
|
||||
return 1, None
|
||||
|
||||
print(f"[A] quantization_config from {cfg_path}:")
|
||||
print(json.dumps(qc, indent=2))
|
||||
|
||||
quant_algo = qc.get("quant_algo")
|
||||
issues = []
|
||||
if qc.get("quant_method") != "modelopt":
|
||||
issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
|
||||
if quant_algo not in SUPPORTED_ALGOS:
|
||||
issues.append(
|
||||
f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
|
||||
)
|
||||
|
||||
if issues:
|
||||
print("[A] WARN — config looks incomplete:")
|
||||
for issue in issues:
|
||||
print(f" - {issue}")
|
||||
return 2, quant_algo
|
||||
print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
|
||||
return 0, quant_algo
|
||||
|
||||
|
||||
def _read_safetensors_header(path: Path) -> dict:
|
||||
"""Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
|
||||
|
||||
Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
|
||||
Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
|
||||
"""
|
||||
import struct
|
||||
|
||||
with open(path, "rb") as f:
|
||||
header_len = struct.unpack("<Q", f.read(8))[0]
|
||||
header = json.loads(f.read(header_len))
|
||||
header.pop("__metadata__", None)
|
||||
return header
|
||||
|
||||
|
||||
def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
|
||||
"""Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
|
||||
|
||||
ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
|
||||
interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
|
||||
and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
|
||||
"""
|
||||
if not weight_scale_shapes:
|
||||
return "no weight_scale tensors found"
|
||||
|
||||
def meaningful_dims(shape: list[int]) -> int:
|
||||
return sum(1 for d in shape if d > 1)
|
||||
|
||||
per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
|
||||
per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
|
||||
per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
|
||||
total = len(weight_scale_shapes)
|
||||
if per_tensor == total:
|
||||
return "per-tensor (all scalar scales)"
|
||||
if per_channel == total:
|
||||
return "per-channel (1 meaningful dim)"
|
||||
if per_block == total:
|
||||
return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
|
||||
return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
|
||||
|
||||
|
||||
def _check_safetensors(transformer_dir: Path) -> int:
|
||||
"""Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
|
||||
files = sorted(transformer_dir.glob("*.safetensors"))
|
||||
if not files:
|
||||
print(f"[FAIL] No *.safetensors in {transformer_dir}.")
|
||||
return 1
|
||||
|
||||
header_dtype_counts: Counter[str] = Counter()
|
||||
sample_quant_weight_keys: list[str] = []
|
||||
sample_scale_keys: list[str] = []
|
||||
weight_scale_shapes: list[list[int]] = []
|
||||
sample_weight_scale_entries: list[tuple[str, list[int]]] = []
|
||||
|
||||
for f in files:
|
||||
try:
|
||||
header = _read_safetensors_header(f)
|
||||
except Exception as exc:
|
||||
print(f"[B] WARN — could not parse header of {f}: {exc}")
|
||||
continue
|
||||
for k, info in header.items():
|
||||
dtype = info.get("dtype", "?")
|
||||
header_dtype_counts[dtype] += 1
|
||||
# FP8 stores weights as F8_E4M3 directly.
|
||||
if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
|
||||
sample_quant_weight_keys.append(k)
|
||||
if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
|
||||
sample_scale_keys.append(k)
|
||||
if k.endswith(".weight_scale"):
|
||||
weight_scale_shapes.append(info.get("shape", []))
|
||||
if len(sample_weight_scale_entries) < 5:
|
||||
sample_weight_scale_entries.append((k, info.get("shape", [])))
|
||||
|
||||
print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
|
||||
for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
|
||||
marker = " <-- FP8" if dtype.startswith("F8") else ""
|
||||
print(f" {dtype:10s} {count:>6d}{marker}")
|
||||
|
||||
quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
|
||||
if quant_count == 0:
|
||||
print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
|
||||
return 1
|
||||
|
||||
print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
|
||||
if sample_quant_weight_keys:
|
||||
print(f" sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
|
||||
if sample_scale_keys:
|
||||
print(f" sample scale tensors: {sample_scale_keys[:3]}")
|
||||
print(" (Note: torch's get_tensor() may return these as bf16 views on some versions —")
|
||||
print(" irrelevant; vLLM's loader uses native FP8 ops.)")
|
||||
|
||||
# Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
|
||||
print(f"\n weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
|
||||
for key, shape in sample_weight_scale_entries[:3]:
|
||||
print(f" {key}: shape {shape}")
|
||||
return 0
|
||||
|
||||
|
||||
def _disk_size_gib(p: Path) -> float:
|
||||
return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
|
||||
|
||||
|
||||
def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
|
||||
"""Returns 0 always (informational only)."""
|
||||
quant_size = _disk_size_gib(transformer_dir)
|
||||
print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
|
||||
|
||||
if baseline is None:
|
||||
print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
|
||||
return 0
|
||||
|
||||
baseline_path = Path(baseline)
|
||||
if not baseline_path.exists():
|
||||
# Try HF download.
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except ImportError:
|
||||
print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
|
||||
return 0
|
||||
print(f" Downloading baseline transformer from HF: {baseline}")
|
||||
baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
|
||||
|
||||
bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
|
||||
bf16_size = _disk_size_gib(bf16_dir)
|
||||
if bf16_size == 0:
|
||||
print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
|
||||
return 0
|
||||
|
||||
# Expected reduction for FP8: ~50%.
|
||||
min_reduction = 30
|
||||
reduction = (1 - quant_size / bf16_size) * 100
|
||||
print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
|
||||
print(f"[C] Disk reduction: {reduction:.1f}% (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
|
||||
if reduction < min_reduction:
|
||||
print(
|
||||
f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
|
||||
p.add_argument(
|
||||
"--baseline",
|
||||
default=None,
|
||||
help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
|
||||
)
|
||||
args = p.parse_args()
|
||||
|
||||
out_root = Path(args.output).expanduser().resolve()
|
||||
transformer_dir = out_root / "transformer"
|
||||
if not transformer_dir.exists():
|
||||
print(f"[FAIL] {transformer_dir} does not exist.")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Checking: {out_root}\n")
|
||||
|
||||
fail = 0
|
||||
config_status, _quant_algo = _check_config(transformer_dir)
|
||||
fail |= config_status
|
||||
fail |= _check_safetensors(transformer_dir)
|
||||
_check_size_vs_baseline(transformer_dir, args.baseline)
|
||||
|
||||
print()
|
||||
if fail == 0:
|
||||
print("=" * 60)
|
||||
print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
|
||||
elif fail == 1:
|
||||
print("=" * 60)
|
||||
print("FAILURES detected — calibration may need to be re-run.")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("=" * 60)
|
||||
print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,448 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
|
||||
|
||||
Calibrates the DiT transformer using a small video prompt set and exports a
|
||||
diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
|
||||
The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
|
||||
(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
|
||||
|
||||
Layers kept full precision match the #2728 / #2795 pattern: modulation,
|
||||
AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
|
||||
and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
|
||||
empirically degrades under FP8 (see #2920 ablation).
|
||||
|
||||
Example:
|
||||
python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
|
||||
--output ./hv15-480p-modelopt-fp8 \\
|
||||
--overwrite
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
DEFAULT_PROMPTS = [
|
||||
"A dog running across a field of golden wheat.",
|
||||
"An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
|
||||
"A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
|
||||
"A crackling campfire at night under a starry sky, sparks rising into the dark.",
|
||||
"An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
|
||||
"A close-up of a blooming rose covered in morning dew, soft natural light.",
|
||||
"A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
|
||||
"A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
|
||||
]
|
||||
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
|
||||
p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
|
||||
p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
|
||||
p.add_argument("--height", type=int, default=480)
|
||||
p.add_argument("--width", type=int, default=832)
|
||||
p.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=33,
|
||||
help="Frames per calibration sample. 33 matches the typical short benchmark.",
|
||||
)
|
||||
p.add_argument("--guidance-scale", type=float, default=6.0)
|
||||
p.add_argument(
|
||||
"--calib-steps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
|
||||
)
|
||||
p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument(
|
||||
"--prompt",
|
||||
action="append",
|
||||
default=[],
|
||||
help="Custom calibration prompt. Repeat to provide multiple.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--quantize-mha",
|
||||
action="store_true",
|
||||
help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--weight-block-size",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
|
||||
"Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
|
||||
"negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
|
||||
)
|
||||
p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
|
||||
return p
|
||||
|
||||
|
||||
def _parse_block_size(spec: str | None) -> list[int] | None:
|
||||
if spec is None:
|
||||
return None
|
||||
parts = [int(x) for x in spec.split(",") if x.strip()]
|
||||
if len(parts) != 2:
|
||||
raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
|
||||
return parts
|
||||
|
||||
|
||||
def _require_modelopt() -> Any:
|
||||
try:
|
||||
import modelopt.torch.quantization as mtq
|
||||
except ModuleNotFoundError as exc:
|
||||
raise SystemExit(
|
||||
"NVIDIA ModelOpt is not installed. Install with:\n"
|
||||
" pip install 'nvidia-modelopt[all]'\n"
|
||||
f"Original error: {exc}"
|
||||
) from exc
|
||||
return mtq
|
||||
|
||||
|
||||
def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
|
||||
model_path = args.model
|
||||
output_dir = Path(args.output).expanduser().resolve()
|
||||
if output_dir.exists():
|
||||
if not args.overwrite:
|
||||
raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
|
||||
shutil.rmtree(output_dir)
|
||||
return model_path, output_dir
|
||||
|
||||
|
||||
def _select_dtype(name: str) -> torch.dtype:
|
||||
return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
|
||||
|
||||
|
||||
def _build_prompts(args: argparse.Namespace) -> list[str]:
|
||||
prompts = args.prompt or DEFAULT_PROMPTS
|
||||
if args.calib_size <= 0:
|
||||
raise SystemExit("--calib-size must be positive.")
|
||||
if len(prompts) < args.calib_size:
|
||||
repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
|
||||
prompts = (prompts * repeats)[: args.calib_size]
|
||||
return prompts[: args.calib_size]
|
||||
|
||||
|
||||
# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
|
||||
# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
|
||||
# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
|
||||
# - proj_out: final output projection
|
||||
# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
|
||||
def _filter_func_hv15(name: str) -> bool:
|
||||
pattern = re.compile(
|
||||
r"(proj_out.*|"
|
||||
r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
|
||||
r"time_embed|cond_type_embed|"
|
||||
r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
|
||||
r"token_refiner).*)"
|
||||
)
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def _mha_filter_func(name: str) -> bool:
|
||||
pattern = re.compile(
|
||||
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
|
||||
)
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
|
||||
if not hasattr(mtq, "disable_quantizer"):
|
||||
return
|
||||
mtq.disable_quantizer(backbone, _filter_func_hv15)
|
||||
if not quantize_mha:
|
||||
mtq.disable_quantizer(backbone, _mha_filter_func)
|
||||
|
||||
|
||||
def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
|
||||
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
|
||||
if hasattr(pipe, "set_progress_bar_config"):
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
|
||||
generator = torch.Generator(device="cuda")
|
||||
|
||||
# Try to set guidance on the pipeline's guider object up front (modern
|
||||
# diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
|
||||
# back silently — calibration uses whatever default the pipeline ships with.
|
||||
guider = getattr(pipe, "guider", None)
|
||||
if guider is not None and hasattr(guider, "guidance_scale"):
|
||||
try:
|
||||
guider.guidance_scale = args.guidance_scale
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
base_kwargs = dict(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.calib_steps,
|
||||
output_type="latent",
|
||||
)
|
||||
|
||||
def forward_loop(*_unused_args, **_unused_kwargs) -> None:
|
||||
with torch.inference_mode():
|
||||
for idx, prompt in enumerate(prompts):
|
||||
generator.manual_seed(args.seed + idx)
|
||||
# Try with guidance_scale first; fall back without on TypeError
|
||||
# for pipelines (like HV-1.5) that take CFG via guider config.
|
||||
try:
|
||||
pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
|
||||
except TypeError as exc:
|
||||
if "guidance_scale" not in str(exc):
|
||||
raise
|
||||
pipe(prompt=prompt, generator=generator, **base_kwargs)
|
||||
|
||||
return forward_loop
|
||||
|
||||
|
||||
def _summarize_export(output_dir: Path) -> None:
|
||||
cfg_path = output_dir / "transformer" / "config.json"
|
||||
if not cfg_path.exists():
|
||||
print(f"[warn] {cfg_path} missing.", file=sys.stderr)
|
||||
return
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
qc = cfg.get("quantization_config")
|
||||
if not isinstance(qc, dict):
|
||||
print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
|
||||
return
|
||||
print("Export summary:")
|
||||
print(f" quant_method: {qc.get('quant_method')}")
|
||||
print(f" quant_algo: {qc.get('quant_algo')}")
|
||||
producer = qc.get("producer")
|
||||
if isinstance(producer, dict):
|
||||
print(f" producer: {producer.get('name')} {producer.get('version')}")
|
||||
print(f" config path: {cfg_path}")
|
||||
|
||||
|
||||
def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
|
||||
"""Convert in-memory weights of quantized modules to actual FP8 storage.
|
||||
|
||||
`export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
|
||||
in ModelOpt's recognized-model registry), so we must call the per-weight
|
||||
export helper ourselves. Same workaround as the HunyuanImage-3 calibration
|
||||
helper.
|
||||
"""
|
||||
from modelopt.torch.export.quant_utils import (
|
||||
QUANTIZATION_NONE,
|
||||
get_quantization_format,
|
||||
quantizer_attr_names,
|
||||
weight_attr_names,
|
||||
)
|
||||
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
|
||||
|
||||
exported = 0
|
||||
for name, module in backbone.named_modules():
|
||||
try:
|
||||
quantization_format = get_quantization_format(module)
|
||||
except Exception as exc:
|
||||
print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
|
||||
continue
|
||||
if quantization_format == QUANTIZATION_NONE:
|
||||
continue
|
||||
for weight_name in weight_attr_names(module):
|
||||
quantizer_attrs = quantizer_attr_names(weight_name)
|
||||
weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
|
||||
if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
|
||||
continue
|
||||
_export_quantized_weight(module, dtype, weight_name)
|
||||
exported += 1
|
||||
return exported
|
||||
|
||||
|
||||
def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
|
||||
"""Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
|
||||
|
||||
Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
|
||||
`weight_block_size` is given, advertise block-wise weight quantization in
|
||||
the saved metadata (so consumers know to expect multi-element scale tensors).
|
||||
"""
|
||||
weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
|
||||
if weight_block_size is not None:
|
||||
weights_cfg["strategy"] = "block"
|
||||
weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
|
||||
return {
|
||||
"config_groups": {
|
||||
"group_0": {
|
||||
"input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
|
||||
"weights": weights_cfg,
|
||||
"targets": ["Linear"],
|
||||
}
|
||||
},
|
||||
"ignore": [
|
||||
"context_embedder*",
|
||||
"context_embedder_2*",
|
||||
"cond_type_embed*",
|
||||
"image_embedder*",
|
||||
"norm1.linear*",
|
||||
"norm1_context.linear*",
|
||||
"norm2*",
|
||||
"norm2_context*",
|
||||
"norm_out*",
|
||||
"proj_out*",
|
||||
"time_embed*",
|
||||
"token_refiner*",
|
||||
"x_embedder*",
|
||||
],
|
||||
"producer": {"name": "modelopt"},
|
||||
"quant_algo": "FP8",
|
||||
"quant_method": "modelopt",
|
||||
}
|
||||
|
||||
|
||||
def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
|
||||
"""Inject quant_algo: FP8 + config_groups into transformer/config.json so
|
||||
vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
|
||||
cfg_path = output_dir / "transformer" / "config.json"
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
|
||||
new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
|
||||
existing = cfg.get("quantization_config")
|
||||
if isinstance(existing, dict):
|
||||
producer = existing.get("producer")
|
||||
if isinstance(producer, dict):
|
||||
new_qc["producer"] = producer
|
||||
|
||||
cfg["quantization_config"] = new_qc
|
||||
with cfg_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(cfg, f, indent=2)
|
||||
|
||||
|
||||
def _save_pipeline_with_fp8_transformer(
|
||||
pipe: DiffusionPipeline,
|
||||
model_path: str,
|
||||
output_dir: Path,
|
||||
max_shard_size: str = "5GB",
|
||||
) -> None:
|
||||
"""Save the pipeline with the (now FP8) transformer.
|
||||
|
||||
Copies the source directory verbatim except for `transformer/`, then
|
||||
saves the transformer with quantizers hidden so the state dict contains
|
||||
only the FP8 weights + scale tensors.
|
||||
"""
|
||||
from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
|
||||
|
||||
src = Path(model_path)
|
||||
if not src.exists():
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
src = Path(snapshot_download(model_path))
|
||||
|
||||
if output_dir.exists():
|
||||
shutil.rmtree(output_dir)
|
||||
shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
|
||||
|
||||
transformer_out = output_dir / "transformer"
|
||||
# `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
|
||||
# nn.Module (transformer), not the diffusers Pipeline wrapper.
|
||||
with hide_quantizers_from_state_dict(pipe.transformer):
|
||||
pipe.transformer.save_pretrained(
|
||||
str(transformer_out),
|
||||
safe_serialization=True,
|
||||
max_shard_size=max_shard_size,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _build_parser().parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
|
||||
|
||||
mtq = _require_modelopt()
|
||||
model_path, output_dir = _ensure_paths(args)
|
||||
dtype = _select_dtype(args.dtype)
|
||||
prompts = _build_prompts(args)
|
||||
|
||||
print("Quantization plan:")
|
||||
weight_block_size = _parse_block_size(args.weight_block_size)
|
||||
|
||||
print(f" input: {args.model}")
|
||||
print(f" output: {output_dir}")
|
||||
print(f" dtype: {dtype}")
|
||||
print(f" height/width: {args.height}x{args.width}")
|
||||
print(f" num_frames: {args.num_frames}")
|
||||
print(f" calib_size: {len(prompts)}")
|
||||
print(f" calib_steps: {args.calib_steps}")
|
||||
print(f" quantize_mha: {args.quantize_mha}")
|
||||
print(
|
||||
f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
|
||||
)
|
||||
|
||||
pipe = _load_pipeline(model_path, dtype)
|
||||
backbone = pipe.transformer
|
||||
|
||||
quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
|
||||
if weight_block_size is not None:
|
||||
# Switch from per-tensor (default) to block-wise weight quantization.
|
||||
# ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
|
||||
quant_config["quant_cfg"]["*weight_quantizer"] = {
|
||||
"num_bits": (4, 3), # E4M3 (FP8 weights, same as default)
|
||||
"block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
|
||||
}
|
||||
print(
|
||||
f" -> overriding weight quantizer with block_sizes={weight_block_size} "
|
||||
f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
|
||||
)
|
||||
|
||||
forward_loop = _build_forward_loop(pipe, args, prompts)
|
||||
quantized = mtq.quantize(backbone, quant_config, forward_loop)
|
||||
if quantized is not None:
|
||||
pipe.transformer = quantized
|
||||
backbone = quantized
|
||||
|
||||
_disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
|
||||
|
||||
print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
|
||||
print("so we have to call the per-weight export helper ourselves)...")
|
||||
exported = _force_export_quantized_weights(backbone, dtype)
|
||||
print(f" -> {exported} weights converted to FP8 in memory")
|
||||
if exported == 0:
|
||||
raise SystemExit(
|
||||
"No quantized weights were exported. Calibration may have skipped every layer "
|
||||
"(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
|
||||
"weight quantizers."
|
||||
)
|
||||
|
||||
print("\nSaving pipeline with FP8 transformer...")
|
||||
_save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
|
||||
_patch_quant_config(output_dir, weight_block_size=weight_block_size)
|
||||
print(f"Saved to: {output_dir}")
|
||||
_summarize_export(output_dir)
|
||||
|
||||
print("\nNext: validate the checkpoint with vllm-omni:")
|
||||
print(
|
||||
" python examples/offline_inference/text_to_video/text_to_video.py \\\n"
|
||||
f" --model {output_dir} \\\n"
|
||||
" --quantization fp8 \\\n"
|
||||
" --prompt 'A dog running across a field of golden wheat.' \\\n"
|
||||
f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
|
||||
" --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
|
||||
" --output outputs/hv15_modelopt_fp8.mp4 \\\n"
|
||||
" --enforce-eager"
|
||||
)
|
||||
print(
|
||||
"\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
|
||||
"checkpoint's config.json has modelopt metadata.)"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,471 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
|
||||
|
||||
Calibrates the DiT transformer using a small video prompt set and exports a
|
||||
diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
|
||||
The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
|
||||
(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
|
||||
|
||||
Layers kept full precision match the #2728 / #2795 pattern: condition embedder
|
||||
(time/text/image), patch embedding, modulation (scale_shift_table), final
|
||||
norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
|
||||
are quantized — static calibration handles the numerics that online FP8
|
||||
couldn't (see #2920 ablation).
|
||||
|
||||
Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
|
||||
fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
|
||||
|
||||
Example:
|
||||
python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
|
||||
--output ./wan22-ti2v-modelopt-fp8 \\
|
||||
--overwrite
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
DEFAULT_PROMPTS = [
|
||||
"A dog running across a field of golden wheat.",
|
||||
"An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
|
||||
"A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
|
||||
"A crackling campfire at night under a starry sky, sparks rising into the dark.",
|
||||
"An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
|
||||
"A close-up of a blooming rose covered in morning dew, soft natural light.",
|
||||
"A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
|
||||
"A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
|
||||
]
|
||||
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
|
||||
p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
|
||||
p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
|
||||
p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
|
||||
p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
|
||||
p.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=49,
|
||||
help="Frames per calibration sample. 49 matches the typical short benchmark; "
|
||||
"use 17 to reduce memory pressure during calibration.",
|
||||
)
|
||||
p.add_argument("--guidance-scale", type=float, default=5.0)
|
||||
p.add_argument(
|
||||
"--calib-steps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
|
||||
)
|
||||
p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument(
|
||||
"--prompt",
|
||||
action="append",
|
||||
default=[],
|
||||
help="Custom calibration prompt. Repeat to provide multiple.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--quantize-mha",
|
||||
action="store_true",
|
||||
help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
|
||||
"sequences amplified FP8 drift in the online ablation (see #2920).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--weight-block-size",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
|
||||
"Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
|
||||
"for the HV-1.5 investigation status before relying on this.",
|
||||
)
|
||||
p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
|
||||
return p
|
||||
|
||||
|
||||
def _parse_block_size(spec: str | None) -> list[int] | None:
|
||||
if spec is None:
|
||||
return None
|
||||
parts = [int(x) for x in spec.split(",") if x.strip()]
|
||||
if len(parts) != 2:
|
||||
raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
|
||||
return parts
|
||||
|
||||
|
||||
def _require_modelopt() -> Any:
|
||||
try:
|
||||
import modelopt.torch.quantization as mtq
|
||||
except ModuleNotFoundError as exc:
|
||||
raise SystemExit(
|
||||
"NVIDIA ModelOpt is not installed. Install with:\n"
|
||||
" pip install 'nvidia-modelopt[all]'\n"
|
||||
f"Original error: {exc}"
|
||||
) from exc
|
||||
return mtq
|
||||
|
||||
|
||||
def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
|
||||
model_path = args.model
|
||||
output_dir = Path(args.output).expanduser().resolve()
|
||||
if output_dir.exists():
|
||||
if not args.overwrite:
|
||||
raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
|
||||
shutil.rmtree(output_dir)
|
||||
return model_path, output_dir
|
||||
|
||||
|
||||
def _select_dtype(name: str) -> torch.dtype:
|
||||
return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
|
||||
|
||||
|
||||
def _build_prompts(args: argparse.Namespace) -> list[str]:
|
||||
prompts = args.prompt or DEFAULT_PROMPTS
|
||||
if args.calib_size <= 0:
|
||||
raise SystemExit("--calib-size must be positive.")
|
||||
if len(prompts) < args.calib_size:
|
||||
repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
|
||||
prompts = (prompts * repeats)[: args.calib_size]
|
||||
return prompts[: args.calib_size]
|
||||
|
||||
|
||||
# Layers to KEEP at full precision. Wan2.2's module naming:
|
||||
# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
|
||||
# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
|
||||
# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
|
||||
# - norm_out: AdaLayerNorm final
|
||||
# - proj_out: final nn.Linear
|
||||
# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
|
||||
def _filter_func_wan22(name: str) -> bool:
|
||||
pattern = re.compile(
|
||||
r"(proj_out.*|"
|
||||
r".*(condition_embedder|patch_embedding|"
|
||||
r"norm_out|scale_shift_table|"
|
||||
r"timestep_proj_prepare|output_scale_shift_prepare).*)"
|
||||
)
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def _mha_filter_func(name: str) -> bool:
|
||||
pattern = re.compile(
|
||||
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
|
||||
)
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
|
||||
if not hasattr(mtq, "disable_quantizer"):
|
||||
return
|
||||
mtq.disable_quantizer(backbone, _filter_func_wan22)
|
||||
if not quantize_mha:
|
||||
mtq.disable_quantizer(backbone, _mha_filter_func)
|
||||
|
||||
|
||||
def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
|
||||
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
|
||||
if hasattr(pipe, "set_progress_bar_config"):
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
|
||||
generator = torch.Generator(device="cuda")
|
||||
|
||||
# Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
|
||||
guider = getattr(pipe, "guider", None)
|
||||
if guider is not None and hasattr(guider, "guidance_scale"):
|
||||
try:
|
||||
guider.guidance_scale = args.guidance_scale
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
base_kwargs = dict(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.calib_steps,
|
||||
output_type="latent",
|
||||
)
|
||||
|
||||
def forward_loop(*_unused_args, **_unused_kwargs) -> None:
|
||||
with torch.inference_mode():
|
||||
for idx, prompt in enumerate(prompts):
|
||||
generator.manual_seed(args.seed + idx)
|
||||
# Try with guidance_scale first; fall back without on TypeError
|
||||
# for pipelines that take CFG via guider config only.
|
||||
try:
|
||||
pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
|
||||
except TypeError as exc:
|
||||
if "guidance_scale" not in str(exc):
|
||||
raise
|
||||
pipe(prompt=prompt, generator=generator, **base_kwargs)
|
||||
|
||||
return forward_loop
|
||||
|
||||
|
||||
def _summarize_export(output_dir: Path) -> None:
|
||||
cfg_path = output_dir / "transformer" / "config.json"
|
||||
if not cfg_path.exists():
|
||||
print(f"[warn] {cfg_path} missing.", file=sys.stderr)
|
||||
return
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
qc = cfg.get("quantization_config")
|
||||
if not isinstance(qc, dict):
|
||||
print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
|
||||
return
|
||||
print("Export summary:")
|
||||
print(f" quant_method: {qc.get('quant_method')}")
|
||||
print(f" quant_algo: {qc.get('quant_algo')}")
|
||||
producer = qc.get("producer")
|
||||
if isinstance(producer, dict):
|
||||
print(f" producer: {producer.get('name')} {producer.get('version')}")
|
||||
print(f" config path: {cfg_path}")
|
||||
|
||||
|
||||
def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
|
||||
"""Convert in-memory weights of quantized modules to actual FP8 storage.
|
||||
|
||||
`export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
|
||||
in ModelOpt's recognized-model registry), so we must call the per-weight
|
||||
export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
|
||||
calibration helpers.
|
||||
"""
|
||||
from modelopt.torch.export.quant_utils import (
|
||||
QUANTIZATION_NONE,
|
||||
get_quantization_format,
|
||||
quantizer_attr_names,
|
||||
weight_attr_names,
|
||||
)
|
||||
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
|
||||
|
||||
exported = 0
|
||||
for name, module in backbone.named_modules():
|
||||
try:
|
||||
quantization_format = get_quantization_format(module)
|
||||
except Exception as exc:
|
||||
print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
|
||||
continue
|
||||
if quantization_format == QUANTIZATION_NONE:
|
||||
continue
|
||||
for weight_name in weight_attr_names(module):
|
||||
quantizer_attrs = quantizer_attr_names(weight_name)
|
||||
weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
|
||||
if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
|
||||
continue
|
||||
_export_quantized_weight(module, dtype, weight_name)
|
||||
exported += 1
|
||||
return exported
|
||||
|
||||
|
||||
def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
|
||||
"""Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
|
||||
weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
|
||||
if weight_block_size is not None:
|
||||
weights_cfg["strategy"] = "block"
|
||||
weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
|
||||
return {
|
||||
"config_groups": {
|
||||
"group_0": {
|
||||
"input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
|
||||
"weights": weights_cfg,
|
||||
"targets": ["Linear"],
|
||||
}
|
||||
},
|
||||
"ignore": [
|
||||
"condition_embedder*",
|
||||
"norm_out*",
|
||||
"output_scale_shift_prepare*",
|
||||
"patch_embedding*",
|
||||
"proj_out*",
|
||||
"scale_shift_table*",
|
||||
"timestep_proj_prepare*",
|
||||
],
|
||||
"producer": {"name": "modelopt"},
|
||||
"quant_algo": "FP8",
|
||||
"quant_method": "modelopt",
|
||||
}
|
||||
|
||||
|
||||
def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
|
||||
"""Inject quant_algo: FP8 + config_groups into each transformer's config.json
|
||||
so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
|
||||
|
||||
Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
|
||||
for sub in ("transformer", "transformer_2"):
|
||||
cfg_path = output_dir / sub / "config.json"
|
||||
if not cfg_path.exists():
|
||||
continue
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
|
||||
new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
|
||||
existing = cfg.get("quantization_config")
|
||||
if isinstance(existing, dict):
|
||||
producer = existing.get("producer")
|
||||
if isinstance(producer, dict):
|
||||
new_qc["producer"] = producer
|
||||
|
||||
cfg["quantization_config"] = new_qc
|
||||
with cfg_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(cfg, f, indent=2)
|
||||
|
||||
|
||||
def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
|
||||
"""Return [(attr_name, module)] for every DiT on the pipeline.
|
||||
|
||||
Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
|
||||
(high-noise). TI2V-5B has only ``transformer``.
|
||||
"""
|
||||
out = []
|
||||
for attr in ("transformer", "transformer_2"):
|
||||
mod = getattr(pipe, attr, None)
|
||||
if mod is not None:
|
||||
out.append((attr, mod))
|
||||
return out
|
||||
|
||||
|
||||
def _save_pipeline_with_fp8_transformer(
|
||||
pipe: DiffusionPipeline,
|
||||
model_path: str,
|
||||
output_dir: Path,
|
||||
max_shard_size: str = "5GB",
|
||||
) -> None:
|
||||
"""Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
|
||||
|
||||
Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
|
||||
as the single-DiT TI2V-5B.
|
||||
"""
|
||||
from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
|
||||
|
||||
src = Path(model_path)
|
||||
if not src.exists():
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
src = Path(snapshot_download(model_path))
|
||||
|
||||
if output_dir.exists():
|
||||
shutil.rmtree(output_dir)
|
||||
shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
|
||||
|
||||
for attr, backbone in _list_transformers(pipe):
|
||||
transformer_out = output_dir / attr
|
||||
# Pass the nn.Module (transformer), not the Pipeline wrapper.
|
||||
with hide_quantizers_from_state_dict(backbone):
|
||||
backbone.save_pretrained(
|
||||
str(transformer_out),
|
||||
safe_serialization=True,
|
||||
max_shard_size=max_shard_size,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _build_parser().parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
|
||||
|
||||
mtq = _require_modelopt()
|
||||
model_path, output_dir = _ensure_paths(args)
|
||||
dtype = _select_dtype(args.dtype)
|
||||
prompts = _build_prompts(args)
|
||||
weight_block_size = _parse_block_size(args.weight_block_size)
|
||||
|
||||
print("Quantization plan:")
|
||||
print(f" input: {args.model}")
|
||||
print(f" output: {output_dir}")
|
||||
print(f" dtype: {dtype}")
|
||||
print(f" height/width: {args.height}x{args.width}")
|
||||
print(f" num_frames: {args.num_frames}")
|
||||
print(f" calib_size: {len(prompts)}")
|
||||
print(f" calib_steps: {args.calib_steps}")
|
||||
print(f" quantize_mha: {args.quantize_mha}")
|
||||
print(
|
||||
f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
|
||||
)
|
||||
|
||||
pipe = _load_pipeline(model_path, dtype)
|
||||
transformers = _list_transformers(pipe)
|
||||
if not transformers:
|
||||
raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
|
||||
print(f" found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
|
||||
|
||||
quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
|
||||
if weight_block_size is not None:
|
||||
quant_config["quant_cfg"]["*weight_quantizer"] = {
|
||||
"num_bits": (4, 3),
|
||||
"block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
|
||||
}
|
||||
print(
|
||||
f" -> overriding weight quantizer with block_sizes={weight_block_size} "
|
||||
f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
|
||||
)
|
||||
|
||||
forward_loop = _build_forward_loop(pipe, args, prompts)
|
||||
|
||||
# Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
|
||||
# on every attached transformer (A14B: transformer + transformer_2) get
|
||||
# calibrated during the same passes.
|
||||
for attr, backbone in transformers:
|
||||
print(f"\nQuantizing {attr}...")
|
||||
quantized = mtq.quantize(backbone, quant_config, forward_loop)
|
||||
if quantized is not None:
|
||||
setattr(pipe, attr, quantized)
|
||||
_disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
|
||||
|
||||
# Export weights to FP8 only AFTER all transformers are calibrated. Exporting
|
||||
# inside the loop would convert transformer's weights to real FP8, and the
|
||||
# next transformer's forward_loop (full pipeline) would then fake-quantize an
|
||||
# already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
|
||||
print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
|
||||
print("so we have to call the per-weight export helper ourselves)...")
|
||||
total_exported = 0
|
||||
for attr, backbone in _list_transformers(pipe):
|
||||
exported = _force_export_quantized_weights(backbone, dtype)
|
||||
print(f" -> {exported} weights converted to FP8 in {attr}")
|
||||
total_exported += exported
|
||||
|
||||
if total_exported == 0:
|
||||
raise SystemExit(
|
||||
"No quantized weights were exported. Calibration may have skipped every layer "
|
||||
"(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
|
||||
"weight quantizers."
|
||||
)
|
||||
|
||||
print("\nSaving pipeline with FP8 transformer(s)...")
|
||||
_save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
|
||||
_patch_quant_config(output_dir, weight_block_size=weight_block_size)
|
||||
print(f"Saved to: {output_dir}")
|
||||
_summarize_export(output_dir)
|
||||
|
||||
print("\nNext: validate the checkpoint with vllm-omni:")
|
||||
print(
|
||||
" python examples/offline_inference/text_to_video/text_to_video.py \\\n"
|
||||
f" --model {output_dir} \\\n"
|
||||
" --quantization fp8 \\\n"
|
||||
" --prompt 'A dog running across a field of golden wheat.' \\\n"
|
||||
f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
|
||||
" --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
|
||||
" --output outputs/wan22_modelopt_fp8.mp4"
|
||||
)
|
||||
print(
|
||||
"\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
|
||||
"checkpoint's config.json has modelopt metadata.)"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -90,12 +90,6 @@ minicpmo = [
|
||||
"stepaudio2-minicpmo",
|
||||
]
|
||||
|
||||
# Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported
|
||||
# as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md.
|
||||
quack = [
|
||||
"quack-kernels>=0.3.11",
|
||||
]
|
||||
|
||||
docs = [
|
||||
"mkdocs>=1.5.0",
|
||||
"mkdocs-api-autonav",
|
||||
|
||||
@@ -1,110 +0,0 @@
|
||||
# Higgs-Audio V3 TTS
|
||||
|
||||
> Multilingual text-to-speech with voice cloning on 1×H100
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Boson AI
|
||||
- Model: `bosonai/higgs-audio-v3-tts-4b`
|
||||
- Task: Text-to-speech synthesis with optional voice cloning (100+ languages)
|
||||
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API; offline batch inference
|
||||
- Maintainer: @yuekaiz
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to serve `bosonai/higgs-audio-v3-tts-4b` for high-quality
|
||||
multilingual TTS. The model generates 24 kHz speech, supports zero-shot voice
|
||||
cloning from a reference clip, and handles 100+ languages with inline control
|
||||
tokens for emotion, style, and prosody. The architecture is a ~4B Qwen3
|
||||
backbone with fused multi-codebook embedding/head (8 codebooks × 1026 vocab,
|
||||
MusicGen-style delay pattern).
|
||||
|
||||
## References
|
||||
|
||||
- Model card: [bosonai/higgs-audio-v3-tts-4b](https://huggingface.co/bosonai/higgs-audio-v3-tts-4b)
|
||||
- Offline example: [`examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py`](../../examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py)
|
||||
- Online example: [`examples/online_serving/text_to_speech/higgs_audio_v3/`](../../examples/online_serving/text_to_speech/higgs_audio_v3/)
|
||||
- Benchmark results: see Performance section below
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1×H100 80GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.12+
|
||||
- CUDA: 12.x
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version or commit: `36e048fd` (branch `higgs-v3`)
|
||||
|
||||
#### Command
|
||||
|
||||
**Online serving:**
|
||||
|
||||
```bash
|
||||
vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
|
||||
--host 0.0.0.0 --port 8095 \
|
||||
--trust-remote-code --omni
|
||||
```
|
||||
|
||||
The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is
|
||||
loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`).
|
||||
|
||||
**Offline batch inference:**
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
|
||||
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
|
||||
--output-dir results/higgs_v3_wavs
|
||||
```
|
||||
|
||||
**Offline voice clone:**
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
|
||||
--texts "Text to synthesize in the cloned voice." \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference clip." \
|
||||
--output-dir results/higgs_v3_clone
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
Basic TTS via curl:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8095/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "bosonai/higgs-audio-v3-tts-4b",
|
||||
"input": "Hello, how are you?"
|
||||
}' --output hello.wav
|
||||
```
|
||||
|
||||
Voice clone via Python client:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--model bosonai/higgs-audio-v3-tts-4b \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference." \
|
||||
--prompts "Text to clone."
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: Stage 0 (talker, ~4B) uses ~60% GPU memory; Stage 1 (codec decoder) uses ~25%.
|
||||
- Key flags: `--trust-remote-code` and `--omni` are required.
|
||||
- Output: 24 kHz mono WAV.
|
||||
- Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity.
|
||||
- Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`).
|
||||
- `max_num_seqs=16` for both stages.
|
||||
- `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1).
|
||||
- Known limitations:
|
||||
- CUDA graph for stage 0 helps batch=1 (-13% RTF) but not batch>1 due to model-owned sampler running outside graph.
|
||||
- Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture).
|
||||
- Async streaming (chunk-based) not yet implemented; sync-only pipeline.
|
||||
+2
-3
@@ -36,8 +36,8 @@ recipes/
|
||||
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
|
||||
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
|
||||
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
|
||||
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
|
||||
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
|
||||
@@ -50,7 +50,6 @@ recipes/
|
||||
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
|
||||
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
|
||||
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
|
||||
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
|
||||
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
|
||||
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
|
||||
|
|
||||
|
||||
@@ -1,139 +0,0 @@
|
||||
# Stable Diffusion 3.5
|
||||
|
||||
> Text-to-image serving with Stability AI's SD 3.5 family
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Stability AI
|
||||
- Model: `stabilityai/stable-diffusion-3.5-medium`, `stabilityai/stable-diffusion-3.5-large`
|
||||
- Task: Text-to-image generation
|
||||
- Mode: Online serving with the OpenAI-compatible images API
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a lightweight, fast text-to-image model that
|
||||
fits on a single consumer or workstation GPU. SD 3.5-medium requires only
|
||||
~17 GiB VRAM and generates 1024×1024 images in ~4 seconds, making it one of
|
||||
the most accessible diffusion models in vLLM-Omni. SD 3.5-large produces
|
||||
noticeably sharper images at the cost of ~34 GiB VRAM and ~21 seconds per
|
||||
image.
|
||||
|
||||
## References
|
||||
|
||||
- Model (medium): <https://huggingface.co/stabilityai/stable-diffusion-3.5-medium>
|
||||
- Model (large): <https://huggingface.co/stabilityai/stable-diffusion-3.5-large>
|
||||
- Related example under `examples/`:
|
||||
[`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
|
||||
- User guide:
|
||||
[`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents two CUDA GPU configurations on a single RTX A6000 48 GB,
|
||||
one for each model variant. Extend it with more hardware sections as community
|
||||
validation lands.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-medium
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04
|
||||
- Python: 3.12
|
||||
- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve stabilityai/stable-diffusion-3.5-medium --omni --port 8091
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "stabilityai/stable-diffusion-3.5-medium",
|
||||
"prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
|
||||
"negative_prompt": "blurry, out of focus",
|
||||
"guidance_scale": 4.5,
|
||||
"n": 1,
|
||||
"size": "1024x1024"
|
||||
}' | python3 -c "
|
||||
import json, base64, sys
|
||||
data = json.load(sys.stdin)
|
||||
with open('output.png', 'wb') as f:
|
||||
f.write(base64.b64decode(data['data'][0]['b64_json']))
|
||||
print('saved output.png')
|
||||
"
|
||||
```
|
||||
|
||||
Or use the example client:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_image/openai_chat_client.py \
|
||||
--server http://localhost:8091 \
|
||||
--prompt "a cute cat sitting on a windowsill" \
|
||||
--output output.png
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: Model weights ~15.6 GiB, peak VRAM ~17.0 GiB (no CFG), ~20.9 GiB with `guidance_scale > 1` (CFG enabled).
|
||||
- Generation time: ~4 seconds without CFG, ~8 seconds with `guidance_scale=4.5`.
|
||||
- Model loading time: ~55 seconds (including weight download on first run).
|
||||
- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
|
||||
|
||||
### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-large
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04
|
||||
- Python: 3.12
|
||||
- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve stabilityai/stable-diffusion-3.5-large --omni --port 8091
|
||||
```
|
||||
|
||||
> **Note**: `stable-diffusion-3.5-large` is a gated model. Accept the license
|
||||
> on [Hugging Face](https://huggingface.co/stabilityai/stable-diffusion-3.5-large)
|
||||
> and set `HF_TOKEN` before running.
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "stabilityai/stable-diffusion-3.5-large",
|
||||
"prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
|
||||
"negative_prompt": "blurry, out of focus",
|
||||
"guidance_scale": 4.5,
|
||||
"n": 1,
|
||||
"size": "1024x1024"
|
||||
}' | python3 -c "
|
||||
import json, base64, sys
|
||||
data = json.load(sys.stdin)
|
||||
with open('output.png', 'wb') as f:
|
||||
f.write(base64.b64decode(data['data'][0]['b64_json']))
|
||||
print('saved output.png')
|
||||
"
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: Peak VRAM ~31.6 GiB at 1024×1024, ~45.5 GiB at 2048×2048 with `guidance_scale=4.5`.
|
||||
- Generation time: ~19 seconds at 1024×1024, ~102 seconds at 2048×2048.
|
||||
- Image quality is noticeably better than medium, with sharper details.
|
||||
- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
|
||||
- 2048×2048 fits within 48 GB but leaves little headroom (~45.5 GiB used out of 49 GiB).
|
||||
@@ -1,43 +1,34 @@
|
||||
# Cosmos3-Nano
|
||||
|
||||
> Text-to-image, text-to-video, image-to-video, and video-to-video serving
|
||||
> Text-to-image, text-to-video, and image-to-video serving
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: NVIDIA
|
||||
- Model: `nvidia/Cosmos3-Nano`
|
||||
- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy
|
||||
- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound), action policy
|
||||
- Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation.
|
||||
A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves these modes; the
|
||||
mode is selected per request:
|
||||
A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes;
|
||||
the mode is selected per request:
|
||||
|
||||
- **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`).
|
||||
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image/video.
|
||||
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image.
|
||||
- **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
|
||||
upload, or `image_reference` JSON).
|
||||
- **V2V** — `POST /v1/videos/sync` with a reference video (`input_reference`
|
||||
video upload, or `video_reference` JSON). Cosmos3 conditions on selected
|
||||
reference-video latent frames; use `extra_params.condition_frame_indexes_vision`
|
||||
and `extra_params.condition_video_keep` to choose which prefix/tail frames guide
|
||||
generation.
|
||||
- **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
|
||||
T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
|
||||
the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
|
||||
- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
|
||||
- `forward_dynamics` — given a first frame or video **and** an action trajectory,
|
||||
roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
|
||||
- `policy` — given a first frame or video and a language instruction,
|
||||
**predict** the action trajectory (and a rollout video). Use the async
|
||||
`POST /v1/videos` endpoint and read the predicted action from the top-level
|
||||
`action` field.
|
||||
- `inverse_dynamics` — given a video, **recover** the action trajectory. Use
|
||||
the async `POST /v1/videos` endpoint and read the recovered action from
|
||||
the top-level `action` field
|
||||
- `forward_dynamics` — given a first frame **and** an action trajectory, roll out
|
||||
the resulting video. Synchronous: `POST /v1/videos/sync`.
|
||||
- `policy` — given a first frame and a language instruction, **predict** the action
|
||||
trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and
|
||||
read the predicted action from the top-level `action` field
|
||||
(`{data, shape, dtype, raw_action_dim, domain_id}`).
|
||||
|
||||
Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`,
|
||||
@@ -47,9 +38,8 @@ mode is selected per request:
|
||||
**`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
|
||||
(`domain_name=droid_lerobot`).
|
||||
|
||||
Action requests can use `input_reference` or `video_reference` for video input.
|
||||
`policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics`
|
||||
requires a video reference.
|
||||
`inverse_dynamics` (recover the action from a given video) is supported by the
|
||||
pipeline; **online inference of inverse dynamics will be added in a follow-up MR.**
|
||||
|
||||
## References
|
||||
|
||||
@@ -152,36 +142,6 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
||||
-o cosmos3_i2v.mp4
|
||||
|
||||
# Video-to-video -> /v1/videos/sync with an uploaded reference video.
|
||||
# By default Cosmos3 conditions on latent indexes [0, 1]. For the default
|
||||
# temporal VAE stride this decodes only the first 5 input frames.
|
||||
# The model works best when the prompt describes the actual situation happening in the video.
|
||||
# Generic prompts may create sub-standard generations.
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Nano" \
|
||||
-F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
|
||||
-F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
|
||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" \
|
||||
-F "num_inference_steps=35" -F "guidance_scale=6.0" \
|
||||
-F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
||||
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true,"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
|
||||
-F "seed=2222" \
|
||||
-F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
|
||||
-o cosmos3_v2v.mp4
|
||||
|
||||
# V2V can also use a JSON-safe URL/data-URL video reference. Do not combine
|
||||
# video_reference with input_reference or image_reference.
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Nano" \
|
||||
-F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
|
||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" \
|
||||
-F "num_inference_steps=35" -F "guidance_scale=6.0" \
|
||||
-F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
||||
-F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"last"}' \
|
||||
-F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \
|
||||
-o cosmos3_v2v_from_url.mp4
|
||||
|
||||
# Text-to-video-with-sound
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
@@ -230,21 +190,6 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
|
||||
# poll until status == completed, then:
|
||||
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
|
||||
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4
|
||||
|
||||
# Action — inverse dynamics (video -> recovered action trajectory).
|
||||
# Asynchronous: use the job metadata to read the recovered action.
|
||||
VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
--form-string "model=nvidia/Cosmos3-Nano" \
|
||||
--form-string "prompt=Recover the robot action trajectory from this clip." \
|
||||
-F "input_reference=@motion_clip.mp4;type=video/mp4" \
|
||||
-F "size=640x480" -F "num_frames=17" -F "fps=5" \
|
||||
-F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
|
||||
--form-string 'extra_params={"action_mode":"inverse_dynamics","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
|
||||
-F "seed=0" | jq -r '.id')
|
||||
# poll until status == completed, then:
|
||||
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
|
||||
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4
|
||||
```
|
||||
|
||||
#### Notes
|
||||
@@ -261,27 +206,24 @@ curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inver
|
||||
- **Determinism:** identical seed reproduces identical output on the same
|
||||
hardware; outputs are not bit-identical across different GPU types.
|
||||
- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
|
||||
3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
|
||||
1280×720, 189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
|
||||
3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720,
|
||||
189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
|
||||
- **Key flags / params:** `--no-guardrails` (server) or
|
||||
`extra_params={"guardrails":false}` (per request) toggles safety. The
|
||||
per-request flag only takes effect when the server was launched **with**
|
||||
guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
|
||||
`use_resolution_template` / `use_duration_template` are off by default and only
|
||||
needed when not using upsampled prompts that already encode resolution/duration.
|
||||
For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
|
||||
frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
|
||||
API decodes the first or last needed reference frames (`"first"` by default).
|
||||
- **Known limitations:**
|
||||
- Guardrails-on requires `cosmos-guardrail` **and** access to the gated
|
||||
`nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
|
||||
the server fails at pipeline build with a gated-repo / safety-checker error.
|
||||
- A guardrail-blocked prompt currently returns HTTP 500
|
||||
(`"Guardrail blocked prompt"`).
|
||||
- Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
|
||||
online. Use async `POST /v1/videos` when you need the predicted/recovered
|
||||
action payload under the top-level `action` field; sync `/v1/videos/sync`
|
||||
returns raw MP4 bytes and does not expose action metadata in the response body.
|
||||
- Action `forward_dynamics` (sync `/v1/videos/sync`) and `policy` (async
|
||||
`/v1/videos`, returns the predicted action under the top-level `action`
|
||||
field) are supported online. **Online inference of inverse dynamics will be
|
||||
added in a follow-up MR.**
|
||||
|
||||
### 1x GPU (Offline generation)
|
||||
|
||||
@@ -345,9 +287,8 @@ python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image
|
||||
|
||||
- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
|
||||
`prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
|
||||
`num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and
|
||||
`multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V.
|
||||
For video, frames are returned in `outputs[0].request_output.images` as an
|
||||
`num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video,
|
||||
frames are returned in `outputs[0].request_output.images` as an
|
||||
`(B, F, H, W, 3)` array.
|
||||
- The offline entry must be guarded by `if __name__ == "__main__":` — the engine
|
||||
spawns workers with the `spawn` start method.
|
||||
@@ -1,12 +1,12 @@
|
||||
# Cosmos3-Super
|
||||
|
||||
> Frontier 64B world model: text-to-image, text-to-video, image-to-video, video-to-video (+ optional audio)
|
||||
> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: NVIDIA
|
||||
- Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`)
|
||||
- Task: T2I, T2V, I2V, V2V generation, with optional synchronized audio (video + sound)
|
||||
- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound)
|
||||
- Mode: Online serving with the OpenAI-compatible image/video APIs
|
||||
- Maintainer: Community
|
||||
|
||||
@@ -16,8 +16,8 @@ Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality
|
||||
Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request
|
||||
formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the
|
||||
recommended parallelism differ. Mode is selected per request (T2I →
|
||||
`/v1/images/generations`; T2V/I2V/V2V → `/v1/videos/sync`; add
|
||||
`generate_sound=true` for audio).
|
||||
`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true`
|
||||
for audio).
|
||||
|
||||
## References
|
||||
|
||||
@@ -62,9 +62,8 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs.
|
||||
#### Verification
|
||||
|
||||
Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
|
||||
T2I/T2V/I2V/V2V/T2VS curls); official params: `size=1280x720, num_frames=189,
|
||||
fps=24, num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0,
|
||||
max_sequence_length=4096`.
|
||||
T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24,
|
||||
num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`.
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/v1/models
|
||||
@@ -85,15 +84,6 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
||||
-o cosmos3_super_i2v.mp4
|
||||
|
||||
# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last".
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \
|
||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
|
||||
-F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
||||
-F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
|
||||
-F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
|
||||
-o cosmos3_super_v2v.mp4
|
||||
|
||||
# T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
|
||||
@@ -113,12 +103,10 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
|
||||
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
|
||||
- **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
|
||||
- Same generation defaults, supported sizes, V2V reference-video controls
|
||||
(`condition_frame_indexes_vision`, `condition_video_keep`), and
|
||||
`generate_sound`/`sound_duration`
|
||||
semantics as Nano, including the **action** modality: `forward_dynamics`,
|
||||
`policy`, and `inverse_dynamics` — see the Cosmos3-Nano recipe for the request
|
||||
shapes. Use async `/v1/videos` when you need predicted/recovered action metadata
|
||||
under the top-level `action` field. Verified on the 64B Super under
|
||||
`--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
|
||||
and the rollout video reliably.
|
||||
- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
|
||||
semantics as Nano, including the **action** modality: `forward_dynamics`
|
||||
(sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under
|
||||
the top-level `action` field) — see the Cosmos3-Nano recipe for the request shape.
|
||||
Online inference of inverse dynamics will be added in a follow-up MR. Verified on
|
||||
the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted
|
||||
action (`[16, 10]`) and the rollout video reliably.
|
||||
@@ -1,3 +1,3 @@
|
||||
-r common.txt
|
||||
onnxruntime-cann>=1.23.2
|
||||
torchaudio==2.9.0
|
||||
torchaudio==2.10.0
|
||||
@@ -1,46 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Tests for the DiffusersPipelineLoader.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from huggingface_hub import snapshot_download
|
||||
from vllm.config.load import LoadConfig
|
||||
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.models.helios import HeliosPipeline
|
||||
from vllm_omni.diffusion.registry import initialize_model
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
model_path = "hf-internal-testing/tiny-helios-modular-pipe"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def prefetch_helios_model():
|
||||
"""Downloads the tiny helios model prior to running a test."""
|
||||
snapshot_download(model_path)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def mock_tp_group(mocker):
|
||||
"""Mocks the tensor parallel group; this is needed to initialize the Helios model."""
|
||||
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
|
||||
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
|
||||
mock_group = mocker.MagicMock()
|
||||
mock_group.world_size = 1
|
||||
mock_group.rank_in_group = 0
|
||||
mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
|
||||
|
||||
|
||||
class _DummyPipelineModel(nn.Module):
|
||||
def __init__(self, *, source_prefix: str):
|
||||
@@ -69,13 +42,7 @@ class _DummyPipelineModel(nn.Module):
|
||||
|
||||
|
||||
def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
quantization_config=None,
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
|
||||
loader = object.__new__(DiffusersPipelineLoader)
|
||||
loader.counter_before_loading_weights = 0.0
|
||||
loader.counter_after_loading_weights = 0.0
|
||||
|
||||
@@ -183,15 +150,16 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
|
||||
quantization_config=None,
|
||||
)
|
||||
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
loader = object.__new__(DiffusersPipelineLoader)
|
||||
loader.load_weights = lambda model: None # type: ignore[assignment]
|
||||
loader._process_weights_after_loading = lambda model, target_device: None # type: ignore[assignment]
|
||||
loader._is_gguf_quantization = lambda: False # type: ignore[assignment]
|
||||
loader._is_gguf_quantization = lambda _od_config: False # type: ignore[assignment]
|
||||
|
||||
monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
|
||||
monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
|
||||
|
||||
model = loader.load_model(
|
||||
od_config,
|
||||
load_device="cpu",
|
||||
load_format="custom_pipeline",
|
||||
custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
|
||||
@@ -200,33 +168,3 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
|
||||
assert model.captured_config is od_config
|
||||
assert model.seen_config_during_init is od_config
|
||||
assert get_current_diffusion_config_or_none() is None
|
||||
|
||||
|
||||
def test_get_all_weights(prefetch_helios_model, mock_tp_group):
|
||||
"""Ensure that get all weights on a tiny model resolves to nonempty weights."""
|
||||
od_config = OmniDiffusionConfig(
|
||||
model_class_name="HeliosPipeline",
|
||||
model=model_path,
|
||||
)
|
||||
loader = DiffusersPipelineLoader(
|
||||
load_config=LoadConfig(),
|
||||
od_config=od_config,
|
||||
)
|
||||
pipeline = HeliosPipeline(od_config=od_config)
|
||||
|
||||
weights = list(loader.get_all_weights(pipeline))
|
||||
assert len(weights) > 0
|
||||
|
||||
|
||||
def test_load_model(prefetch_helios_model, mock_tp_group):
|
||||
"""Ensure that load model creates an instance of the expected pipeline class."""
|
||||
od_config = OmniDiffusionConfig(
|
||||
model_class_name="HeliosPipeline",
|
||||
model=model_path,
|
||||
)
|
||||
loader = DiffusersPipelineLoader(
|
||||
load_config=LoadConfig(),
|
||||
od_config=od_config,
|
||||
)
|
||||
model = loader.load_model(load_device="cpu")
|
||||
assert isinstance(model, HeliosPipeline)
|
||||
@@ -62,7 +62,6 @@ def _make_loader() -> DiffusersPipelineLoader:
|
||||
ignore_patterns=["*.tmp"],
|
||||
)
|
||||
loader.od_config = None
|
||||
loader.quant_config = None
|
||||
loader.counter_before_loading_weights = 0.0
|
||||
loader.counter_after_loading_weights = 0.0
|
||||
return loader
|
||||
@@ -70,47 +69,53 @@ def _make_loader() -> DiffusersPipelineLoader:
|
||||
|
||||
def test_is_gguf_quantization_accepts_dict_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = {
|
||||
"method": "gguf",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
od_config = SimpleNamespace(
|
||||
quantization_config={
|
||||
"method": "gguf",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
)
|
||||
|
||||
assert loader._is_gguf_quantization() is True
|
||||
assert loader._is_gguf_quantization(od_config) is True
|
||||
|
||||
|
||||
def test_is_gguf_quantization_rejects_non_gguf_dict_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = {
|
||||
"method": "fp8",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
od_config = SimpleNamespace(
|
||||
quantization_config={
|
||||
"method": "fp8",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
)
|
||||
|
||||
assert loader._is_gguf_quantization() is False
|
||||
assert loader._is_gguf_quantization(od_config) is False
|
||||
|
||||
|
||||
def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = {"method": "gguf"}
|
||||
od_config = SimpleNamespace(quantization_config={"method": "gguf"})
|
||||
|
||||
with pytest.raises(ValueError, match="gguf_model"):
|
||||
loader._is_gguf_quantization()
|
||||
loader._is_gguf_quantization(od_config)
|
||||
|
||||
|
||||
def test_is_gguf_quantization_accepts_object_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(
|
||||
get_name=lambda: "gguf",
|
||||
gguf_model="weights.gguf",
|
||||
od_config = SimpleNamespace(
|
||||
quantization_config=SimpleNamespace(
|
||||
get_name=lambda: "gguf",
|
||||
gguf_model="weights.gguf",
|
||||
)
|
||||
)
|
||||
|
||||
assert loader._is_gguf_quantization() is True
|
||||
assert loader._is_gguf_quantization(od_config) is True
|
||||
|
||||
|
||||
def test_is_gguf_quantization_uses_fallback_object_without_get_name():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
|
||||
od_config = SimpleNamespace(quantization_config=SimpleNamespace(gguf_model="weights.gguf"))
|
||||
|
||||
assert loader._is_gguf_quantization() is True
|
||||
assert loader._is_gguf_quantization(od_config) is True
|
||||
|
||||
|
||||
def test_get_model_loadable_names_collects_parameters_and_buffers():
|
||||
@@ -166,14 +171,16 @@ def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.Mon
|
||||
|
||||
def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
|
||||
loader.od_config = SimpleNamespace(revision="main")
|
||||
source = DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path="dummy",
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
)
|
||||
od_config = SimpleNamespace(
|
||||
revision="main",
|
||||
quantization_config=SimpleNamespace(gguf_model="weights.gguf"),
|
||||
)
|
||||
|
||||
class _Adapter:
|
||||
def weights_iterator(self):
|
||||
@@ -183,7 +190,7 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
|
||||
monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
|
||||
monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
|
||||
|
||||
weights = list(loader._get_gguf_weights_iterator(source, object()))
|
||||
weights = list(loader._get_gguf_weights_iterator(source, object(), od_config))
|
||||
|
||||
assert weights[0][0] == "transformer.weight"
|
||||
assert torch.equal(weights[0][1], torch.ones((2, 2)))
|
||||
@@ -194,12 +201,13 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
|
||||
def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
model = _DummyModel()
|
||||
od_config = SimpleNamespace()
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_gguf_weights_iterator",
|
||||
lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
|
||||
lambda source, model, od_config: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
|
||||
)
|
||||
|
||||
def _hf_iter(source):
|
||||
@@ -214,7 +222,7 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
|
||||
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
|
||||
)
|
||||
|
||||
loaded = loader._load_weights_with_gguf(model)
|
||||
loaded = loader._load_weights_with_gguf(model, od_config)
|
||||
|
||||
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
|
||||
assert model.calls[0] == ["transformer.weight"]
|
||||
@@ -225,12 +233,13 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
|
||||
def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
model = _DummyModel()
|
||||
od_config = SimpleNamespace()
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_gguf_weights_iterator",
|
||||
lambda source, model: iter(
|
||||
lambda source, model, od_config: iter(
|
||||
[
|
||||
("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
|
||||
("transformer.bias", torch.zeros(2, dtype=torch.float32)),
|
||||
@@ -253,7 +262,7 @@ def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete
|
||||
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
|
||||
)
|
||||
|
||||
loaded = loader._load_weights_with_gguf(model)
|
||||
loaded = loader._load_weights_with_gguf(model, od_config)
|
||||
|
||||
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
|
||||
assert hf_calls == ["vae"]
|
||||
@@ -374,19 +374,6 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None:
|
||||
assert tuple(additional["preprocessed_image"].shape) == (1, 3, 16, 32)
|
||||
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 3, 16, 32)
|
||||
|
||||
frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "purple", "black")]
|
||||
v2v = SimpleNamespace(
|
||||
prompts=[{"prompt": "Continue.", "multi_modal_data": {"video": frames}}],
|
||||
sampling_params=SimpleNamespace(
|
||||
height=16,
|
||||
width=32,
|
||||
extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"},
|
||||
),
|
||||
)
|
||||
additional = preprocess(v2v).prompts[0]["additional_information"]
|
||||
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
|
||||
assert additional["condition_frame_indexes_vision"] == [0, 1]
|
||||
|
||||
|
||||
def test_postprocess_handles_image_video_audio_and_validation() -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
|
||||
@@ -477,19 +464,6 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
|
||||
assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]]
|
||||
assert image_latent.shape == (1, 2, 1, 2, 3)
|
||||
|
||||
pipeline._encode_video_tensor = lambda *args, **kwargs: torch.full((1, 2, 3, 2, 3), 6.0)
|
||||
v2v_latents, v2v_velocity_mask, v2v_condition = pipeline._prepare_latents_v2v(
|
||||
torch.zeros(1, 3, 5, 16, 24),
|
||||
16,
|
||||
24,
|
||||
9,
|
||||
torch.Generator(device="cpu").manual_seed(0),
|
||||
[0, 1],
|
||||
)
|
||||
torch.testing.assert_close(v2v_latents[:, :, 0:2], torch.full((1, 2, 2, 2, 3), 6.0))
|
||||
assert v2v_velocity_mask.tolist() == [[[[[0.0]], [[0.0]], [[1.0]]]]]
|
||||
assert v2v_condition.shape == (1, 2, 3, 2, 3)
|
||||
|
||||
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
||||
pipeline._sound_tokenizer = SimpleNamespace(
|
||||
sample_rate=10,
|
||||
@@ -696,34 +670,6 @@ class TestForwardRouting:
|
||||
)
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
|
||||
|
||||
video_tensor = torch.zeros(1, 3, 5, 16, 16)
|
||||
v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
|
||||
v2v_mask = torch.tensor([[[[[0.0]], [[1.0]]]]])
|
||||
pipeline._prepare_latents_v2v = lambda *args, **kwargs: (
|
||||
torch.zeros(1, 2, 2, 1, 1),
|
||||
v2v_mask,
|
||||
v2v_condition,
|
||||
)
|
||||
pipeline.forward(
|
||||
SimpleNamespace(
|
||||
prompts=[
|
||||
{
|
||||
"prompt": "continue",
|
||||
"modalities": ["video"],
|
||||
"additional_information": {
|
||||
"preprocessed_video": video_tensor,
|
||||
"condition_frame_indexes_vision": [0],
|
||||
},
|
||||
}
|
||||
],
|
||||
sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
|
||||
)
|
||||
)
|
||||
assert captured["flow_shifts"][-1] == 10.0
|
||||
assert captured["format"]["negative_prompt"] == ""
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
|
||||
assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
|
||||
|
||||
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
||||
sound_latents = torch.zeros(1, 3, 4)
|
||||
pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)
|
||||
|
||||
@@ -1,2 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -1,160 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for HunyuanVideo-1.5 quant_config propagation through transformer creation.
|
||||
|
||||
Tests cover:
|
||||
- HunyuanVideo15Pipeline passes quant_config to HunyuanVideo15Transformer3DModel
|
||||
- HunyuanVideo15I2VPipeline passes quant_config to HunyuanVideo15Transformer3DModel
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 as t2v_module
|
||||
import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v as i2v_module
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class TestHunyuanVideoQuantConfigPropagation:
|
||||
"""Verify quant_config is propagated to the transformer model in HunyuanVideo-1.5 pipelines."""
|
||||
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_local_device",
|
||||
return_value="cpu",
|
||||
)
|
||||
@patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.prefetch_subfolders")
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.from_pretrained_with_prefetch",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.Qwen2Tokenizer.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.ByT5Tokenizer.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.AutoConfig.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.T5EncoderModel",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.FlowMatchEulerDiscreteScheduler.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_transformer_config_kwargs",
|
||||
return_value={},
|
||||
)
|
||||
def test_t2v_quant_config_passed_through(
|
||||
self,
|
||||
mock_get_kwargs,
|
||||
mock_scheduler,
|
||||
mock_t5_encoder,
|
||||
mock_auto_config,
|
||||
mock_byt5_tok,
|
||||
mock_qwen2_tok,
|
||||
mock_from_pretrained,
|
||||
mock_prefetch,
|
||||
mock_get_device,
|
||||
):
|
||||
captured_kwargs = {}
|
||||
|
||||
class FakeTransformer:
|
||||
def __init__(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
with patch.object(t2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
|
||||
fake_qc = MagicMock()
|
||||
od_config = SimpleNamespace(
|
||||
model="fake-t2v-model",
|
||||
dtype=torch.bfloat16,
|
||||
flow_shift=None,
|
||||
tf_model_config=SimpleNamespace(),
|
||||
quantization_config=fake_qc,
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
|
||||
t2v_module.HunyuanVideo15Pipeline(od_config=od_config)
|
||||
|
||||
assert captured_kwargs.get("quant_config") is fake_qc
|
||||
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_local_device",
|
||||
return_value="cpu",
|
||||
)
|
||||
@patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.prefetch_subfolders")
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.from_pretrained_with_prefetch",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.Qwen2Tokenizer.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.ByT5Tokenizer.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.AutoConfig.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.T5EncoderModel",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.SiglipImageProcessor.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.FlowMatchEulerDiscreteScheduler.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_transformer_config_kwargs",
|
||||
return_value={},
|
||||
)
|
||||
def test_i2v_quant_config_passed_through(
|
||||
self,
|
||||
mock_get_kwargs,
|
||||
mock_scheduler,
|
||||
mock_siglip_processor,
|
||||
mock_t5_encoder,
|
||||
mock_auto_config,
|
||||
mock_byt5_tok,
|
||||
mock_qwen2_tok,
|
||||
mock_from_pretrained,
|
||||
mock_prefetch,
|
||||
mock_get_device,
|
||||
):
|
||||
captured_kwargs = {}
|
||||
|
||||
class FakeTransformer:
|
||||
def __init__(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
with patch.object(i2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
|
||||
fake_qc = MagicMock()
|
||||
od_config = SimpleNamespace(
|
||||
model="fake-i2v-model",
|
||||
dtype=torch.bfloat16,
|
||||
flow_shift=None,
|
||||
tf_model_config=SimpleNamespace(),
|
||||
quantization_config=fake_qc,
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
|
||||
i2v_module.HunyuanVideo15I2VPipeline(od_config=od_config)
|
||||
|
||||
assert captured_kwargs.get("quant_config") is fake_qc
|
||||
@@ -72,31 +72,6 @@ class TestPipelineIndependence:
|
||||
|
||||
assert issubclass(LTX23Pipeline, CFGParallelMixin)
|
||||
|
||||
def test_ltx23_pipeline_declares_offload_components(self):
|
||||
"""LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
||||
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
||||
|
||||
pipe = object.__new__(LTX23Pipeline)
|
||||
torch.nn.Module.__init__(pipe)
|
||||
pipe.transformer = torch.nn.Linear(1, 1)
|
||||
pipe.text_encoder = torch.nn.Linear(1, 1)
|
||||
pipe.connectors = torch.nn.Linear(1, 1)
|
||||
pipe.vae = torch.nn.Linear(1, 1)
|
||||
pipe.audio_vae = torch.nn.Linear(1, 1)
|
||||
pipe.vocoder = torch.nn.Linear(1, 1)
|
||||
|
||||
modules = ModuleDiscovery.discover(pipe)
|
||||
|
||||
assert LTX23Pipeline._dit_modules == ["transformer"]
|
||||
assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"]
|
||||
assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"]
|
||||
assert LTX23Pipeline._resident_modules == ["vocoder"]
|
||||
assert modules.dit_names == ["transformer"]
|
||||
assert modules.encoder_names == ["text_encoder", "connectors"]
|
||||
assert modules.resident_names == ["vocoder"]
|
||||
assert len(modules.vaes) == 2
|
||||
|
||||
|
||||
class TestCFGParallelHelpers:
|
||||
"""Test LTX-2.3 CFG helper math without loading model weights."""
|
||||
|
||||
@@ -1,22 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm
|
||||
|
||||
|
||||
def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
|
||||
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
|
||||
|
||||
assert "weight" not in dict(norm.named_parameters())
|
||||
assert "weight" in dict(norm.named_buffers())
|
||||
assert "weight" not in norm.state_dict()
|
||||
|
||||
|
||||
def test_ltx_rms_norm_affine_weight_remains_parameter():
|
||||
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=True)
|
||||
|
||||
assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter)
|
||||
assert "weight" not in dict(norm.named_buffers())
|
||||
assert "weight" in norm.state_dict()
|
||||
@@ -41,7 +41,7 @@ def mock_od_config(mocker: MockerFixture):
|
||||
config.cache_backend = None
|
||||
config.cache_config = None
|
||||
config.model = "test-model"
|
||||
config.diffusion_load_format = "default"
|
||||
config.diffusion_load_format = None
|
||||
config.dtype = "float32"
|
||||
config.max_cpu_loras = 0
|
||||
config.lora_path = None
|
||||
|
||||
@@ -1,273 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end ModelOpt FP8 quality accuracy for video-gen DiTs.
|
||||
|
||||
Modeled on ``test_hunyuan_image3.py``'s ``test_quantized_dit_matches_bf16_accuracy``:
|
||||
for each calibrated FP8 checkpoint the test generates one video with BF16 and one
|
||||
with the FP8 DiT under the same seed, scores prompt-faithfulness with CLIP on the
|
||||
middle frame, then asserts both ``CLIP >= absolute_floor`` and
|
||||
``CLIP_drop <= drop_threshold``. The test is gated by an opt-in env var (model
|
||||
paths are local) so it never runs in generic CI.
|
||||
|
||||
To run::
|
||||
|
||||
VIDEOGEN_RUN_QUANT_ACCURACY=1 \\
|
||||
WAN22_A14B_BF16_MODEL=Wan-AI/Wan2.2-T2V-A14B-Diffusers \\
|
||||
WAN22_A14B_FP8_MODEL=/path/to/wan22-a14b-modelopt-fp8 \\
|
||||
HV15_BF16_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
|
||||
HV15_FP8_MODEL=/path/to/hv15-modelopt-fp8 \\
|
||||
pytest -s -v tests/e2e/accuracy/test_videogen_modelopt_quant.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from tests.e2e.accuracy.helpers import CLIPScorer, model_output_dir
|
||||
|
||||
pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tunables (override via env to suit local hardware / calibration budget)
|
||||
# ---------------------------------------------------------------------------
|
||||
SEED = 42
|
||||
HEIGHT = int(os.environ.get("VIDEOGEN_QUANT_HEIGHT", "480"))
|
||||
WIDTH = int(os.environ.get("VIDEOGEN_QUANT_WIDTH", "832"))
|
||||
NUM_FRAMES = int(os.environ.get("VIDEOGEN_QUANT_NUM_FRAMES", "17"))
|
||||
NUM_INFERENCE_STEPS = int(os.environ.get("VIDEOGEN_QUANT_STEPS", "20"))
|
||||
|
||||
WAN_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
|
||||
HV_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
|
||||
|
||||
# CLIP gates: absolute floor + max drop from the BF16 reference.
|
||||
CLIP_ABSOLUTE_FLOOR = float(os.environ.get("VIDEOGEN_QUANT_CLIP_FLOOR", "20.0"))
|
||||
CLIP_DROP_FP8 = float(os.environ.get("VIDEOGEN_QUANT_CLIP_DROP_FP8", "7.0"))
|
||||
|
||||
RUN_ENV = "VIDEOGEN_RUN_QUANT_ACCURACY"
|
||||
_TRUE = {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@dataclass(frozen=True)
|
||||
class _Spec:
|
||||
"""A model under test (BF16 + its FP8 variant)."""
|
||||
|
||||
family: str
|
||||
bf16_env: str
|
||||
fp8_env: str
|
||||
prompt: str
|
||||
guidance_scale: float
|
||||
guidance_scale_2: float | None = None
|
||||
boundary_ratio: float | None = None
|
||||
|
||||
|
||||
WAN_SPEC = _Spec(
|
||||
family="wan22_a14b_t2v",
|
||||
bf16_env="WAN22_A14B_BF16_MODEL",
|
||||
fp8_env="WAN22_A14B_FP8_MODEL",
|
||||
prompt=WAN_PROMPT,
|
||||
guidance_scale=3.5,
|
||||
guidance_scale_2=4.0,
|
||||
boundary_ratio=0.875,
|
||||
)
|
||||
|
||||
HV_SPEC = _Spec(
|
||||
family="hunyuanvideo15_t2v",
|
||||
bf16_env="HV15_BF16_MODEL",
|
||||
fp8_env="HV15_FP8_MODEL",
|
||||
prompt=HV_PROMPT,
|
||||
guidance_scale=6.0,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _QuantCase:
|
||||
name: str
|
||||
spec: _Spec
|
||||
clip_drop: float
|
||||
|
||||
|
||||
def _cases() -> list[pytest.ParameterSet]:
|
||||
cases = [
|
||||
_QuantCase("wan_fp8", WAN_SPEC, CLIP_DROP_FP8),
|
||||
_QuantCase("hv_fp8", HV_SPEC, CLIP_DROP_FP8),
|
||||
]
|
||||
params: list[pytest.ParameterSet] = []
|
||||
opted_in = os.environ.get(RUN_ENV, "").lower() in _TRUE
|
||||
for case in cases:
|
||||
marks: list[Any] = []
|
||||
if not opted_in:
|
||||
marks.append(pytest.mark.skip(reason=f"Set {RUN_ENV}=1 to run video-gen quant accuracy."))
|
||||
if not os.environ.get(case.spec.bf16_env):
|
||||
marks.append(pytest.mark.skip(reason=f"Set {case.spec.bf16_env} to run {case.name}."))
|
||||
if not os.environ.get(case.spec.fp8_env):
|
||||
marks.append(pytest.mark.skip(reason=f"Set {case.spec.fp8_env} to run {case.name}."))
|
||||
params.append(pytest.param(case, id=case.name, marks=marks))
|
||||
return params
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
def _middle_frame(frames: list[Image.Image] | list[np.ndarray] | np.ndarray) -> Image.Image:
|
||||
if isinstance(frames, np.ndarray):
|
||||
idx = frames.shape[0] // 2
|
||||
arr = frames[idx]
|
||||
if arr.dtype != np.uint8:
|
||||
arr = (arr * 255).clip(0, 255).astype(np.uint8)
|
||||
return Image.fromarray(arr).convert("RGB")
|
||||
if not frames:
|
||||
raise AssertionError("Empty video output")
|
||||
mid = frames[len(frames) // 2]
|
||||
if isinstance(mid, np.ndarray):
|
||||
if mid.dtype != np.uint8:
|
||||
mid = (mid * 255).clip(0, 255).astype(np.uint8)
|
||||
return Image.fromarray(mid).convert("RGB")
|
||||
return mid.convert("RGB")
|
||||
|
||||
|
||||
def _extract_frames(result) -> list[Image.Image]:
|
||||
"""Pull a frame list out of an Omni text-to-video result."""
|
||||
from vllm_omni.entrypoints.omni import OmniRequestOutput
|
||||
|
||||
if isinstance(result, list):
|
||||
result = result[0] if result else None
|
||||
if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None:
|
||||
result = result.request_output
|
||||
images = getattr(result, "images", None)
|
||||
if not images:
|
||||
raise AssertionError("Pipeline output had no images")
|
||||
if isinstance(images, list) and images and isinstance(images[0], tuple):
|
||||
# (frames, fps)
|
||||
images = images[0][0]
|
||||
return list(images)
|
||||
|
||||
|
||||
def _generate_video(
|
||||
*,
|
||||
model: str,
|
||||
spec: _Spec,
|
||||
quantization: str | None,
|
||||
output_path: Path,
|
||||
) -> tuple[Image.Image, float]:
|
||||
"""Generate one video and return its middle frame + elapsed wall time."""
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
kwargs: dict[str, Any] = dict(
|
||||
model=model,
|
||||
parallel_config=DiffusionParallelConfig(),
|
||||
enforce_eager=True,
|
||||
vae_use_tiling=True,
|
||||
)
|
||||
if quantization:
|
||||
kwargs["quantization"] = quantization
|
||||
if spec.boundary_ratio is not None:
|
||||
kwargs["boundary_ratio"] = spec.boundary_ratio
|
||||
omni = Omni(**kwargs)
|
||||
|
||||
try:
|
||||
gen = torch.Generator(device=current_omni_platform.device_type).manual_seed(SEED)
|
||||
sampling = dict(
|
||||
height=HEIGHT,
|
||||
width=WIDTH,
|
||||
generator=gen,
|
||||
guidance_scale=spec.guidance_scale,
|
||||
num_inference_steps=NUM_INFERENCE_STEPS,
|
||||
num_frames=NUM_FRAMES,
|
||||
seed=SEED,
|
||||
)
|
||||
if spec.guidance_scale_2 is not None:
|
||||
sampling["guidance_scale_2"] = spec.guidance_scale_2
|
||||
|
||||
import time as _time
|
||||
|
||||
t0 = _time.perf_counter()
|
||||
out = omni.generate({"prompt": spec.prompt}, OmniDiffusionSamplingParams(**sampling))
|
||||
elapsed = _time.perf_counter() - t0
|
||||
|
||||
frames = _extract_frames(out)
|
||||
mid = _middle_frame(frames)
|
||||
mid.save(output_path)
|
||||
return mid, elapsed
|
||||
finally:
|
||||
del omni
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", _cases())
|
||||
@pytest.mark.skipif(
|
||||
not torch.cuda.is_available() or torch.accelerator.device_count() < 1,
|
||||
reason="Needs at least 1 GPU.",
|
||||
)
|
||||
def test_quantized_videogen_dit_matches_bf16_clip(case: _QuantCase, accuracy_artifact_root: Path) -> None:
|
||||
"""FP8 DiT should preserve prompt-aligned video quality vs BF16."""
|
||||
out_dir = model_output_dir(accuracy_artifact_root, case.spec.family + "-quant")
|
||||
|
||||
with tempfile.TemporaryDirectory():
|
||||
bf16_model = os.environ[case.spec.bf16_env]
|
||||
quant_model = os.environ[case.spec.fp8_env]
|
||||
|
||||
bf16_frame, bf16_time = _generate_video(
|
||||
model=bf16_model,
|
||||
spec=case.spec,
|
||||
quantization=None,
|
||||
output_path=out_dir / "bf16.png",
|
||||
)
|
||||
quant_frame, quant_time = _generate_video(
|
||||
model=quant_model,
|
||||
spec=case.spec,
|
||||
quantization="fp8",
|
||||
output_path=out_dir / f"{case.name}.png",
|
||||
)
|
||||
|
||||
clip = CLIPScorer()
|
||||
bf16_clip = clip.score(bf16_frame, case.spec.prompt)
|
||||
quant_clip = clip.score(quant_frame, case.spec.prompt)
|
||||
clip_drop = bf16_clip - quant_clip
|
||||
speedup = (bf16_time / quant_time) if quant_time else float("nan")
|
||||
|
||||
metrics = {
|
||||
"case": case.name,
|
||||
"family": case.spec.family,
|
||||
"quantization": "fp8",
|
||||
"bf16_model": bf16_model,
|
||||
"quant_model": quant_model,
|
||||
"prompt": case.spec.prompt,
|
||||
"height": HEIGHT,
|
||||
"width": WIDTH,
|
||||
"num_frames": NUM_FRAMES,
|
||||
"num_inference_steps": NUM_INFERENCE_STEPS,
|
||||
"bf16_elapsed_s": bf16_time,
|
||||
"quant_elapsed_s": quant_time,
|
||||
"speedup_vs_bf16": speedup,
|
||||
"bf16_clip_score": bf16_clip,
|
||||
"quant_clip_score": quant_clip,
|
||||
"clip_score_drop": clip_drop,
|
||||
}
|
||||
metrics_path = out_dir / f"{case.name}_metrics.json"
|
||||
metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
|
||||
|
||||
print(f"\n[{case.name}] BF16 {bf16_time:.2f}s | FP8 {quant_time:.2f}s (x{speedup:.2f})")
|
||||
print(f" CLIP BF16={bf16_clip:.2f} FP8={quant_clip:.2f} drop={clip_drop:.2f}")
|
||||
print(f" metrics={metrics_path}")
|
||||
|
||||
assert quant_clip >= CLIP_ABSOLUTE_FLOOR, (
|
||||
f"{case.name} CLIP below floor: got {quant_clip:.2f}, expected >= {CLIP_ABSOLUTE_FLOOR:.2f}"
|
||||
)
|
||||
assert clip_drop <= case.clip_drop, (
|
||||
f"{case.name} CLIP drop too large: got {clip_drop:.2f}, expected <= {case.clip_drop:.2f}"
|
||||
)
|
||||
@@ -1,388 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end online tests for higgs-audio v3 against /v1/audio/speech.
|
||||
|
||||
Mirrors the higgs_audio_v2 test layout. Covers the plain-text-in / audio-out
|
||||
happy path, a small concurrent burst (Stage 0 prefix caching has to coexist
|
||||
with batching), the voice-clone path via ``ref_audio`` + ``ref_text``, the
|
||||
upstream ``references[]`` cookbook alias added by
|
||||
``normalize_references_alias`` in ``protocol/audio.py``, and the few
|
||||
validator rejections that should remain 4xx.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import load_test_audio_data_url
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
MODEL = "bosonai/higgs-audio-v3-tts-4b"
|
||||
STAGE_CONFIG = get_deploy_config_path("higgs_multimodal_qwen3.yaml")
|
||||
SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
|
||||
SERVER_ENV = {"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"}
|
||||
|
||||
TEST_PARAMS = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=STAGE_CONFIG,
|
||||
server_args=SERVER_ARGS,
|
||||
env_dict=SERVER_ENV,
|
||||
),
|
||||
id="higgs_audio_v3_plain_text",
|
||||
)
|
||||
]
|
||||
|
||||
DEFAULT_SPEECH_TIMEOUT_S = 180.0
|
||||
# Floor for ~0.5 s of 24 kHz mono PCM_16: 24000 * 0.5 * 2 bytes ~= 24 KiB.
|
||||
# A WAV header adds 44 bytes; conservative floor catches truncated /
|
||||
# silence-only outputs without flagging short legitimate clips.
|
||||
_MIN_AUDIO_BYTES = 20_000
|
||||
|
||||
# Reuse the shared TTS reference clip (clean ~5 s 24 kHz mono human speech)
|
||||
# vendored under tests/assets/qwen3_tts/. Keeps a single WAV across TTS
|
||||
# tests rather than duplicating asset bytes.
|
||||
_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
|
||||
_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestHiggsAudioV3OnlineHappyPath:
|
||||
"""Plain-text -> audio happy paths against the live HTTP server."""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_plain_text_wav(self, omni_server, openai_client) -> None:
|
||||
"""Single non-streaming WAV request - canonical TTS happy path."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
|
||||
"""``max_new_tokens`` is one of the few extra fields the v3 validator accepts."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Innovation distinguishes between a leader and a follower.",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"max_new_tokens": 500,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
|
||||
"""Three concurrent non-streaming requests - guards per-slot audio state and Stage-0 PC under batching."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "It was the night before my birthday.",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
},
|
||||
request_num=3,
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
|
||||
"""Streaming PCM via talker2code2wav_async_chunk + the connector's codec_streaming path.
|
||||
|
||||
The deploy YAML pins ``async_chunk: true`` and ``codec_streaming: true`` so chunks
|
||||
flow Stage 0 -> Stage 1 per AR loop. The Stage-1 codec honors
|
||||
``meta.left_context_size`` and ``meta.right_holdback_size`` so per-chunk codec
|
||||
windows stitch into a coherent PCM stream. The byte-count gate is the same as
|
||||
the sync paths; per-chunk audio content is verified offline against Whisper.
|
||||
|
||||
NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
|
||||
check still catches catastrophic codec failure (silence, white noise, sample
|
||||
scramble all give HNR << 0) while allowing for the sliding-window codec's
|
||||
slightly-noisier-than-sync output. The default 1.0 dB threshold has only
|
||||
~0.16 dB margin over measured single-request output (1.16 dB on L4), so
|
||||
flake risk is too high there.
|
||||
"""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "The quick brown fox jumps over the lazy dog.",
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
"min_hnr_db": 0.0,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None:
|
||||
"""Three concurrent streaming requests - guards per-request frame cursors
|
||||
in ``talker2code2wav_async_chunk`` and per-slot delay-pattern state under batched AR.
|
||||
|
||||
NOTE: ``min_hnr_db=-2.0`` is looser than the single-request streaming test
|
||||
because batched AR adds per-request codec quality variance (DAC boundary
|
||||
artifacts compounded across the 3-way batch); -2.0 dB gives ~1.6 dB margin
|
||||
over the measured worst-of-3 on L4 (-0.42 dB) while still well above the
|
||||
catastrophic-failure region.
|
||||
"""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "She sells seashells by the seashore.",
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
"min_hnr_db": -2.0,
|
||||
},
|
||||
request_num=3,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestHiggsAudioV3OnlineInlineControlTokens:
|
||||
"""Inline control-token surface from the bosonai cookbook.
|
||||
|
||||
The model exposes four categories of ``<|category:value|>`` tokens
|
||||
(emotion, style, prosody, sfx). Delivery tokens (emotion / style /
|
||||
prosody speed-pitch-expressive) go at the START of input; positional
|
||||
tokens (pause, sfx) go inline. SFX tokens must pair with their written
|
||||
onomatopoeia.
|
||||
|
||||
These tests check the *serving surface* - the validator accepts the
|
||||
payload, the engine produces audio, and the WAV has non-trivial size.
|
||||
They do not assert audio quality (out of scope for CI).
|
||||
"""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None:
|
||||
"""Delivery tokens (emotion + expressive_high) at the start of input."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": ("<|emotion:amusement|><|prosody:expressive_high|>Wait, that was actually hilarious."),
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_inline_style_whispering(self, omni_server, openai_client) -> None:
|
||||
"""Style token at the start - ``whispering``."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "<|style:whispering|>It is just between you and me, alright?",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> None:
|
||||
"""Two prosody tokens at the start - slow speed plus low pitch."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": (
|
||||
"<|prosody:speed_slow|><|prosody:pitch_low|>The radar shows a storm approaching from the east."
|
||||
),
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_inline_pause_mid_text(self, omni_server, openai_client) -> None:
|
||||
"""Positional pause token placed inline between two clauses."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hold on a moment <|prosody:pause|> let me think about it.",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_inline_sfx_with_onomatopoeia(self, omni_server, openai_client) -> None:
|
||||
"""SFX token paired with its written onomatopoeia (``<|sfx:laughter|>Hehe``)."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": (
|
||||
"<|emotion:amusement|>I cannot believe that just happened. "
|
||||
"<|sfx:laughter|>Hehe, I am still recovering from it."
|
||||
),
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestHiggsAudioV3OnlineVoiceClone:
|
||||
"""Voice clone via the two payload shapes the model serves."""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_voice_clone_ref_audio_ref_text(self, omni_server, openai_client) -> None:
|
||||
"""Canonical vllm-omni voice clone via ``ref_audio`` + ``ref_text``."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"ref_audio": _REF_AUDIO_URL,
|
||||
"ref_text": _REF_TEXT,
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_voice_clone_references_alias(self, omni_server, openai_client) -> None:
|
||||
"""BosonAI cookbook payload: ``references=[{audio_path, text}]``.
|
||||
|
||||
``normalize_references_alias`` (``protocol/audio.py``) is supposed to
|
||||
translate the cookbook field into ``ref_audio`` / ``ref_text``; without
|
||||
it the request would silently fall through to zero-shot synthesis.
|
||||
"""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"references": [{"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT}],
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestHiggsAudioV3OnlineValidatorRejections:
|
||||
"""Out-of-scope shapes must come back as 4xx.
|
||||
|
||||
These cases use ``send_audio_speech_http_request`` (raw HTTP POST against
|
||||
``/v1/audio/speech``) rather than the OpenAI SDK helper because:
|
||||
|
||||
- We need the failure-path matcher; the SDK helper asserts success.
|
||||
- The cookbook ``references`` field is not in the OpenAI SDK schema, so
|
||||
the SDK strips it before the request leaves the client and the server
|
||||
sees a plain TTS request that returns 200. Raw HTTP keeps the JSON body
|
||||
intact so the validator actually fires.
|
||||
"""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_empty_input(self, omni_server, openai_client) -> None:
|
||||
"""Empty text input - validator should reject before reaching the engine."""
|
||||
openai_client.send_audio_speech_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"input": "",
|
||||
"response_format": "wav",
|
||||
},
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"err_code": (400, 422),
|
||||
"err_message": ("input", "empty"),
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_multi_reference_payload(self, omni_server, openai_client) -> None:
|
||||
"""``references[]`` with more than one entry - multi-shot voice clone is not supported."""
|
||||
openai_client.send_audio_speech_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"references": [
|
||||
{"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT},
|
||||
{"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT},
|
||||
],
|
||||
"response_format": "wav",
|
||||
},
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"err_code": (400, 422),
|
||||
"err_message": "references",
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_conflicting_ref_audio_and_references(self, omni_server, openai_client) -> None:
|
||||
"""``ref_audio`` and ``references`` cannot both be set to different values."""
|
||||
openai_client.send_audio_speech_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"ref_audio": _REF_AUDIO_URL,
|
||||
"references": [{"audio_path": "https://example.com/other.wav"}],
|
||||
"response_format": "wav",
|
||||
},
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"err_code": (400, 422),
|
||||
"err_message": "mutually exclusive",
|
||||
}
|
||||
)
|
||||
@@ -33,7 +33,6 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
|
||||
SpeechBatchItem,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.serving_speech import (
|
||||
_TTS_LANGUAGES,
|
||||
OmniOpenAIServingSpeech,
|
||||
_create_wav_header,
|
||||
)
|
||||
@@ -1258,67 +1257,6 @@ class TestTTSMethods:
|
||||
# Verify speakers are normalized to lowercase
|
||||
assert server.supported_speakers == {"ryan", "vivian", "aiden"}
|
||||
|
||||
def test_load_supported_languages_from_config(self, speech_server):
|
||||
"""Languages/dialects from codec_language_id are loaded title-cased; 'Auto' is added."""
|
||||
speech_server._tts_model_type = "qwen3_tts"
|
||||
speech_server.engine_client.model_config = SimpleNamespace(
|
||||
hf_config=SimpleNamespace(
|
||||
talker_config=SimpleNamespace(
|
||||
codec_language_id={"chinese": 2055, "english": 2050, "beijing_dialect": 2074}
|
||||
)
|
||||
)
|
||||
)
|
||||
assert speech_server._load_supported_languages() == {
|
||||
"Chinese",
|
||||
"English",
|
||||
"Beijing_Dialect",
|
||||
"Auto",
|
||||
}
|
||||
|
||||
def test_load_supported_languages_from_dict_config(self, speech_server):
|
||||
"""talker_config provided as a plain dict is handled the same as an object."""
|
||||
speech_server._tts_model_type = "qwen3_tts"
|
||||
speech_server.engine_client.model_config = SimpleNamespace(
|
||||
hf_config=SimpleNamespace(talker_config={"codec_language_id": {"chinese": 2055, "english": 2050}})
|
||||
)
|
||||
assert speech_server._load_supported_languages() == {"Chinese", "English", "Auto"}
|
||||
|
||||
def test_validate_language_custom_dialect_accepted(self, speech_server):
|
||||
"""A language present in the model config passes validation, case-insensitively."""
|
||||
speech_server.supported_languages = {"Chinese", "English", "Beijing_Dialect", "Auto"}
|
||||
for language in ("Beijing_Dialect", "beijing_dialect", "English", "english", "Auto", "AUTO"):
|
||||
req = OpenAICreateSpeechRequest(input="Hello", language=language)
|
||||
result = speech_server._validate_tts_request(req)
|
||||
assert result is None or "Invalid language" not in result
|
||||
# Language is normalized to the title-cased config form.
|
||||
assert req.language == language.title()
|
||||
|
||||
def test_validate_language_unknown_rejected(self, speech_server):
|
||||
"""A language not in the configured set is rejected."""
|
||||
speech_server.supported_languages = {"Chinese", "English", "Auto"}
|
||||
for language in ("Klingon", "klingon"):
|
||||
req = OpenAICreateSpeechRequest(input="Hello", language=language)
|
||||
assert "Invalid language" in speech_server._validate_tts_request(req)
|
||||
|
||||
def test_load_supported_languages_default_when_no_config(self, speech_server):
|
||||
"""Empty/missing codec_language_id on a Qwen3-TTS model falls back to the default list."""
|
||||
speech_server._tts_model_type = "qwen3_tts"
|
||||
speech_server.engine_client.model_config = SimpleNamespace(
|
||||
hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_language_id={}))
|
||||
)
|
||||
assert speech_server._load_supported_languages() == _TTS_LANGUAGES
|
||||
|
||||
def test_load_supported_languages_default_on_config_error(self, speech_server):
|
||||
"""If the model config cannot be read, fall back to the default list."""
|
||||
speech_server._tts_model_type = "qwen3_tts"
|
||||
speech_server.engine_client = SimpleNamespace() # no model_config -> AttributeError
|
||||
assert speech_server._load_supported_languages() == _TTS_LANGUAGES
|
||||
|
||||
def test_load_supported_languages_default_for_non_qwen(self, speech_server):
|
||||
"""Non-qwen3_tts model types get the default language set."""
|
||||
speech_server._tts_model_type = None
|
||||
assert speech_server._load_supported_languages() == _TTS_LANGUAGES
|
||||
|
||||
def test_build_tts_params_with_uploaded_voice(self, speech_server, mocker: MockerFixture):
|
||||
"""Test _build_tts_params auto-sets ref_audio for uploaded voices (x_vector only)."""
|
||||
speech_server.uploaded_speakers = {
|
||||
|
||||
@@ -13,14 +13,12 @@ import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from PIL import Image
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.entrypoints.openai import api_server
|
||||
from vllm_omni.entrypoints.openai.api_server import router
|
||||
from vllm_omni.entrypoints.openai.protocol.videos import (
|
||||
@@ -86,8 +84,7 @@ class BlockingVideoHandler:
|
||||
if self.stage_configs is None:
|
||||
self.stage_configs = stage_configs
|
||||
|
||||
async def generate_video_bytes(self, request, reference_id, *, reference_image=None, reference_video=None):
|
||||
del request, reference_id, reference_image, reference_video
|
||||
async def generate_video_bytes(self, request, reference_id, *, reference_image=None):
|
||||
self.started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
@@ -133,30 +130,6 @@ def _make_test_image_data_url(size=(64, 64)) -> str:
|
||||
return f"data:image/png;base64,{encoded}"
|
||||
|
||||
|
||||
def _make_test_video_bytes(size=(32, 24), num_frames=3) -> bytes:
|
||||
width, height = size
|
||||
frames = np.zeros((num_frames, height, width, 3), dtype=np.uint8)
|
||||
for idx in range(num_frames):
|
||||
frames[idx, :, :, 0] = idx * 40
|
||||
frames[idx, :, :, 1] = 128
|
||||
frames[idx, :, :, 2] = 255 - idx * 40
|
||||
return mux_video_audio_bytes(frames, fps=8, video_codec_options={"preset": "ultrafast", "threads": "0"})
|
||||
|
||||
|
||||
def _make_test_video_data_url(size=(32, 24), num_frames=3) -> str:
|
||||
encoded = base64.b64encode(_make_test_video_bytes(size, num_frames)).decode("utf-8")
|
||||
return f"data:video/mp4;base64,{encoded}"
|
||||
|
||||
|
||||
def _cosmos3_stage_configs():
|
||||
return [
|
||||
SimpleNamespace(
|
||||
stage_type="diffusion",
|
||||
engine_args=SimpleNamespace(model_class_name="Cosmos3OmniDiffusersPipeline"),
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _wait_for_status(client: TestClient, video_id: str, status: str, timeout_s: float = 2.0):
|
||||
deadline = time.time() + timeout_s
|
||||
last_payload = None
|
||||
@@ -359,104 +332,6 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
|
||||
assert input_image.size == (40, 24)
|
||||
|
||||
|
||||
def test_v2v_video_generation_form(test_client, mocker: MockerFixture):
|
||||
video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
|
||||
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||
return_value=b"fake-video",
|
||||
)
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={"prompt": "Continue this motion."},
|
||||
files={"input_reference": ("input.mp4", video_bytes, "video/mp4")},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
video_id = response.json()["id"]
|
||||
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
|
||||
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
prompt = engine.captured_prompt
|
||||
assert "multi_modal_data" in prompt
|
||||
assert "video" in prompt["multi_modal_data"]
|
||||
input_video = prompt["multi_modal_data"]["video"]
|
||||
assert len(input_video) == 3
|
||||
assert all(isinstance(frame, Image.Image) for frame in input_video)
|
||||
assert input_video[0].size == (32, 24)
|
||||
|
||||
|
||||
def test_v2v_video_generation_with_video_reference_form(test_client, mocker: MockerFixture):
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||
return_value=b"fake-video",
|
||||
)
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={
|
||||
"prompt": "Continue this motion.",
|
||||
"video_reference": json.dumps({"video_url": _make_test_video_data_url((32, 24), 2)}),
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
video_id = response.json()["id"]
|
||||
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
|
||||
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
input_video = engine.captured_prompt["multi_modal_data"]["video"]
|
||||
assert len(input_video) == 2
|
||||
assert input_video[0].size == (32, 24)
|
||||
|
||||
|
||||
def test_decode_video_bytes_can_keep_last_frames():
|
||||
from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
|
||||
|
||||
frames = _decode_video_bytes(
|
||||
_make_test_video_bytes((32, 24), num_frames=6),
|
||||
source="input_reference",
|
||||
max_frames=2,
|
||||
keep="last",
|
||||
)
|
||||
|
||||
assert len(frames) == 2
|
||||
red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames]
|
||||
assert red_means[0] > 100
|
||||
assert red_means[1] > red_means[0]
|
||||
|
||||
|
||||
def test_cosmos3_reference_video_limit_uses_v2v_condition_frames():
|
||||
request = VideoGenerationRequest(
|
||||
prompt="Continue this motion.",
|
||||
num_frames=189,
|
||||
extra_params={"condition_frame_indexes_vision": [0, 2]},
|
||||
)
|
||||
|
||||
spec = api_server._reference_video_decode_spec(request, _cosmos3_stage_configs())
|
||||
assert spec.max_frames == 9
|
||||
assert spec.keep == "first"
|
||||
|
||||
|
||||
def test_cosmos3_reference_video_limit_preserves_action_frames():
|
||||
request = VideoGenerationRequest(
|
||||
prompt="Predict the action.",
|
||||
num_frames=17,
|
||||
extra_params={"action_mode": "inverse_dynamics", "action_chunk_size": 16},
|
||||
)
|
||||
|
||||
assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 17
|
||||
|
||||
|
||||
def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames():
|
||||
request = VideoGenerationRequest(
|
||||
prompt="Continue this motion.",
|
||||
num_frames=5,
|
||||
extra_params={"condition_frame_indexes_vision": [0, 20]},
|
||||
)
|
||||
|
||||
assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5
|
||||
|
||||
|
||||
def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
|
||||
fps_values = []
|
||||
|
||||
@@ -919,20 +794,7 @@ def test_rejects_input_reference_and_image_reference_together(test_client):
|
||||
files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
||||
|
||||
|
||||
def test_rejects_image_reference_and_video_reference_together(test_client):
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={
|
||||
"prompt": "bad refs",
|
||||
"image_reference": '{"image_url": "https://example.com/cat.png"}',
|
||||
"video_reference": '{"video_url": "https://example.com/cat.mp4"}',
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
||||
assert "either input_reference or image_reference" in response.json()["detail"].lower()
|
||||
|
||||
|
||||
def test_invalid_seconds_returns_422(test_client):
|
||||
@@ -994,18 +856,6 @@ def test_unsupported_image_reference_file_id_returns_400(test_client):
|
||||
assert response.json()["detail"] == "Invalid image_reference: file_id is not supported yet."
|
||||
|
||||
|
||||
def test_unsupported_video_reference_file_id_returns_400(test_client):
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={
|
||||
"prompt": "unsupported ref",
|
||||
"video_reference": '{"file_id": "file-123"}',
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.json()["detail"] == "Invalid video_reference: file_id is not supported yet."
|
||||
|
||||
|
||||
def test_invalid_uploaded_input_reference_returns_400(test_client):
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
@@ -1013,7 +863,7 @@ def test_invalid_uploaded_input_reference_returns_400(test_client):
|
||||
files={"input_reference": ("input.png", b"not-an-image", "image/png")},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image or video."
|
||||
assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image."
|
||||
|
||||
|
||||
def test_video_request_validation():
|
||||
@@ -1030,8 +880,6 @@ def test_video_request_validation():
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"})
|
||||
with pytest.raises(ValueError):
|
||||
VideoGenerationRequest(prompt="test", video_reference={"file_id": "file-1", "video_url": "https://example.com"})
|
||||
with pytest.raises(ValueError):
|
||||
VideoGenerationRequest(prompt="test", frame_interpolation_exp=0)
|
||||
with pytest.raises(ValueError):
|
||||
@@ -1417,23 +1265,6 @@ def test_sync_i2v_with_image_reference(test_client, mocker: MockerFixture):
|
||||
assert response.content == b"ref-video"
|
||||
|
||||
|
||||
def test_sync_v2v_returns_video_bytes(test_client, mocker: MockerFixture):
|
||||
video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
|
||||
_mock_encode_video_bytes(mocker, b"v2v-video-data")
|
||||
response = test_client.post(
|
||||
"/v1/videos/sync",
|
||||
data={"prompt": "Continue this motion."},
|
||||
files={"input_reference": ("input.mp4", video_bytes, "video/mp4")},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.content == b"v2v-video-data"
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
input_video = engine.captured_prompt["multi_modal_data"]["video"]
|
||||
assert len(input_video) == 3
|
||||
assert input_video[0].size == (32, 24)
|
||||
|
||||
|
||||
def test_sync_missing_handler_returns_503():
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
@@ -1466,7 +1297,7 @@ def test_sync_rejects_both_references(test_client):
|
||||
files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
|
||||
assert "either input_reference or image_reference" in response.json()["detail"].lower()
|
||||
|
||||
|
||||
def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture):
|
||||
|
||||
@@ -532,10 +532,7 @@ async def test_pure_diffusion_scenario(tp_size: int):
|
||||
try:
|
||||
await engine.sleep(level=1)
|
||||
await engine.wake_up()
|
||||
async for _ in engine.generate(
|
||||
"test",
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),
|
||||
):
|
||||
async for _ in engine.generate("test", sampling_params=SamplingParams()):
|
||||
pass
|
||||
logger.info("Pure diffusion OK (TP=%s)", tp_size)
|
||||
finally:
|
||||
|
||||
@@ -17,6 +17,7 @@ from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.config.stage_config import (
|
||||
_PIPELINE_REGISTRY,
|
||||
DeployConfig,
|
||||
ModelPipeline,
|
||||
PipelineConfig,
|
||||
StageConfig,
|
||||
StageConfigFactory,
|
||||
@@ -350,6 +351,93 @@ class TestStageConfig:
|
||||
assert "parallel_config" not in omega_config.engine_args
|
||||
|
||||
|
||||
class TestModelPipeline:
|
||||
"""Tests for ModelPipeline class."""
|
||||
|
||||
def test_valid_linear_dag(self):
|
||||
"""Test validation of a valid linear DAG."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
|
||||
StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]),
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert errors == [], f"Unexpected errors: {errors}"
|
||||
|
||||
def test_valid_branching_dag(self):
|
||||
"""Test validation of a valid branching DAG."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="input", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]),
|
||||
StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]),
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert errors == [], f"Unexpected errors: {errors}"
|
||||
|
||||
def test_missing_entry_point(self):
|
||||
"""Test that missing entry point is detected."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]),
|
||||
StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]),
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("entry point" in e.lower() for e in errors)
|
||||
|
||||
def test_missing_dependency(self):
|
||||
"""Test that missing stage reference is detected."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="input", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="output", input_sources=[99]), # Invalid
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("non-existent" in e.lower() for e in errors)
|
||||
|
||||
def test_duplicate_stage_ids(self):
|
||||
"""Test that duplicate stage IDs are detected."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]),
|
||||
StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]), # Duplicate
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("duplicate" in e.lower() for e in errors)
|
||||
|
||||
def test_self_reference(self):
|
||||
"""Test that self-references are detected."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="entry", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]), # Self
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("itself" in e.lower() for e in errors)
|
||||
|
||||
def test_get_stage_by_id(self):
|
||||
"""Test getting stage by ID."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
|
||||
stage = pipeline.get_stage(1)
|
||||
assert stage is not None
|
||||
assert stage.model_stage == "talker"
|
||||
|
||||
missing = pipeline.get_stage(99)
|
||||
assert missing is None
|
||||
|
||||
def test_empty_pipeline(self):
|
||||
"""Test validation of empty pipeline."""
|
||||
pipeline = ModelPipeline(model_type="test", stages=[])
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("no stages" in e.lower() for e in errors)
|
||||
|
||||
|
||||
class TestStageConfigFactory:
|
||||
"""Tests for StageConfigFactory class."""
|
||||
|
||||
@@ -510,6 +598,304 @@ class TestStageResolutionHelpers:
|
||||
assert overridden == []
|
||||
|
||||
|
||||
class TestPipelineYamlParsing:
|
||||
"""Tests for pipeline YAML file parsing (@ZJY0516)."""
|
||||
|
||||
def test_parse_qwen3_omni_moe_yaml(self, tmp_path):
|
||||
"""Test parsing the qwen3_omni_moe pipeline YAML."""
|
||||
yaml_content = """\
|
||||
model_type: qwen3_omni_moe
|
||||
async_chunk: false
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: thinker
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
hf_config_name: thinker_config
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
|
||||
- stage_id: 1
|
||||
model_stage: talker
|
||||
stage_type: llm
|
||||
input_sources: [0]
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
hf_config_name: talker_config
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
|
||||
|
||||
- stage_id: 2
|
||||
model_stage: code2wav
|
||||
stage_type: llm
|
||||
input_sources: [1]
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
hf_config_name: thinker_config
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
"""
|
||||
yaml_file = tmp_path / "qwen3_omni_moe.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe")
|
||||
|
||||
assert pipeline.model_type == "qwen3_omni_moe"
|
||||
assert len(pipeline.stages) == 3
|
||||
|
||||
# Stage 0: thinker
|
||||
s0 = pipeline.stages[0]
|
||||
assert s0.stage_id == 0
|
||||
assert s0.model_stage == "thinker"
|
||||
assert s0.stage_type == StageType.LLM
|
||||
assert s0.input_sources == []
|
||||
assert s0.worker_type == "ar"
|
||||
assert s0.final_output is True
|
||||
assert s0.final_output_type == "text"
|
||||
assert s0.is_comprehension is True
|
||||
|
||||
# Stage 1: talker
|
||||
s1 = pipeline.stages[1]
|
||||
assert s1.stage_id == 1
|
||||
assert s1.input_sources == [0]
|
||||
assert s1.custom_process_input_func == (
|
||||
"vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker"
|
||||
)
|
||||
assert s1.final_output is False
|
||||
|
||||
# Stage 2: code2wav
|
||||
s2 = pipeline.stages[2]
|
||||
assert s2.stage_id == 2
|
||||
assert s2.input_sources == [1]
|
||||
assert s2.worker_type == "generation"
|
||||
assert s2.final_output_type == "audio"
|
||||
|
||||
def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path):
|
||||
"""Test backward compatibility with engine_input_source field."""
|
||||
yaml_content = """\
|
||||
model_type: legacy_model
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
- stage_id: 1
|
||||
model_stage: downstream
|
||||
stage_type: llm
|
||||
engine_input_source: [0]
|
||||
"""
|
||||
yaml_file = tmp_path / "legacy.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model")
|
||||
assert pipeline.stages[1].input_sources == [0]
|
||||
|
||||
def test_parse_yaml_with_connectors_and_edges(self, tmp_path):
|
||||
"""Test parsing pipeline with optional connectors and edges."""
|
||||
yaml_content = """\
|
||||
model_type: test_model
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
|
||||
connectors:
|
||||
type: ray
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
"""
|
||||
yaml_file = tmp_path / "with_connectors.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model")
|
||||
assert pipeline.connectors == {"type": "ray"}
|
||||
assert pipeline.edges == [{"from": 0, "to": 1}]
|
||||
|
||||
def test_parsed_pipeline_passes_validation(self, tmp_path):
|
||||
"""Test that a well-formed YAML produces a valid pipeline."""
|
||||
yaml_content = """\
|
||||
model_type: valid_model
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
- stage_id: 1
|
||||
model_stage: next
|
||||
stage_type: llm
|
||||
input_sources: [0]
|
||||
"""
|
||||
yaml_file = tmp_path / "valid.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model")
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert errors == [], f"Unexpected validation errors: {errors}"
|
||||
|
||||
def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path):
|
||||
"""Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs."""
|
||||
yaml_content = """\
|
||||
model_type: legacy_batch
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 16
|
||||
engine_args:
|
||||
model_arch: SomeModel
|
||||
"""
|
||||
yaml_file = tmp_path / "legacy_batch.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch")
|
||||
s0 = pipeline.stages[0]
|
||||
assert "max_batch_size" not in s0.yaml_runtime
|
||||
assert s0.yaml_engine_args.get("max_num_seqs") == 16
|
||||
|
||||
def test_parse_diffusion_stage_type(self, tmp_path):
|
||||
"""Test parsing a diffusion stage type from YAML."""
|
||||
yaml_content = """\
|
||||
model_type: diff_model
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: dit
|
||||
stage_type: diffusion
|
||||
input_sources: []
|
||||
final_output: true
|
||||
final_output_type: image
|
||||
"""
|
||||
yaml_file = tmp_path / "diffusion.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model")
|
||||
assert pipeline.stages[0].stage_type == StageType.DIFFUSION
|
||||
|
||||
def test_parse_mimo_audio_yaml(self, tmp_path):
|
||||
"""Test parsing the MiMo Audio pipeline YAML."""
|
||||
yaml_content = """\
|
||||
model_type: mimo_audio
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: fused_thinker_talker
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
|
||||
- stage_id: 1
|
||||
model_stage: code2wav
|
||||
stage_type: llm
|
||||
input_sources: [0]
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
"""
|
||||
yaml_file = tmp_path / "mimo_audio.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio")
|
||||
|
||||
assert pipeline.model_type == "mimo_audio"
|
||||
assert len(pipeline.stages) == 2
|
||||
|
||||
# Stage 0: fused_thinker_talker
|
||||
s0 = pipeline.stages[0]
|
||||
assert s0.stage_id == 0
|
||||
assert s0.model_stage == "fused_thinker_talker"
|
||||
assert s0.input_sources == []
|
||||
assert s0.worker_type == "ar"
|
||||
assert s0.is_comprehension is True
|
||||
assert s0.final_output is True
|
||||
assert s0.final_output_type == "text"
|
||||
|
||||
# Stage 1: code2wav
|
||||
s1 = pipeline.stages[1]
|
||||
assert s1.stage_id == 1
|
||||
assert s1.input_sources == [0]
|
||||
assert s1.worker_type == "generation"
|
||||
assert s1.custom_process_input_func == (
|
||||
"vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav"
|
||||
)
|
||||
assert s1.final_output is True
|
||||
assert s1.final_output_type == "audio"
|
||||
|
||||
# Pipeline validation
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert errors == [], f"Unexpected errors: {errors}"
|
||||
|
||||
|
||||
class TestAsyncChunk:
|
||||
"""Tests for async_chunk pipeline flag."""
|
||||
|
||||
def test_model_pipeline_async_chunk_default(self):
|
||||
"""Test ModelPipeline defaults async_chunk to False."""
|
||||
stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
assert pipeline.async_chunk is False
|
||||
|
||||
def test_model_pipeline_async_chunk_set(self):
|
||||
"""Test ModelPipeline with async_chunk=True."""
|
||||
stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True)
|
||||
assert pipeline.async_chunk is True
|
||||
|
||||
def test_parse_async_chunk_from_yaml(self, tmp_path):
|
||||
"""Test that async_chunk is parsed from pipeline YAML."""
|
||||
yaml_content = """\
|
||||
model_type: test_async
|
||||
async_chunk: true
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
"""
|
||||
yaml_file = tmp_path / "async.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async")
|
||||
assert pipeline.async_chunk is True
|
||||
|
||||
def test_parse_missing_async_chunk_defaults_false(self, tmp_path):
|
||||
"""Test that missing async_chunk defaults to False."""
|
||||
yaml_content = """\
|
||||
model_type: test_no_async
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
"""
|
||||
yaml_file = tmp_path / "no_async.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async")
|
||||
assert pipeline.async_chunk is False
|
||||
|
||||
|
||||
class TestPipelineDiscovery:
|
||||
"""Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``)."""
|
||||
|
||||
|
||||
Whitespace-only changes.
File diff suppressed because it is too large.
Load diff
@@ -6,6 +6,7 @@ from vllm_omni.config.lora import LoRAConfig
|
||||
from vllm_omni.config.model import OmniModelConfig
|
||||
from vllm_omni.config.stage_config import (
|
||||
DeployConfig,
|
||||
ModelPipeline,
|
||||
PipelineConfig,
|
||||
StageConfig,
|
||||
StageConfigFactory,
|
||||
@@ -29,6 +30,7 @@ __all__ = [
|
||||
"LoRAConfig",
|
||||
"StageConfig",
|
||||
"StageConfigFactory",
|
||||
"ModelPipeline",
|
||||
"StageType",
|
||||
"StageExecutionType",
|
||||
"StagePipelineConfig",
|
||||
|
||||
@@ -149,10 +149,6 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
|
||||
"vllm_omni.model_executor.models.higgs_audio_v2.pipeline",
|
||||
"HIGGS_AUDIO_V2_PIPELINE",
|
||||
),
|
||||
"higgs_multimodal_qwen3": (
|
||||
"vllm_omni.model_executor.models.higgs_audio_v3.pipeline",
|
||||
"HIGGS_AUDIO_V3_PIPELINE",
|
||||
),
|
||||
"dynin_omni": (
|
||||
"vllm_omni.model_executor.models.dynin_omni.pipeline",
|
||||
"DYNIN_OMNI_PIPELINE",
|
||||
|
||||
@@ -20,6 +20,13 @@ from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
_MODELS_DIR = Path(__file__).resolve().parent.parent / "model_executor" / "models"
|
||||
|
||||
|
||||
def get_pipeline_path(model_dir: str, filename: str) -> Path:
|
||||
return _MODELS_DIR / model_dir / filename
|
||||
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@@ -1054,6 +1061,77 @@ class StageConfig:
|
||||
return create_config(config_dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelPipeline:
|
||||
"""Complete pipeline definition for a multi-stage model (legacy).
|
||||
|
||||
TODO(@lishunyang12): remove once all models migrate to PipelineConfig.
|
||||
"""
|
||||
|
||||
model_type: str
|
||||
stages: list[StageConfig]
|
||||
|
||||
# Pipeline-wide behavior flags
|
||||
async_chunk: bool = False
|
||||
|
||||
# Optional distributed configuration
|
||||
connectors: dict[str, Any] | None = None
|
||||
edges: list[dict[str, Any]] | None = None
|
||||
|
||||
def get_stage(self, stage_id: int) -> StageConfig | None:
|
||||
"""Look up a stage by its ID.
|
||||
|
||||
Args:
|
||||
stage_id: The stage ID to search for.
|
||||
|
||||
Returns:
|
||||
The matching StageConfig, or None if not found.
|
||||
"""
|
||||
for stage in self.stages:
|
||||
if stage.stage_id == stage_id:
|
||||
return stage
|
||||
return None
|
||||
|
||||
def validate_pipeline(self) -> list[str]:
|
||||
"""Validate pipeline topology at model integration time (not runtime).
|
||||
|
||||
Checks:
|
||||
- All stage IDs are unique
|
||||
- All input_sources reference valid stage IDs
|
||||
- At least one entry point (stage with empty input_sources)
|
||||
|
||||
Returns:
|
||||
List of validation error messages. Empty list if valid.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if not self.stages:
|
||||
errors.append("Topology has no stages defined")
|
||||
return errors
|
||||
|
||||
# Check for unique stage IDs
|
||||
stage_ids = [s.stage_id for s in self.stages]
|
||||
if len(stage_ids) != len(set(stage_ids)):
|
||||
errors.append("Duplicate stage IDs found")
|
||||
|
||||
stage_id_set = set(stage_ids)
|
||||
|
||||
# Check input_sources reference valid stages
|
||||
for stage in self.stages:
|
||||
for source_id in stage.input_sources:
|
||||
if source_id not in stage_id_set:
|
||||
errors.append(f"Stage {stage.stage_id} references non-existent input source {source_id}")
|
||||
if source_id == stage.stage_id:
|
||||
errors.append(f"Stage {stage.stage_id} references itself as input source")
|
||||
|
||||
# Check for at least one entry point
|
||||
entry_points = [s for s in self.stages if not s.input_sources]
|
||||
if not entry_points:
|
||||
errors.append("No entry point found (stage with empty input_sources)")
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
class StageConfigFactory:
|
||||
"""Factory that loads pipeline YAML and merges CLI overrides.
|
||||
|
||||
@@ -1128,9 +1206,33 @@ class StageConfigFactory:
|
||||
continue
|
||||
return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path)
|
||||
|
||||
# Not in the pipeline registry — let the caller fall back to the
|
||||
# legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
|
||||
return None
|
||||
# --- Legacy path: load from pipeline YAML ---
|
||||
pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code)
|
||||
|
||||
if pipeline is None:
|
||||
return None
|
||||
|
||||
errors = pipeline.validate_pipeline()
|
||||
if errors:
|
||||
logger.warning(f"Pipeline validation warnings for {model}: {errors}")
|
||||
|
||||
# Materialize the resolved pipeline-wide async_chunk value into every
|
||||
# stage so build_engine_args_dict() can inject the stage connector
|
||||
# spec and explicit False overrides are preserved.
|
||||
resolved_async_chunk = cli_overrides.get("async_chunk")
|
||||
if resolved_async_chunk is None:
|
||||
resolved_async_chunk = bool(pipeline.async_chunk)
|
||||
for stage in pipeline.stages:
|
||||
stage.yaml_engine_args["async_chunk"] = bool(resolved_async_chunk)
|
||||
|
||||
# Apply CLI overrides
|
||||
result: list[StageConfig] = []
|
||||
for stage in pipeline.stages:
|
||||
# Merge global CLI overrides
|
||||
stage.runtime_overrides = cls._merge_cli_overrides(stage, cli_overrides)
|
||||
result.append(stage)
|
||||
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def _create_from_registry(
|
||||
@@ -1253,6 +1355,193 @@ class StageConfigFactory:
|
||||
|
||||
return [config_dict]
|
||||
|
||||
@classmethod
|
||||
def _load_pipeline(cls, model: str, trust_remote_code: bool = True) -> ModelPipeline | None:
|
||||
"""Load a legacy ``pipeline.yaml`` for the model.
|
||||
|
||||
Searches ``model_executor/models/<dir>/pipeline.yaml`` by trying
|
||||
(a) the raw ``model_type`` as the directory name, then
|
||||
(b) ``model_type`` with hyphens replaced by underscores,
|
||||
and finally (c) scanning every ``pipeline.yaml`` for one that
|
||||
declares a matching ``model_type`` or ``hf_architectures``.
|
||||
|
||||
Returns None if no pipeline.yaml is found — caller handles the
|
||||
``resolve_model_config_path`` fallback via stage_configs/ YAMLs.
|
||||
"""
|
||||
model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
|
||||
if model_type is None:
|
||||
return None
|
||||
|
||||
# Direct lookups by convention
|
||||
candidates = [model_type, model_type.replace("-", "_")]
|
||||
for dir_name in candidates:
|
||||
pipeline_path = get_pipeline_path(dir_name, "pipeline.yaml")
|
||||
if pipeline_path.exists():
|
||||
return cls._parse_pipeline_yaml(pipeline_path, model_type)
|
||||
|
||||
# Scan fallback: read every pipeline.yaml and match on declared fields
|
||||
hf_archs = set(getattr(hf_config, "architectures", []) or []) if hf_config else set()
|
||||
if _MODELS_DIR.exists():
|
||||
for subdir in sorted(_MODELS_DIR.iterdir()):
|
||||
if not subdir.is_dir():
|
||||
continue
|
||||
pipeline_path = subdir / "pipeline.yaml"
|
||||
if not pipeline_path.exists():
|
||||
continue
|
||||
try:
|
||||
cfg = load_yaml_config(pipeline_path)
|
||||
except Exception as exc:
|
||||
logger.debug("Skip %s: %s", pipeline_path, exc)
|
||||
continue
|
||||
declared_type = getattr(cfg, "model_type", None)
|
||||
declared_archs = set(getattr(cfg, "hf_architectures", None) or [])
|
||||
if declared_type == model_type or (hf_archs and hf_archs.intersection(declared_archs)):
|
||||
return cls._parse_pipeline_yaml(pipeline_path, declared_type or model_type)
|
||||
|
||||
logger.debug("No pipeline.yaml found for model_type %s (archs=%s)", model_type, sorted(hf_archs))
|
||||
return None
|
||||
|
||||
# Keys consumed as explicit StageConfig fields — everything else is
|
||||
# passed through via yaml_extras.
|
||||
_KNOWN_STAGE_KEYS: set[str] = {
|
||||
"stage_id",
|
||||
"model_stage",
|
||||
"stage_type",
|
||||
"input_sources",
|
||||
"engine_input_source",
|
||||
"custom_process_input_func",
|
||||
"final_output",
|
||||
"final_output_type",
|
||||
"worker_type",
|
||||
"scheduler_cls",
|
||||
"hf_config_name",
|
||||
"is_comprehension",
|
||||
"engine_args",
|
||||
"runtime",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline:
|
||||
"""Parse a pipeline YAML file.
|
||||
|
||||
Args:
|
||||
path: Path to the YAML file.
|
||||
model_type: Model type identifier.
|
||||
|
||||
Returns:
|
||||
ModelPipeline object.
|
||||
"""
|
||||
config_data = load_yaml_config(path)
|
||||
|
||||
stages: list[StageConfig] = []
|
||||
for stage_data in config_data.stages:
|
||||
# Use .get() for optional fields — idiomatic for OmegaConf DictConfig
|
||||
stage_type_str = stage_data.get("stage_type", "llm")
|
||||
stage_type = StageType(stage_type_str) if stage_type_str else StageType.LLM
|
||||
|
||||
# Handle both 'input_sources' (new) and 'engine_input_source' (legacy)
|
||||
input_sources = stage_data.get("input_sources", None)
|
||||
if input_sources is None:
|
||||
input_sources = stage_data.get("engine_input_source", [])
|
||||
if input_sources is None:
|
||||
input_sources = []
|
||||
input_sources = list(input_sources)
|
||||
|
||||
# Extract per-stage engine_args and runtime dicts
|
||||
raw_ea = stage_data.get("engine_args", None)
|
||||
yaml_engine_args = to_dict(raw_ea) if raw_ea is not None else {}
|
||||
raw_rt = stage_data.get("runtime", None)
|
||||
yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}
|
||||
|
||||
# Migrate legacy runtime.max_batch_size → engine_args.max_num_seqs
|
||||
if "max_batch_size" in yaml_runtime:
|
||||
mbs = yaml_runtime.pop("max_batch_size")
|
||||
yaml_engine_args.setdefault("max_num_seqs", int(mbs))
|
||||
logger.debug(
|
||||
"Stage %s: migrated runtime.max_batch_size=%s to engine_args.max_num_seqs",
|
||||
stage_data.get("stage_id", "?"),
|
||||
mbs,
|
||||
)
|
||||
|
||||
# Topology-level fields that also live inside engine_args in legacy
|
||||
# YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
|
||||
worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
|
||||
scheduler_cls = stage_data.get("scheduler_cls", None) or yaml_engine_args.pop("scheduler_cls", None)
|
||||
if scheduler_cls:
|
||||
async_sched = yaml_engine_args.get("async_scheduling")
|
||||
if async_sched is not None:
|
||||
logger.warning(
|
||||
"Stage %s: async_scheduling=%r and scheduler_cls=%r "
|
||||
"should not be set together. scheduler_cls will take "
|
||||
"precedence for which scheduler is used.",
|
||||
stage_data.stage_id,
|
||||
async_sched,
|
||||
scheduler_cls,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Stage %s: scheduler_cls=%r is deprecated. Use async_scheduling instead.",
|
||||
stage_data.stage_id,
|
||||
scheduler_cls,
|
||||
)
|
||||
hf_config_name = stage_data.get("hf_config_name", None) or yaml_engine_args.pop("hf_config_name", None)
|
||||
model_stage = getattr(stage_data, "model_stage", None) or yaml_engine_args.pop("model_stage", None)
|
||||
|
||||
# Collect pass-through fields (default_sampling_params,
|
||||
# output_connectors, input_connectors, tts_args, etc.)
|
||||
yaml_extras: dict[str, Any] = {}
|
||||
for key in stage_data:
|
||||
if key not in cls._KNOWN_STAGE_KEYS:
|
||||
val = stage_data[key]
|
||||
try:
|
||||
yaml_extras[key] = to_dict(val)
|
||||
except ValueError:
|
||||
yaml_extras[key] = val
|
||||
|
||||
stage = StageConfig(
|
||||
stage_id=stage_data.stage_id,
|
||||
model_stage=model_stage or "",
|
||||
stage_type=stage_type,
|
||||
input_sources=input_sources,
|
||||
custom_process_input_func=stage_data.get("custom_process_input_func", None),
|
||||
final_output=stage_data.get("final_output", False),
|
||||
final_output_type=stage_data.get("final_output_type", None),
|
||||
worker_type=worker_type,
|
||||
scheduler_cls=scheduler_cls,
|
||||
hf_config_name=hf_config_name,
|
||||
is_comprehension=stage_data.get("is_comprehension", False),
|
||||
yaml_engine_args=yaml_engine_args,
|
||||
yaml_runtime=yaml_runtime,
|
||||
yaml_extras=yaml_extras,
|
||||
)
|
||||
stages.append(stage)
|
||||
|
||||
# Get pipeline-wide flags
|
||||
async_chunk = config_data.get("async_chunk", False)
|
||||
|
||||
# Get optional connector config — check both top-level and nested
|
||||
# under ``runtime`` (legacy stage_configs format).
|
||||
connectors = None
|
||||
edges = None
|
||||
if hasattr(config_data, "connectors"):
|
||||
connectors = to_dict(config_data.connectors)
|
||||
if hasattr(config_data, "edges"):
|
||||
edges = to_dict(config_data.edges)
|
||||
if hasattr(config_data, "runtime") and config_data.runtime is not None:
|
||||
top_runtime = config_data.runtime
|
||||
if connectors is None and hasattr(top_runtime, "connectors"):
|
||||
connectors = to_dict(top_runtime.connectors)
|
||||
if edges is None and hasattr(top_runtime, "edges"):
|
||||
edges = to_dict(top_runtime.edges)
|
||||
|
||||
return ModelPipeline(
|
||||
model_type=getattr(config_data, "model_type", model_type),
|
||||
stages=stages,
|
||||
async_chunk=async_chunk,
|
||||
connectors=connectors,
|
||||
edges=edges,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
|
||||
"""Auto-detect model_type from model directory.
|
||||
|
||||
@@ -391,33 +391,27 @@ class OmniTensorPrefixCache:
|
||||
self,
|
||||
query_start_loc: torch.Tensor,
|
||||
input_batch: InputBatch,
|
||||
multimodal_outputs: dict | None,
|
||||
multimodal_outputs: dict,
|
||||
num_scheduled_tokens: dict[str, int],
|
||||
):
|
||||
"""Get the merged multimodal states if hidden state prefix caching is enabled."""
|
||||
combined_multimodal_outputs = {}
|
||||
# Talkers that produce multimodal outputs only at decode (e.g. Higgs
|
||||
# Audio v3's audio codes) have ``multimodal_outputs is None`` at the
|
||||
# initial prefill step. Treat it as an empty mapping so the merger
|
||||
# short-circuits cleanly and the cache merge still runs once codes
|
||||
# start arriving.
|
||||
mm_outputs = multimodal_outputs if multimodal_outputs is not None else {}
|
||||
# First get the prefix cached tensors that are present in the mm data
|
||||
for mm_key in self.mm_cache_keys:
|
||||
if mm_key in mm_outputs:
|
||||
if mm_key in multimodal_outputs:
|
||||
combined_multimodal_outputs[mm_key] = self._get_merged_tensors(
|
||||
query_start_loc=query_start_loc,
|
||||
input_batch=input_batch,
|
||||
cache=self.mm_outputs_cache[mm_key],
|
||||
hidden_states=mm_outputs[mm_key],
|
||||
hidden_states=multimodal_outputs[mm_key],
|
||||
num_scheduled_tokens=num_scheduled_tokens,
|
||||
)
|
||||
|
||||
# Then, get everything else (passthrough data); first, convert to CPU
|
||||
# tensors similarly to the non prefix cached path, and then populate
|
||||
# the subdicts mapping request IDs -> payload objects
|
||||
passthrough_keys = set(mm_outputs.keys()) - self.mm_cache_keys
|
||||
passthrough_mm_data = {k: v for k, v in mm_outputs.items() if k in passthrough_keys}
|
||||
passthrough_keys = set(multimodal_outputs.keys()) - self.mm_cache_keys
|
||||
passthrough_mm_data = {k: v for k, v in multimodal_outputs.items() if k in passthrough_keys}
|
||||
mm_cpu = build_mm_cpu(multimodal_outputs=passthrough_mm_data)
|
||||
|
||||
for mm_key, mm_val in mm_cpu.items():
|
||||
|
||||
@@ -63,7 +63,6 @@ class OmniPayloadMeta(TypedDict, total=False):
|
||||
stream_finished: torch.Tensor
|
||||
req_id: list[str]
|
||||
left_context_size: int
|
||||
right_holdback_size: int
|
||||
override_keys: list[tuple[str, str]]
|
||||
num_processed_tokens: int
|
||||
next_stage_prompt_len: int
|
||||
@@ -151,7 +150,6 @@ class MetaStruct(_StructBase):
|
||||
stream_finished: torch.Tensor | None = None
|
||||
req_id: list[str] | None = None
|
||||
left_context_size: int | None = None
|
||||
right_holdback_size: int | None = None
|
||||
override_keys: list[tuple[str, str]] | None = None
|
||||
num_processed_tokens: int | None = None
|
||||
next_stage_prompt_len: int | None = None
|
||||
|
||||
@@ -1,74 +0,0 @@
|
||||
# Higgs-audio v3 deploy: Stage 0 (talker) -> Stage 1 (code2wav) via shared memory.
|
||||
#
|
||||
# Plain TTS only (no voice cloning). async_chunk streams codec chunks from
|
||||
# Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's
|
||||
# worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size
|
||||
# so the per-chunk codec windows stitch into a coherent PCM stream.
|
||||
async_chunk: true
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
# Stream PCM bytes back to the client as Stage 1 emits each chunk.
|
||||
codec_streaming: true
|
||||
# Streaming chunk shape for talker2code2wav_async_chunk:
|
||||
# codec_chunk_frames steady-state new frames per emit
|
||||
# codec_left_context_frames sliding past context to avoid leading-edge artifacts
|
||||
# codec_right_holdback_frames trailing frames held back so the codec
|
||||
# has future context for emit-region tail samples
|
||||
# initial_codec_chunk_frames smaller first chunk to minimize TTFA
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
codec_right_holdback_frames: 4
|
||||
initial_codec_chunk_frames: 1
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 5000
|
||||
connector_get_max_wait: 500
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 16
|
||||
gpu_memory_utilization: 0.6
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
# Stage-0 prefix caching is on by default. The talker's HS opt-out and
|
||||
# deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py)
|
||||
# keep per-step bookkeeping negligible; any shared text prefix across
|
||||
# concurrent requests is a net throughput win. Stage 1 (codec decoder)
|
||||
# has no KV cache and stays disabled.
|
||||
enable_prefix_caching: true
|
||||
async_scheduling: false
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 8192
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 1.0
|
||||
top_p: 0.95
|
||||
top_k: 50
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
|
||||
- stage_id: 1
|
||||
max_num_seqs: 16
|
||||
gpu_memory_utilization: 0.25
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: false
|
||||
max_num_batched_tokens: 65536
|
||||
max_model_len: 65536
|
||||
devices: "0"
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
repetition_penalty: 1.0
|
||||
@@ -81,7 +81,7 @@ class DefaultAdapter:
|
||||
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config)
|
||||
# load_model will handle dtypes / device placement, put in .eval() mode
|
||||
return loader.load_model(load_device=device)
|
||||
return loader.load_model(od_config=od_config, load_device=device)
|
||||
|
||||
@staticmethod
|
||||
def get_transformer(pipeline: Any) -> tuple[Any, str]:
|
||||
|
||||
@@ -610,6 +610,9 @@ class OmniDiffusionConfig:
|
||||
# Step mode settings
|
||||
step_execution: bool = False
|
||||
|
||||
# sleep mode
|
||||
enable_sleep_mode: bool = False
|
||||
|
||||
# Maximum number of sequences to generate in a batch
|
||||
max_num_seqs: int = 1
|
||||
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import glob
|
||||
import os
|
||||
@@ -31,9 +30,8 @@ from vllm.transformers_utils.repo_utils import file_exists
|
||||
from vllm.utils.import_utils import resolve_obj_by_qualname
|
||||
from vllm.utils.torch_utils import set_default_torch_dtype
|
||||
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp_to_model
|
||||
from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
|
||||
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
|
||||
get_checkpoint_adapter,
|
||||
)
|
||||
@@ -102,20 +100,18 @@ class DiffusersPipelineLoader:
|
||||
counter_before_loading_weights: float = 0.0
|
||||
counter_after_loading_weights: float = 0.0
|
||||
|
||||
def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig):
|
||||
def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig | None = None):
|
||||
self.load_config = load_config
|
||||
self.od_config = od_config
|
||||
self.quant_config = od_config.quantization_config
|
||||
self.parallel_config = od_config.parallel_config
|
||||
|
||||
def _prepare_weights(
|
||||
self,
|
||||
model_name_or_path: Path | str,
|
||||
model_name_or_path: Path,
|
||||
subfolder: str | None,
|
||||
revision: str | None,
|
||||
fall_back_to_pt: bool,
|
||||
allow_patterns_overrides: list[str] | None,
|
||||
) -> tuple[Path | str, list[str], bool]:
|
||||
) -> tuple[str, list[str], bool]:
|
||||
"""Prepare weights for the model.
|
||||
|
||||
If the model is not local, it will be downloaded."""
|
||||
@@ -212,13 +208,15 @@ class DiffusersPipelineLoader:
|
||||
source.allow_patterns_overrides,
|
||||
)
|
||||
|
||||
od_config = self.od_config
|
||||
use_multithread = (
|
||||
use_safetensors
|
||||
and getattr(self.od_config, "enable_multithread_weight_load", False)
|
||||
and od_config is not None
|
||||
and getattr(od_config, "enable_multithread_weight_load", False)
|
||||
and self.load_config.safetensors_load_strategy != "torchao"
|
||||
)
|
||||
if use_multithread:
|
||||
num_threads = getattr(self.od_config, "num_weight_load_threads", 4)
|
||||
num_threads = getattr(od_config, "num_weight_load_threads", 4)
|
||||
# Keep deterministic shard order before passing to vLLM helper.
|
||||
sorted_hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key)
|
||||
weights_iterator = multi_thread_safetensors_weights_iterator(
|
||||
@@ -244,7 +242,10 @@ class DiffusersPipelineLoader:
|
||||
return prefixed_weights_iterator
|
||||
|
||||
def _get_source_quant_config(self, source: "ComponentSource") -> object | None:
|
||||
quant_config = self.quant_config
|
||||
if self.od_config is None:
|
||||
return None
|
||||
|
||||
quant_config = self.od_config.quantization_config
|
||||
if hasattr(quant_config, "resolve"):
|
||||
return quant_config.resolve(source.prefix.rstrip("."))
|
||||
return quant_config
|
||||
@@ -307,6 +308,7 @@ class DiffusersPipelineLoader:
|
||||
|
||||
def load_model(
|
||||
self,
|
||||
od_config: OmniDiffusionConfig,
|
||||
load_device: str,
|
||||
load_format: str | None = "default",
|
||||
custom_pipeline_name: str | type[nn.Module] | None = None,
|
||||
@@ -315,13 +317,14 @@ class DiffusersPipelineLoader:
|
||||
"""Load a model with the given configurations."""
|
||||
if load_format is None:
|
||||
load_format = "default"
|
||||
self.od_config = od_config
|
||||
# CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8),
|
||||
# weights are already quantized in the checkpoint — load directly on CPU.
|
||||
# For online quantization, load on device so quantization can run on accelerator,
|
||||
# then move back to CPU afterward.
|
||||
offload_after_quant = False
|
||||
if load_device == "cpu" and self.quant_config is not None and device is not None:
|
||||
quant_cfg = self.quant_config
|
||||
if load_device == "cpu" and od_config.quantization_config is not None:
|
||||
quant_cfg = od_config.quantization_config
|
||||
is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr(
|
||||
quant_cfg, "is_checkpoint_quantized", False
|
||||
)
|
||||
@@ -336,21 +339,48 @@ class DiffusersPipelineLoader:
|
||||
logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU")
|
||||
|
||||
target_device = torch.device(load_device)
|
||||
with set_default_torch_dtype(self.od_config.dtype):
|
||||
if self.parallel_config.use_hsdp:
|
||||
with set_default_torch_dtype(od_config.dtype):
|
||||
if od_config.parallel_config.use_hsdp:
|
||||
model = self._load_model_with_hsdp(
|
||||
target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name
|
||||
od_config, target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name
|
||||
)
|
||||
else:
|
||||
model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
|
||||
if load_format == "custom_pipeline":
|
||||
# NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)`
|
||||
# internally. If we construct them under `with target_device:` (CUDA),
|
||||
# safetensors takes a direct-to-GPU fast path that calls `cudaMalloc`
|
||||
# via the driver API and BYPASSES PyTorch's caching allocator.
|
||||
# That makes those bytes invisible to CuMemAllocator, so `sleep()`
|
||||
# cannot offload/unmap them and GPU memory stays pinned.
|
||||
#
|
||||
# Fix: build the custom pipeline on CPU first (no default device
|
||||
# context), then explicitly move it to the target device. The
|
||||
# subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)`
|
||||
# + `copy_`, which goes through the caching allocator and is fully
|
||||
# tracked by CuMemAllocator.
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
|
||||
model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
|
||||
with set_current_diffusion_config(od_config):
|
||||
model = model_cls(od_config=od_config)
|
||||
if target_device.type != "cpu":
|
||||
model.to(target_device)
|
||||
else:
|
||||
with target_device:
|
||||
if load_format == "default":
|
||||
model = initialize_model(od_config)
|
||||
elif load_format == "diffusers":
|
||||
model = DiffusersAdapterPipeline(od_config=od_config, device=target_device)
|
||||
else:
|
||||
raise ValueError(f"Unknown load_format: {load_format}")
|
||||
logger.debug("Loading weights on %s ...", load_device)
|
||||
if load_format == "diffusers":
|
||||
# DiffusersAdapterPipeline.load_weights() calls
|
||||
# DiffusionPipeline.from_pretrained() internally; it does
|
||||
# not use our native customized pipeline classes.
|
||||
cast(DiffusersAdapterPipeline, model).load_weights()
|
||||
elif self._is_gguf_quantization():
|
||||
self._load_weights_with_gguf(model)
|
||||
elif self._is_gguf_quantization(od_config):
|
||||
self._load_weights_with_gguf(model, od_config)
|
||||
else:
|
||||
# Quantization does not happen in `load_weights` but after it
|
||||
self.load_weights(model)
|
||||
@@ -373,7 +403,7 @@ class DiffusersPipelineLoader:
|
||||
"""
|
||||
for _, module in model.named_modules():
|
||||
quant_method = getattr(module, "quant_method", None)
|
||||
if quant_method is not None and isinstance(quant_method, QuantizeMethodBase):
|
||||
if isinstance(quant_method, QuantizeMethodBase):
|
||||
# Move module to target device for processing if needed
|
||||
module_device = next(module.parameters(), None)
|
||||
if module_device is not None:
|
||||
@@ -403,13 +433,8 @@ class DiffusersPipelineLoader:
|
||||
# that have loaded weights tracking currently.
|
||||
if loaded_weights is not None:
|
||||
weights_not_loaded = weights_to_load - loaded_weights
|
||||
# NOTE: if the model is quantized, ignore not_loaded check for scale
|
||||
# weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a
|
||||
# static activation `input_scale`, which the quant method may
|
||||
# fold/track differently than plain parameters.
|
||||
weights_scale_not_loaded = {
|
||||
name for name in weights_not_loaded if name.endswith(("weight_scale", "input_scale"))
|
||||
}
|
||||
# NOTE: if the model is quantized, ignore not_loaded check for scale weights
|
||||
weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")}
|
||||
weights_not_loaded = weights_not_loaded - weights_scale_not_loaded
|
||||
if weights_not_loaded:
|
||||
self._check_unloaded_weights(weights_not_loaded)
|
||||
@@ -451,11 +476,9 @@ class DiffusersPipelineLoader:
|
||||
are logged as a warning. Any *other* missing weight raises
|
||||
``ValueError`` regardless of quantization.
|
||||
"""
|
||||
if self.quant_config is None:
|
||||
raise ValueError(
|
||||
"The quantization config is None, and the following weights "
|
||||
f"were not initialized from checkpoint: {weights_not_loaded}"
|
||||
)
|
||||
od_config = getattr(self, "od_config", None)
|
||||
if od_config is None or od_config.quantization_config is None:
|
||||
raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
|
||||
|
||||
expected_missing = {w for w in weights_not_loaded if self._is_expected_quantized_weight(w)}
|
||||
unexpected_missing = weights_not_loaded - expected_missing
|
||||
@@ -468,33 +491,39 @@ class DiffusersPipelineLoader:
|
||||
if unexpected_missing:
|
||||
raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
|
||||
|
||||
def _get_gguf_model_from_config(self) -> str | None:
|
||||
"""Extract gguf_model from quant_config (handles both dict and object)."""
|
||||
if self.quant_config is None:
|
||||
return None
|
||||
if isinstance(self.quant_config, dict):
|
||||
return self.quant_config.get("gguf_model")
|
||||
return getattr(self.quant_config, "gguf_model", None)
|
||||
def _is_gguf_quantization(self, od_config: OmniDiffusionConfig) -> bool:
|
||||
quant_config = od_config.quantization_config
|
||||
if quant_config is None:
|
||||
return False
|
||||
|
||||
def _is_gguf_quantization(self) -> bool:
|
||||
"""Check whether or not this pipeline loader is pointing at a GGUF config."""
|
||||
# New API: DiffusionGGUFConfig (QuantizationConfig subclass)
|
||||
from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
|
||||
|
||||
maybe_gguf_model = self._get_gguf_model_from_config()
|
||||
is_gguf = (
|
||||
isinstance(self.quant_config, DiffusionGGUFConfig)
|
||||
or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf")
|
||||
or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf")
|
||||
)
|
||||
if is_gguf:
|
||||
if not maybe_gguf_model:
|
||||
if isinstance(quant_config, DiffusionGGUFConfig):
|
||||
if quant_config.gguf_model is None:
|
||||
raise ValueError("GGUF quantization requires gguf_model")
|
||||
return True
|
||||
|
||||
if isinstance(self.quant_config, dict):
|
||||
# Dict-style config: {"method": "gguf", "gguf_model": "..."}
|
||||
if isinstance(quant_config, dict):
|
||||
if quant_config.get("method") == "gguf":
|
||||
if not quant_config.get("gguf_model"):
|
||||
raise ValueError("GGUF quantization requires gguf_model")
|
||||
return True
|
||||
return False
|
||||
|
||||
return bool(maybe_gguf_model)
|
||||
# Check by name for any config that reports as "gguf"
|
||||
if hasattr(quant_config, "get_name") and quant_config.get_name() == "gguf":
|
||||
gguf_model = getattr(quant_config, "gguf_model", None)
|
||||
if gguf_model is None:
|
||||
raise ValueError("GGUF quantization requires gguf_model")
|
||||
return True
|
||||
|
||||
# Fallback: object with gguf_model attribute but no get_name
|
||||
if hasattr(quant_config, "gguf_model") and quant_config.gguf_model:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _is_transformer_source(self, source: "ComponentSource") -> bool:
|
||||
if source.subfolder == "transformer":
|
||||
@@ -537,23 +566,25 @@ class DiffusersPipelineLoader:
|
||||
self,
|
||||
source: "ComponentSource",
|
||||
model: nn.Module,
|
||||
od_config: OmniDiffusionConfig,
|
||||
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
||||
gguf_model = self._get_gguf_model_from_config()
|
||||
quant_config = od_config.quantization_config
|
||||
gguf_model = getattr(quant_config, "gguf_model", None)
|
||||
if gguf_model is None:
|
||||
raise ValueError("GGUF quantization requires quantization_config.gguf_model")
|
||||
gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision)
|
||||
adapter = get_gguf_adapter(gguf_file, model, source, self.od_config)
|
||||
gguf_file = self._resolve_gguf_model_path(gguf_model, od_config.revision)
|
||||
adapter = get_gguf_adapter(gguf_file, model, source, od_config)
|
||||
weights_iter = adapter.weights_iterator()
|
||||
return ((source.prefix + name, tensor) for (name, tensor) in weights_iter)
|
||||
|
||||
def _load_weights_with_gguf(self, model: nn.Module) -> set[str]:
|
||||
def _load_weights_with_gguf(self, model: nn.Module, od_config: OmniDiffusionConfig) -> set[str]:
|
||||
sources = self._get_weight_sources(model)
|
||||
loaded: set[str] = set()
|
||||
loadable_names: set[str] | None = None
|
||||
|
||||
for source in sources:
|
||||
if self._is_transformer_source(source):
|
||||
loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model))
|
||||
loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model, od_config))
|
||||
|
||||
# GGUF checkpoints can be transformer-only or partially quantized.
|
||||
# Only fall back to HF if this source still has missing loadable weights.
|
||||
@@ -578,45 +609,9 @@ class DiffusersPipelineLoader:
|
||||
raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
|
||||
return loaded
|
||||
|
||||
def _init_from_load_format(
|
||||
self,
|
||||
load_format: str,
|
||||
target_device: torch.device,
|
||||
custom_pipeline_name: str | type[nn.Module] | None = None,
|
||||
is_hsdp: bool = False,
|
||||
) -> nn.Module:
|
||||
"""Initialize the model from a specified load format."""
|
||||
if load_format == "custom_pipeline":
|
||||
# NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)`
|
||||
# internally. If we construct them under `with target_device:` (CUDA),
|
||||
# safetensors takes a direct-to-GPU fast path that calls `cudaMalloc`
|
||||
# via the driver API and BYPASSES PyTorch's caching allocator.
|
||||
# That makes those bytes invisible to CuMemAllocator, so `sleep()`
|
||||
# cannot offload/unmap them and GPU memory stays pinned.
|
||||
#
|
||||
# Fix: build the custom pipeline on CPU first (no default device
|
||||
# context), then explicitly move it to the target device. The
|
||||
# subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)`
|
||||
# + `copy_`, which goes through the caching allocator and is fully
|
||||
# tracked by CuMemAllocator.
|
||||
model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
|
||||
with set_current_diffusion_config(self.od_config):
|
||||
model = model_cls(od_config=self.od_config)
|
||||
if not is_hsdp and target_device.type != "cpu":
|
||||
model.to(target_device)
|
||||
else:
|
||||
device_ctx = target_device if not is_hsdp else contextlib.nullcontext()
|
||||
with device_ctx:
|
||||
if load_format == "default":
|
||||
model = initialize_model(self.od_config)
|
||||
elif load_format == "diffusers":
|
||||
model = DiffusersAdapterPipeline(od_config=self.od_config, device=target_device)
|
||||
else:
|
||||
raise ValueError(f"Unknown load_format: {load_format}")
|
||||
return model
|
||||
|
||||
def _load_model_with_hsdp(
|
||||
self,
|
||||
od_config: OmniDiffusionConfig,
|
||||
target_device: torch.device,
|
||||
load_format: str = "default",
|
||||
custom_pipeline_name: str | type[nn.Module] | None = None,
|
||||
@@ -629,11 +624,14 @@ class DiffusersPipelineLoader:
|
||||
Approach: Load weights first using model's load_weights (handles QKV fusion etc.),
|
||||
then apply HSDP sharding to redistribute weights across GPUs.
|
||||
"""
|
||||
from vllm_omni.diffusion.distributed.hsdp import apply_hsdp_to_model
|
||||
|
||||
parallel_config = od_config.parallel_config
|
||||
hsdp_config = HSDPInferenceConfig(
|
||||
enabled=True,
|
||||
hsdp_replicate_size=self.parallel_config.hsdp_replicate_size,
|
||||
hsdp_shard_size=self.parallel_config.hsdp_shard_size,
|
||||
param_dtype=self.od_config.dtype,
|
||||
hsdp_replicate_size=parallel_config.hsdp_replicate_size,
|
||||
hsdp_shard_size=parallel_config.hsdp_shard_size,
|
||||
param_dtype=od_config.dtype,
|
||||
)
|
||||
|
||||
# Initialize model WITHOUT device context (weights start on CPU).
|
||||
@@ -641,9 +639,14 @@ class DiffusersPipelineLoader:
|
||||
# directly on GPU, HSDP needs weights on CPU first so they can be redistributed
|
||||
# across GPUs by apply_hsdp_to_model. The model's load_weights handles weight
|
||||
# mapping (QKV fusion, etc.).
|
||||
if load_format == "diffusers":
|
||||
raise ValueError("HSDP is not supported with the diffusers adapter load format")
|
||||
model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
|
||||
if load_format == "default":
|
||||
model = initialize_model(od_config)
|
||||
elif load_format == "custom_pipeline":
|
||||
model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
|
||||
with set_current_diffusion_config(od_config):
|
||||
model = model_cls(od_config=od_config)
|
||||
self.load_weights(model)
|
||||
|
||||
# Collect all transformers to shard (some models have transformer_2 for MoE)
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Cosmos3 text/image/video-to-video and text-to-image pipeline for vllm-omni.
|
||||
"""Cosmos3 text/image-to-video and text-to-image pipeline for vllm-omni.
|
||||
|
||||
Single pipeline class supports T2V, I2V, V2V, and T2I; the mode is selected at
|
||||
Single pipeline class supports T2V, I2V, and T2I; the mode is selected at
|
||||
runtime by:
|
||||
|
||||
* ``prompt["modalities"]`` contains ``"image"``: **T2I** (text-to-image).
|
||||
@@ -10,8 +10,6 @@ runtime by:
|
||||
(text-to-video).
|
||||
* ``multi_modal_data['image']`` present on the prompt: **I2V**
|
||||
(handled by :func:`get_cosmos3_pre_process_func`)
|
||||
* ``multi_modal_data['video']`` present on the prompt with no action mode:
|
||||
**V2V** (handled by :func:`get_cosmos3_pre_process_func`)
|
||||
|
||||
"""
|
||||
|
||||
@@ -42,10 +40,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
)
|
||||
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 (
|
||||
ReferenceVideoDecodeSpec,
|
||||
SupportImageInput,
|
||||
)
|
||||
from vllm_omni.diffusion.models.interface import SupportImageInput
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
|
||||
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
@@ -68,17 +63,6 @@ from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION = (0, 1)
|
||||
COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first"
|
||||
# Mirrors the WAN VAE's temporal compression. Authoritative value is
|
||||
# ``self.vae.config.scale_factor_temporal`` at runtime; this constant exists so
|
||||
# off-line / API code that runs before the pipeline is constructed can compute
|
||||
# pixel-frame budgets without instantiating the VAE.
|
||||
COSMOS3_VAE_TEMPORAL_COMPRESSION = 4
|
||||
COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = (
|
||||
max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1
|
||||
)
|
||||
COSMOS3_V2V_DEFAULT_FLOW_SHIFT = 10.0
|
||||
COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS."
|
||||
COSMOS3_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution."
|
||||
COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution."
|
||||
@@ -107,55 +91,16 @@ COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0)
|
||||
COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
|
||||
|
||||
|
||||
def _normalize_condition_frame_indexes_vision(value: Any) -> tuple[int, ...]:
|
||||
"""Normalize Cosmos3 vision-conditioning latent frame indexes."""
|
||||
if value is None:
|
||||
return COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION
|
||||
if isinstance(value, str):
|
||||
value = [item.strip() for item in value.split(",") if item.strip()]
|
||||
elif isinstance(value, int):
|
||||
value = [value]
|
||||
|
||||
if not isinstance(value, Iterable):
|
||||
raise TypeError(
|
||||
"Cosmos3 condition_frame_indexes_vision must be an int, comma-separated string, "
|
||||
f"or iterable of ints; got {type(value)!r}."
|
||||
)
|
||||
|
||||
indexes = tuple(sorted({int(index) for index in value}))
|
||||
if not indexes:
|
||||
raise ValueError("Cosmos3 condition_frame_indexes_vision must contain at least one index.")
|
||||
if any(index < 0 for index in indexes):
|
||||
raise ValueError(f"Cosmos3 condition_frame_indexes_vision must be non-negative, got {indexes}.")
|
||||
return indexes
|
||||
|
||||
|
||||
def _condition_pixel_frame_count(
|
||||
condition_frame_indexes_vision: Iterable[int],
|
||||
temporal_compression: int = COSMOS3_VAE_TEMPORAL_COMPRESSION,
|
||||
) -> int:
|
||||
return max(condition_frame_indexes_vision) * int(temporal_compression) + 1
|
||||
|
||||
|
||||
def _normalize_condition_video_keep(value: Any) -> str:
|
||||
keep = str(value or COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP).strip().lower()
|
||||
if keep not in {"first", "last"}:
|
||||
raise ValueError("Cosmos3 condition_video_keep must be either 'first' or 'last'.")
|
||||
return keep
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Post-process function (registered in registry.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
"""Pre-process function for T2V, I2V, V2V, and action inputs.
|
||||
"""Pre-process function for both T2V and I2V.
|
||||
|
||||
For T2V (no image in ``multi_modal_data``), the request is returned
|
||||
unchanged after the optional guardrails check. For I2V (image present),
|
||||
the conditioning image is loaded, aspect-resized + center-cropped, and
|
||||
stored back on the prompt as ``additional_information.preprocessed_image``.
|
||||
For V2V (video present without action mode), source frames are cropped to
|
||||
the target size and stored as ``additional_information.preprocessed_video``.
|
||||
"""
|
||||
from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled
|
||||
|
||||
@@ -192,27 +137,7 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
return PIL.Image.open(value).convert("RGB")
|
||||
if isinstance(value, PIL.Image.Image):
|
||||
return value.convert("RGB")
|
||||
if isinstance(value, np.ndarray):
|
||||
array = value
|
||||
if array.ndim == 3 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
|
||||
array = np.transpose(array, (1, 2, 0))
|
||||
if np.issubdtype(array.dtype, np.floating):
|
||||
if array.min() < 0.0 or array.max() > 1.0:
|
||||
array = np.clip(array, -1.0, 1.0) * 0.5 + 0.5
|
||||
array = (np.clip(array, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||||
return PIL.Image.fromarray(array).convert("RGB")
|
||||
if isinstance(value, torch.Tensor):
|
||||
tensor = value.detach().cpu()
|
||||
if tensor.ndim == 3 and tensor.shape[0] in (3, 4):
|
||||
tensor = tensor.permute(1, 2, 0)
|
||||
if tensor.is_floating_point():
|
||||
if tensor.min().item() < 0.0 or tensor.max().item() > 1.0:
|
||||
tensor = tensor.clamp(-1.0, 1.0) * 0.5 + 0.5
|
||||
tensor = (tensor.clamp(0.0, 1.0) * 255.0).round().to(torch.uint8)
|
||||
return PIL.Image.fromarray(tensor.numpy()).convert("RGB")
|
||||
raise TypeError(
|
||||
f"Cosmos3 preprocessing expected PIL image, numpy array, torch tensor, or path, got {type(value)!r}."
|
||||
)
|
||||
raise TypeError(f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}.")
|
||||
|
||||
def _resize_and_pad_action_image(image: PIL.Image.Image, target_h: int, target_w: int) -> PIL.Image.Image:
|
||||
scale = min(target_w / image.width, target_h / image.height, 1.0)
|
||||
@@ -245,61 +170,6 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
processed = [_preprocess_action_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in frames]
|
||||
return torch.stack(processed, dim=1).unsqueeze(0).contiguous()
|
||||
|
||||
def _preprocess_condition_image(image: PIL.Image.Image, target_h: int, target_w: int) -> torch.Tensor:
|
||||
scale = max(target_w / image.width, target_h / image.height)
|
||||
resize_w = int(np.ceil(scale * image.width))
|
||||
resize_h = int(np.ceil(scale * image.height))
|
||||
image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
|
||||
left = (resize_w - target_w) // 2
|
||||
top = (resize_h - target_h) // 2
|
||||
image = image.crop((left, top, left + target_w, top + target_h))
|
||||
return video_processor.preprocess(image, height=target_h, width=target_w)
|
||||
|
||||
def _video_payload_to_frames(video: Any) -> list[Any]:
|
||||
if isinstance(video, list):
|
||||
return video
|
||||
if isinstance(video, torch.Tensor):
|
||||
tensor = video.detach().cpu()
|
||||
if tensor.ndim == 5:
|
||||
if tensor.shape[0] != 1:
|
||||
raise TypeError("Cosmos3 video preprocessing supports only batch size 1.")
|
||||
tensor = tensor[0]
|
||||
if tensor.ndim == 4 and tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4):
|
||||
return [tensor[:, i] for i in range(tensor.shape[1])]
|
||||
if tensor.ndim == 4 and tensor.shape[-1] in (3, 4):
|
||||
return [tensor[i] for i in range(tensor.shape[0])]
|
||||
if isinstance(video, np.ndarray):
|
||||
array = video
|
||||
if array.ndim == 5:
|
||||
if array.shape[0] != 1:
|
||||
raise TypeError("Cosmos3 video preprocessing supports only batch size 1.")
|
||||
array = array[0]
|
||||
if array.ndim == 4 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
|
||||
return [array[:, i] for i in range(array.shape[1])]
|
||||
if array.ndim == 4 and array.shape[-1] in (3, 4):
|
||||
return [array[i] for i in range(array.shape[0])]
|
||||
raise TypeError("Cosmos3 video input must be a non-empty list of frames or a single video tensor/array.")
|
||||
|
||||
def _select_video_frames(frames: list[Any], max_frames: int, keep: str) -> list[Any]:
|
||||
if not frames:
|
||||
raise ValueError("Cosmos3 video input must contain at least one frame.")
|
||||
if keep == "last":
|
||||
return frames[-max_frames:]
|
||||
return frames[:max_frames]
|
||||
|
||||
def _preprocess_condition_video(
|
||||
frames: list[Any],
|
||||
target_h: int,
|
||||
target_w: int,
|
||||
max_frames: int,
|
||||
keep: str,
|
||||
) -> torch.Tensor:
|
||||
selected = _select_video_frames(frames, max_frames, keep)
|
||||
processed = [
|
||||
_preprocess_condition_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in selected
|
||||
]
|
||||
return torch.stack(processed, dim=1).unsqueeze(0).contiguous()
|
||||
|
||||
def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
|
||||
action_mode = _request_action_mode(request)
|
||||
if is_guardrails_enabled(od_config, request.sampling_params):
|
||||
@@ -313,23 +183,16 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
multi_modal_data = prompt.get("multi_modal_data", {}) or {}
|
||||
raw_image = multi_modal_data.get("image")
|
||||
raw_video = multi_modal_data.get("video")
|
||||
if raw_image is None and raw_video is None:
|
||||
if raw_image is None and not (action_mode is not None and raw_video is not None):
|
||||
continue
|
||||
if raw_image is not None and raw_video is not None and action_mode is None:
|
||||
raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
|
||||
|
||||
if "additional_information" not in prompt:
|
||||
prompt["additional_information"] = {}
|
||||
|
||||
raw_video_frames: list[Any] | None = None
|
||||
if raw_video is not None:
|
||||
raw_video_frames = _video_payload_to_frames(raw_video)
|
||||
if not raw_video_frames:
|
||||
raise TypeError("Cosmos3 video input must be a non-empty list of PIL images or image paths.")
|
||||
|
||||
if raw_image is None:
|
||||
assert raw_video_frames is not None # raw_image and raw_video can't both be None here
|
||||
image = _pil_to_rgb(raw_video_frames[0])
|
||||
if not isinstance(raw_video, list) or not raw_video:
|
||||
raise TypeError("Cosmos3 action video input must be a non-empty list of PIL images or image paths.")
|
||||
image = _pil_to_rgb(raw_video[0])
|
||||
else:
|
||||
image = _pil_to_rgb(raw_image)
|
||||
|
||||
@@ -356,38 +219,23 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
|
||||
int(target_h),
|
||||
int(target_w),
|
||||
)
|
||||
elif raw_video is None:
|
||||
prompt["additional_information"]["preprocessed_image"] = _preprocess_condition_image(
|
||||
image,
|
||||
int(target_h),
|
||||
int(target_w),
|
||||
)
|
||||
else:
|
||||
assert raw_video_frames is not None
|
||||
extra = _extra_args(request)
|
||||
condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision(
|
||||
extra.get(
|
||||
"condition_frame_indexes_vision",
|
||||
prompt.get("condition_frame_indexes_vision"),
|
||||
)
|
||||
scale = max(target_w / image.width, target_h / image.height)
|
||||
resize_w = int(np.ceil(scale * image.width))
|
||||
resize_h = int(np.ceil(scale * image.height))
|
||||
image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
|
||||
left = (resize_w - target_w) // 2
|
||||
top = (resize_h - target_h) // 2
|
||||
image = image.crop((left, top, left + target_w, top + target_h))
|
||||
|
||||
prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
|
||||
image, height=target_h, width=target_w
|
||||
)
|
||||
keep = _normalize_condition_video_keep(
|
||||
extra.get("condition_video_keep", prompt.get("condition_video_keep"))
|
||||
)
|
||||
max_frames = _condition_pixel_frame_count(condition_frame_indexes_vision)
|
||||
prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video(
|
||||
raw_video_frames,
|
||||
int(target_h),
|
||||
int(target_w),
|
||||
max_frames,
|
||||
keep,
|
||||
)
|
||||
prompt["additional_information"]["condition_frame_indexes_vision"] = list(
|
||||
condition_frame_indexes_vision
|
||||
)
|
||||
if action_mode is not None and raw_video_frames is not None:
|
||||
if action_mode is not None and raw_video is not None:
|
||||
if not isinstance(raw_video, list):
|
||||
raise TypeError("Cosmos3 action video input must be a list of PIL images or image paths.")
|
||||
prompt["additional_information"]["preprocessed_video"] = _preprocess_action_video(
|
||||
raw_video_frames,
|
||||
raw_video,
|
||||
int(target_h),
|
||||
int(target_w),
|
||||
)
|
||||
@@ -497,13 +345,13 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
|
||||
class Cosmos3OmniDiffusersPipeline(
|
||||
nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
|
||||
):
|
||||
"""Cosmos3 text/image/video-to-video / text-to-image pipeline.
|
||||
"""Cosmos3 text/image-to-video / text-to-image pipeline.
|
||||
|
||||
Architecture: Mixture-of-Transformers with Qwen3-VL backbone.
|
||||
- Understanding pathway: causal self-attention on text (runs once, K/V cached)
|
||||
- Generation pathway: cross-attention on noisy visual latents (runs each step)
|
||||
|
||||
Supports T2V, I2V, V2V, and T2I from the same class. Mode is selected at
|
||||
Supports T2V, I2V, and T2I from the same class. Mode is selected at
|
||||
runtime:
|
||||
|
||||
* **T2I** when ``prompt["modalities"]`` contains ``"image"``. Latent
|
||||
@@ -517,45 +365,12 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
Frame 0 of the initial latent is set to the VAE-encoded conditioning
|
||||
image, frame-0 noise predictions are masked to zero, and the clean
|
||||
image latent is re-injected at frame 0 after each scheduler step.
|
||||
* **V2V** when the request supplies a preprocessed video via
|
||||
``multi_modal_data['video']`` without an action mode. Explicit latent
|
||||
frame indexes are kept clean with ``noisy_frame_mask`` and re-injected
|
||||
after each scheduler step.
|
||||
* **T2V** otherwise (default video generation).
|
||||
"""
|
||||
|
||||
support_image_input: ClassVar[bool] = True
|
||||
color_format: ClassVar[str] = "RGB"
|
||||
|
||||
@classmethod
|
||||
def reference_video_decode_spec(
|
||||
cls,
|
||||
*,
|
||||
num_frames: int | None = None,
|
||||
extra_args: dict[str, Any] | None = None,
|
||||
) -> ReferenceVideoDecodeSpec:
|
||||
extra_args = extra_args if isinstance(extra_args, dict) else {}
|
||||
action_mode = normalize_action_mode(extra_args.get("action_mode"))
|
||||
if action_mode is not None:
|
||||
if num_frames is not None:
|
||||
return ReferenceVideoDecodeSpec(max_frames=int(num_frames), keep="first")
|
||||
action_chunk_size = extra_args.get("action_chunk_size")
|
||||
if action_chunk_size is not None:
|
||||
try:
|
||||
max_frames = int(action_chunk_size) + 1
|
||||
except (TypeError, ValueError):
|
||||
max_frames = None
|
||||
if max_frames is not None and max_frames > 0:
|
||||
return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first")
|
||||
return ReferenceVideoDecodeSpec(max_frames=None, keep="first")
|
||||
|
||||
condition_indexes = _normalize_condition_frame_indexes_vision(extra_args.get("condition_frame_indexes_vision"))
|
||||
max_frames = _condition_pixel_frame_count(condition_indexes)
|
||||
if num_frames is not None:
|
||||
max_frames = min(max_frames, int(num_frames))
|
||||
keep = _normalize_condition_video_keep(extra_args.get("condition_video_keep"))
|
||||
return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=keep)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
@@ -1373,70 +1188,6 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
velocity_mask = 1.0 - condition_mask
|
||||
return latents, velocity_mask, image_latent
|
||||
|
||||
def _prepare_latents_v2v(
|
||||
self,
|
||||
video_tensor: torch.Tensor,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
generator: torch.Generator,
|
||||
condition_frame_indexes_vision: Iterable[int] | int | str | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Prepare V2V latents with explicit clean conditioned latent frames."""
|
||||
del height, width
|
||||
if video_tensor.ndim == 4:
|
||||
video_tensor = video_tensor.unsqueeze(0)
|
||||
if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
|
||||
raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
|
||||
if video_tensor.shape[2] < 1:
|
||||
raise ValueError("Cosmos3 V2V video tensor must contain at least one frame.")
|
||||
|
||||
C = self.transformer.latent_channel_size
|
||||
T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
||||
H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
|
||||
W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
|
||||
indexes = _normalize_condition_frame_indexes_vision(condition_frame_indexes_vision)
|
||||
out_of_range = [index for index in indexes if index >= T_lat]
|
||||
if out_of_range:
|
||||
raise ValueError(
|
||||
"Cosmos3 condition_frame_indexes_vision contains indexes outside the latent video: "
|
||||
f"indexes={indexes}, latent_frames={T_lat}."
|
||||
)
|
||||
|
||||
noise = randn_tensor(
|
||||
(1, C, T_lat, H_lat, W_lat),
|
||||
generator=generator,
|
||||
device=self.device,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
condition_pixel_frames = _condition_pixel_frame_count(indexes, self.vae_scale_factor_temporal)
|
||||
condition_video = video_tensor[:, :, :condition_pixel_frames]
|
||||
if condition_video.shape[2] < condition_pixel_frames:
|
||||
pad = condition_video[:, :, -1:].repeat(1, 1, condition_pixel_frames - condition_video.shape[2], 1, 1)
|
||||
condition_video = torch.cat([condition_video, pad], dim=2)
|
||||
|
||||
cond_prefix_latent = self._encode_video_tensor(condition_video)
|
||||
expected_prefix = (1, C, max(indexes) + 1, H_lat, W_lat)
|
||||
if (
|
||||
cond_prefix_latent.shape[0] != expected_prefix[0]
|
||||
or cond_prefix_latent.shape[1] != expected_prefix[1]
|
||||
or cond_prefix_latent.shape[2] < expected_prefix[2]
|
||||
or cond_prefix_latent.shape[3:] != expected_prefix[3:]
|
||||
):
|
||||
raise ValueError(
|
||||
"Cosmos3 V2V condition latent shape mismatch: "
|
||||
f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}."
|
||||
)
|
||||
|
||||
condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype)
|
||||
condition_latents = torch.zeros_like(noise)
|
||||
for index in indexes:
|
||||
condition_mask[:, :, index, :, :] = 1.0
|
||||
condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1]
|
||||
latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
|
||||
velocity_mask = 1.0 - condition_mask
|
||||
return latents, velocity_mask, condition_latents
|
||||
|
||||
def _prepare_latents_action_video(
|
||||
self,
|
||||
video_tensor: torch.Tensor,
|
||||
@@ -1816,20 +1567,19 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
prompt = prompt_data
|
||||
negative_prompt = None
|
||||
image_tensor = None
|
||||
video_tensor = None
|
||||
action_video_tensor = None
|
||||
else:
|
||||
prompt = prompt_data.get("prompt", "")
|
||||
negative_prompt = prompt_data.get("negative_prompt")
|
||||
additional_info = prompt_data.get("additional_information", {}) or {}
|
||||
image_tensor = additional_info.get("preprocessed_image")
|
||||
video_tensor = additional_info.get("preprocessed_video")
|
||||
action_video_tensor = additional_info.get("preprocessed_video")
|
||||
|
||||
sp = req.sampling_params
|
||||
is_t2i = self._is_t2i_request(req)
|
||||
sound_enabled = self._is_sound_request(prompt_data, sp)
|
||||
action_mode = self._get_action_mode(prompt_data, sp)
|
||||
action_enabled = action_mode is not None
|
||||
action_video_tensor = video_tensor if action_enabled else None
|
||||
if action_enabled and is_t2i:
|
||||
raise ValueError("Cosmos3 action generation is supported only for video outputs.")
|
||||
if action_enabled and sound_enabled:
|
||||
@@ -1853,11 +1603,6 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
)
|
||||
if negative_prompt is None:
|
||||
negative_prompt = ""
|
||||
if image_tensor is not None and video_tensor is not None and not action_enabled:
|
||||
raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
|
||||
if video_tensor is not None and is_t2i:
|
||||
raise ValueError("Cosmos3 video-to-video generation is supported only for video outputs.")
|
||||
is_v2v = video_tensor is not None and not is_t2i and not action_enabled
|
||||
|
||||
# T2I and T2V share the same model + forward path; only defaults
|
||||
# differ:
|
||||
@@ -1881,7 +1626,7 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
# Fall back to the engine-init shift, NOT None: passing None
|
||||
# to ``_set_flow_shift`` would leak a prior T2I rebuild
|
||||
# (shift=3.0) into a subsequent video request.
|
||||
default_flow_shift = COSMOS3_V2V_DEFAULT_FLOW_SHIFT if is_v2v else self._engine_init_flow_shift
|
||||
default_flow_shift = self._engine_init_flow_shift
|
||||
default_guidance_interval = None
|
||||
batch_size = 1 # Existing video pipeline assumes B=1.
|
||||
|
||||
@@ -2030,23 +1775,6 @@ class Cosmos3OmniDiffusersPipeline(
|
||||
generator,
|
||||
)
|
||||
image_latent = condition_latents[:, :, 0:1]
|
||||
elif is_v2v:
|
||||
condition_frame_indexes_vision = _normalize_condition_frame_indexes_vision(
|
||||
self._get_sp_param(
|
||||
sp,
|
||||
"condition_frame_indexes_vision",
|
||||
self._get_prompt_param(prompt_data, "condition_frame_indexes_vision", None),
|
||||
)
|
||||
)
|
||||
latents, velocity_mask, condition_latents = self._prepare_latents_v2v(
|
||||
video_tensor,
|
||||
height,
|
||||
width,
|
||||
num_frames,
|
||||
generator,
|
||||
condition_frame_indexes_vision,
|
||||
)
|
||||
image_latent = None
|
||||
elif image_tensor is not None and not is_t2i:
|
||||
latents, velocity_mask, image_latent = self._prepare_latents_i2v(
|
||||
image_tensor,
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Iterable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
@@ -31,9 +31,6 @@ from vllm_omni.diffusion.forward_context import get_forward_context
|
||||
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
|
||||
from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@@ -335,8 +332,6 @@ class HunyuanVideo15Attention(nn.Module):
|
||||
out_bias: bool = True,
|
||||
eps: float = 1e-6,
|
||||
out_dim: int | None = None,
|
||||
quant_config: "QuantizationConfig | None" = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
@@ -355,8 +350,6 @@ class HunyuanVideo15Attention(nn.Module):
|
||||
head_size=self.head_dim,
|
||||
total_num_heads=self.heads,
|
||||
bias=bias,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_qkv",
|
||||
)
|
||||
|
||||
self.to_out = nn.ModuleList(
|
||||
@@ -367,8 +360,6 @@ class HunyuanVideo15Attention(nn.Module):
|
||||
bias=out_bias,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_out.0",
|
||||
),
|
||||
nn.Identity(), # placeholder for dropout (none used)
|
||||
]
|
||||
@@ -383,8 +374,6 @@ class HunyuanVideo15Attention(nn.Module):
|
||||
head_size=self.head_dim,
|
||||
total_num_heads=self.heads,
|
||||
bias=added_proj_bias,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.add_kv_proj",
|
||||
)
|
||||
|
||||
self.to_add_out = RowParallelLinear(
|
||||
@@ -393,8 +382,6 @@ class HunyuanVideo15Attention(nn.Module):
|
||||
bias=out_bias,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_add_out",
|
||||
)
|
||||
|
||||
self.rope = RotaryEmbedding(is_neox_style=False)
|
||||
@@ -414,8 +401,6 @@ class HunyuanVideo15Attention(nn.Module):
|
||||
image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
hidden_states_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
# Ensure contiguous for FP8 quantized linear layers
|
||||
hidden_states = hidden_states.contiguous()
|
||||
qkv, _ = self.to_qkv(hidden_states)
|
||||
q_size = self.to_qkv.num_heads * self.head_dim
|
||||
kv_size = self.to_qkv.num_kv_heads * self.head_dim
|
||||
@@ -436,7 +421,6 @@ class HunyuanVideo15Attention(nn.Module):
|
||||
key = self.rope(key, cos, sin)
|
||||
|
||||
if encoder_hidden_states is not None:
|
||||
encoder_hidden_states = encoder_hidden_states.contiguous()
|
||||
encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
|
||||
add_q_size = self.add_kv_proj.num_heads * self.head_dim
|
||||
add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
|
||||
@@ -508,8 +492,6 @@ class HunyuanVideo15TransformerBlock(nn.Module):
|
||||
attention_head_dim: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qk_norm: str = "rms_norm",
|
||||
quant_config: "QuantizationConfig | None" = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
hidden_size = num_attention_heads * attention_head_dim
|
||||
@@ -525,23 +507,13 @@ class HunyuanVideo15TransformerBlock(nn.Module):
|
||||
out_dim=hidden_size,
|
||||
bias=True,
|
||||
eps=1e-6,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.attn",
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.ff = FeedForward(
|
||||
dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio, quant_config=quant_config, prefix=f"{prefix}.ff"
|
||||
)
|
||||
self.ff = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio)
|
||||
|
||||
self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.ff_context = FeedForward(
|
||||
dim=hidden_size,
|
||||
dim_out=hidden_size,
|
||||
mult=mlp_ratio,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.ff_context",
|
||||
)
|
||||
self.ff_context = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -629,7 +601,6 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
|
||||
target_size: int = 640,
|
||||
task_type: str = "i2v",
|
||||
use_meanflow: bool = False,
|
||||
quant_config: "QuantizationConfig | None" = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
@@ -661,14 +632,9 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
HunyuanVideo15TransformerBlock(
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
mlp_ratio=mlp_ratio,
|
||||
qk_norm=qk_norm,
|
||||
quant_config=quant_config,
|
||||
prefix=f"transformer_blocks.{i}",
|
||||
num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm
|
||||
)
|
||||
for i in range(num_layers)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@@ -143,9 +143,7 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff
|
||||
self.scheduler._shift = od_config.flow_shift
|
||||
|
||||
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel)
|
||||
self.transformer = HunyuanVideo15Transformer3DModel(
|
||||
od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
|
||||
)
|
||||
self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs)
|
||||
|
||||
# Check if model uses meanflow (distilled variants)
|
||||
self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False)
|
||||
|
||||
@@ -186,9 +186,7 @@ class HunyuanVideo15I2VPipeline(
|
||||
self.scheduler._shift = od_config.flow_shift
|
||||
|
||||
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel)
|
||||
self.transformer = HunyuanVideo15Transformer3DModel(
|
||||
od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
|
||||
)
|
||||
self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs)
|
||||
|
||||
self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False)
|
||||
|
||||
|
||||
@@ -2,12 +2,10 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
ClassVar,
|
||||
Literal,
|
||||
Protocol,
|
||||
runtime_checkable,
|
||||
)
|
||||
@@ -25,12 +23,6 @@ class SupportImageInput(Protocol):
|
||||
color_format: ClassVar[str] = "RGB" # Default color format
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ReferenceVideoDecodeSpec:
|
||||
max_frames: int | None = None
|
||||
keep: Literal["first", "last"] = "first"
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SupportAudioInput(Protocol):
|
||||
support_audio_input: ClassVar[bool] = True
|
||||
|
||||
@@ -793,9 +793,17 @@ class LancePipeline(BagelPipeline):
|
||||
curr_rope=cfg_text_context["ropes"],
|
||||
video_shapes=[video_shape],
|
||||
)
|
||||
cfg_img_lat = self.bagel.prepare_video_latent_cfg(
|
||||
curr_kvlens=gen_context["kv_lens"],
|
||||
curr_rope=gen_context["ropes"],
|
||||
video_shapes=[video_shape],
|
||||
)
|
||||
for k, v in cfg_text_lat.items():
|
||||
if torch.is_tensor(v):
|
||||
cfg_text_lat[k] = v.to(self.device)
|
||||
for k, v in cfg_img_lat.items():
|
||||
if torch.is_tensor(v):
|
||||
cfg_img_lat[k] = v.to(self.device)
|
||||
|
||||
self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed)
|
||||
|
||||
@@ -814,7 +822,7 @@ class LancePipeline(BagelPipeline):
|
||||
latents, *_ = self.bagel.generate_image(
|
||||
past_key_values=gen_context["past_key_values"],
|
||||
cfg_text_past_key_values=cfg_text_context["past_key_values"],
|
||||
cfg_img_past_key_values=None, # no img CFG branch
|
||||
cfg_img_past_key_values=gen_context["past_key_values"], # no img CFG branch
|
||||
num_timesteps=num_timesteps,
|
||||
timestep_shift=timestep_shift,
|
||||
cfg_text_scale=cfg_text_scale,
|
||||
@@ -827,7 +835,7 @@ class LancePipeline(BagelPipeline):
|
||||
# ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
|
||||
# / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
|
||||
# (derived from ``cfg_text_past_key_values``).
|
||||
cfg_img_packed_position_ids=None,
|
||||
cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
|
||||
# ``cfg_img_*`` index/lens kwargs removed — same as above.
|
||||
)
|
||||
|
||||
@@ -1030,10 +1038,11 @@ class LancePipeline(BagelPipeline):
|
||||
if torch.is_tensor(v):
|
||||
cfg_text_lat[k] = v.to(self.device)
|
||||
|
||||
# cfg_img branch is off by default (cfg_img_scale=1.0 matches
|
||||
# upstream ``cfg_vit_scale=1.0``). Avoid passing synthetic metadata
|
||||
# in that path; Bagel.generate_image only reads it when scale > 1.0.
|
||||
cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None
|
||||
# cfg_img branch: keep off (cfg_img_scale=1.0 matches upstream
|
||||
# ``cfg_vit_scale=1.0``). generate_image needs the cfg_img_* args
|
||||
# but skips the branch when scale<=1.0; just point them at the
|
||||
# cond's lat metadata so shapes/types are fine.
|
||||
cfg_img_lat = cfg_text_lat
|
||||
|
||||
self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed)
|
||||
|
||||
@@ -1062,9 +1071,7 @@ class LancePipeline(BagelPipeline):
|
||||
# ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
|
||||
# / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
|
||||
# (derived from ``cfg_text_past_key_values``).
|
||||
cfg_img_packed_position_ids=(
|
||||
cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None
|
||||
),
|
||||
cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
|
||||
# ``cfg_img_*`` index/lens kwargs removed — same as above.
|
||||
)
|
||||
|
||||
@@ -1262,7 +1269,7 @@ class LancePipeline(BagelPipeline):
|
||||
for k, v in cfg_text_lat.items():
|
||||
if torch.is_tensor(v):
|
||||
cfg_text_lat[k] = v.to(self.device)
|
||||
cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None
|
||||
cfg_img_lat = cfg_text_lat
|
||||
|
||||
self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed)
|
||||
|
||||
@@ -1351,9 +1358,7 @@ class LancePipeline(BagelPipeline):
|
||||
# ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
|
||||
# / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
|
||||
# (derived from ``cfg_text_past_key_values``).
|
||||
cfg_img_packed_position_ids=(
|
||||
cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None
|
||||
),
|
||||
cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
|
||||
# ``cfg_img_*`` index/lens kwargs removed — same as above.
|
||||
frame_condition_token_indexes=frame_condition_token_indexes,
|
||||
)
|
||||
@@ -1582,9 +1587,9 @@ class LancePipeline(BagelPipeline):
|
||||
if torch.is_tensor(v):
|
||||
cfg_text_lat[k] = v.to(self.device)
|
||||
|
||||
# cfg_img branch is off by default; pass metadata only when an
|
||||
# explicit cfg_img_scale enables the branch.
|
||||
cfg_img_lat = cfg_text_lat if cfg_img_scale > 1.0 else None
|
||||
# cfg_img branch is off (cfg_img_scale=1.0); generate_image needs the
|
||||
# arg surface but skips the branch when scale<=1.0.
|
||||
cfg_img_lat = cfg_text_lat
|
||||
|
||||
self._regen_init_noise_on_device(gen_input_lat, req.sampling_params.seed)
|
||||
|
||||
@@ -1613,9 +1618,7 @@ class LancePipeline(BagelPipeline):
|
||||
# ``cfg_text_packed_query_indexes`` / ``cfg_text_key_values_lens``
|
||||
# / ``cfg_text_packed_key_value_indexes`` removed post-main-merge
|
||||
# (derived from ``cfg_text_past_key_values``).
|
||||
cfg_img_packed_position_ids=(
|
||||
cfg_img_lat["cfg_packed_position_ids"] if cfg_img_lat is not None else None
|
||||
),
|
||||
cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
|
||||
# ``cfg_img_*`` index/lens kwargs removed — same as above.
|
||||
)
|
||||
|
||||
|
||||
@@ -62,17 +62,7 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) ->
|
||||
kwargs["has_weight"] = elementwise_affine
|
||||
elif not elementwise_affine:
|
||||
raise TypeError("RMSNorm backend does not support disabling affine weights.")
|
||||
norm = RMSNorm(hidden_size, **kwargs)
|
||||
weight = getattr(norm, "weight", None)
|
||||
if (
|
||||
not elementwise_affine
|
||||
and isinstance(weight, torch.Tensor)
|
||||
and not isinstance(weight, nn.Parameter)
|
||||
and "weight" not in dict(norm.named_buffers())
|
||||
):
|
||||
delattr(norm, "weight")
|
||||
norm.register_buffer("weight", weight, persistent=False)
|
||||
return norm
|
||||
return RMSNorm(hidden_size, **kwargs)
|
||||
|
||||
|
||||
def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
|
||||
|
||||
@@ -10,6 +10,7 @@ This pipeline does NOT inherit from LTX2Pipeline because:
|
||||
versus LTX-2's per_layer_masked_mean_norm + shared projection path
|
||||
- LTX-2.3 uses a BWE vocoder outputting 48kHz audio (not 16kHz)
|
||||
- LTX-2.3 transformer requires the sigma parameter for prompt modulation
|
||||
- CPU offloading is required for the 22B transformer (~44GB VRAM)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -19,7 +20,7 @@ import json
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from contextlib import nullcontext
|
||||
from typing import Any, ClassVar
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -45,9 +46,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||||
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
|
||||
from .pipeline_ltx2 import (
|
||||
@@ -113,7 +112,7 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
|
||||
return post_process_func
|
||||
|
||||
|
||||
class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsComponentDiscovery):
|
||||
class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
|
||||
"""Fully independent LTX-2.3 pipeline.
|
||||
|
||||
Key differences from LTX2Pipeline:
|
||||
@@ -121,14 +120,11 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
- Connectors: uses padding_side API (not additive_mask)
|
||||
- Vocoder: uses LTX2VocoderWithBWE (48kHz output)
|
||||
- Transformer: passes sigma for prompt_adaln
|
||||
- CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU
|
||||
"""
|
||||
|
||||
# Audio is diffused jointly with video; warmup must size audio tokens.
|
||||
dummy_run_num_frames = 2
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "connectors"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"]
|
||||
_resident_modules: ClassVar[list[str]] = ["vocoder"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -170,7 +166,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
# --- Tokenizer (lightweight, stays wherever) ---
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
|
||||
|
||||
# --- Text encoder ---
|
||||
# --- Text encoder: load on CPU, move to GPU only during encoding ---
|
||||
with torch.device("cpu"):
|
||||
self.text_encoder = from_pretrained_with_prefetch(
|
||||
Gemma3ForConditionalGeneration.from_pretrained,
|
||||
@@ -181,7 +177,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
|
||||
# --- Connectors (LTX-2.3 connectors include caption projection) ---
|
||||
# --- Connectors: CPU (LTX-2.3 connectors include caption projection) ---
|
||||
self.connectors = from_pretrained_with_prefetch(
|
||||
LTX2TextConnectors.from_pretrained,
|
||||
model,
|
||||
@@ -191,7 +187,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
|
||||
# --- VAE, Audio VAE ---
|
||||
# --- VAE, Audio VAE: CPU ---
|
||||
self.vae = from_pretrained_with_prefetch(
|
||||
AutoencoderKLLTX2Video.from_pretrained,
|
||||
model,
|
||||
@@ -224,7 +220,6 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
transformer_config = load_transformer_config(model, "transformer", local_files_only)
|
||||
quant_config = getattr(self.od_config, "quantization_config", None)
|
||||
self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
|
||||
self._place_aux_components()
|
||||
|
||||
# --- Scheduler ---
|
||||
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
||||
@@ -266,20 +261,6 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
self._num_timesteps = None
|
||||
self._current_timestep = None
|
||||
|
||||
def _place_aux_components(self) -> None:
|
||||
parallel_config = getattr(self.od_config, "parallel_config", None)
|
||||
use_managed_placement = bool(
|
||||
getattr(self.od_config, "enable_cpu_offload", False)
|
||||
or getattr(self.od_config, "enable_layerwise_offload", False)
|
||||
or getattr(parallel_config, "use_hsdp", False)
|
||||
)
|
||||
if use_managed_placement:
|
||||
return
|
||||
|
||||
modules = ModuleDiscovery.discover(self)
|
||||
for module in (*modules.encoders, *modules.vaes, *modules.resident_modules):
|
||||
module.to(self.device)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Text Encoding (LTX-2.3 specific)
|
||||
# ------------------------------------------------------------------
|
||||
@@ -322,11 +303,16 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
text_input_ids = text_inputs.input_ids.to(device)
|
||||
prompt_attention_mask = text_inputs.attention_mask.to(device)
|
||||
|
||||
# Move text encoder to GPU for encoding
|
||||
self.text_encoder.to(device)
|
||||
text_encoder_outputs = self.text_encoder(
|
||||
input_ids=text_input_ids,
|
||||
attention_mask=prompt_attention_mask,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
# Move text encoder back to CPU immediately
|
||||
self.text_encoder.to("cpu")
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
hidden_states = text_encoder_outputs.hidden_states
|
||||
|
||||
@@ -972,10 +958,14 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
|
||||
prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
|
||||
|
||||
self.connectors.to(device)
|
||||
tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
|
||||
connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
|
||||
prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side
|
||||
)
|
||||
self.connectors.to("cpu")
|
||||
if torch.cuda.is_available():
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
positive_connector_prompt_embeds = connector_prompt_embeds
|
||||
positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds
|
||||
@@ -1251,14 +1241,22 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
|
||||
]
|
||||
latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise
|
||||
|
||||
# Move VAE, audio_vae, vocoder to GPU for decoding
|
||||
self.vae.to(device)
|
||||
latents = latents.to(self.vae.dtype)
|
||||
video = self.vae.decode(latents, timestep_decode, return_dict=False)[0]
|
||||
video = self.video_processor.postprocess_video(video, output_type=output_type)
|
||||
self.vae.to("cpu")
|
||||
|
||||
self.audio_vae.to(device)
|
||||
audio_latents = audio_latents.to(self.audio_vae.dtype)
|
||||
generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
|
||||
self.audio_vae.to("cpu")
|
||||
|
||||
self.vocoder.to(device)
|
||||
audio = self.vocoder(generated_mel_spectrograms)
|
||||
self.vocoder.to("cpu")
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
return DiffusionOutput(output=(video, audio))
|
||||
|
||||
|
||||
@@ -101,8 +101,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
def load_model(
|
||||
self,
|
||||
memory_pool_context_fn: callable | None = None,
|
||||
load_format: str = "default",
|
||||
custom_pipeline_name: str | None = None,
|
||||
load_format: str | None = None,
|
||||
custom_pipeline_name: str | type | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Load the diffusion model, apply compilation and offloading.
|
||||
@@ -138,6 +138,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
with get_memory_context():
|
||||
with DeviceMemoryProfiler() as m:
|
||||
self.pipeline = model_loader.load_model(
|
||||
od_config=self.od_config,
|
||||
load_device=load_device,
|
||||
load_format=load_format,
|
||||
custom_pipeline_name=custom_pipeline_name,
|
||||
|
||||
@@ -239,6 +239,7 @@ class DiffusionWorker:
|
||||
"Final IR op priority after setting vLLM-Omni overrides: %s", vllm_config.kernel_config.ir_op_priority
|
||||
)
|
||||
self.vllm_config = vllm_config
|
||||
current_omni_platform.init_diffusion_worker_vllm_config(vllm_config)
|
||||
|
||||
# Initialize distributed environment
|
||||
with (
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
async def extract_audio_from_video_async(video_url: str) -> tuple[np.ndarray, int | float]:
|
||||
"""Extract audio from a video URL using vllm's load_audio.
|
||||
|
||||
Returns a (audio_array, sample_rate) tuple compatible with audio format.
|
||||
All blocking I/O operations are run in a thread pool.
|
||||
"""
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed_url = urlparse(video_url)
|
||||
temp_video_file_path = None
|
||||
|
||||
def _download_video_sync(url: str) -> bytes:
|
||||
from urllib.request import urlopen
|
||||
|
||||
return urlopen(url).read()
|
||||
|
||||
def _write_temp_file_sync(data: bytes, suffix: str) -> str:
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
|
||||
temp_file.write(data)
|
||||
return temp_file.name
|
||||
|
||||
def _load_audio_sync(file_path: str) -> tuple[np.ndarray, int | float]:
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
return load_audio(file_path, sr=16000)
|
||||
|
||||
def _cleanup_file_sync(file_path: str) -> None:
|
||||
try:
|
||||
if os.path.exists(file_path):
|
||||
os.unlink(file_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
try:
|
||||
if parsed_url.scheme in ("http", "https"):
|
||||
video_data = await asyncio.to_thread(_download_video_sync, video_url)
|
||||
temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4")
|
||||
elif parsed_url.scheme == "file":
|
||||
from urllib.request import url2pathname
|
||||
|
||||
temp_video_file_path = url2pathname(parsed_url.path)
|
||||
elif parsed_url.scheme == "data":
|
||||
import base64
|
||||
|
||||
header, data = video_url.split(",", 1)
|
||||
video_data = base64.b64decode(data)
|
||||
temp_video_file_path = await asyncio.to_thread(_write_temp_file_sync, video_data, ".mp4")
|
||||
else:
|
||||
# Assume local file path
|
||||
temp_video_file_path = video_url
|
||||
|
||||
audio_array, sample_rate = await asyncio.to_thread(_load_audio_sync, temp_video_file_path)
|
||||
return audio_array, sample_rate
|
||||
finally:
|
||||
if temp_video_file_path and parsed_url.scheme in ("http", "https", "data"):
|
||||
await asyncio.to_thread(_cleanup_file_sync, temp_video_file_path)
|
||||
@@ -91,7 +91,6 @@ from vllm.utils import random_uuid
|
||||
from vllm.utils.system_utils import decorate_logs
|
||||
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
|
||||
|
||||
from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
|
||||
from vllm_omni.entrypoints.openai.image_api_utils import (
|
||||
@@ -126,7 +125,7 @@ from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServin
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
||||
from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
|
||||
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage, ReferenceVideo
|
||||
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage
|
||||
from vllm_omni.entrypoints.openai.serving_video_stream import OmniStreamingVideoHandler
|
||||
from vllm_omni.entrypoints.openai.stage_params import (
|
||||
build_stage_sampling_params_list,
|
||||
@@ -2532,84 +2531,6 @@ def _parse_form_json(value: str | None, expected_type: type | None = None) -> An
|
||||
return parsed
|
||||
|
||||
|
||||
def _config_get(config: Any, key: str, default: Any = None) -> Any:
|
||||
if isinstance(config, dict):
|
||||
return config.get(key, default)
|
||||
if hasattr(config, "get"):
|
||||
try:
|
||||
return config.get(key, default)
|
||||
except Exception:
|
||||
pass
|
||||
return getattr(config, key, default)
|
||||
|
||||
|
||||
def _stage_engine_args(stage_cfg: Any) -> Any:
|
||||
return _config_get(stage_cfg, "engine_args", {}) or {}
|
||||
|
||||
|
||||
def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]:
|
||||
if not stage_configs:
|
||||
return []
|
||||
|
||||
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
||||
|
||||
model_classes: list[type] = []
|
||||
for stage_cfg in stage_configs:
|
||||
if get_stage_type(stage_cfg) != "diffusion":
|
||||
continue
|
||||
model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name")
|
||||
if not model_class_name:
|
||||
continue
|
||||
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
|
||||
if model_cls is not None:
|
||||
model_classes.append(model_cls)
|
||||
return model_classes
|
||||
|
||||
|
||||
def _normalize_reference_video_decode_spec(spec: ReferenceVideoDecodeSpec) -> ReferenceVideoDecodeSpec:
|
||||
max_frames = spec.max_frames
|
||||
if max_frames is not None:
|
||||
try:
|
||||
max_frames = int(max_frames)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid reference video decode spec: max_frames must be an integer.",
|
||||
) from exc
|
||||
if max_frames <= 0:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid reference video decode spec: max_frames must be positive.",
|
||||
)
|
||||
|
||||
keep = str(spec.keep or "first").strip().lower()
|
||||
if keep not in {"first", "last"}:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Invalid reference video decode spec: keep must be either 'first' or 'last'.",
|
||||
)
|
||||
return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=cast(Literal["first", "last"], keep))
|
||||
|
||||
|
||||
def _reference_video_decode_spec(
|
||||
req: VideoGenerationRequest,
|
||||
stage_configs: list[Any] | None,
|
||||
) -> ReferenceVideoDecodeSpec:
|
||||
video_params = req.resolve_video_params()
|
||||
extra_params = req.extra_params if isinstance(req.extra_params, dict) else {}
|
||||
for model_cls in _diffusion_model_classes(stage_configs):
|
||||
resolver = getattr(model_cls, "reference_video_decode_spec", None)
|
||||
if resolver is None:
|
||||
continue
|
||||
try:
|
||||
spec = resolver(num_frames=video_params.num_frames, extra_args=extra_params)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc
|
||||
if spec is not None:
|
||||
return _normalize_reference_video_decode_spec(spec)
|
||||
return ReferenceVideoDecodeSpec(max_frames=video_params.num_frames, keep="first")
|
||||
|
||||
|
||||
def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse:
|
||||
resp = VideoResponse(
|
||||
model=model_name,
|
||||
@@ -2676,7 +2597,6 @@ async def _run_video_generation_job(
|
||||
request: VideoGenerationRequest,
|
||||
video_id: str,
|
||||
reference_image: ReferenceImage | None = None,
|
||||
reference_video: ReferenceVideo | None = None,
|
||||
app_state: Any | None = None,
|
||||
) -> None:
|
||||
job = await VIDEO_STORE.get(video_id)
|
||||
@@ -2689,10 +2609,7 @@ async def _run_video_generation_job(
|
||||
output_path = None
|
||||
try:
|
||||
video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes(
|
||||
request,
|
||||
video_id,
|
||||
reference_image=reference_image,
|
||||
reference_video=reference_video,
|
||||
request, video_id, reference_image=reference_image
|
||||
)
|
||||
|
||||
file_name = f"{video_id}.{job.file_extension}"
|
||||
@@ -2759,7 +2676,6 @@ async def _parse_video_form(
|
||||
prompt: str = Form(...),
|
||||
input_reference: UploadFile | None = File(default=None),
|
||||
image_reference: str | None = Form(default=None),
|
||||
video_reference: str | None = Form(default=None),
|
||||
model: str | None = Form(default=None),
|
||||
seconds: SecondStr | None = Form(default=None),
|
||||
size: SizeStr | None = Form(default=None),
|
||||
@@ -2784,7 +2700,7 @@ async def _parse_video_form(
|
||||
frame_interpolation_model_path: str | None = Form(default=None),
|
||||
lora: str | None = Form(default=None),
|
||||
extra_params: str | None = Form(default=None),
|
||||
) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None, ReferenceVideo | None]:
|
||||
) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None]:
|
||||
"""FastAPI dependency that parses video form data, validates inputs,
|
||||
resolves the handler, and decodes any reference image.
|
||||
|
||||
@@ -2792,15 +2708,11 @@ async def _parse_video_form(
|
||||
"""
|
||||
input_reference_bytes = await input_reference.read() if input_reference is not None else None
|
||||
parsed_image_reference = _parse_form_json(image_reference)
|
||||
parsed_video_reference = _parse_form_json(video_reference)
|
||||
|
||||
provided_references = sum(
|
||||
item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference_bytes)
|
||||
)
|
||||
if provided_references > 1:
|
||||
if parsed_image_reference is not None and input_reference_bytes is not None:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Provide only one of input_reference, image_reference, or video_reference.",
|
||||
detail="Provide either input_reference or image_reference, not both.",
|
||||
)
|
||||
|
||||
request_data: dict[str, Any] = {
|
||||
@@ -2809,7 +2721,6 @@ async def _parse_video_form(
|
||||
"seconds": seconds,
|
||||
"size": size,
|
||||
"image_reference": parsed_image_reference,
|
||||
"video_reference": parsed_video_reference,
|
||||
"user": user,
|
||||
"width": width,
|
||||
"height": height,
|
||||
@@ -2863,28 +2774,13 @@ async def _parse_video_form(
|
||||
detail=f"Video generation setup failed: {str(e)}",
|
||||
)
|
||||
|
||||
decode_spec = ReferenceVideoDecodeSpec()
|
||||
if parsed_video_reference is not None or input_reference_bytes is not None:
|
||||
stage_configs = (
|
||||
handler.stage_configs
|
||||
or app_stage_configs
|
||||
or getattr(getattr(handler, "_engine_client", None), "stage_configs", None)
|
||||
)
|
||||
decode_spec = _reference_video_decode_spec(request, stage_configs)
|
||||
try:
|
||||
media_data = await decode_input_reference(
|
||||
request.image_reference,
|
||||
request.video_reference,
|
||||
input_reference_bytes,
|
||||
max_video_frames=decode_spec.max_frames,
|
||||
video_keep=decode_spec.keep,
|
||||
)
|
||||
image_data = await decode_input_reference(request.image_reference, input_reference_bytes)
|
||||
except InvalidInputReferenceError as exc:
|
||||
raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc
|
||||
|
||||
reference_image = ReferenceImage(data=media_data) if isinstance(media_data, Image.Image) else None
|
||||
reference_video = ReferenceVideo(data=media_data) if isinstance(media_data, list) else None
|
||||
return request, handler, effective_model_name, reference_image, reference_video
|
||||
reference_image = ReferenceImage(data=image_data) if image_data is not None else None
|
||||
return request, handler, effective_model_name, reference_image
|
||||
|
||||
|
||||
@router.post(
|
||||
@@ -2898,31 +2794,18 @@ async def _parse_video_form(
|
||||
)
|
||||
async def create_video(
|
||||
raw_request: Request,
|
||||
ctx: tuple[
|
||||
VideoGenerationRequest,
|
||||
OmniOpenAIServingVideo,
|
||||
str,
|
||||
ReferenceImage | None,
|
||||
ReferenceVideo | None,
|
||||
] = Depends(_parse_video_form),
|
||||
ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
|
||||
) -> VideoResponse:
|
||||
"""Create an asynchronous video generation job.
|
||||
|
||||
Accepts multipart form-data (see ``_parse_video_form`` for parameters),
|
||||
persists a queued job record, and starts generation in the background.
|
||||
"""
|
||||
request, handler, effective_model_name, reference_image, reference_video = ctx
|
||||
request, handler, effective_model_name, reference_image = ctx
|
||||
ref = video_response_from_request(effective_model_name, request)
|
||||
await VIDEO_STORE.upsert(ref.id, ref)
|
||||
task = asyncio.create_task(
|
||||
_run_video_generation_job(
|
||||
handler,
|
||||
request,
|
||||
ref.id,
|
||||
reference_image,
|
||||
reference_video,
|
||||
app_state=raw_request.app.state,
|
||||
)
|
||||
_run_video_generation_job(handler, request, ref.id, reference_image, app_state=raw_request.app.state)
|
||||
)
|
||||
await VIDEO_TASKS.upsert(ref.id, task)
|
||||
return ref
|
||||
@@ -2939,13 +2822,7 @@ async def create_video(
|
||||
)
|
||||
async def create_video_sync(
|
||||
raw_request: Request,
|
||||
ctx: tuple[
|
||||
VideoGenerationRequest,
|
||||
OmniOpenAIServingVideo,
|
||||
str,
|
||||
ReferenceImage | None,
|
||||
ReferenceVideo | None,
|
||||
] = Depends(_parse_video_form),
|
||||
ctx: tuple[VideoGenerationRequest, OmniOpenAIServingVideo, str, ReferenceImage | None] = Depends(_parse_video_form),
|
||||
) -> Response:
|
||||
"""Synchronous video generation endpoint.
|
||||
|
||||
@@ -2956,18 +2833,13 @@ async def create_video_sync(
|
||||
Metadata is returned via response headers ``X-Request-Id``,
|
||||
``X-Model``, and ``X-Inference-Time-S``.
|
||||
"""
|
||||
request, handler, effective_model_name, reference_image, reference_video = ctx
|
||||
request, handler, effective_model_name, reference_image = ctx
|
||||
request_id = f"video_sync-{random_uuid()}"
|
||||
raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
|
||||
started_at = time.perf_counter()
|
||||
try:
|
||||
video_bytes, stage_durations, peak_memory_mb, _action = await asyncio.wait_for(
|
||||
handler.generate_video_bytes(
|
||||
request,
|
||||
request_id,
|
||||
reference_image=reference_image,
|
||||
reference_video=reference_video,
|
||||
),
|
||||
handler.generate_video_bytes(request, request_id, reference_image=reference_image),
|
||||
timeout=VIDEO_SYNC_TIMEOUT_S,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
|
||||
@@ -115,54 +115,6 @@ class OpenAICreateSpeechRequest(BaseModel):
|
||||
raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
|
||||
return v
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def normalize_references_alias(cls, data):
|
||||
"""Map the BosonAI Higgs Audio v3 cookbook ``references`` array onto ``ref_audio`` / ``ref_text``.
|
||||
|
||||
Upstream Higgs Audio v3 examples and the boson.ai cookbook send voice
|
||||
clones as::
|
||||
|
||||
{"references": [{"audio_path": str, "text": str}]}
|
||||
|
||||
vllm-omni's existing convention is ``ref_audio`` + ``ref_text``.
|
||||
Without this alias pydantic silently drops the unknown ``references``
|
||||
field and the request falls through to zero-shot synthesis with no
|
||||
error, which is hard to debug from the client side. The normalizer
|
||||
translates a single reference and rejects multi-reference payloads
|
||||
(vllm-omni does not yet support multi-shot voice clone) or conflicts
|
||||
with the explicit ``ref_audio`` / ``ref_text`` fields.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
refs = data.get("references")
|
||||
if refs is None:
|
||||
return data
|
||||
if not isinstance(refs, list) or len(refs) == 0:
|
||||
raise ValueError("'references' must be a non-empty array of {audio_path, text} objects")
|
||||
if len(refs) > 1:
|
||||
raise ValueError("'references' only supports a single reference; multi-shot voice clone is not supported")
|
||||
first = refs[0]
|
||||
if not isinstance(first, dict):
|
||||
raise ValueError("'references[0]' must be an object with 'audio_path' and optional 'text'")
|
||||
audio_path = first.get("audio_path")
|
||||
if not isinstance(audio_path, str) or not audio_path:
|
||||
raise ValueError(
|
||||
"'references[0].audio_path' is required and must be a string (URL, data: URI, or file path)"
|
||||
)
|
||||
existing_ref_audio = data.get("ref_audio")
|
||||
if existing_ref_audio and existing_ref_audio != audio_path:
|
||||
raise ValueError("'references' and 'ref_audio' are mutually exclusive")
|
||||
data["ref_audio"] = audio_path
|
||||
text = first.get("text")
|
||||
if isinstance(text, str) and text.strip():
|
||||
existing_ref_text = data.get("ref_text")
|
||||
if existing_ref_text and existing_ref_text != text:
|
||||
raise ValueError("'references[0].text' and 'ref_text' conflict; supply only one")
|
||||
data["ref_text"] = text
|
||||
data.pop("references", None)
|
||||
return data
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def reject_higgs_audio_v2_unsupported_aliases(cls, data):
|
||||
|
||||
@@ -73,19 +73,6 @@ class UrlImageReference(BaseModel):
|
||||
ImageReference = UrlImageReference | FileImageReference
|
||||
|
||||
|
||||
class FileVideoReference(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
file_id: str
|
||||
|
||||
|
||||
class UrlVideoReference(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
video_url: str
|
||||
|
||||
|
||||
VideoReference = UrlVideoReference | FileVideoReference
|
||||
|
||||
|
||||
class VideoGenerationRequest(BaseModel):
|
||||
"""
|
||||
OpenAI-style video generation request.
|
||||
@@ -112,10 +99,6 @@ class VideoGenerationRequest(BaseModel):
|
||||
default=None,
|
||||
description="Optional JSON-safe image reference that guides generation. Provide either image_url or file_id.",
|
||||
)
|
||||
video_reference: VideoReference | None = Field(
|
||||
default=None,
|
||||
description="Optional JSON-safe video reference that guides generation. Provide either video_url or file_id.",
|
||||
)
|
||||
|
||||
# Video params block for extensibility
|
||||
video_params: VideoParams | None = Field(default=None, description="Optional video-specific parameters")
|
||||
|
||||
@@ -12,16 +12,8 @@ from typing import Any, Final, cast
|
||||
import jinja2
|
||||
import torch
|
||||
from fastapi import Request
|
||||
from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
|
||||
from PIL import Image
|
||||
from pydantic import TypeAdapter
|
||||
from vllm.entrypoints.chat_utils import (
|
||||
ChatCompletionMessageParam,
|
||||
ChatTemplateContentFormatOption,
|
||||
ConversationMessage,
|
||||
get_history_tool_calls_cnt,
|
||||
make_tool_call_id,
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.diffusion_engine import get_extra_body_params, get_extra_output_params
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
@@ -40,6 +32,14 @@ except ImportError:
|
||||
soundfile = None
|
||||
|
||||
|
||||
from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
|
||||
from vllm.entrypoints.chat_utils import (
|
||||
ChatCompletionMessageParam,
|
||||
ChatTemplateContentFormatOption,
|
||||
ConversationMessage,
|
||||
get_history_tool_calls_cnt,
|
||||
make_tool_call_id,
|
||||
)
|
||||
from vllm.entrypoints.launcher import terminate_if_errored
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionNamedToolChoiceParam,
|
||||
@@ -652,15 +652,19 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
)
|
||||
|
||||
tok_params = request.build_tok_params(self.model_config)
|
||||
mm_config = self.model_config.multimodal_config
|
||||
chat_params = request.build_chat_params(
|
||||
default_template,
|
||||
default_template_content_format,
|
||||
).with_defaults(
|
||||
default_template_kwargs,
|
||||
default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
|
||||
default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
|
||||
)
|
||||
).with_defaults(default_template_kwargs)
|
||||
|
||||
# OMNI: When use_audio_in_video=True, the qwen2_5_omni_thinker mm
|
||||
# processor asserts that audio items are present alongside video items
|
||||
# (inside render_chat_async). We must inject audio_url items into the
|
||||
# messages BEFORE calling render_chat_async so the mm processor can
|
||||
# count them correctly during tokenisation.
|
||||
mm_proc_kw = getattr(request, "mm_processor_kwargs", None) or {}
|
||||
if mm_proc_kw.get("use_audio_in_video", False):
|
||||
messages = await self._inject_audio_from_video_urls(messages)
|
||||
|
||||
(conversation,), (engine_prompt,) = await renderer.render_chat_async(
|
||||
[messages],
|
||||
@@ -731,6 +735,95 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
|
||||
return conversation, [engine_prompt]
|
||||
|
||||
async def _inject_audio_from_video_urls(
|
||||
self,
|
||||
messages: list[ChatCompletionMessageParam],
|
||||
) -> list[ChatCompletionMessageParam]:
|
||||
"""Pre-extract audio from video URLs and inject as audio_url content items.
|
||||
|
||||
When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal
|
||||
processor requires that the number of audio items equals the number of
|
||||
video items (it subtracts mm_counts["video"] from mm_counts["audio"]).
|
||||
The client only sends video_url items; this method adds the matching
|
||||
audio_url items on the server side before the renderer processes them.
|
||||
"""
|
||||
import io
|
||||
|
||||
from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async
|
||||
|
||||
new_messages: list[ChatCompletionMessageParam] = []
|
||||
for msg in messages:
|
||||
content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
|
||||
if not isinstance(content, list):
|
||||
new_messages.append(msg)
|
||||
continue
|
||||
|
||||
video_urls = [
|
||||
part.get("video_url", {}).get("url")
|
||||
for part in content
|
||||
if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url")
|
||||
]
|
||||
|
||||
if not video_urls:
|
||||
new_messages.append(msg)
|
||||
continue
|
||||
|
||||
audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls))
|
||||
|
||||
audio_items: list[dict] = []
|
||||
for audio_array, sample_rate in audios:
|
||||
buf = io.BytesIO()
|
||||
if soundfile is not None:
|
||||
soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV")
|
||||
else:
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
|
||||
audio_np = np.asarray(audio_array, dtype=np.float32)
|
||||
sr = int(sample_rate)
|
||||
num_channels = 1
|
||||
bits_per_sample = 32
|
||||
num_frames = len(audio_np)
|
||||
data_size = num_frames * num_channels * (bits_per_sample // 8)
|
||||
# Write minimal RIFF/WAV header
|
||||
buf.write(b"RIFF")
|
||||
buf.write(struct.pack("<I", 36 + data_size))
|
||||
buf.write(b"WAVE")
|
||||
buf.write(b"fmt ")
|
||||
buf.write(
|
||||
struct.pack(
|
||||
"<IHHIIHH",
|
||||
16,
|
||||
3,
|
||||
num_channels,
|
||||
sr,
|
||||
sr * num_channels * (bits_per_sample // 8),
|
||||
num_channels * (bits_per_sample // 8),
|
||||
bits_per_sample,
|
||||
)
|
||||
)
|
||||
buf.write(b"data")
|
||||
buf.write(struct.pack("<I", data_size))
|
||||
buf.write(audio_np.tobytes())
|
||||
|
||||
audio_b64 = base64.b64encode(buf.getvalue()).decode()
|
||||
audio_items.append(
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"},
|
||||
}
|
||||
)
|
||||
|
||||
new_content = list(content) + audio_items
|
||||
if isinstance(msg, dict):
|
||||
new_msg = {**msg, "content": new_content}
|
||||
else:
|
||||
new_msg = msg.model_copy(update={"content": new_content})
|
||||
new_messages.append(new_msg)
|
||||
|
||||
return new_messages
|
||||
|
||||
def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[Any]:
|
||||
"""Convert request dicts to stage-typed sampling params objects.
|
||||
|
||||
|
||||
@@ -9,7 +9,6 @@ import re
|
||||
import struct
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from http import HTTPStatus
|
||||
from pathlib import Path
|
||||
@@ -77,7 +76,6 @@ _MING_TTS_MODEL_STAGES = {"ming_tts"}
|
||||
_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
|
||||
_MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"}
|
||||
_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
|
||||
_HIGGS_AUDIO_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
|
||||
_GLM_TTS_MODEL_STAGES = {"glm_tts"}
|
||||
_TTS_MODEL_STAGES: set[str] = (
|
||||
_VOXTRAL_TTS_MODEL_STAGES
|
||||
@@ -86,7 +84,6 @@ _TTS_MODEL_STAGES: set[str] = (
|
||||
| _COSYVOICE3_TTS_MODEL_STAGES
|
||||
| _OMNIVOICE_TTS_MODEL_STAGES
|
||||
| _HIGGS_AUDIO_V2_TTS_MODEL_STAGES
|
||||
| _HIGGS_AUDIO_V3_TTS_MODEL_STAGES
|
||||
| _COVO_AUDIO_MODEL_STAGES
|
||||
| _VOXCPM2_TTS_MODEL_STAGES
|
||||
| _MING_TTS_MODEL_STAGES
|
||||
@@ -101,23 +98,20 @@ _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
|
||||
"cosyvoice3",
|
||||
"voxcpm2",
|
||||
"higgs_audio_v2",
|
||||
"higgs_audio_v3",
|
||||
}
|
||||
_TTS_LANGUAGES = frozenset(
|
||||
{
|
||||
"Auto",
|
||||
"Chinese",
|
||||
"English",
|
||||
"Japanese",
|
||||
"Korean",
|
||||
"German",
|
||||
"French",
|
||||
"Russian",
|
||||
"Portuguese",
|
||||
"Spanish",
|
||||
"Italian",
|
||||
}
|
||||
)
|
||||
_TTS_LANGUAGES: set[str] = {
|
||||
"Auto",
|
||||
"Chinese",
|
||||
"English",
|
||||
"Japanese",
|
||||
"Korean",
|
||||
"German",
|
||||
"French",
|
||||
"Russian",
|
||||
"Portuguese",
|
||||
"Spanish",
|
||||
"Italian",
|
||||
}
|
||||
_REF_AUDIO_MIN_DURATION = 1.0 # seconds
|
||||
_REF_AUDIO_MAX_DURATION = 30.0 # seconds
|
||||
_REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256
|
||||
@@ -437,9 +431,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
# Merge built-in speakers into the set initialized by _init_speaker_storage.
|
||||
self.supported_speakers |= self._load_supported_speakers()
|
||||
self.supported_speakers |= set(self.precomputed_speakers)
|
||||
|
||||
self.supported_languages = self._load_supported_languages()
|
||||
|
||||
self._tts_tokenizer = None
|
||||
self._voxcpm2_tokenizer = None
|
||||
self._voxcpm2_split_map: dict[int, list[int]] = {}
|
||||
@@ -603,8 +594,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
return "moss_tts"
|
||||
if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES:
|
||||
return "higgs_audio_v2"
|
||||
if model_stage in _HIGGS_AUDIO_V3_TTS_MODEL_STAGES:
|
||||
return "higgs_audio_v3"
|
||||
if model_stage in _GLM_TTS_MODEL_STAGES:
|
||||
return "glm_tts"
|
||||
return None
|
||||
@@ -700,26 +689,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
|
||||
return set()
|
||||
|
||||
def _load_supported_languages(self) -> frozenset[str]:
|
||||
"""Load supported languages (title-cased) from the model configuration"""
|
||||
if self._tts_model_type != "qwen3_tts":
|
||||
return _TTS_LANGUAGES
|
||||
try:
|
||||
config = self.engine_client.model_config.hf_config.talker_config
|
||||
|
||||
if isinstance(config, dict):
|
||||
codec_language_id = config.get("codec_language_id")
|
||||
else:
|
||||
codec_language_id = getattr(config, "codec_language_id", None)
|
||||
|
||||
if codec_language_id and isinstance(codec_language_id, Mapping):
|
||||
return frozenset(str(language).title() for language in codec_language_id) | {"Auto"}
|
||||
|
||||
logger.warning("No codec_language_id found in talker_config; falling back to default languages")
|
||||
except Exception as e:
|
||||
logger.warning("Could not load languages from model config: %s", e)
|
||||
return _TTS_LANGUAGES
|
||||
|
||||
def _estimate_ref_code_len(self, ref_audio: object) -> int | None:
|
||||
"""Estimate ref_code length from ref_audio waveform without running the codec.
|
||||
|
||||
@@ -965,7 +934,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
"moss_tts_nano",
|
||||
"glm_tts",
|
||||
"higgs_audio_v2",
|
||||
"higgs_audio_v3",
|
||||
):
|
||||
label = {
|
||||
"cosyvoice3": "CosyVoice3",
|
||||
@@ -973,7 +941,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
"omnivoice": "OmniVoice",
|
||||
"moss_tts_nano": "MOSS-TTS-Nano",
|
||||
"higgs_audio_v2": "Higgs-Audio V2",
|
||||
"higgs_audio_v3": "Higgs-Audio V3",
|
||||
"glm_tts": "GLM-TTS",
|
||||
}.get(self._tts_model_type, self._tts_model_type)
|
||||
return (
|
||||
@@ -1446,13 +1413,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
if not request.input or not request.input.strip():
|
||||
return "Input text cannot be empty"
|
||||
|
||||
# Validate language (case-insensitive; normalized to the title-cased config form)
|
||||
if request.language is not None:
|
||||
request.language = request.language.title()
|
||||
if request.language not in self.supported_languages:
|
||||
return (
|
||||
f"Invalid language '{request.language}'. Supported: {', '.join(sorted(self.supported_languages))}"
|
||||
)
|
||||
# Validate language
|
||||
if request.language is not None and request.language not in _TTS_LANGUAGES:
|
||||
return f"Invalid language '{request.language}'. Supported: {', '.join(sorted(_TTS_LANGUAGES))}"
|
||||
|
||||
# Validate speaker for CustomVoice task
|
||||
if task_type == "CustomVoice":
|
||||
@@ -1933,90 +1896,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
self._higgs_audio_v2_processor = processor
|
||||
return processor
|
||||
|
||||
# ---- higgs-audio v3 ----
|
||||
|
||||
def _validate_higgs_audio_v3_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
"""Validate higgs_audio_v3 request parameters."""
|
||||
if not request.input or not request.input.strip():
|
||||
return "higgs_audio_v3: input text cannot be empty"
|
||||
if request.ref_audio is not None and not request.ref_text:
|
||||
# Voice clone ref_text is optional for v3 (improves fidelity but not required)
|
||||
pass
|
||||
if request.max_new_tokens is not None:
|
||||
if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
|
||||
return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
|
||||
return None
|
||||
|
||||
async def _build_higgs_audio_v3_params(self, request: OpenAICreateSpeechRequest):
|
||||
"""Build prompt_token_ids for higgs_audio_v3.
|
||||
|
||||
Plain-text path: builds ``[tts, text, tokens, audio]``.
|
||||
Voice-clone path: encodes reference audio, applies delay pattern,
|
||||
builds ``[tts, (ref_text, tokens,) ref_audio, -100xN, text, tokens, audio]``.
|
||||
"""
|
||||
adapter = await self._resolve_higgs_audio_v3_adapter()
|
||||
|
||||
if request.ref_audio is None:
|
||||
prompt_ids = adapter.build_prompt(request.input)
|
||||
return tokens_input(prompt_token_ids=prompt_ids)
|
||||
|
||||
# Voice clone
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
|
||||
apply_delay_pattern,
|
||||
encode_reference_audio,
|
||||
)
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.ref_audio_cache import (
|
||||
cached_encode_reference_audio,
|
||||
)
|
||||
|
||||
wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
|
||||
wav = np.asarray(wav_list, dtype=np.float32)
|
||||
# Cache the codec encode by content hash so identical ref_audio
|
||||
# payloads skip the encode and reuse the same [num_frames,
|
||||
# num_codebooks] codes tensor across requests.
|
||||
ref_codes_raw = await asyncio.to_thread(cached_encode_reference_audio, wav, int(sr), encode_reference_audio)
|
||||
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
|
||||
|
||||
prompt_ids = adapter.build_prompt(
|
||||
request.input,
|
||||
num_ref_tokens=int(ref_codes_delayed.shape[0]),
|
||||
reference_text=request.ref_text or None,
|
||||
)
|
||||
prompt = tokens_input(prompt_token_ids=prompt_ids)
|
||||
import torch
|
||||
|
||||
prompt["additional_information"] = {
|
||||
"audio_input_ids": ref_codes_delayed.to(torch.long),
|
||||
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
|
||||
}
|
||||
return prompt
|
||||
|
||||
async def _resolve_higgs_audio_v3_adapter(self):
|
||||
"""Lazy-load the tokenizer adapter for higgs_audio_v3."""
|
||||
cached = getattr(self, "_higgs_audio_v3_adapter", None)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
|
||||
HiggsAudioV3TokenizerAdapter,
|
||||
)
|
||||
|
||||
model_path = None
|
||||
for stage in self.engine_client.stage_configs:
|
||||
model_path = getattr(getattr(stage, "engine_args", None), "model", None)
|
||||
if model_path:
|
||||
break
|
||||
if model_path is None:
|
||||
model_path = getattr(self.engine_client, "model", None)
|
||||
if model_path is None:
|
||||
raise RuntimeError("higgs_audio_v3 serving could not resolve model path")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
|
||||
self._higgs_audio_v3_adapter = adapter
|
||||
return adapter
|
||||
|
||||
def _validate_fish_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
|
||||
"""Validate Fish Speech request parameters. Returns error message or None."""
|
||||
if not request.input or not request.input.strip():
|
||||
@@ -2917,15 +2796,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
additional["voice_name"] = voice_lower
|
||||
additional["voice_created_at"] = self._voice_created_at(voice_lower)
|
||||
tts_params = {}
|
||||
elif self._tts_model_type == "higgs_audio_v3":
|
||||
err = self._apply_uploaded_speaker(request)
|
||||
if err:
|
||||
raise ValueError(err)
|
||||
validation_error = self._validate_higgs_audio_v3_request(request)
|
||||
if validation_error:
|
||||
raise ValueError(validation_error)
|
||||
prompt = await self._build_higgs_audio_v3_params(request)
|
||||
tts_params = {}
|
||||
elif self._tts_model_type == "voxcpm2":
|
||||
# voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the
|
||||
# raw waveform tuple for prefill-length accounting, not a base64 data URL.
|
||||
|
||||
@@ -39,13 +39,6 @@ class ReferenceImage:
|
||||
data: Image.Image
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReferenceVideo:
|
||||
"""Reference video frames for video-conditioned generation."""
|
||||
|
||||
data: list[Image.Image]
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoGenerationArtifacts:
|
||||
"""Normalized outputs and profiler metadata extracted from one request."""
|
||||
@@ -103,7 +96,6 @@ class OmniOpenAIServingVideo:
|
||||
reference_id: str,
|
||||
*,
|
||||
reference_image: ReferenceImage | None = None,
|
||||
reference_video: ReferenceVideo | None = None,
|
||||
) -> VideoGenerationArtifacts:
|
||||
"""Run the generation pipeline and extract video/audio/profiler outputs."""
|
||||
prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt, modalities=["video"])
|
||||
@@ -113,12 +105,6 @@ class OmniOpenAIServingVideo:
|
||||
gen_params = self._resolve_default_sampling_params()
|
||||
|
||||
input_image = None if reference_image is None else reference_image.data
|
||||
input_video = None if reference_video is None else reference_video.data
|
||||
if input_image is not None and input_video is not None:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST.value,
|
||||
detail="Provide either an image reference or a video reference, not both.",
|
||||
)
|
||||
vp = request.resolve_video_params()
|
||||
if input_image is not None and vp.width is not None and vp.height is not None:
|
||||
target_size = (vp.width, vp.height)
|
||||
@@ -126,8 +112,6 @@ class OmniOpenAIServingVideo:
|
||||
input_image = input_image.resize(target_size, Image.Resampling.LANCZOS)
|
||||
if input_image is not None:
|
||||
prompt["multi_modal_data"] = {"image": input_image}
|
||||
elif input_video is not None:
|
||||
prompt["multi_modal_data"] = {"video": input_video}
|
||||
if vp.width is not None and vp.height is not None:
|
||||
gen_params.width = vp.width
|
||||
gen_params.height = vp.height
|
||||
@@ -227,14 +211,8 @@ class OmniOpenAIServingVideo:
|
||||
reference_id: str,
|
||||
*,
|
||||
reference_image: ReferenceImage | None = None,
|
||||
reference_video: ReferenceVideo | None = None,
|
||||
) -> VideoGenerationResponse:
|
||||
artifacts = await self._run_and_extract(
|
||||
request,
|
||||
reference_id,
|
||||
reference_image=reference_image,
|
||||
reference_video=reference_video,
|
||||
)
|
||||
artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
|
||||
|
||||
video_codec_options = {"preset": "ultrafast", "threads": "0"}
|
||||
if request.extra_params is not None and isinstance(request.extra_params, dict):
|
||||
@@ -274,15 +252,9 @@ class OmniOpenAIServingVideo:
|
||||
reference_id: str,
|
||||
*,
|
||||
reference_image: ReferenceImage | None = None,
|
||||
reference_video: ReferenceVideo | None = None,
|
||||
) -> tuple[bytes, dict[str, float], float, VideoAction | None]:
|
||||
"""Generate a video and return raw MP4 bytes, bypassing base64 encoding."""
|
||||
artifacts = await self._run_and_extract(
|
||||
request,
|
||||
reference_id,
|
||||
reference_image=reference_image,
|
||||
reference_video=reference_video,
|
||||
)
|
||||
artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
|
||||
if len(artifacts.videos) > 1:
|
||||
logger.warning(
|
||||
"Video request %s generated %d outputs; returning only the first.",
|
||||
|
||||
@@ -8,9 +8,8 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
from collections import deque
|
||||
from io import BytesIO
|
||||
from typing import Any, Literal
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import numpy as np
|
||||
@@ -20,11 +19,8 @@ from PIL import Image, UnidentifiedImageError
|
||||
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
|
||||
from vllm_omni.entrypoints.openai.protocol.videos import (
|
||||
FileImageReference,
|
||||
FileVideoReference,
|
||||
ImageReference,
|
||||
UrlImageReference,
|
||||
UrlVideoReference,
|
||||
VideoReference,
|
||||
)
|
||||
|
||||
|
||||
@@ -35,70 +31,6 @@ def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc
|
||||
|
||||
|
||||
def _decode_video_bytes(
|
||||
video_bytes: bytes,
|
||||
*,
|
||||
source: str,
|
||||
max_frames: int | None = None,
|
||||
keep: Literal["first", "last"] = "first",
|
||||
) -> list[Image.Image]:
|
||||
try:
|
||||
import av
|
||||
except ImportError as exc: # pragma: no cover - av is a serving dependency via media_utils
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video decoding requires PyAV.") from exc
|
||||
|
||||
if keep not in {"first", "last"}:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video frame selection must be 'first' or 'last'.")
|
||||
if max_frames is not None and max_frames <= 0:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: max video frames must be positive.")
|
||||
|
||||
frames: list[Image.Image] = []
|
||||
tail_frames: deque[Image.Image] | None = (
|
||||
deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None
|
||||
)
|
||||
try:
|
||||
with av.open(BytesIO(video_bytes)) as container:
|
||||
for frame in container.decode(video=0):
|
||||
image = frame.to_image().convert("RGB")
|
||||
if tail_frames is not None:
|
||||
tail_frames.append(image)
|
||||
else:
|
||||
frames.append(image)
|
||||
if keep == "first" and max_frames is not None and len(frames) >= max_frames:
|
||||
break
|
||||
except Exception as exc:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") from exc
|
||||
|
||||
if tail_frames is not None:
|
||||
frames = list(tail_frames)
|
||||
if not frames:
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.")
|
||||
return frames
|
||||
|
||||
|
||||
def _decode_media_bytes(
|
||||
media_bytes: bytes,
|
||||
*,
|
||||
source: str,
|
||||
max_video_frames: int | None = None,
|
||||
video_keep: Literal["first", "last"] = "first",
|
||||
) -> Image.Image | list[Image.Image]:
|
||||
try:
|
||||
return _decode_image_bytes(media_bytes, source=source)
|
||||
except InvalidInputReferenceError:
|
||||
try:
|
||||
return _decode_video_bytes(
|
||||
media_bytes,
|
||||
source=source,
|
||||
max_frames=max_video_frames,
|
||||
keep=video_keep,
|
||||
)
|
||||
except InvalidInputReferenceError as video_exc:
|
||||
raise InvalidInputReferenceError(
|
||||
f"Invalid {source}: provided content is not a valid image or video."
|
||||
) from video_exc
|
||||
|
||||
|
||||
def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image:
|
||||
if input_reference:
|
||||
if input_reference.startswith("data:image"):
|
||||
@@ -132,96 +64,23 @@ async def decode_image_url(image_url: str) -> Image.Image:
|
||||
raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.")
|
||||
|
||||
|
||||
def _decode_base64_video(
|
||||
video_reference: str,
|
||||
*,
|
||||
source: str,
|
||||
max_frames: int | None = None,
|
||||
keep: Literal["first", "last"] = "first",
|
||||
) -> list[Image.Image]:
|
||||
if video_reference:
|
||||
if video_reference.startswith("data:video"):
|
||||
_, b64_data = video_reference.split(",", 1)
|
||||
else:
|
||||
b64_data = video_reference
|
||||
|
||||
try:
|
||||
video_bytes = base64.b64decode(b64_data)
|
||||
except (binascii.Error, ValueError) as exc: # pragma: no cover - malformed base64
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video data is not valid base64.") from exc
|
||||
return _decode_video_bytes(video_bytes, source=source, max_frames=max_frames, keep=keep)
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.")
|
||||
|
||||
|
||||
async def decode_video_url(
|
||||
video_url: str,
|
||||
*,
|
||||
max_frames: int | None = None,
|
||||
keep: Literal["first", "last"] = "first",
|
||||
) -> list[Image.Image]:
|
||||
if video_url.startswith("data:video"):
|
||||
return _decode_base64_video(
|
||||
video_url,
|
||||
source="video_reference.video_url",
|
||||
max_frames=max_frames,
|
||||
keep=keep,
|
||||
)
|
||||
|
||||
if video_url.startswith(("http://", "https://")):
|
||||
async with httpx.AsyncClient(timeout=60) as client:
|
||||
try:
|
||||
response = await client.get(video_url)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as exc:
|
||||
raise InvalidInputReferenceError(
|
||||
"Invalid video_reference.video_url: failed to download video."
|
||||
) from exc
|
||||
return _decode_video_bytes(
|
||||
response.content,
|
||||
source="video_reference.video_url",
|
||||
max_frames=max_frames,
|
||||
keep=keep,
|
||||
)
|
||||
|
||||
raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.")
|
||||
|
||||
|
||||
async def decode_input_reference(
|
||||
image_reference: ImageReference | None,
|
||||
video_reference: VideoReference | None,
|
||||
input_reference_bytes: bytes | None,
|
||||
*,
|
||||
max_video_frames: int | None = None,
|
||||
video_keep: Literal["first", "last"] = "first",
|
||||
) -> Image.Image | list[Image.Image] | None:
|
||||
"""Decode media input from multipart bytes, data URLs, or typed references."""
|
||||
) -> Image.Image | None:
|
||||
"""Decode image input from multipart bytes, base64/data URL, or image_reference."""
|
||||
|
||||
provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference))
|
||||
if provided > 1:
|
||||
raise InvalidInputReferenceError("Provide only one of input_reference, image_reference, or video_reference.")
|
||||
if input_reference_bytes is not None and image_reference is not None:
|
||||
raise InvalidInputReferenceError("Provide either input_reference or image_reference, not both.")
|
||||
|
||||
if isinstance(input_reference_bytes, bytes):
|
||||
return _decode_media_bytes(
|
||||
input_reference_bytes,
|
||||
source="input_reference",
|
||||
max_video_frames=max_video_frames,
|
||||
video_keep=video_keep,
|
||||
)
|
||||
return _decode_image_bytes(input_reference_bytes, source="input_reference")
|
||||
|
||||
if isinstance(image_reference, UrlImageReference):
|
||||
return await decode_image_url(image_reference.image_url)
|
||||
elif isinstance(image_reference, FileImageReference):
|
||||
raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.")
|
||||
|
||||
if isinstance(video_reference, UrlVideoReference):
|
||||
return await decode_video_url(
|
||||
video_reference.video_url,
|
||||
max_frames=max_video_frames,
|
||||
keep=video_keep,
|
||||
)
|
||||
elif isinstance(video_reference, FileVideoReference):
|
||||
raise InvalidInputReferenceError("Invalid video_reference: file_id is not supported yet.")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -1,9 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""vllm-omni integration for boson-ai's higgs-audio v3 (two-stage TTS)."""
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
|
||||
__all__ = ["HiggsAudioV3Config"]
|
||||
@@ -1,182 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Configuration class for higgs-audio v3 (HiggsMultimodalQwen3) in vllm-omni.
|
||||
|
||||
``HiggsAudioV3Config.from_pretrained(model_path)`` returns a config with
|
||||
``tts_token_id``, ``text_token_id``, ``audio_continuation_id``, and
|
||||
``eos_token_id`` already resolved from the checkpoint tokenizer. If the
|
||||
tokenizer is unavailable or missing required specials, the load raises.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import transformers
|
||||
from transformers import AutoConfig, PretrainedConfig
|
||||
|
||||
__all__ = ["HiggsAudioV3Config"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_QWEN3_ROPE_THETA = 1_000_000
|
||||
|
||||
_REQUIRED_SPECIALS = ("<|tts|>", "<|text|>", "<|audio|>")
|
||||
|
||||
|
||||
def _build_text_config(raw: Any) -> PretrainedConfig:
|
||||
if isinstance(raw, PretrainedConfig):
|
||||
return raw
|
||||
cfg = dict(raw or {})
|
||||
model_type = cfg.get("model_type", "qwen3")
|
||||
if model_type == "qwen3" and cfg.get("rope_theta") is None:
|
||||
cfg["rope_theta"] = _QWEN3_ROPE_THETA
|
||||
try:
|
||||
cfg_cls = transformers.CONFIG_MAPPING[model_type]
|
||||
except KeyError as exc:
|
||||
raise ValueError(f"Unknown text backbone model_type {model_type!r}") from exc
|
||||
return cfg_cls(**cfg)
|
||||
|
||||
|
||||
def _resolve_model_dir(pretrained_model_name_or_path: str) -> str | None:
|
||||
"""Resolve a model name/path to a local directory containing tokenizer files."""
|
||||
if os.path.isdir(pretrained_model_name_or_path):
|
||||
return pretrained_model_name_or_path
|
||||
try:
|
||||
from huggingface_hub import try_to_load_from_cache
|
||||
|
||||
cached = try_to_load_from_cache(repo_id=pretrained_model_name_or_path, filename="tokenizer_config.json")
|
||||
if isinstance(cached, str) and os.path.isfile(cached):
|
||||
return os.path.dirname(cached)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
from huggingface_hub.constants import HF_HUB_CACHE
|
||||
|
||||
safe = pretrained_model_name_or_path.replace("/", "--")
|
||||
snapshots = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots")
|
||||
if os.path.isdir(snapshots):
|
||||
for rev in os.listdir(snapshots):
|
||||
candidate = os.path.join(snapshots, rev)
|
||||
if os.path.isfile(os.path.join(candidate, "tokenizer_config.json")):
|
||||
return candidate
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
class HiggsAudioV3Config(PretrainedConfig):
|
||||
"""Typed config for higgs-audio v3 (HiggsMultimodalQwen3).
|
||||
|
||||
``from_pretrained()`` automatically resolves ``<|tts|>``, ``<|text|>``,
|
||||
``<|audio|>`` and ``eos_token_id`` from the checkpoint tokenizer.
|
||||
"""
|
||||
|
||||
model_type: str = "higgs_multimodal_qwen3"
|
||||
is_composition = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
audio_encoder_config: dict[str, Any] | None = None,
|
||||
text_config: dict[str, Any] | PretrainedConfig | None = None,
|
||||
audio_token_id: int = -100,
|
||||
mel_per_sample: int = 8,
|
||||
num_codebooks: int = 8,
|
||||
codebook_size: int = 1026,
|
||||
audio_stream_bos_id: int = 1024,
|
||||
audio_stream_eos_id: int = 1025,
|
||||
sample_rate: int = 24000,
|
||||
frame_rate: int = 25,
|
||||
tts_token_id: int | None = None,
|
||||
text_token_id: int | None = None,
|
||||
audio_continuation_id: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self.audio_token_id = audio_token_id
|
||||
self.mel_per_sample = mel_per_sample
|
||||
|
||||
if audio_encoder_config is None:
|
||||
audio_encoder_config = {
|
||||
"encoder_type": "discrete",
|
||||
"num_codebooks": num_codebooks,
|
||||
"vocab_size": codebook_size,
|
||||
"tie_word_embeddings": True,
|
||||
}
|
||||
self.audio_encoder_config = audio_encoder_config
|
||||
|
||||
self.num_codebooks = int(audio_encoder_config.get("num_codebooks", num_codebooks))
|
||||
if self.num_codebooks <= 0:
|
||||
raise ValueError(f"num_codebooks must be > 0, got {self.num_codebooks}")
|
||||
self.codebook_size = int(audio_encoder_config.get("vocab_size", codebook_size))
|
||||
self.tie_modality_embeddings = bool(audio_encoder_config.get("tie_word_embeddings", True))
|
||||
|
||||
self.audio_stream_bos_id = audio_stream_bos_id
|
||||
self.audio_stream_eos_id = audio_stream_eos_id
|
||||
self.sample_rate = sample_rate
|
||||
self.frame_rate = frame_rate
|
||||
|
||||
self.text_config = _build_text_config(text_config)
|
||||
self.audio_hidden_size = int(audio_encoder_config.get("out_dim", self.text_config.hidden_size))
|
||||
|
||||
self.tts_token_id = tts_token_id
|
||||
self.text_token_id = text_token_id
|
||||
self.audio_continuation_id = audio_continuation_id
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs: Any) -> HiggsAudioV3Config:
|
||||
"""Load config and resolve special token IDs from the checkpoint tokenizer.
|
||||
|
||||
Passes the original ``pretrained_model_name_or_path`` (local dir or
|
||||
HF repo ID) directly to ``AutoTokenizer.from_pretrained()`` so it can
|
||||
handle cache hits, downloads, and local paths uniformly. Raises if the
|
||||
tokenizer is missing required specials.
|
||||
"""
|
||||
config = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
|
||||
config.resolve_special_tokens(pretrained_model_name_or_path)
|
||||
return config
|
||||
|
||||
def get_text_config(self, decoder: bool = False) -> PretrainedConfig:
|
||||
del decoder
|
||||
return self.text_config
|
||||
|
||||
@property
|
||||
def num_real_codes(self) -> int:
|
||||
return self.audio_stream_bos_id
|
||||
|
||||
@property
|
||||
def hidden_size(self) -> int:
|
||||
return self.text_config.hidden_size
|
||||
|
||||
def resolve_special_tokens(self, model_path: str) -> None:
|
||||
"""Resolve <|tts|>, <|text|>, <|audio|> and eos from the HF tokenizer.
|
||||
|
||||
Raises ``ValueError`` if any of the 3 required specials is missing
|
||||
from the tokenizer's added vocabulary.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
vocab = dict(tokenizer.get_added_vocab())
|
||||
|
||||
missing = [t for t in _REQUIRED_SPECIALS if t not in vocab]
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"Tokenizer at {model_path} is missing required Higgs TTS v3 "
|
||||
f"special tokens: {missing}. Available added tokens: "
|
||||
f"{sorted(vocab.keys())[:20]}"
|
||||
)
|
||||
|
||||
self.tts_token_id = vocab["<|tts|>"]
|
||||
self.text_token_id = vocab["<|text|>"]
|
||||
self.audio_continuation_id = vocab["<|audio|>"]
|
||||
|
||||
if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
|
||||
self.eos_token_id = int(tokenizer.eos_token_id)
|
||||
|
||||
|
||||
AutoConfig.register("higgs_multimodal_qwen3", HiggsAudioV3Config)
|
||||
@@ -1,524 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Stage 1 (codec decoder) for higgs-audio v3.
|
||||
|
||||
Reuses higgs-audio-v2's RVQ + DAC codec decoder but loads weights from the
|
||||
v3 checkpoint's bundled codec (``tied.embedding.modality_embeddings.0.model.*``
|
||||
prefix) rather than from a standalone tokenizer repo.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import (
|
||||
HiggsAudioRVQ,
|
||||
build_boson_dac_decoder,
|
||||
build_higgs_audio_acoustic_decoder,
|
||||
load_higgs_audio_codec,
|
||||
)
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
__all__ = [
|
||||
"HiggsAudioV3Code2Wav",
|
||||
"HiggsAudioV3Code2WavForConditionalGeneration",
|
||||
]
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Prefix under which codec weights are stored in the v3 checkpoint.
|
||||
_CODEC_PREFIX = "tied.embedding.modality_embeddings.0.model."
|
||||
|
||||
|
||||
class HiggsAudioV3Code2Wav(nn.Module):
|
||||
"""Stage-1 codec decoder for higgs-audio v3.
|
||||
|
||||
The codec architecture is identical to v2 (HiggsAudioRVQ + fc2 +
|
||||
BosonDacDecoder); only the weight loading path differs. V3 bundles
|
||||
codec weights inside the main checkpoint under a known prefix, while
|
||||
v2 uses a standalone tokenizer repo.
|
||||
|
||||
For backward compatibility and simplicity, we delegate to the v2
|
||||
``load_higgs_audio_codec`` function when the audio tokenizer repo
|
||||
is available. When not (or as primary path), we extract codec weights
|
||||
from the v3 safetensors via the engine's weight iterator.
|
||||
"""
|
||||
|
||||
input_modalities = "audio"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: HiggsAudioV3Config | None = None,
|
||||
*,
|
||||
vllm_config: VllmConfig | None = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
if vllm_config is not None:
|
||||
hf_config = vllm_config.model_config.hf_config
|
||||
if isinstance(hf_config, HiggsAudioV3Config):
|
||||
self.config = hf_config
|
||||
else:
|
||||
self.config = HiggsAudioV3Config(**hf_config.to_dict())
|
||||
self._model_path: str | None = vllm_config.model_config.model
|
||||
self.vllm_config: VllmConfig | None = vllm_config
|
||||
else:
|
||||
if config is None:
|
||||
raise TypeError("HiggsAudioV3Code2Wav: provide either `config` or `vllm_config`.")
|
||||
self.config = config
|
||||
self._model_path = None
|
||||
self.vllm_config = None
|
||||
|
||||
self.sample_rate: int = int(self.config.sample_rate)
|
||||
self.num_codebooks: int = int(self.config.num_codebooks)
|
||||
self.num_real_codes: int = int(self.config.num_real_codes)
|
||||
self.hop_length: int = 960
|
||||
|
||||
# Engine-runner hooks
|
||||
self.have_multimodal_outputs = True
|
||||
self.has_preprocess = False
|
||||
self.has_postprocess = False
|
||||
self.enable_update_additional_information = True
|
||||
self.requires_raw_input_tokens = True
|
||||
|
||||
# Codec modules (populated by load_weights)
|
||||
self.quantizer: HiggsAudioRVQ | None = None
|
||||
self.fc2: nn.Linear | None = None
|
||||
self.acoustic_decoder: nn.Module | None = None
|
||||
self._loaded: bool = False
|
||||
# Do NOT eagerly load from standalone tokenizer here.
|
||||
# load_weights() will try bundled V3 codec first, then fall back.
|
||||
|
||||
def _load_from_tokenizer_repo(self) -> None:
|
||||
"""Load codec from the standalone higgs-audio-v2-tokenizer repo."""
|
||||
tokenizer_id = "bosonai/higgs-audio-v2-tokenizer"
|
||||
# Try to resolve from HF cache
|
||||
from huggingface_hub import try_to_load_from_cache
|
||||
|
||||
cached = try_to_load_from_cache(repo_id=tokenizer_id, filename="config.json")
|
||||
if isinstance(cached, str) and os.path.isfile(cached):
|
||||
tokenizer_dir = os.path.dirname(cached)
|
||||
else:
|
||||
from huggingface_hub.constants import HF_HUB_CACHE
|
||||
|
||||
safe = tokenizer_id.replace("/", "--")
|
||||
snapshots_dir = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots")
|
||||
if os.path.isdir(snapshots_dir):
|
||||
tokenizer_dir = None
|
||||
for rev in os.listdir(snapshots_dir):
|
||||
candidate = os.path.join(snapshots_dir, rev)
|
||||
if os.path.isfile(os.path.join(candidate, "config.json")):
|
||||
tokenizer_dir = candidate
|
||||
break
|
||||
if tokenizer_dir is None:
|
||||
raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}")
|
||||
else:
|
||||
raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}")
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
quantizer, fc2, acoustic_decoder, _cfg = load_higgs_audio_codec(tokenizer_dir, device)
|
||||
self.quantizer = quantizer
|
||||
self.fc2 = fc2
|
||||
self.acoustic_decoder = acoustic_decoder
|
||||
self._loaded = True
|
||||
logger.info("Loaded HiggsAudioV3Code2Wav from standalone tokenizer repo.")
|
||||
|
||||
# ------------------------------------------------------------------ engine hooks
|
||||
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
||||
if input_ids.numel() == 0:
|
||||
return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
|
||||
return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
|
||||
|
||||
def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------ load
|
||||
def load_weights(self, weights_or_model_dir, device: torch.device | None = None):
|
||||
if self._loaded:
|
||||
if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)):
|
||||
for _ in weights_or_model_dir:
|
||||
pass
|
||||
return {name for name, _ in self.named_parameters()}
|
||||
|
||||
if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)):
|
||||
# Engine path: try to extract V3 bundled codec weights first
|
||||
codec_state: dict[str, torch.Tensor] = {}
|
||||
for name, tensor in weights_or_model_dir:
|
||||
if name.startswith(_CODEC_PREFIX):
|
||||
stripped = name[len(_CODEC_PREFIX) :]
|
||||
codec_state[stripped] = tensor
|
||||
|
||||
if codec_state:
|
||||
# Bundled keys found — load strictly, no fallback on error.
|
||||
self._load_from_bundled_state(codec_state, device)
|
||||
logger.info(
|
||||
"Loaded HiggsAudioV3Code2Wav from bundled V3 checkpoint (%d codec keys consumed).",
|
||||
len(codec_state),
|
||||
)
|
||||
return {name for name, _ in self.named_parameters()}
|
||||
|
||||
# No bundled keys — fallback to standalone tokenizer repo
|
||||
if not self._loaded:
|
||||
try:
|
||||
self._load_from_tokenizer_repo()
|
||||
except (FileNotFoundError, OSError) as exc2:
|
||||
raise RuntimeError(
|
||||
"HiggsAudioV3Code2Wav: could not load codec weights. "
|
||||
"Neither V3 bundled codec keys (tied.embedding."
|
||||
"modality_embeddings.0.model.*) were found in the "
|
||||
"checkpoint, nor is bosonai/higgs-audio-v2-tokenizer "
|
||||
f"cached in HF_HOME. Last error: {exc2}"
|
||||
) from exc2
|
||||
return {name for name, _ in self.named_parameters()}
|
||||
|
||||
# File-based path
|
||||
if not self._loaded:
|
||||
self._load_from_tokenizer_repo()
|
||||
return None
|
||||
|
||||
def _load_from_bundled_state(
|
||||
self,
|
||||
codec_state: dict[str, torch.Tensor],
|
||||
device: torch.device | None = None,
|
||||
) -> None:
|
||||
"""Load codec from V3 checkpoint's bundled weights.
|
||||
|
||||
Keys arrive with the ``tied.embedding.modality_embeddings.0.model.``
|
||||
prefix already stripped. We remap to the same layout that
|
||||
``load_higgs_audio_codec`` expects and build the modules.
|
||||
"""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import (
|
||||
_remap_boson_model_pth_state_dict,
|
||||
)
|
||||
|
||||
if device is None:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
# The bundled codec may use boson-ai key naming. Try remapping.
|
||||
needs_remap = any(k.startswith("quantizer.vq.layers.") for k in codec_state)
|
||||
if needs_remap:
|
||||
codec_state = _remap_boson_model_pth_state_dict(codec_state)
|
||||
|
||||
# Discover parameters from the state dict
|
||||
proj_key = "quantizer.quantizers.0.project_out.weight"
|
||||
if proj_key not in codec_state:
|
||||
raise KeyError(
|
||||
f"Bundled codec state dict missing '{proj_key}'. "
|
||||
f"Available keys (first 10): {list(codec_state.keys())[:10]}"
|
||||
)
|
||||
hidden_size = codec_state[proj_key].shape[0]
|
||||
codebook_dim = 64
|
||||
codebook_size = 1024
|
||||
required_quantizers = self.num_codebooks
|
||||
|
||||
# Build RVQ — require exactly num_codebooks quantizers
|
||||
quantizer = HiggsAudioRVQ(
|
||||
num_quantizers=required_quantizers,
|
||||
codebook_size=codebook_size,
|
||||
codebook_dim=codebook_dim,
|
||||
hidden_size=hidden_size,
|
||||
).to(device)
|
||||
for i in range(required_quantizers):
|
||||
prefix = f"quantizer.quantizers.{i}"
|
||||
# Require codebook weight
|
||||
embed_key = None
|
||||
for key_suffix in [".codebook.embed", ".codebook.weight"]:
|
||||
candidate = f"{prefix}{key_suffix}"
|
||||
if candidate in codec_state:
|
||||
embed_key = candidate
|
||||
break
|
||||
if embed_key is None:
|
||||
raise KeyError(
|
||||
f"Bundled codec missing codebook weight for quantizer {i}. "
|
||||
f"Expected {prefix}.codebook.embed or {prefix}.codebook.weight"
|
||||
)
|
||||
quantizer.quantizers[i].codebook.weight.data.copy_(codec_state[embed_key].to(device))
|
||||
# Require projector weights
|
||||
proj_w = f"{prefix}.project_out.weight"
|
||||
proj_b = f"{prefix}.project_out.bias"
|
||||
if proj_w not in codec_state:
|
||||
raise KeyError(f"Bundled codec missing {proj_w}")
|
||||
if proj_b not in codec_state:
|
||||
raise KeyError(f"Bundled codec missing {proj_b}")
|
||||
quantizer.quantizers[i].project_out.weight.data.copy_(codec_state[proj_w].to(device))
|
||||
quantizer.quantizers[i].project_out.bias.data.copy_(codec_state[proj_b].to(device))
|
||||
|
||||
# Build fc2
|
||||
fc2_w = codec_state.get("fc2.weight")
|
||||
fc2_b = codec_state.get("fc2.bias")
|
||||
if fc2_w is None or fc2_b is None:
|
||||
raise KeyError(
|
||||
"Bundled codec state dict missing fc2.weight and/or fc2.bias. "
|
||||
f"Available keys (first 10): {list(codec_state.keys())[:10]}"
|
||||
)
|
||||
fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device)
|
||||
fc2.weight.data.copy_(fc2_w.to(device))
|
||||
fc2.bias.data.copy_(fc2_b.to(device))
|
||||
|
||||
# Build acoustic decoder — detect layout from key names.
|
||||
# V3 checkpoint uses OmniVoice layout (block.N/conv1/conv2/snake1),
|
||||
# not boson-ai standalone layout (model.0/model.1).
|
||||
ad_keys = {
|
||||
k[len("acoustic_decoder.") :]: v for k, v in codec_state.items() if k.startswith("acoustic_decoder.")
|
||||
}
|
||||
is_boson_layout = any(k.startswith("model.") for k in ad_keys)
|
||||
if is_boson_layout:
|
||||
acoustic_decoder = build_boson_dac_decoder(device)
|
||||
else:
|
||||
# OmniVoice layout — need the v2 tokenizer config for DacConfig
|
||||
tokenizer_config = {
|
||||
"acoustic_model_config": {
|
||||
"codebook_dim": 8,
|
||||
"codebook_size": 1024,
|
||||
"decoder_hidden_size": 1024,
|
||||
"downsampling_ratios": [8, 5, 4, 2, 3],
|
||||
"encoder_hidden_size": 64,
|
||||
"hidden_size": 256,
|
||||
"hop_length": 960,
|
||||
"model_type": "dac",
|
||||
"n_codebooks": 9,
|
||||
"sampling_rate": 16000,
|
||||
"upsampling_ratios": [8, 5, 4, 2, 3],
|
||||
}
|
||||
}
|
||||
acoustic_decoder = build_higgs_audio_acoustic_decoder(tokenizer_config, device)
|
||||
if not ad_keys:
|
||||
raise KeyError(
|
||||
"Bundled codec state dict has no acoustic_decoder.* keys. "
|
||||
f"Available keys (first 10): {list(codec_state.keys())[:10]}"
|
||||
)
|
||||
load_report = acoustic_decoder.load_state_dict(ad_keys, strict=False)
|
||||
if load_report.missing_keys:
|
||||
raise RuntimeError(
|
||||
f"Bundled DAC decoder missing {len(load_report.missing_keys)} keys: {load_report.missing_keys[:5]}..."
|
||||
)
|
||||
if load_report.unexpected_keys:
|
||||
raise RuntimeError(
|
||||
f"Bundled DAC decoder has {len(load_report.unexpected_keys)} unexpected keys: "
|
||||
f"{load_report.unexpected_keys[:5]}..."
|
||||
)
|
||||
acoustic_decoder.eval()
|
||||
|
||||
# Verify required decoder-side codec keys were consumed. The V3
|
||||
# checkpoint also bundles encoder-side keys (acoustic_encoder.*,
|
||||
# encoder_semantic.*, decoder_semantic.*, semantic_model.*, fc.*, fc1.*)
|
||||
# that the DAC decoder does not need — those are allowed but ignored.
|
||||
consumed: set[str] = set()
|
||||
for i in range(required_quantizers):
|
||||
prefix = f"quantizer.quantizers.{i}"
|
||||
for suffix in [".codebook.embed", ".codebook.weight"]:
|
||||
if f"{prefix}{suffix}" in codec_state:
|
||||
consumed.add(f"{prefix}{suffix}")
|
||||
consumed.add(f"{prefix}.project_out.weight")
|
||||
consumed.add(f"{prefix}.project_out.bias")
|
||||
consumed.add("fc2.weight")
|
||||
consumed.add("fc2.bias")
|
||||
for k in codec_state:
|
||||
if k.startswith("acoustic_decoder."):
|
||||
consumed.add(k)
|
||||
|
||||
# Known unused prefixes: encoder-side modules and quantizer training
|
||||
# artifacts (cluster_size, embed_avg, inited, project_in) that are
|
||||
# bundled in the checkpoint but not needed for inference decoding.
|
||||
_KNOWN_UNUSED_PREFIXES = (
|
||||
"acoustic_encoder.",
|
||||
"encoder_semantic.",
|
||||
"decoder_semantic.",
|
||||
"semantic_model.",
|
||||
"fc.",
|
||||
"fc1.",
|
||||
)
|
||||
# Quantizer training-only suffixes (per-codebook)
|
||||
_QUANTIZER_TRAINING_SUFFIXES = (
|
||||
".codebook.cluster_size",
|
||||
".codebook.embed_avg",
|
||||
".codebook.inited",
|
||||
".project_in.weight",
|
||||
".project_in.bias",
|
||||
)
|
||||
unconsumed = set()
|
||||
for k in codec_state:
|
||||
if k in consumed:
|
||||
continue
|
||||
if any(k.startswith(p) for p in _KNOWN_UNUSED_PREFIXES):
|
||||
continue
|
||||
if any(k.endswith(s) for s in _QUANTIZER_TRAINING_SUFFIXES):
|
||||
continue
|
||||
unconsumed.add(k)
|
||||
if unconsumed:
|
||||
raise RuntimeError(
|
||||
f"Bundled codec has {len(unconsumed)} unexpected decoder-side keys: {sorted(unconsumed)[:10]}..."
|
||||
)
|
||||
|
||||
self.quantizer = quantizer
|
||||
self.fc2 = fc2
|
||||
self.acoustic_decoder = acoustic_decoder
|
||||
self._loaded = True
|
||||
|
||||
# ------------------------------------------------------------------ decode
|
||||
@torch.inference_mode()
|
||||
def decode_codes(self, audio_codes: torch.Tensor) -> torch.Tensor:
|
||||
"""Decode [B, num_codebooks=8, T] codes to PCM [B, 1, T*960]."""
|
||||
if not self._loaded:
|
||||
raise RuntimeError("HiggsAudioV3Code2Wav not loaded.")
|
||||
|
||||
codes = self._validate_codes(audio_codes)
|
||||
rvq_codes = codes.transpose(0, 1).long()
|
||||
quantized = self.quantizer.decode(rvq_codes)
|
||||
quantized = quantized.to(dtype=self.fc2.weight.dtype)
|
||||
quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)
|
||||
first_param = next(self.acoustic_decoder.parameters(), None)
|
||||
if first_param is not None and quantized.dtype != first_param.dtype:
|
||||
quantized = quantized.to(dtype=first_param.dtype)
|
||||
audio = self.acoustic_decoder(quantized)
|
||||
if audio.dim() == 2:
|
||||
audio = audio.unsqueeze(1)
|
||||
return audio.to(codes.device)
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward_chunk(
|
||||
self,
|
||||
audio_codes: torch.Tensor,
|
||||
*,
|
||||
left_context_size: int = 0,
|
||||
right_holdback_size: int = 0,
|
||||
hop_length: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
# ``left_context_size`` frames trimmed off the front (already-emitted
|
||||
# prefix). ``right_holdback_size`` frames trimmed off the end
|
||||
# (deferred to the next chunk so the codec has future context for
|
||||
# the emit region's trailing samples).
|
||||
hop = int(hop_length) if hop_length is not None else self.hop_length
|
||||
pcm = self.decode_codes(audio_codes)
|
||||
right_trim = right_holdback_size * hop
|
||||
left_trim = left_context_size * hop
|
||||
if right_trim > 0:
|
||||
if pcm.shape[-1] <= right_trim:
|
||||
return pcm[..., :0]
|
||||
pcm = pcm[..., :-right_trim]
|
||||
if left_trim == 0:
|
||||
return pcm
|
||||
if pcm.shape[-1] <= left_trim:
|
||||
return pcm[..., :0]
|
||||
return pcm[..., left_trim:]
|
||||
|
||||
# ------------------------------------------------------------------ runtime forward
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: Any = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
runtime_additional_information: list[dict[str, Any]] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor | OmniOutput:
|
||||
if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 3:
|
||||
return self.decode_codes(input_ids)
|
||||
|
||||
sr_val = int(self.sample_rate)
|
||||
sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
|
||||
empty = torch.zeros((0,), dtype=torch.float32)
|
||||
|
||||
if input_ids is None or input_ids.numel() == 0:
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]},
|
||||
)
|
||||
|
||||
ids = input_ids.reshape(-1).to(dtype=torch.long)
|
||||
request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
|
||||
|
||||
left_context_size = [0] * len(request_ids_list)
|
||||
right_holdback_size = [0] * len(request_ids_list)
|
||||
if runtime_additional_information is not None:
|
||||
for i, info in enumerate(runtime_additional_information):
|
||||
if i >= len(left_context_size):
|
||||
break
|
||||
meta = info.get("meta", {}) if isinstance(info, dict) else {}
|
||||
if "left_context_size" in meta:
|
||||
left_context_size[i] = int(meta["left_context_size"])
|
||||
if "right_holdback_size" in meta:
|
||||
right_holdback_size[i] = int(meta["right_holdback_size"])
|
||||
|
||||
wavs: list[torch.Tensor] = []
|
||||
for i, req_ids in enumerate(request_ids_list):
|
||||
n = int(req_ids.numel())
|
||||
if n == 0:
|
||||
wavs.append(empty)
|
||||
continue
|
||||
if n % self.num_codebooks != 0:
|
||||
logger.warning(
|
||||
"HiggsAudioV3Code2Wav: flat code length %d not divisible by %d",
|
||||
n,
|
||||
self.num_codebooks,
|
||||
)
|
||||
wavs.append(empty)
|
||||
continue
|
||||
frames = n // self.num_codebooks
|
||||
codes_qf = req_ids.reshape(self.num_codebooks, frames)
|
||||
codes_bqf = codes_qf.unsqueeze(0)
|
||||
try:
|
||||
pcm = self.forward_chunk(
|
||||
codes_bqf,
|
||||
left_context_size=left_context_size[i],
|
||||
right_holdback_size=right_holdback_size[i],
|
||||
hop_length=self.hop_length,
|
||||
)
|
||||
except ValueError as exc:
|
||||
logger.warning("HiggsAudioV3Code2Wav: decode skipped (%s)", exc)
|
||||
wavs.append(empty)
|
||||
continue
|
||||
wavs.append(pcm.squeeze(0).squeeze(0).to(torch.float32).cpu())
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"model_outputs": wavs,
|
||||
"sr": [sr_tensor] * len(wavs),
|
||||
},
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ helpers
|
||||
def _validate_codes(self, audio_codes: torch.Tensor) -> torch.Tensor:
|
||||
if audio_codes.ndim != 3:
|
||||
raise ValueError(
|
||||
f"audio_codes must have shape [B, {self.num_codebooks}, T]; got {tuple(audio_codes.shape)}"
|
||||
)
|
||||
if int(audio_codes.shape[1]) != self.num_codebooks:
|
||||
raise ValueError(f"dim 1 must equal num_codebooks={self.num_codebooks}; got {int(audio_codes.shape[1])}")
|
||||
# ``.item()`` triggers a GPU->CPU sync and is illegal inside CUDA-graph
|
||||
# capture (the dummy run before warmup hits this path). Skip the
|
||||
# range check while a graph is being captured; the per-step validator
|
||||
# still fires at real inference time.
|
||||
if audio_codes.numel() > 0 and not torch.cuda.is_current_stream_capturing():
|
||||
max_val = int(audio_codes.max().item())
|
||||
min_val = int(audio_codes.min().item())
|
||||
if max_val >= self.num_real_codes or min_val < 0:
|
||||
raise ValueError(
|
||||
f"audio_codes out of range: min={min_val}, max={max_val}; expected [0, {self.num_real_codes - 1}]"
|
||||
)
|
||||
return audio_codes
|
||||
|
||||
@staticmethod
|
||||
def _split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]:
|
||||
if seq_token_counts is not None and len(seq_token_counts) > 1:
|
||||
boundaries = [0]
|
||||
for count in seq_token_counts:
|
||||
boundaries.append(boundaries[-1] + int(count))
|
||||
n = int(ids.numel())
|
||||
return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))]
|
||||
return [ids]
|
||||
|
||||
|
||||
HiggsAudioV3Code2WavForConditionalGeneration = HiggsAudioV3Code2Wav
|
||||
@@ -1,970 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Stage-0 talker for higgs-audio v3 (Qwen3 backbone, fused multi-codebook).
|
||||
|
||||
Architecture:
|
||||
- Backbone: Qwen3 (~4B, 36 layers, 2560 hidden, GQA 32/8). No DualFFN.
|
||||
- Fused multi-codebook embedding: [N*V, D] weight, offset lookup, sum across N
|
||||
- Fused multi-codebook head: same weight (tied), reshape to [L, N, V]
|
||||
- MusicGen-style delay pattern [0,1,...,7] with BOC/EOC
|
||||
- Audio feedback: replace continuation-token embedding with fused codebook embed
|
||||
|
||||
Weight loading maps from the HF checkpoint's prefixes:
|
||||
tied.embedding.text_embedding. -> model.embed_tokens.
|
||||
body.layers. -> model.layers.
|
||||
body.norm. -> model.norm.
|
||||
tied.head.text_head. -> lm_head.
|
||||
tied.embedding.modality_embeddings.0.embedding. -> multimodal_embedding.
|
||||
tied.embedding.modality_embeddings.0.model.* -> skipped (codec for code2wav)
|
||||
tied.head.modality_heads.0.* -> skipped when tied
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
|
||||
from vllm.model_executor.models.qwen3 import Qwen3Model
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
|
||||
HiggsAudioV3Config,
|
||||
)
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
__all__ = ["HiggsAudioV3TalkerForConditionalGeneration"]
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Delay pattern constants
|
||||
BOC_ID = 1024 # beginning of codebook
|
||||
EOC_ID = 1025 # end of codebook
|
||||
|
||||
# Checkpoint prefix mapping
|
||||
_BACKBONE_PREFIX_MAP = {
|
||||
"tied.embedding.text_embedding.": "model.embed_tokens.",
|
||||
"body.layers.": "model.layers.",
|
||||
"body.norm.": "model.norm.",
|
||||
"tied.head.text_head.": "lm_head.",
|
||||
}
|
||||
_MODALITY_EMBEDDING_PREFIX = "tied.embedding.modality_embeddings.0.embedding."
|
||||
_MODALITY_HEAD_PREFIX = "tied.head.modality_heads.0."
|
||||
_CODEC_PREFIX = "tied.embedding.modality_embeddings.0.model."
|
||||
|
||||
|
||||
class HiggsFusedMultiTextEmbedding(nn.Module):
|
||||
"""Fused multi-codebook embedding: [N*V, D] weight + offset lookup."""
|
||||
|
||||
def __init__(self, num_codebooks: int, vocab_size: int, hidden_size: int):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty(num_codebooks * vocab_size, hidden_size))
|
||||
self.num_codebooks = num_codebooks
|
||||
self.vocab_size = vocab_size
|
||||
|
||||
def forward(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
N = self.num_codebooks
|
||||
V = self.vocab_size
|
||||
offsets = torch.arange(N, device=codes.device, dtype=codes.dtype) * V
|
||||
fused_ids = codes + offsets
|
||||
return F.embedding(fused_ids, self.weight).sum(dim=-2)
|
||||
|
||||
|
||||
class HiggsFusedMultiTextHead(nn.Module):
|
||||
"""Fused multi-codebook head: [L, D] -> [L, N, V] via one linear."""
|
||||
|
||||
def __init__(self, num_codebooks: int, vocab_size: int, hidden_size: int):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty(num_codebooks * vocab_size, hidden_size))
|
||||
self.num_codebooks = num_codebooks
|
||||
self.vocab_size = vocab_size
|
||||
|
||||
def generate(self, hidden: torch.Tensor) -> torch.Tensor:
|
||||
logits = F.linear(hidden, self.weight)
|
||||
return logits.reshape(hidden.shape[0], self.num_codebooks, self.vocab_size)
|
||||
|
||||
|
||||
class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
|
||||
"""Stage-0 talker for higgs-audio v3.
|
||||
|
||||
Wraps Qwen3Model backbone + fused multi-codebook modules for TTS generation
|
||||
with MusicGen-style delay pattern sampling and audio feedback embedding.
|
||||
"""
|
||||
|
||||
# Tell the AR runner to call model.sample() instead of the stock sampler.
|
||||
prefer_model_sampler: bool = True
|
||||
# Tell the runner to call postprocess() to emit per-step audio codes.
|
||||
have_multimodal_outputs: bool = True
|
||||
has_postprocess: bool = True
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
hf_config = vllm_config.model_config.hf_config
|
||||
if isinstance(hf_config, HiggsAudioV3Config):
|
||||
self.config = hf_config
|
||||
else:
|
||||
self.config = HiggsAudioV3Config(**hf_config.to_dict())
|
||||
|
||||
self.vllm_config = vllm_config
|
||||
self.num_codebooks = int(self.config.num_codebooks)
|
||||
self.codebook_size = int(self.config.codebook_size)
|
||||
hidden_size = int(self.config.audio_hidden_size)
|
||||
self.tie_modality = self.config.tie_modality_embeddings
|
||||
|
||||
# Fused multi-codebook modules
|
||||
self.multimodal_embedding = HiggsFusedMultiTextEmbedding(self.num_codebooks, self.codebook_size, hidden_size)
|
||||
self.modality_head = HiggsFusedMultiTextHead(self.num_codebooks, self.codebook_size, hidden_size)
|
||||
if self.tie_modality:
|
||||
self.modality_head.weight = self.multimodal_embedding.weight
|
||||
|
||||
# Qwen3 backbone
|
||||
self._backbone_config = self.config.text_config
|
||||
backbone_vllm_config = copy.copy(vllm_config)
|
||||
backbone_model_config = copy.copy(vllm_config.model_config)
|
||||
backbone_model_config.hf_config = self._backbone_config
|
||||
backbone_vllm_config.model_config = backbone_model_config
|
||||
|
||||
self.model = Qwen3Model(
|
||||
vllm_config=backbone_vllm_config,
|
||||
prefix=f"{prefix}.model" if prefix else "model",
|
||||
)
|
||||
|
||||
if self._backbone_config.tie_word_embeddings:
|
||||
self.lm_head = self.model.embed_tokens
|
||||
else:
|
||||
self.lm_head = ParallelLMHead(
|
||||
self._backbone_config.vocab_size,
|
||||
self._backbone_config.hidden_size,
|
||||
prefix=f"{prefix}.lm_head" if prefix else "lm_head",
|
||||
)
|
||||
|
||||
self.logits_processor = LogitsProcessor(self._backbone_config.vocab_size)
|
||||
|
||||
# Audio continuation token ID — resolved lazily from tokenizer.
|
||||
# This is the <|audio|> token that serves as the LM-level continuation
|
||||
# marker during audio generation (equivalent to v2's audio_token_id).
|
||||
self._audio_continuation_id: int | None = None
|
||||
self._eos_token_id: int | None = None
|
||||
self._resolved_tokens = False
|
||||
|
||||
# Per-request audio state keyed by batch row index.
|
||||
# Reset per slot via _slot_output_len tracking (same pattern as v2).
|
||||
self._audio_state: dict[int, dict[str, Any]] = {}
|
||||
self._slot_output_len: dict[int, int] = {}
|
||||
self._last_logits_hidden: torch.Tensor | None = None
|
||||
self._last_step_input_ids: torch.Tensor | None = None
|
||||
self._last_step_query_start_loc: torch.Tensor | None = None
|
||||
self._last_first_audio_after_start: torch.Tensor | None = None
|
||||
self._last_audio_codes: torch.Tensor | None = None
|
||||
self._postprocess_cursor: int = 0
|
||||
|
||||
# Pre-allocated decode-step audio feedback buffers (CUDA-graph safe).
|
||||
# Populated by sample(), read by forward() via torch.where (no dict).
|
||||
max_bs = 64 # safe upper bound; will grow if needed
|
||||
self._decode_last_codes = torch.zeros(max_bs, self.num_codebooks, dtype=torch.long)
|
||||
self._decode_has_codes = torch.zeros(max_bs, dtype=torch.bool)
|
||||
|
||||
# PrefixCache opt-outs (mirror qwen3_tts pattern):
|
||||
# 1. The talker only consumes the last token's hidden state, so the
|
||||
# runner can skip the per-step full hidden-state GPU->CPU merge
|
||||
# that PrefixCache otherwise does.
|
||||
# 2. Per-step ``codes.audio`` rows stay GPU-resident; defer the CPU
|
||||
# write of the prefix-cache mm-output copy to request finish so
|
||||
# the per-step bookkeeping does not block batching. Stage 0 can
|
||||
# then set ``enable_prefix_caching: true`` without the regression
|
||||
# observed in qwen3_tts (#3665).
|
||||
self.requires_full_prefix_cached_hidden_states = False
|
||||
self.deferred_prefix_cache_mm_keys = {"codes.audio"}
|
||||
|
||||
def _resolve_token_ids(self) -> None:
|
||||
"""Resolve <|audio|> and eos token IDs.
|
||||
|
||||
Prefers config's pre-resolved IDs (from ``resolve_special_tokens()``),
|
||||
falls back to loading the HF tokenizer directly.
|
||||
"""
|
||||
if self._resolved_tokens:
|
||||
return
|
||||
self._resolved_tokens = True
|
||||
|
||||
# Try config first (populated by resolve_special_tokens or from_pretrained)
|
||||
cfg_audio = getattr(self.config, "audio_continuation_id", None)
|
||||
cfg_eos = getattr(self.config, "eos_token_id", None)
|
||||
if cfg_audio is not None:
|
||||
self._audio_continuation_id = int(cfg_audio)
|
||||
if cfg_eos is not None:
|
||||
self._eos_token_id = int(cfg_eos)
|
||||
|
||||
if self._audio_continuation_id is not None:
|
||||
logger.info(
|
||||
"Resolved v3 token IDs from config: audio_continuation=%s, eos=%s",
|
||||
self._audio_continuation_id,
|
||||
self._eos_token_id,
|
||||
)
|
||||
return
|
||||
|
||||
# Fallback: load tokenizer directly
|
||||
model_path = getattr(self.vllm_config.model_config, "model", None)
|
||||
if model_path is None:
|
||||
return
|
||||
try:
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
vocab = dict(tokenizer.get_added_vocab())
|
||||
if "<|audio|>" in vocab:
|
||||
self._audio_continuation_id = vocab["<|audio|>"]
|
||||
if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
|
||||
self._eos_token_id = int(tokenizer.eos_token_id)
|
||||
logger.info(
|
||||
"Resolved v3 token IDs from tokenizer: audio_continuation=%s, eos=%s",
|
||||
self._audio_continuation_id,
|
||||
self._eos_token_id,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to resolve token IDs from tokenizer: %s", exc)
|
||||
|
||||
# ------------------------------------------------------------------ forward
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: Any | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor:
|
||||
if inputs_embeds is None:
|
||||
# Mask -100 placeholders to 0 before embedding. Use torch.where
|
||||
# (no Python data-dependent branch) so this is CUDA-graph safe.
|
||||
safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids)
|
||||
hidden_states = self.model.embed_tokens(safe_ids)
|
||||
else:
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
if input_ids is not None:
|
||||
self._last_step_input_ids = input_ids
|
||||
|
||||
# Stash query_start_loc and max_query_len for prefill detection
|
||||
_max_query_len = None
|
||||
try:
|
||||
from vllm.forward_context import get_forward_context
|
||||
|
||||
attn_metadata = get_forward_context().attn_metadata
|
||||
if isinstance(attn_metadata, dict) and attn_metadata:
|
||||
attn = next(iter(attn_metadata.values()))
|
||||
else:
|
||||
attn = attn_metadata
|
||||
qsl = getattr(attn, "query_start_loc", None)
|
||||
if isinstance(qsl, torch.Tensor):
|
||||
self._last_step_query_start_loc = qsl.detach().clone()
|
||||
else:
|
||||
self._last_step_query_start_loc = None
|
||||
_max_query_len = getattr(attn, "max_query_len", None)
|
||||
except Exception:
|
||||
self._last_step_query_start_loc = None
|
||||
|
||||
# Detect prefill vs decode using attn_metadata.max_query_len.
|
||||
# max_query_len == 1 means pure decode (even with N concurrent
|
||||
# requests, each contributes exactly 1 token). numel > 1 is NOT
|
||||
# reliable because decode with N>1 concurrent requests gives
|
||||
# input_ids.numel() == N.
|
||||
if _max_query_len is not None:
|
||||
is_prefill = int(_max_query_len) > 1
|
||||
else:
|
||||
is_prefill = input_ids is not None and inputs_embeds is None and int(input_ids.numel()) > 1
|
||||
if is_prefill:
|
||||
# Voice clone: replace -100 placeholder positions with ref audio embeddings
|
||||
info_dicts = kwargs.get("model_intermediate_buffer")
|
||||
if info_dicts is None:
|
||||
info_dicts = kwargs.get("runtime_additional_information")
|
||||
hidden_states = self._apply_ref_audio_substitution(hidden_states, input_ids, info_dicts)
|
||||
|
||||
# Audio feedback: replace continuation token embeddings with audio
|
||||
# embeddings from the last decoded frame (CUDA-graph safe).
|
||||
if input_ids is not None and inputs_embeds is None:
|
||||
hidden_states = self._apply_audio_feedback(hidden_states, input_ids)
|
||||
|
||||
residual: torch.Tensor | None = None
|
||||
for layer in self.model.layers:
|
||||
hidden_states, residual = layer(positions, hidden_states, residual)
|
||||
|
||||
norm_out = self.model.norm(hidden_states, residual)
|
||||
if isinstance(norm_out, tuple):
|
||||
norm_out = norm_out[0]
|
||||
return norm_out
|
||||
|
||||
def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor:
|
||||
self._last_logits_hidden = hidden_states
|
||||
return self.logits_processor(self.lm_head, hidden_states, sampling_metadata)
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
||||
safe_ids = input_ids
|
||||
if input_ids is not None and (input_ids < 0).any():
|
||||
safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids)
|
||||
text_embed = self.model.embed_tokens(safe_ids)
|
||||
return self._apply_audio_feedback(text_embed, input_ids)
|
||||
|
||||
# ------------------------------------------------------------------ ref audio substitution
|
||||
def _apply_ref_audio_substitution(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
input_ids: torch.Tensor,
|
||||
info_dicts: list[dict[str, Any]] | None,
|
||||
) -> torch.Tensor:
|
||||
"""Replace -100 placeholder positions with fused multi-codebook embeddings
|
||||
of the delay-pattern-encoded reference audio codes.
|
||||
|
||||
Called at prefill to inject voice clone reference. ``info_dicts`` is a
|
||||
list of per-request dicts from ``model_intermediate_buffer``, each
|
||||
containing ``audio_input_ids`` ([T, N] delayed codes) and
|
||||
``audio_input_ids_mask`` ([T] bool mask).
|
||||
"""
|
||||
if not info_dicts:
|
||||
return hidden_states
|
||||
|
||||
PLACEHOLDER = -100
|
||||
flat_ids = input_ids.reshape(-1)
|
||||
placeholder_mask = flat_ids == PLACEHOLDER
|
||||
if not placeholder_mask.any():
|
||||
return hidden_states
|
||||
|
||||
# Use query_start_loc to map placeholders to per-request spans
|
||||
q_start = self._last_step_query_start_loc
|
||||
if not isinstance(q_start, torch.Tensor) or q_start.numel() < 2:
|
||||
# Fallback: single-request batch
|
||||
q_start_list = [0, int(flat_ids.numel())]
|
||||
else:
|
||||
q_start_list = q_start.detach().to("cpu").tolist()
|
||||
|
||||
new_hidden: torch.Tensor | None = None
|
||||
num_requests = min(len(info_dicts), len(q_start_list) - 1)
|
||||
|
||||
for i in range(num_requests):
|
||||
info = info_dicts[i]
|
||||
if not isinstance(info, dict):
|
||||
continue
|
||||
|
||||
codes = info.get("audio_input_ids")
|
||||
mask = info.get("audio_input_ids_mask")
|
||||
|
||||
# Handle msgspec serialization (may be list-wrapped)
|
||||
if isinstance(codes, list):
|
||||
codes = codes[0] if codes else None
|
||||
if isinstance(mask, list):
|
||||
mask = mask[0] if mask else None
|
||||
if not isinstance(codes, torch.Tensor):
|
||||
continue
|
||||
|
||||
# codes shape: [T, num_codebooks] delayed reference codes
|
||||
if codes.ndim == 3:
|
||||
codes = codes[0]
|
||||
if codes.ndim != 2:
|
||||
continue
|
||||
|
||||
if isinstance(mask, torch.Tensor):
|
||||
if mask.ndim == 2:
|
||||
mask = mask[0]
|
||||
codes = codes[mask.to(dtype=torch.bool)]
|
||||
|
||||
if codes.numel() == 0:
|
||||
continue
|
||||
|
||||
# Find placeholder positions in this request's span
|
||||
s = int(q_start_list[i])
|
||||
e = int(q_start_list[i + 1])
|
||||
if e - s <= 1:
|
||||
continue # Decode step, skip
|
||||
|
||||
span_mask = placeholder_mask[s:e]
|
||||
placeholders = span_mask.nonzero(as_tuple=True)[0]
|
||||
n_codes = int(codes.shape[0])
|
||||
|
||||
if int(placeholders.numel()) < n_codes:
|
||||
continue # Mismatch
|
||||
|
||||
# Embed delayed codes via fused multi-codebook embedding
|
||||
target = placeholders[:n_codes] + s
|
||||
codes_device = codes.to(device=hidden_states.device, dtype=torch.long)
|
||||
embeds = self.multimodal_embedding(codes_device) # [n_codes, hidden]
|
||||
|
||||
if new_hidden is None:
|
||||
new_hidden = hidden_states.clone()
|
||||
flat_hidden = new_hidden.reshape(-1, new_hidden.shape[-1])
|
||||
flat_hidden[target] = embeds.to(new_hidden.dtype)
|
||||
|
||||
return new_hidden if new_hidden is not None else hidden_states
|
||||
|
||||
# ------------------------------------------------------------------ audio feedback
|
||||
def _apply_audio_feedback(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""Replace decode-step embeddings with audio feedback from pre-allocated buffers.
|
||||
|
||||
CUDA-graph safe: reads from pre-allocated _decode_last_codes and
|
||||
_decode_has_codes tensors using torch.where (no Python dict lookup).
|
||||
The buffers are populated by sample() after each step.
|
||||
|
||||
For decode steps (1 token per request): position i maps to row i.
|
||||
For prefill: audio feedback is not needed (ref audio substitution
|
||||
handles the prefill path separately).
|
||||
"""
|
||||
bs = hidden_states.shape[0]
|
||||
|
||||
# Ensure buffers are on the right device and large enough
|
||||
if self._decode_last_codes.device != hidden_states.device:
|
||||
self._decode_last_codes = self._decode_last_codes.to(hidden_states.device)
|
||||
self._decode_has_codes = self._decode_has_codes.to(hidden_states.device)
|
||||
if bs > self._decode_last_codes.shape[0]:
|
||||
old_size = self._decode_last_codes.shape[0]
|
||||
new_size = max(bs, old_size * 2)
|
||||
new_codes = torch.zeros(new_size, self.num_codebooks, dtype=torch.long, device=hidden_states.device)
|
||||
new_has = torch.zeros(new_size, dtype=torch.bool, device=hidden_states.device)
|
||||
new_codes[:old_size] = self._decode_last_codes
|
||||
new_has[:old_size] = self._decode_has_codes
|
||||
self._decode_last_codes = new_codes
|
||||
self._decode_has_codes = new_has
|
||||
|
||||
# Compute audio embeddings from last_codes for ALL rows (graph-safe)
|
||||
codes_slice = self._decode_last_codes[:bs] # [bs, N]
|
||||
has_codes = self._decode_has_codes[:bs].unsqueeze(-1) # [bs, 1]
|
||||
audio_embeds = self.multimodal_embedding(codes_slice) # [bs, D]
|
||||
audio_embeds = audio_embeds.to(dtype=hidden_states.dtype)
|
||||
|
||||
# Select: where has_codes, use audio embed; else keep text embed
|
||||
return torch.where(has_codes, audio_embeds, hidden_states)
|
||||
|
||||
# ------------------------------------------------------------------ sampling
|
||||
def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> Any:
|
||||
"""Model-owned sampler with delay-pattern audio dispatch.
|
||||
|
||||
Mirrors v2's pattern: bias LM logits to force audio continuation,
|
||||
sample multi-codebook codes via the fused head, apply delay pattern,
|
||||
and accumulate per-request state.
|
||||
"""
|
||||
self._resolve_token_ids()
|
||||
|
||||
sampler = getattr(self, "_stock_sampler", None)
|
||||
if sampler is None:
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
sampler = Sampler()
|
||||
self._stock_sampler = sampler
|
||||
|
||||
audio_id = self._audio_continuation_id
|
||||
num_codebooks = self.num_codebooks
|
||||
|
||||
# Slot reuse detection: check output_token_ids length for each row.
|
||||
# If a slot's output length decreased (fresh request reusing a finished
|
||||
# slot), drop stale _audio_state for that row.
|
||||
output_ids = getattr(sampling_metadata, "output_token_ids", None)
|
||||
if output_ids is not None:
|
||||
for i in range(len(output_ids)):
|
||||
cur_len = len(output_ids[i]) if output_ids[i] else 0
|
||||
prev_len = self._slot_output_len.get(i, 0)
|
||||
if cur_len < prev_len:
|
||||
# Fresh request reusing this slot — clear stale state
|
||||
self._audio_state.pop(i, None)
|
||||
self._slot_output_len[i] = 0
|
||||
self._slot_output_len[i] = cur_len
|
||||
|
||||
# Bias LM logits for audio continuation
|
||||
self._apply_audio_mode_bias(logits, sampling_metadata)
|
||||
sampler_output = sampler(logits=logits, sampling_metadata=sampling_metadata)
|
||||
|
||||
hidden = self._last_logits_hidden
|
||||
self._last_logits_hidden = None
|
||||
if hidden is None or audio_id is None:
|
||||
self._last_audio_codes = None
|
||||
return sampler_output
|
||||
|
||||
sampled = getattr(sampler_output, "sampled_token_ids", None)
|
||||
if sampled is None:
|
||||
self._last_audio_codes = None
|
||||
return sampler_output
|
||||
sampled_flat = sampled.reshape(-1)
|
||||
if int(sampled_flat.numel()) != int(hidden.shape[0]):
|
||||
self._last_audio_codes = None
|
||||
return sampler_output
|
||||
|
||||
is_audio = sampled_flat == audio_id
|
||||
|
||||
# Handle first-after-<|audio|> transition: initialize state but do NOT
|
||||
# skip sampling. Unlike v2 (which seeds an all-BOC frame and skips the
|
||||
# first audio step), v3/sglang samples codebook 0 on the very first
|
||||
# step with delay_count=0 (codebooks 1-7 masked to BOC).
|
||||
first_after_start = self._last_first_audio_after_start
|
||||
self._last_first_audio_after_start = None
|
||||
|
||||
if isinstance(first_after_start, torch.Tensor) and first_after_start.numel() == is_audio.shape[0]:
|
||||
first_after_start = first_after_start.to(is_audio.device)
|
||||
init_rows = first_after_start & is_audio
|
||||
if bool(init_rows.any()):
|
||||
boc_frame = torch.full((num_codebooks,), BOC_ID, dtype=torch.long, device=hidden.device)
|
||||
for bi in torch.nonzero(init_rows, as_tuple=False).reshape(-1).tolist():
|
||||
bi = int(bi)
|
||||
self._slot_output_len[bi] = 0
|
||||
self._audio_state[bi] = {
|
||||
"num_delay": 0,
|
||||
"num_remaining_delays": None,
|
||||
"audio_out_ids": None,
|
||||
"last_codes": boc_frame.clone(),
|
||||
"should_terminate": False,
|
||||
}
|
||||
# Update CUDA-graph-safe feedback buffers
|
||||
self._decode_last_codes[bi] = boc_frame.to(self._decode_last_codes.device)
|
||||
self._decode_has_codes[bi] = True
|
||||
# Do NOT remove init_rows from is_audio — let them be sampled
|
||||
|
||||
if not bool(is_audio.any()):
|
||||
self._last_audio_codes = None
|
||||
return sampler_output
|
||||
|
||||
audio_row_indices = torch.nonzero(is_audio, as_tuple=False).reshape(-1).tolist()
|
||||
|
||||
# Per-codebook logits at audio positions
|
||||
cb_logits = self._audio_codebook_logits(hidden, is_audio)
|
||||
|
||||
# Apply delay pattern masking BEFORE sampling
|
||||
self._apply_delay_pattern_masking(cb_logits, audio_row_indices)
|
||||
|
||||
# Sample per-codebook
|
||||
cb_logits_2d = cb_logits.reshape(-1, cb_logits.shape[-1])
|
||||
codes_2d = self._sample_audio_codes(cb_logits_2d)
|
||||
codes_flat = codes_2d.view(cb_logits.shape[0], cb_logits.shape[1]).to(torch.long)
|
||||
|
||||
# Update delay pattern state
|
||||
eos_stream = EOC_ID
|
||||
bos = BOC_ID
|
||||
new_codes_flat: list[torch.Tensor] = []
|
||||
for local_i, batch_i in enumerate(audio_row_indices):
|
||||
state = self._audio_state.setdefault(
|
||||
int(batch_i),
|
||||
{
|
||||
"num_delay": 0,
|
||||
"num_remaining_delays": None,
|
||||
"audio_out_ids": None,
|
||||
"last_codes": torch.full((num_codebooks,), bos, dtype=torch.long, device=hidden.device),
|
||||
"should_terminate": False,
|
||||
},
|
||||
)
|
||||
num_delay: int = state["num_delay"]
|
||||
num_remaining_delays: int | None = state["num_remaining_delays"]
|
||||
this_codes = codes_flat[local_i].clone()
|
||||
|
||||
# Leading delay-pattern BOS pad
|
||||
if num_delay + 1 < num_codebooks:
|
||||
this_codes[num_delay + 1 :] = bos
|
||||
num_delay += 1
|
||||
|
||||
# Trailing eos ramp-down
|
||||
if num_remaining_delays is not None:
|
||||
this_codes[: num_codebooks - num_remaining_delays] = eos_stream
|
||||
num_remaining_delays -= 1
|
||||
else:
|
||||
eos_positions = (this_codes == eos_stream).nonzero(as_tuple=False).reshape(-1)
|
||||
if eos_positions.numel() > 0:
|
||||
last_eos_idx = int(eos_positions[-1].item())
|
||||
this_codes[: last_eos_idx + 1] = eos_stream
|
||||
num_remaining_delays = num_codebooks - last_eos_idx - 1
|
||||
|
||||
if num_remaining_delays is not None and num_remaining_delays <= 0:
|
||||
# Ramp-down complete — terminate without emitting this frame.
|
||||
state["num_delay"] = 0
|
||||
state["num_remaining_delays"] = None
|
||||
state["should_terminate"] = True
|
||||
new_codes_flat.append(torch.full_like(this_codes, -1))
|
||||
continue
|
||||
|
||||
state["num_delay"] = num_delay
|
||||
state["num_remaining_delays"] = num_remaining_delays
|
||||
state["last_codes"] = this_codes.clone()
|
||||
# Update CUDA-graph-safe feedback buffers
|
||||
self._decode_last_codes[int(batch_i)] = this_codes.to(self._decode_last_codes.device)
|
||||
self._decode_has_codes[int(batch_i)] = True
|
||||
if state["audio_out_ids"] is None:
|
||||
state["audio_out_ids"] = this_codes.unsqueeze(-1).clone()
|
||||
else:
|
||||
state["audio_out_ids"] = torch.cat([state["audio_out_ids"], this_codes.unsqueeze(-1)], dim=-1)
|
||||
|
||||
new_codes_flat.append(this_codes)
|
||||
|
||||
# Build full codes tensor [batch_size, num_codebooks]
|
||||
# -1 marks "no audio code at this position"
|
||||
codes_full = torch.full(
|
||||
(int(sampled_flat.numel()), num_codebooks),
|
||||
-1,
|
||||
dtype=torch.long,
|
||||
device=hidden.device,
|
||||
)
|
||||
if new_codes_flat:
|
||||
stacked = torch.stack(new_codes_flat, dim=0).to(hidden.device)
|
||||
codes_full[is_audio] = stacked
|
||||
|
||||
self._last_audio_codes = codes_full
|
||||
self._postprocess_cursor = 0
|
||||
return sampler_output
|
||||
|
||||
# ------------------------------------------------------------------ postprocess
|
||||
def postprocess(
|
||||
self,
|
||||
hidden_states_slice: torch.Tensor,
|
||||
multimodal_outputs: Any = None,
|
||||
**req_infos: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Publish per-request audio codes into model_intermediate_buffer.
|
||||
|
||||
Called once per request in batch order. Indexes _last_audio_codes
|
||||
by a running cursor (one row per request per step).
|
||||
"""
|
||||
_ = multimodal_outputs
|
||||
codes_full = self._last_audio_codes
|
||||
if codes_full is None:
|
||||
return {}
|
||||
|
||||
cursor = int(self._postprocess_cursor)
|
||||
if cursor >= int(codes_full.shape[0]):
|
||||
self._postprocess_cursor = 0
|
||||
return {}
|
||||
slice_codes = codes_full[cursor : cursor + 1]
|
||||
self._postprocess_cursor = cursor + 1
|
||||
|
||||
# Drop placeholder rows (-1)
|
||||
audio_rows = slice_codes[:, 0] >= 0
|
||||
if not bool(audio_rows.any()):
|
||||
return {}
|
||||
new_codes = slice_codes[audio_rows].to(torch.int32)
|
||||
return {"codes": {"audio": new_codes}}
|
||||
|
||||
# ------------------------------------------------------------------ helpers
|
||||
def _audio_codebook_logits(self, hidden_states: torch.Tensor, audio_mask: torch.Tensor) -> torch.Tensor:
|
||||
mask = audio_mask.reshape(-1).to(hidden_states.device)
|
||||
hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1])
|
||||
if not mask.any():
|
||||
return torch.empty(
|
||||
(0, self.num_codebooks, self.codebook_size),
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
return self.modality_head.generate(hidden_flat[mask])
|
||||
|
||||
def _apply_delay_pattern_masking(self, cb_logits: torch.Tensor, audio_row_indices: list[int]) -> None:
|
||||
"""Mask per-codebook logits according to delay pattern state, in-place.
|
||||
|
||||
During delay phase: codebooks beyond delay_count only allow BOC.
|
||||
During ramp-down: locked codebooks only allow EOC.
|
||||
Normal generation: BOC disallowed; only cb0 allows EOC.
|
||||
"""
|
||||
bos_pre = BOC_ID
|
||||
eos_pre = EOC_ID
|
||||
num_codebooks = self.num_codebooks
|
||||
for local_i, batch_i in enumerate(audio_row_indices):
|
||||
state = self._audio_state.get(int(batch_i))
|
||||
num_delay = int(state["num_delay"]) if state else 0
|
||||
num_rem = state.get("num_remaining_delays") if state else None
|
||||
|
||||
if num_rem is not None:
|
||||
lock_until = num_codebooks - int(num_rem)
|
||||
for q in range(num_codebooks):
|
||||
row = cb_logits[local_i, q]
|
||||
if q < lock_until:
|
||||
mask = torch.full_like(row, float("-inf"))
|
||||
mask[eos_pre] = row[eos_pre]
|
||||
cb_logits[local_i, q] = mask
|
||||
else:
|
||||
cb_logits[local_i, q, bos_pre] = float("-inf")
|
||||
cb_logits[local_i, q, eos_pre] = float("-inf")
|
||||
else:
|
||||
for q in range(num_codebooks):
|
||||
row = cb_logits[local_i, q]
|
||||
if q > num_delay:
|
||||
mask = torch.full_like(row, float("-inf"))
|
||||
mask[bos_pre] = row[bos_pre]
|
||||
cb_logits[local_i, q] = mask
|
||||
else:
|
||||
cb_logits[local_i, q, bos_pre] = float("-inf")
|
||||
if q != 0:
|
||||
cb_logits[local_i, q, eos_pre] = float("-inf")
|
||||
|
||||
def _sample_audio_codes(self, logits_2d: torch.Tensor) -> torch.Tensor:
|
||||
"""Replicate upstream sampling: temperature → top-k → top-p → multinomial."""
|
||||
x = logits_2d.float()
|
||||
top_k = 50
|
||||
top_p = 0.95
|
||||
if 0 < top_k < x.shape[-1]:
|
||||
kth = x.topk(top_k, dim=-1).values[..., -1:]
|
||||
x = x.masked_fill(x < kth, float("-inf"))
|
||||
if 0.0 < top_p < 1.0:
|
||||
sorted_logits, sorted_idx = x.sort(dim=-1, descending=True)
|
||||
cumprobs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
|
||||
sorted_mask = cumprobs > top_p
|
||||
sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
|
||||
sorted_mask[..., 0] = False
|
||||
mask = torch.zeros_like(x, dtype=torch.bool)
|
||||
mask.scatter_(-1, sorted_idx, sorted_mask)
|
||||
x = x.masked_fill(mask, float("-inf"))
|
||||
# Detect all-masked rows BEFORE softmax (softmax of all-inf yields NaN).
|
||||
has_finite = torch.isfinite(x).any(dim=-1)
|
||||
all_masked = ~has_finite
|
||||
if all_masked.any():
|
||||
# For all-masked rows, fall back to argmax (picks least-negative).
|
||||
fallback = x.argmax(dim=-1)
|
||||
if has_finite.any():
|
||||
probs = x[has_finite].softmax(dim=-1)
|
||||
sampled = torch.multinomial(probs, num_samples=1).squeeze(-1)
|
||||
result = fallback.clone()
|
||||
result[has_finite] = sampled
|
||||
return result
|
||||
return fallback
|
||||
probs = x.softmax(dim=-1)
|
||||
return torch.multinomial(probs, num_samples=1).squeeze(-1)
|
||||
|
||||
def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> None:
|
||||
"""Detect <|audio|> transition, force continuation, force eos at ramp-down.
|
||||
|
||||
Mirrors v2's _apply_audio_mode_bias: walks per-request to find the
|
||||
previous token, and if it was <|audio|> or the continuation token,
|
||||
forces the next emit to <|audio|>. On ramp-down completion, forces eos.
|
||||
"""
|
||||
if logits is None or logits.ndim != 2:
|
||||
return
|
||||
|
||||
audio_id = self._audio_continuation_id
|
||||
eos_id = self._eos_token_id
|
||||
if audio_id is None:
|
||||
return
|
||||
|
||||
num_rows = int(logits.shape[0])
|
||||
prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None)
|
||||
output_ids = getattr(sampling_metadata, "output_token_ids", None)
|
||||
|
||||
# Fallback prev-token source from stashed input_ids
|
||||
stash_ids = self._last_step_input_ids
|
||||
stash_tail: list[int] | None = None
|
||||
if isinstance(stash_ids, torch.Tensor) and stash_ids.numel() > 0:
|
||||
q_start = self._last_step_query_start_loc
|
||||
if isinstance(q_start, torch.Tensor) and int(q_start.numel()) == num_rows + 1:
|
||||
q_start_cpu = q_start.detach().to("cpu").tolist()
|
||||
tail_idx = [max(0, int(q_start_cpu[i + 1]) - 1) for i in range(num_rows)]
|
||||
flat_ids = stash_ids.detach().to("cpu").tolist()
|
||||
stash_tail = [int(flat_ids[idx]) if idx < len(flat_ids) else -1 for idx in tail_idx]
|
||||
elif int(stash_ids.numel()) >= num_rows:
|
||||
stash_tail = stash_ids[-num_rows:].detach().to("cpu").tolist()
|
||||
|
||||
first_after_start = torch.zeros(num_rows, dtype=torch.bool, device=logits.device)
|
||||
|
||||
for i in range(num_rows):
|
||||
prev: int | None = None
|
||||
if output_ids is not None and i < len(output_ids):
|
||||
hist = output_ids[i]
|
||||
if hist:
|
||||
prev = int(hist[-1])
|
||||
if prev is None and prompt_ids is not None:
|
||||
try:
|
||||
p_i = prompt_ids[i]
|
||||
if hasattr(p_i, "tolist"):
|
||||
p_i = p_i.tolist()
|
||||
if p_i:
|
||||
prev = int(p_i[-1])
|
||||
except (IndexError, TypeError):
|
||||
prev = None
|
||||
if prev is None and stash_tail is not None and i < len(stash_tail):
|
||||
prev = int(stash_tail[i])
|
||||
if prev is None:
|
||||
continue
|
||||
|
||||
# Only bias if previous token was <|audio|> (the continuation token)
|
||||
if prev != audio_id:
|
||||
continue
|
||||
|
||||
# Check if this is the FIRST step after <|audio|> appears
|
||||
# (i.e., transitioning from prompt to audio generation)
|
||||
audio_state = self._audio_state.get(i)
|
||||
if audio_state is None or audio_state.get("should_terminate"):
|
||||
# No state yet, or stale state from a finished prior request
|
||||
# reusing this slot — treat as first audio step
|
||||
self._audio_state.pop(i, None)
|
||||
# Clear CUDA-graph feedback buffers for this slot
|
||||
self._decode_has_codes[i] = False
|
||||
first_after_start[i] = True
|
||||
|
||||
# Check for ramp-down termination
|
||||
should_terminate = bool(isinstance(audio_state, dict) and audio_state.get("should_terminate"))
|
||||
if should_terminate and eos_id is not None and 0 <= eos_id < int(logits.shape[-1]):
|
||||
row = logits[i]
|
||||
mask = torch.full_like(row, float("-inf"))
|
||||
mask[eos_id] = row[eos_id]
|
||||
logits[i].copy_(mask)
|
||||
audio_state["should_terminate"] = False
|
||||
continue
|
||||
|
||||
# Force audio continuation token
|
||||
row = logits[i]
|
||||
mask = torch.full_like(row, float("-inf"))
|
||||
if 0 <= audio_id < row.shape[-1]:
|
||||
mask[audio_id] = row[audio_id]
|
||||
logits[i].copy_(mask)
|
||||
|
||||
self._last_first_audio_after_start = first_after_start
|
||||
|
||||
# ------------------------------------------------------------------ omni output
|
||||
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
|
||||
if isinstance(model_outputs, OmniOutput):
|
||||
return model_outputs
|
||||
hidden = model_outputs
|
||||
|
||||
info_dicts = kwargs.get("model_intermediate_buffer")
|
||||
if info_dicts is None:
|
||||
info_dicts = kwargs.get("runtime_additional_information")
|
||||
if info_dicts is None:
|
||||
info_dicts = []
|
||||
|
||||
audio_codes_list: list[torch.Tensor] = []
|
||||
any_nonempty = False
|
||||
for info in info_dicts:
|
||||
ac: torch.Tensor | None = None
|
||||
if isinstance(info, dict):
|
||||
codes_field = info.get("codes")
|
||||
if isinstance(codes_field, dict):
|
||||
ac = codes_field.get("audio")
|
||||
else:
|
||||
ac = info.get("audio_codes")
|
||||
if isinstance(ac, torch.Tensor) and ac.numel() > 0:
|
||||
audio_codes_list.append(ac)
|
||||
any_nonempty = True
|
||||
else:
|
||||
audio_codes_list.append(torch.empty(0, dtype=torch.long))
|
||||
|
||||
if any_nonempty:
|
||||
return OmniOutput(
|
||||
text_hidden_states=hidden,
|
||||
multimodal_outputs={"codes": {"audio": audio_codes_list}},
|
||||
)
|
||||
return OmniOutput(text_hidden_states=hidden, multimodal_outputs=None)
|
||||
|
||||
# ------------------------------------------------------------------ weight loading
|
||||
|
||||
# Per-layer suffixes from the actual V3 checkpoint (results/higgs_v3_checkpoint_analysis.txt)
|
||||
_V3_LAYER_SUFFIXES = (
|
||||
"input_layernorm.weight",
|
||||
"mlp.down_proj.weight",
|
||||
"mlp.gate_proj.weight",
|
||||
"mlp.up_proj.weight",
|
||||
"post_attention_layernorm.weight",
|
||||
"self_attn.k_norm.weight",
|
||||
"self_attn.k_proj.weight",
|
||||
"self_attn.o_proj.weight",
|
||||
"self_attn.q_norm.weight",
|
||||
"self_attn.q_proj.weight",
|
||||
"self_attn.v_proj.weight",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _build_required_keys(cls, num_layers: int) -> set[str]:
|
||||
"""Build the exact set of required V3 checkpoint keys."""
|
||||
keys = {
|
||||
"tied.embedding.text_embedding.weight",
|
||||
"body.norm.weight",
|
||||
f"{_MODALITY_EMBEDDING_PREFIX}weight",
|
||||
}
|
||||
for i in range(num_layers):
|
||||
for suffix in cls._V3_LAYER_SUFFIXES:
|
||||
keys.add(f"body.layers.{i}.{suffix}")
|
||||
return keys
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
backbone_weights: list[tuple[str, torch.Tensor]] = []
|
||||
loaded_params: set[str] = set()
|
||||
own_params = dict(self.named_parameters())
|
||||
seen_checkpoint_keys: set[str] = set()
|
||||
|
||||
for name, tensor in weights:
|
||||
seen_checkpoint_keys.add(name)
|
||||
|
||||
mapped = self._map_weight_name(name)
|
||||
if mapped is None:
|
||||
continue
|
||||
|
||||
if mapped.startswith("model.") or mapped.startswith("lm_head."):
|
||||
backbone_weights.append((mapped, tensor))
|
||||
elif mapped in own_params:
|
||||
param = own_params[mapped]
|
||||
if param.shape != tensor.shape:
|
||||
raise ValueError(
|
||||
f"Shape mismatch for {mapped}: expected {param.shape}, "
|
||||
f"got {tensor.shape} (checkpoint key: {name})"
|
||||
)
|
||||
param.data.copy_(tensor.to(param.dtype))
|
||||
loaded_params.add(mapped)
|
||||
|
||||
if backbone_weights:
|
||||
backbone_module = _BackboneWrapper(self.model, self.lm_head, self._backbone_config)
|
||||
loaded = backbone_module.load_weights(iter(backbone_weights))
|
||||
loaded_params.update(loaded)
|
||||
|
||||
# Resolve special token IDs from the tokenizer
|
||||
model_path = getattr(self.vllm_config.model_config, "model", None)
|
||||
if model_path:
|
||||
self.config.resolve_special_tokens(model_path)
|
||||
self._resolve_token_ids()
|
||||
|
||||
# Verify every required checkpoint key was seen.
|
||||
num_layers = int(self._backbone_config.num_hidden_layers)
|
||||
required = self._build_required_keys(num_layers)
|
||||
missing = required - seen_checkpoint_keys
|
||||
if missing:
|
||||
raise RuntimeError(
|
||||
f"HiggsAudioV3Talker: {len(missing)} required checkpoint keys missing: {sorted(missing)[:5]}..."
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"HiggsAudioV3Talker: loaded %d params, modality_embedding=%s, tied=%s",
|
||||
len(loaded_params),
|
||||
tuple(self.multimodal_embedding.weight.shape),
|
||||
self.tie_modality,
|
||||
)
|
||||
return loaded_params
|
||||
|
||||
def _map_weight_name(self, name: str) -> str | None:
|
||||
if name.startswith(_CODEC_PREFIX):
|
||||
return None
|
||||
if name.startswith(_MODALITY_HEAD_PREFIX):
|
||||
if self.tie_modality:
|
||||
return None
|
||||
return name.replace(_MODALITY_HEAD_PREFIX, "modality_head.")
|
||||
if name.startswith(_MODALITY_EMBEDDING_PREFIX):
|
||||
return name.replace(_MODALITY_EMBEDDING_PREFIX, "multimodal_embedding.")
|
||||
for ckpt_prefix, model_prefix in _BACKBONE_PREFIX_MAP.items():
|
||||
if name.startswith(ckpt_prefix):
|
||||
return name.replace(ckpt_prefix, model_prefix, 1)
|
||||
# Reject unexpected non-codec Higgs checkpoint keys
|
||||
raise ValueError(
|
||||
f"Unexpected checkpoint key with no known mapping: {name!r}. "
|
||||
f"Known prefixes: {list(_BACKBONE_PREFIX_MAP.keys())}, "
|
||||
f"{_MODALITY_EMBEDDING_PREFIX!r}, {_MODALITY_HEAD_PREFIX!r}, {_CODEC_PREFIX!r}"
|
||||
)
|
||||
|
||||
|
||||
class _BackboneWrapper(nn.Module):
|
||||
"""Wrapper to use AutoWeightsLoader for Qwen3 backbone."""
|
||||
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
}
|
||||
|
||||
def __init__(self, model, lm_head, config):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.lm_head = lm_head
|
||||
self.config = config
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
|
||||
skip = ["lm_head."] if getattr(self.config, "tie_word_embeddings", False) else None
|
||||
loader = AutoWeightsLoader(self, skip_prefixes=skip)
|
||||
return loader.load_weights(weights)
|
||||
@@ -1,135 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Prompt builder for higgs-audio v3 TTS.
|
||||
|
||||
Prompt formats:
|
||||
Zero-shot:
|
||||
<|tts|> <|text|> {text tokens} <|audio|>
|
||||
Voice clone (no ref text):
|
||||
<|tts|> <|ref_audio|> [-100]×N <|text|> {text tokens} <|audio|>
|
||||
Voice clone (with ref text):
|
||||
<|tts|> <|ref_text|> {ref text tokens} <|ref_audio|> [-100]×N <|text|> {text tokens} <|audio|>
|
||||
|
||||
-100 placeholders are replaced at prefill time with fused multi-codebook
|
||||
embeddings of the delay-pattern-encoded reference audio codes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"HiggsAudioV3TokenizerAdapter",
|
||||
"AUDIO_PLACEHOLDER_ID",
|
||||
"encode_reference_audio",
|
||||
"apply_delay_pattern",
|
||||
]
|
||||
|
||||
AUDIO_PLACEHOLDER_ID = -100
|
||||
|
||||
_REQUIRED_SPECIALS: tuple[str, ...] = (
|
||||
"<|tts|>",
|
||||
"<|text|>",
|
||||
"<|audio|>",
|
||||
)
|
||||
|
||||
# Codec stream special token IDs (inside the codebook vocab, not the LM vocab)
|
||||
BOC_ID = 1024
|
||||
EOC_ID = 1025
|
||||
|
||||
|
||||
def apply_delay_pattern(codes_tn: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply MusicGen-style delay pattern to raw codes.
|
||||
|
||||
Input: [T, N] raw codes (T frames, N codebooks).
|
||||
Output: [T + N - 1, N] delayed codes with BOC/EOC padding.
|
||||
|
||||
Codebook c is delayed by c positions: rows 0..c-1 get BOC,
|
||||
rows c..c+T-1 get real codes, rows c+T..T+N-2 get EOC.
|
||||
"""
|
||||
T, N = codes_tn.shape
|
||||
L = T + N - 1
|
||||
delayed = torch.full((L, N), EOC_ID, dtype=codes_tn.dtype)
|
||||
for c in range(N):
|
||||
delayed[:c, c] = BOC_ID
|
||||
delayed[c : c + T, c] = codes_tn[:, c]
|
||||
return delayed
|
||||
|
||||
|
||||
def encode_reference_audio(
|
||||
wav: np.ndarray,
|
||||
sr: int,
|
||||
) -> torch.Tensor:
|
||||
"""Encode a reference audio clip to codec codes [T, num_codebooks].
|
||||
|
||||
Uses the same HiggsAudioV2TokenizerModel as v2 (same codec).
|
||||
Returns raw codes before delay pattern application.
|
||||
"""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
|
||||
_encode_ref_audio_codes,
|
||||
)
|
||||
|
||||
# _encode_ref_audio_codes returns [num_codebooks, T_raw]
|
||||
codes_qt = _encode_ref_audio_codes(wav, sr)
|
||||
# Transpose to [T_raw, num_codebooks]
|
||||
return codes_qt.transpose(0, 1).contiguous()
|
||||
|
||||
|
||||
class HiggsAudioV3TokenizerAdapter:
|
||||
"""Wraps the HF tokenizer and builds TTS prompts."""
|
||||
|
||||
def __init__(self, tokenizer: Any) -> None:
|
||||
self._tok = tokenizer
|
||||
vocab = dict(tokenizer.get_added_vocab())
|
||||
missing = [t for t in _REQUIRED_SPECIALS if t not in vocab]
|
||||
if missing:
|
||||
raise ValueError(f"Tokenizer is missing Higgs TTS v3 specials: {missing}")
|
||||
self.tts_id: int = vocab["<|tts|>"]
|
||||
self.text_id: int = vocab["<|text|>"]
|
||||
self.audio_id: int = vocab["<|audio|>"]
|
||||
self.ref_audio_id: int | None = vocab.get("<|ref_audio|>")
|
||||
self.ref_text_id: int | None = vocab.get("<|ref_text|>")
|
||||
|
||||
@property
|
||||
def tokenizer(self) -> Any:
|
||||
return self._tok
|
||||
|
||||
def build_prompt(
|
||||
self,
|
||||
text: str,
|
||||
*,
|
||||
num_ref_tokens: int = 0,
|
||||
reference_text: str | None = None,
|
||||
) -> list[int]:
|
||||
"""Build a TTS prompt.
|
||||
|
||||
Args:
|
||||
text: Target text to synthesize.
|
||||
num_ref_tokens: Number of delay-pattern reference code rows.
|
||||
0 means zero-shot (no voice clone).
|
||||
reference_text: Optional transcript of the reference audio.
|
||||
"""
|
||||
if not text or not text.strip():
|
||||
raise ValueError("Text input must be non-empty for TTS")
|
||||
if num_ref_tokens < 0:
|
||||
raise ValueError(f"num_ref_tokens must be >= 0, got {num_ref_tokens}")
|
||||
|
||||
ids: list[int] = [self.tts_id]
|
||||
|
||||
# Voice clone: optional ref text + ref audio placeholders
|
||||
if reference_text and num_ref_tokens > 0 and self.ref_text_id is not None:
|
||||
ids.append(self.ref_text_id)
|
||||
ids.extend(self._tok.encode(reference_text, add_special_tokens=False))
|
||||
if num_ref_tokens > 0:
|
||||
if self.ref_audio_id is None:
|
||||
raise ValueError("Tokenizer missing <|ref_audio|> for voice clone")
|
||||
ids.append(self.ref_audio_id)
|
||||
ids.extend([AUDIO_PLACEHOLDER_ID] * num_ref_tokens)
|
||||
|
||||
ids.append(self.text_id)
|
||||
ids.extend(self._tok.encode(text, add_special_tokens=False))
|
||||
ids.append(self.audio_id)
|
||||
return ids
|
||||
@@ -1,60 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""higgs-audio v3 pipeline: Talker (text -> 8-codebook codec) -> Code2Wav (codec -> 24 kHz PCM).
|
||||
|
||||
Two delivery modes are wired here:
|
||||
|
||||
* ``sync_process_input_func`` runs when the deploy YAML has
|
||||
``async_chunk: false``. The orchestrator collects the entire Stage-0 emit,
|
||||
reverts the delay pattern once, and hands a single payload to Stage 1.
|
||||
* ``async_chunk_process_next_stage_input_func`` runs when the deploy YAML
|
||||
has ``async_chunk: true``. Stage 0 dispatches per AR step; the streaming
|
||||
adapter buffers raw delay-pattern rows, slides a window with left context
|
||||
and right holdback, and emits codec-ready frames per chunk. Stage 1
|
||||
trims the overlap on both ends so the client sees a coherent PCM stream.
|
||||
"""
|
||||
|
||||
from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
)
|
||||
|
||||
_PROC = "vllm_omni.model_executor.stage_input_processors.higgs_audio_v3"
|
||||
|
||||
HIGGS_AUDIO_V3_PIPELINE = PipelineConfig(
|
||||
model_type="higgs_multimodal_qwen3",
|
||||
model_arch="HiggsAudioV3TalkerForConditionalGeneration",
|
||||
hf_architectures=("HiggsMultimodalQwen3ForConditionalGeneration",),
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="higgs_audio_v3",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
input_sources=(),
|
||||
owns_tokenizer=True,
|
||||
engine_output_type="latent",
|
||||
# stop_token_ids: the model-owned sampler forces eos at ramp-down
|
||||
# completion. Safety stops from the actual V3 checkpoint:
|
||||
# 151643 = <|endoftext|> (eos_token_id from config.json)
|
||||
# 151671 = <|audio_end|> (audio generation end marker)
|
||||
sampling_constraints={
|
||||
"detokenize": False,
|
||||
"stop_token_ids": [151643, 151671],
|
||||
},
|
||||
async_chunk_process_next_stage_input_func=f"{_PROC}.talker2code2wav_async_chunk",
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=1,
|
||||
model_stage="code2wav",
|
||||
execution_type=StageExecutionType.LLM_GENERATION,
|
||||
input_sources=(0,),
|
||||
final_output=True,
|
||||
final_output_type="audio",
|
||||
engine_output_type="audio",
|
||||
model_arch="HiggsAudioV3Code2WavForConditionalGeneration",
|
||||
sync_process_input_func=f"{_PROC}.talker2code2wav",
|
||||
sampling_constraints={"detokenize": True},
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -1,111 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Bounded LRU cache for Higgs Audio v3 reference-audio codec encoding.
|
||||
|
||||
Voice-clone requests that re-send the same ref_audio are otherwise re-encoded
|
||||
through the DAC encoder on every call. The cache keys by a stable content
|
||||
hash (size + sample rate + head/mid/tail sentinel bytes) so the same audio
|
||||
served via different URLs or repeated base64 payloads also hits warm.
|
||||
Tensors are cloned on get/put so callers and the cache cannot alias storage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
__all__ = ["cached_encode_reference_audio", "cache_stats", "clear_cache"]
|
||||
|
||||
_DEFAULT_MAX_ENTRIES = 128
|
||||
_SENTINEL_BYTES = 4096
|
||||
|
||||
_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
|
||||
_cache_lock = threading.Lock()
|
||||
_hits = 0
|
||||
_misses = 0
|
||||
|
||||
|
||||
def _content_hash(wav: np.ndarray, sr: int) -> str:
|
||||
"""Stable content hash. Head/mid/tail sentinel keeps cost O(1)."""
|
||||
raw = wav.tobytes()
|
||||
n = len(raw)
|
||||
h = hashlib.sha256()
|
||||
h.update(n.to_bytes(8, "little"))
|
||||
h.update(int(sr).to_bytes(4, "little"))
|
||||
if n == 0:
|
||||
return h.hexdigest()
|
||||
if n <= _SENTINEL_BYTES * 3:
|
||||
h.update(raw)
|
||||
return h.hexdigest()
|
||||
h.update(raw[:_SENTINEL_BYTES])
|
||||
mid = (n - _SENTINEL_BYTES) // 2
|
||||
h.update(raw[mid : mid + _SENTINEL_BYTES])
|
||||
h.update(raw[-_SENTINEL_BYTES:])
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def cached_encode_reference_audio(
|
||||
wav: np.ndarray,
|
||||
sr: int,
|
||||
encode_fn: Callable[[np.ndarray, int], torch.Tensor],
|
||||
max_entries: int = _DEFAULT_MAX_ENTRIES,
|
||||
) -> torch.Tensor:
|
||||
"""LRU-cached wrapper around ``encode_reference_audio``.
|
||||
|
||||
``encode_fn`` is the actual codec encoder; passed in so this module
|
||||
stays decoupled from the tokenizer import (avoids circular imports
|
||||
during stage-input-processor discovery).
|
||||
|
||||
Returns a CPU tensor clone — never the cache's interior storage.
|
||||
"""
|
||||
global _hits, _misses
|
||||
key = _content_hash(wav, sr)
|
||||
with _cache_lock:
|
||||
cached = _cache.get(key)
|
||||
if cached is not None:
|
||||
_hits += 1
|
||||
_cache.move_to_end(key)
|
||||
return cached.clone()
|
||||
|
||||
# Miss: encode outside the lock so concurrent misses on different
|
||||
# refs don't serialize on this single mutex.
|
||||
codes = encode_fn(wav, sr)
|
||||
codes_cpu = codes.detach().to("cpu").contiguous()
|
||||
with _cache_lock:
|
||||
if key in _cache:
|
||||
# Concurrent miss winner already stored it; refresh LRU position
|
||||
# and return that copy to avoid divergent tensors per request.
|
||||
_cache.move_to_end(key)
|
||||
_misses += 1 # we still paid the encode cost
|
||||
return _cache[key].clone()
|
||||
_cache[key] = codes_cpu
|
||||
if len(_cache) > max_entries:
|
||||
_cache.popitem(last=False)
|
||||
_misses += 1
|
||||
return codes_cpu.clone()
|
||||
|
||||
|
||||
def cache_stats() -> dict:
|
||||
"""Return (hits, misses, size, hit_rate) — for /metrics-style reporting."""
|
||||
with _cache_lock:
|
||||
total = _hits + _misses
|
||||
return {
|
||||
"hits": _hits,
|
||||
"misses": _misses,
|
||||
"size": len(_cache),
|
||||
"hit_rate": (_hits / total) if total > 0 else 0.0,
|
||||
}
|
||||
|
||||
|
||||
def clear_cache() -> None:
|
||||
"""Drop all cached entries. Used by tests."""
|
||||
global _hits, _misses
|
||||
with _cache_lock:
|
||||
_cache.clear()
|
||||
_hits = 0
|
||||
_misses = 0
|
||||
@@ -1457,7 +1457,6 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
|
||||
HunyuanImage3Inputs: TypeAlias = HunyuanImage3PixelInputs
|
||||
|
||||
prefer_model_sampler = True
|
||||
logitsprocs_need_output_token_ids = True
|
||||
|
||||
# Siglip2 ViT supports data-parallel encoding (mm_encoder_tp_mode="data"):
|
||||
# weights are replicated and the image batch is sharded across TP ranks.
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker with minimal overrides."""
|
||||
"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker (v0.14) with minimal overrides."""
|
||||
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
|
||||
@@ -50,7 +49,9 @@ from vllm.model_executor.models.utils import (
|
||||
WeightsMapper,
|
||||
init_vllm_registered_model,
|
||||
maybe_prefix,
|
||||
split_list_into_ranges,
|
||||
)
|
||||
from vllm.model_executor.models.vision import get_llm_pos_ids_for_vision
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
@@ -91,15 +92,20 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
|
||||
self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
|
||||
self._validate_mm_updates(mm_prompt_updates, mm_item_counts)
|
||||
|
||||
# Detect use_audio_in_video from mm_kwargs
|
||||
use_audio_in_video = False
|
||||
if "video" in mm_kwargs:
|
||||
for item in mm_kwargs["video"]:
|
||||
if item and item.get("use_audio_in_video"):
|
||||
use_audio_in_video_tensor = item["use_audio_in_video"].data
|
||||
if use_audio_in_video_tensor.numel() > 0:
|
||||
use_audio_in_video = bool(use_audio_in_video_tensor.item())
|
||||
break
|
||||
non_none_items = [item for item in mm_kwargs["video"] if item is not None]
|
||||
if non_none_items:
|
||||
for item in non_none_items:
|
||||
if item.get("use_audio_in_video"):
|
||||
uaiv_tensor = item["use_audio_in_video"].data
|
||||
if uaiv_tensor.numel() > 0:
|
||||
use_audio_in_video = bool(uaiv_tensor.item())
|
||||
break
|
||||
elif "audio" in mm_prompt_updates:
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
|
||||
use_audio_in_video = audio_pad_id not in prompt_ids
|
||||
# for mutilmodality cache
|
||||
if any(item is None for item in mm_kwargs["video"]):
|
||||
video_token_id = self.info.get_hf_config().video_token_id
|
||||
@@ -415,193 +421,167 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
|
||||
def get_language_model(self) -> torch.nn.Module:
|
||||
return self.language_model
|
||||
|
||||
def _get_audio_for_video_mapping(self, mm_features: list[MultiModalFeatureSpec]) -> tuple[dict[int, int], set[int]]:
|
||||
"""
|
||||
Map video offset -> paired audio_feature_length for use_audio_in_video.
|
||||
|
||||
When use_audio_in_video=True, audio is interleaved within video chunks.
|
||||
The pairing is based on feature order in mm_features.
|
||||
"""
|
||||
videos_with_audio = [
|
||||
f
|
||||
for f in mm_features
|
||||
if f.modality == "video" and f.data.get("use_audio_in_video") and f.data["use_audio_in_video"].data.item()
|
||||
]
|
||||
audios = [f for f in mm_features if f.modality == "audio"]
|
||||
|
||||
mapping: dict[int, int] = {}
|
||||
paired_audio_offsets: set[int] = set()
|
||||
for i, video_f in enumerate(videos_with_audio):
|
||||
if i < len(audios):
|
||||
audio_len = audios[i].data["audio_feature_lengths"].data.item()
|
||||
mapping[video_f.mm_position.offset] = audio_len
|
||||
paired_audio_offsets.add(audios[i].mm_position.offset)
|
||||
return mapping, paired_audio_offsets
|
||||
|
||||
def _compute_audio_token_count(self, audio_feature_length: int) -> int:
|
||||
return ((audio_feature_length - 1) // 2 + 1 - 2) // 2 + 1
|
||||
|
||||
def iter_mm_features(self, mm_features: list[MultiModalFeatureSpec]) -> Iterator[tuple[int, str, dict[str, Any]]]:
|
||||
thinker_config = self.config
|
||||
spatial_merge_size = thinker_config.vision_config.spatial_merge_size
|
||||
tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
|
||||
|
||||
sorted_features = sorted(mm_features, key=lambda f: f.mm_position.offset)
|
||||
audio_for_video, paired_audio_offsets = self._get_audio_for_video_mapping(sorted_features)
|
||||
|
||||
for mm_feature in sorted_features:
|
||||
offset = mm_feature.mm_position.offset
|
||||
modality = mm_feature.modality
|
||||
|
||||
if modality == "image":
|
||||
t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
|
||||
yield (
|
||||
offset,
|
||||
"image",
|
||||
{
|
||||
"grid_t": t,
|
||||
"grid_h": h // spatial_merge_size,
|
||||
"grid_w": w // spatial_merge_size,
|
||||
"t_factor": 1.0 * tokens_per_second,
|
||||
},
|
||||
)
|
||||
elif modality == "video":
|
||||
t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
|
||||
second_per_grid_ts = 1.0
|
||||
if mm_feature.data.get("second_per_grid_ts"):
|
||||
second_per_grid_ts = mm_feature.data["second_per_grid_ts"].data.item()
|
||||
use_audio_in_video = False
|
||||
if mm_feature.data.get("use_audio_in_video"):
|
||||
use_audio_in_video = bool(mm_feature.data["use_audio_in_video"].data.item())
|
||||
|
||||
yield (
|
||||
offset,
|
||||
"video",
|
||||
{
|
||||
"grid_t": t,
|
||||
"grid_h": h // spatial_merge_size,
|
||||
"grid_w": w // spatial_merge_size,
|
||||
"t_factor": second_per_grid_ts * tokens_per_second,
|
||||
"use_audio_in_video": use_audio_in_video,
|
||||
"audio_feature_length": audio_for_video.get(offset),
|
||||
},
|
||||
)
|
||||
elif modality == "audio":
|
||||
if offset not in paired_audio_offsets:
|
||||
audio_len = mm_feature.data["audio_feature_lengths"].data.item()
|
||||
yield offset, "audio", {"audio_feature_length": audio_len}
|
||||
|
||||
def _compute_interleaved_positions(self, start_idx: int, data: dict[str, Any]) -> tuple[np.ndarray, int]:
|
||||
grid_t = data["grid_t"]
|
||||
grid_h = data["grid_h"]
|
||||
grid_w = data["grid_w"]
|
||||
t_factor = data["t_factor"]
|
||||
audio_len = data["audio_feature_length"]
|
||||
|
||||
thinker_config = self.config
|
||||
tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
|
||||
seconds_per_chunk = thinker_config.seconds_per_chunk
|
||||
t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
|
||||
|
||||
t_index = (np.arange(grid_t) * t_factor).astype(np.int64)
|
||||
t_index_split_chunk: list[list[int]] = [[] for _ in range((int(t_index.max()) // t_ntoken_per_chunk) + 1)]
|
||||
for t_val in t_index:
|
||||
idx = int(t_val) // t_ntoken_per_chunk
|
||||
t_index_split_chunk[idx].append(int(t_val))
|
||||
|
||||
pure_audio_len = self._compute_audio_token_count(audio_len)
|
||||
added_audio_len = 0
|
||||
pos_ids_list: list[np.ndarray] = []
|
||||
audio_start_idx = start_idx
|
||||
|
||||
for t_chunk in t_index_split_chunk:
|
||||
if not t_chunk:
|
||||
continue
|
||||
|
||||
chunk_t = len(t_chunk)
|
||||
|
||||
h_indices = np.tile(np.arange(grid_h).reshape(1, -1, 1), (chunk_t, 1, grid_w)).flatten()
|
||||
w_indices = np.tile(np.arange(grid_w).reshape(1, 1, -1), (chunk_t, grid_h, 1)).flatten()
|
||||
t_indices = np.repeat(np.array(t_chunk), grid_h * grid_w)
|
||||
|
||||
vision_pos = np.stack([t_indices, h_indices, w_indices]) + start_idx
|
||||
pos_ids_list.append(vision_pos)
|
||||
|
||||
audio_chunk_size = min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
|
||||
if audio_chunk_size > 0:
|
||||
audio_pos = np.broadcast_to(np.arange(audio_chunk_size), (3, audio_chunk_size)) + audio_start_idx
|
||||
pos_ids_list.append(audio_pos)
|
||||
audio_start_idx = audio_start_idx + audio_chunk_size
|
||||
added_audio_len += audio_chunk_size
|
||||
|
||||
if added_audio_len < pure_audio_len:
|
||||
remaining = pure_audio_len - added_audio_len
|
||||
remaining_audio_pos = np.broadcast_to(np.arange(remaining), (3, remaining)) + audio_start_idx
|
||||
pos_ids_list.append(remaining_audio_pos)
|
||||
|
||||
vision_tokens = grid_t * grid_h * grid_w
|
||||
total_tokens = vision_tokens + pure_audio_len
|
||||
|
||||
return np.concatenate(pos_ids_list, axis=1), total_tokens
|
||||
|
||||
def get_mrope_input_positions(
|
||||
self,
|
||||
input_tokens: list[int],
|
||||
mm_features: list[MultiModalFeatureSpec],
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
llm_pos_ids_list: list[np.ndarray] = []
|
||||
st = 0
|
||||
"""
|
||||
Example:
|
||||
|
||||
for offset, modality, data in self.iter_mm_features(mm_features):
|
||||
text_len = offset - st
|
||||
st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0
|
||||
if text_len > 0:
|
||||
llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx)
|
||||
st_idx += text_len
|
||||
(V_i are vision position ids, A_i are audio position ids)
|
||||
|
||||
if modality == "audio":
|
||||
audio_tokens = self._compute_audio_token_count(data["audio_feature_length"])
|
||||
llm_pos_ids_list.append(np.broadcast_to(np.arange(audio_tokens), (3, audio_tokens)) + st_idx)
|
||||
st = offset + audio_tokens
|
||||
|V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
|
||||
|vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
|
||||
"""
|
||||
kwargs = MultiModalFeatureSpec.gather_kwargs(
|
||||
mm_features,
|
||||
{
|
||||
"image_grid_thw",
|
||||
"video_grid_thw",
|
||||
"second_per_grid_ts",
|
||||
"audio_feature_lengths",
|
||||
"use_audio_in_video",
|
||||
},
|
||||
)
|
||||
image_grid_thw = kwargs.get("image_grid_thw", [])
|
||||
video_grid_thw = kwargs.get("video_grid_thw", [])
|
||||
second_per_grid_ts = kwargs.get("second_per_grid_ts", [])
|
||||
audio_feature_lengths = kwargs.get("audio_feature_lengths", [])
|
||||
use_audio_in_video = any(kwargs.get("use_audio_in_video", []))
|
||||
|
||||
elif modality == "image":
|
||||
grid_t = data["grid_t"]
|
||||
grid_h = data["grid_h"]
|
||||
grid_w = data["grid_w"]
|
||||
t_factor = data["t_factor"]
|
||||
image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(image_grid_thw)
|
||||
video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(video_grid_thw)
|
||||
|
||||
grid_indices = np.indices((grid_t, grid_h, grid_w))
|
||||
if t_factor != 1.0:
|
||||
grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
|
||||
llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx)
|
||||
st = offset + grid_t * grid_h * grid_w
|
||||
# TODO(fyabc): refactor and share more code with
|
||||
# _vl_get_input_positions_tensor.
|
||||
|
||||
elif modality == "video":
|
||||
grid_t = data["grid_t"]
|
||||
grid_h = data["grid_h"]
|
||||
grid_w = data["grid_w"]
|
||||
t_factor = data["t_factor"]
|
||||
thinker_config = self.config
|
||||
audio_token_id = thinker_config.audio_token_index
|
||||
image_token_id = thinker_config.image_token_index
|
||||
video_token_id = thinker_config.video_token_index
|
||||
audio_start_token_id = thinker_config.audio_start_token_id
|
||||
audio_end_token_id = thinker_config.audio_end_token_id
|
||||
vision_start_token_id = thinker_config.vision_start_token_id
|
||||
vision_end_token_id = thinker_config.vision_end_token_id
|
||||
seconds_per_chunk = thinker_config.seconds_per_chunk
|
||||
spatial_merge_size = thinker_config.vision_config.spatial_merge_size
|
||||
tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25)
|
||||
|
||||
if not data["use_audio_in_video"]:
|
||||
grid_indices = np.indices((grid_t, grid_h, grid_w))
|
||||
if t_factor != 1.0:
|
||||
grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
|
||||
llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx)
|
||||
st = offset + grid_t * grid_h * grid_w
|
||||
else:
|
||||
pos_ids, token_count = self._compute_interleaved_positions(st_idx, data)
|
||||
llm_pos_ids_list.append(pos_ids)
|
||||
st = offset + token_count
|
||||
src_item = input_tokens
|
||||
audio_seqlens = audio_feature_lengths
|
||||
if not second_per_grid_ts:
|
||||
second_per_grid_ts = [1] * video_grid_thw.shape[0]
|
||||
audio_idx = 0
|
||||
video_idx = 0
|
||||
image_idx = 0
|
||||
new_src_item: list[int] = []
|
||||
llm_pos_ids_list: list[torch.Tensor] = []
|
||||
|
||||
if st < len(input_tokens):
|
||||
st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0
|
||||
text_len = len(input_tokens) - st
|
||||
llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx)
|
||||
idx = 0
|
||||
while idx < len(src_item):
|
||||
new_src_item_len = len(new_src_item)
|
||||
start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
|
||||
if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]:
|
||||
if use_audio_in_video and idx > 0:
|
||||
if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id:
|
||||
# processing the <|audio_eos|> before <|vision_eos|>
|
||||
start_idx -= 1
|
||||
elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id:
|
||||
# processing the <|audio_bos|> after <|vision_eos|>
|
||||
start_idx -= 1
|
||||
new_src_item.append(src_item[idx])
|
||||
llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1)
|
||||
llm_pos_ids_list.append(llm_pos_ids)
|
||||
elif src_item[idx] == audio_token_id:
|
||||
assert audio_seqlens is not None
|
||||
audio_seqlen = audio_seqlens[audio_idx]
|
||||
place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1
|
||||
new_src_item.extend([audio_token_id] * place_num)
|
||||
llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
|
||||
llm_pos_ids_list.append(llm_pos_ids)
|
||||
audio_idx += 1
|
||||
elif src_item[idx] == image_token_id:
|
||||
grid_t = image_grid_thw[image_idx][0]
|
||||
grid_hs = image_grid_thw[:, 1]
|
||||
grid_ws = image_grid_thw[:, 2]
|
||||
t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()
|
||||
llm_pos_ids = get_llm_pos_ids_for_vision(
|
||||
start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
|
||||
)
|
||||
llm_pos_ids_list.append(llm_pos_ids)
|
||||
vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
|
||||
new_src_item.extend([image_token_id] * vision_seqlen)
|
||||
image_idx += 1
|
||||
elif src_item[idx] == video_token_id and not use_audio_in_video:
|
||||
grid_t = video_grid_thw[video_idx][0]
|
||||
grid_hs = video_grid_thw[:, 1]
|
||||
grid_ws = video_grid_thw[:, 2]
|
||||
t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
|
||||
llm_pos_ids = get_llm_pos_ids_for_vision(
|
||||
start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
|
||||
)
|
||||
llm_pos_ids_list.append(llm_pos_ids)
|
||||
vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
|
||||
new_src_item.extend([video_token_id] * vision_seqlen)
|
||||
video_idx += 1
|
||||
else:
|
||||
# read audio from video
|
||||
assert audio_seqlens is not None
|
||||
audio_seqlen = audio_seqlens[audio_idx]
|
||||
vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
|
||||
grid_t = video_grid_thw[video_idx][0]
|
||||
grid_h = video_grid_thw[video_idx][1]
|
||||
grid_w = video_grid_thw[video_idx][2]
|
||||
grid_hs = video_grid_thw[:, 1]
|
||||
grid_ws = video_grid_thw[:, 2]
|
||||
t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
|
||||
t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long()
|
||||
t_index_split_chunk = split_list_into_ranges(t_index, t_ntoken_per_chunk)
|
||||
place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
|
||||
pure_audio_len = place_num - 2
|
||||
added_audio_len = 0
|
||||
audio_llm_pos_ids_list: list[torch.Tensor] = []
|
||||
for t_chunk in t_index_split_chunk:
|
||||
vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
|
||||
new_src_item.extend([video_token_id] * vision_ntoken_per_chunk)
|
||||
vision_llm_pos_ids_list = get_llm_pos_ids_for_vision(
|
||||
start_idx,
|
||||
video_idx,
|
||||
spatial_merge_size,
|
||||
t_chunk,
|
||||
grid_hs,
|
||||
grid_ws,
|
||||
).split(1, dim=1)
|
||||
llm_pos_ids_list.extend(vision_llm_pos_ids_list)
|
||||
new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id])
|
||||
audio_start_idx = (
|
||||
start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
|
||||
)
|
||||
if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0:
|
||||
audio_llm_pos_ids_list = (
|
||||
torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1)
|
||||
+ audio_start_idx
|
||||
).split(1, dim=1)
|
||||
else:
|
||||
audio_llm_pos_ids_list = []
|
||||
added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)
|
||||
llm_pos_ids_list.extend(audio_llm_pos_ids_list)
|
||||
if added_audio_len < pure_audio_len:
|
||||
new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id])
|
||||
audio_llm_pos_ids_list = (
|
||||
torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1
|
||||
).split(1, dim=1)
|
||||
llm_pos_ids_list.extend(audio_llm_pos_ids_list)
|
||||
audio_idx += 1
|
||||
video_idx += 1
|
||||
# move to the next token
|
||||
idx += len(new_src_item) - new_src_item_len
|
||||
|
||||
llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
|
||||
mrope_position_delta = int(llm_positions.max()) + 1 - len(input_tokens)
|
||||
llm_positions = torch.cat(llm_pos_ids_list, dim=1)
|
||||
mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item)
|
||||
|
||||
return torch.from_numpy(llm_positions), mrope_position_delta
|
||||
return llm_positions, mrope_position_delta
|
||||
|
||||
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
|
||||
mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
|
||||
@@ -649,12 +629,14 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
|
||||
# Check for audio-in-video: interleaved video and audio tokens
|
||||
# in the multimodal region. Only use the interleaved path when
|
||||
# needed; otherwise fall back to the default parent implementation.
|
||||
# vLLM's _gather_mm_embeddings builds is_multimodal on CPU for merge
|
||||
# indexing; bitwise ops with input_ids require the same device.
|
||||
is_mm_device = is_multimodal.to(device=input_ids.device, non_blocking=True)
|
||||
video_token_id = self.config.video_token_index
|
||||
audio_token_id = self.config.audio_token_index
|
||||
|
||||
input_ids_cpu = input_ids.cpu()
|
||||
is_video = is_multimodal & (input_ids_cpu == video_token_id)
|
||||
is_audio = is_multimodal & (input_ids_cpu == audio_token_id)
|
||||
is_video = is_mm_device & (input_ids == video_token_id)
|
||||
is_audio = is_mm_device & (input_ids == audio_token_id)
|
||||
|
||||
num_video = is_video.sum().item()
|
||||
num_audio = is_audio.sum().item()
|
||||
@@ -670,7 +652,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
|
||||
multimodal_embeddings,
|
||||
is_video,
|
||||
is_audio,
|
||||
is_multimodal,
|
||||
is_mm_device,
|
||||
num_video,
|
||||
num_audio,
|
||||
)
|
||||
|
||||
@@ -758,11 +758,17 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
|
||||
|
||||
use_audio_in_video = False
|
||||
if "video" in mm_kwargs:
|
||||
for item in mm_kwargs["video"]:
|
||||
if item and item["use_audio_in_video"].data:
|
||||
use_audio_in_video = True
|
||||
else:
|
||||
use_audio_in_video = False
|
||||
non_none_items = [item for item in mm_kwargs["video"] if item is not None]
|
||||
if non_none_items:
|
||||
# Normal case: at least one non-cached item, read flag directly
|
||||
use_audio_in_video = any(item["use_audio_in_video"].data for item in non_none_items)
|
||||
elif "audio" in mm_prompt_updates:
|
||||
# All video items are from cache (None); infer from prompt:
|
||||
# use_audio_in_video=True means the prompt has no <|audio_pad|>
|
||||
# placeholder (audio is embedded in video tokens instead)
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
|
||||
use_audio_in_video = audio_pad_id not in prompt_ids
|
||||
# for mutilmodality cache
|
||||
if any(item is None for item in mm_kwargs["video"]):
|
||||
video_token_id = self.info.get_hf_config().video_token_id
|
||||
|
||||
@@ -119,22 +119,6 @@ _OMNI_MODELS = {
|
||||
"higgs_audio_v2_code2wav",
|
||||
"HiggsAudioV2Code2WavForConditionalGeneration",
|
||||
),
|
||||
## higgs-audio v3
|
||||
"HiggsMultimodalQwen3ForConditionalGeneration": (
|
||||
"higgs_audio_v3",
|
||||
"higgs_audio_v3_talker",
|
||||
"HiggsAudioV3TalkerForConditionalGeneration",
|
||||
),
|
||||
"HiggsAudioV3TalkerForConditionalGeneration": (
|
||||
"higgs_audio_v3",
|
||||
"higgs_audio_v3_talker",
|
||||
"HiggsAudioV3TalkerForConditionalGeneration",
|
||||
),
|
||||
"HiggsAudioV3Code2WavForConditionalGeneration": (
|
||||
"higgs_audio_v3",
|
||||
"higgs_audio_v3_code2wav",
|
||||
"HiggsAudioV3Code2WavForConditionalGeneration",
|
||||
),
|
||||
## mimo_audio
|
||||
"MiMoAudioModel": (
|
||||
"mimo_audio",
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
# Stage config for running HunyuanVideo-1.5 DiT with ModelOpt FP8 auto-detect.
|
||||
# Single GPU. Bump `tensor_parallel_size` and `devices` for multi-GPU TP.
|
||||
#
|
||||
# Use with a ModelOpt FP8 checkpoint (e.g. produced by
|
||||
# examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py).
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: diffusion
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
model_class_name: HunyuanVideo15Pipeline
|
||||
max_num_seqs: 1
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
distributed_executor_backend: "mp"
|
||||
# HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of
|
||||
# the full video (the DiT contributes little), so tiling the decode is the
|
||||
# main memory lever — far more than FP8 (e.g. 480x832x33: ~117 -> ~47 GiB).
|
||||
# Default it on here; override with --vae-use-tiling/false if you have headroom.
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 1
|
||||
|
||||
final_output: true
|
||||
final_output_type: video
|
||||
is_comprehension: false
|
||||
default_sampling_params:
|
||||
seed: 42
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
@@ -1,41 +0,0 @@
|
||||
# Stage config for running Wan2.2 TI2V-5B DiT with ModelOpt FP8 auto-detect.
|
||||
# Single GPU (TI2V-5B fits 80GB BF16; FP8 drops by ~half).
|
||||
# For the A14B MoE variants, bump `tensor_parallel_size` and `devices`.
|
||||
#
|
||||
# Use with a ModelOpt FP8 checkpoint (e.g. produced by
|
||||
# examples/quantization/quantize_wan2_2_modelopt_fp8.py).
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: diffusion
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: dit
|
||||
# Wan2.2 TI2V-5B is served by the base Wan22Pipeline, which auto-detects
|
||||
# expand_timesteps (TI2V-5B mode) from model_index.json. (The standalone
|
||||
# Wan22TI2VPipeline was consolidated into Wan22Pipeline upstream.)
|
||||
model_class_name: Wan22Pipeline
|
||||
max_num_seqs: 1
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
distributed_executor_backend: "mp"
|
||||
# Tile the VAE decode by default. Wan2.2's 480p peak is DiT-dominated (so
|
||||
# FP8 is the main lever here), but tiling adds headroom for higher
|
||||
# resolutions / longer clips at negligible cost. Override with --vae-use-tiling/false.
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 1
|
||||
|
||||
final_output: true
|
||||
final_output_type: video
|
||||
is_comprehension: false
|
||||
default_sampling_params:
|
||||
seed: 42
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
@@ -1,404 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Stage-input processor for higgs-audio v3: Talker -> Code2Wav.
|
||||
|
||||
Two adapters:
|
||||
|
||||
* ``talker2code2wav`` is the sync (non-async-chunk) adapter that runs once,
|
||||
after Stage 0 has finished, collecting every audio frame into a single
|
||||
``OmniTokensPrompt`` for Stage 1.
|
||||
* ``talker2code2wav_async_chunk`` is the streaming adapter invoked once per
|
||||
emitted talker step. It accumulates raw delay-pattern rows per request and
|
||||
flushes a sliding-window chunk to Stage 1 once enough new frames are
|
||||
available, holding back ``codec_right_holdback_frames`` trailing frames
|
||||
so the codec has future context for the emit region's tail samples. Stage 1
|
||||
trims ``left_context_size * hop`` samples off the front and
|
||||
``right_holdback_size * hop`` off the end before emitting PCM.
|
||||
|
||||
Key difference from v2: BOC/EOC filtering happens AFTER delay pattern reversal
|
||||
to avoid corrupting valid tail content.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.data_entry_keys import (
|
||||
CodesStruct,
|
||||
MetaStruct,
|
||||
OmniPayload,
|
||||
OmniPayloadStruct,
|
||||
)
|
||||
|
||||
__all__ = ["talker2code2wav", "talker2code2wav_async_chunk"]
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_NUM_CODEBOOKS = 8
|
||||
_NUM_REAL_CODES = 1024 # codes in [0, 1023] are real
|
||||
|
||||
# Defaults for async_chunk streaming when the deploy YAML connector config
|
||||
# does not override them. Mirrors the qwen3_tts canonical setting at 25 fps
|
||||
# codec rate (25 frames = ~1.0 s of audio).
|
||||
_DEFAULT_CODEC_CHUNK_FRAMES = 25
|
||||
_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 25
|
||||
_DEFAULT_CODEC_RIGHT_HOLDBACK_FRAMES = 4
|
||||
|
||||
|
||||
def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor:
|
||||
"""Reverse the MusicGen-style delay pattern.
|
||||
|
||||
Input shape: [num_codebooks, seq_len + num_codebooks - 1].
|
||||
Output shape: [num_codebooks, seq_len].
|
||||
|
||||
For each codebook i, extract delayed[i, i : i + seq_len] to remove
|
||||
the i leading BOC pads and Q-1-i trailing EOC entries.
|
||||
"""
|
||||
if audio_codes_qt.ndim != 2:
|
||||
raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}")
|
||||
q, t = audio_codes_qt.shape
|
||||
if q != _NUM_CODEBOOKS:
|
||||
raise ValueError(f"Expected exactly {_NUM_CODEBOOKS} codebook rows, got {q}. Input shape: [{q}, {t}]")
|
||||
if t < q:
|
||||
raise ValueError(f"Not enough frames to revert delay pattern: T={t} < Q={q}")
|
||||
seq_len = t - q + 1
|
||||
out_l = []
|
||||
for i in range(q):
|
||||
out_l.append(audio_codes_qt[i : i + 1, i : seq_len + i])
|
||||
return torch.cat(out_l, dim=0)
|
||||
|
||||
|
||||
def _filter_real_code_frames(audio_codes_qt: torch.Tensor) -> torch.Tensor:
|
||||
"""Keep only frames where ALL codebook values are in [0, 1023].
|
||||
|
||||
Input shape: [num_codebooks, num_frames].
|
||||
Called AFTER delay pattern reversal.
|
||||
"""
|
||||
if audio_codes_qt.numel() == 0:
|
||||
return audio_codes_qt
|
||||
# Transpose to [num_frames, num_codebooks] for per-frame filtering
|
||||
frames = audio_codes_qt.t()
|
||||
valid = (frames >= 0).all(dim=1) & (frames < _NUM_REAL_CODES).all(dim=1)
|
||||
return frames[valid].t().contiguous()
|
||||
|
||||
|
||||
def talker2code2wav(
|
||||
source_outputs: list[Any],
|
||||
prompt: Any = None,
|
||||
_requires_multimodal_data: bool = False,
|
||||
) -> list[Any]:
|
||||
"""Sync: collect all talker codes, revert delay pattern, filter, pass to code2wav."""
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
|
||||
code2wav_inputs: list[OmniTokensPrompt] = []
|
||||
for talker_output in source_outputs:
|
||||
if not talker_output.finished:
|
||||
continue
|
||||
output = talker_output.outputs[0]
|
||||
mm = output.multimodal_output
|
||||
mm_codes = mm.get("codes", {})
|
||||
|
||||
audio_codes = mm_codes.get("audio")
|
||||
if audio_codes is None or not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
|
||||
code2wav_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=[],
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
additional_information=None,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
audio_codes = audio_codes.to(torch.long)
|
||||
if audio_codes.ndim == 1:
|
||||
if audio_codes.numel() % _NUM_CODEBOOKS != 0:
|
||||
raise ValueError(
|
||||
f"flat audio_codes length {audio_codes.numel()} not divisible by num_codebooks={_NUM_CODEBOOKS}"
|
||||
)
|
||||
audio_codes = audio_codes.reshape(-1, _NUM_CODEBOOKS)
|
||||
|
||||
if audio_codes.ndim != 2:
|
||||
raise ValueError(f"audio_codes must be 1D or 2D; got shape {tuple(audio_codes.shape)}")
|
||||
|
||||
# Validate codebook count
|
||||
if audio_codes.shape[1] != _NUM_CODEBOOKS:
|
||||
raise ValueError(
|
||||
f"Expected {_NUM_CODEBOOKS} codebooks per frame, "
|
||||
f"got {audio_codes.shape[1]}. Audio codes shape: {tuple(audio_codes.shape)}"
|
||||
)
|
||||
|
||||
# Transpose to [Q, T] for delay pattern reversal
|
||||
codes_qt = audio_codes.transpose(0, 1).contiguous().cpu()
|
||||
|
||||
# Step 1: Revert delay pattern
|
||||
try:
|
||||
codes_qt = _revert_delay_pattern(codes_qt)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"Skipping request: insufficient frames for delay pattern "
|
||||
"reversal (T=%d < Q=%d). This is a request-level error, "
|
||||
"not a server-fatal condition.",
|
||||
codes_qt.shape[1],
|
||||
codes_qt.shape[0],
|
||||
)
|
||||
code2wav_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=[],
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
additional_information=None,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# Step 2: Replace out-of-range codes (BOC=1024, EOC=1025, -1) with 0.
|
||||
# Must use torch.where, NOT clamp: clamp(max=1023) turns 1025→1023
|
||||
# which is a valid codec code and decodes to audio artifacts.
|
||||
# Matches sglang's: torch.where(codes >= codec_vocab, 0, codes)
|
||||
codes_qt = torch.where(
|
||||
(codes_qt >= _NUM_REAL_CODES) | (codes_qt < 0),
|
||||
torch.zeros_like(codes_qt),
|
||||
codes_qt,
|
||||
)
|
||||
|
||||
# Step 3: Trim the last frame. After de-delay, the final frame
|
||||
# contains residual ramp-down codes (EOC→0 substituted) that
|
||||
# decode to a brief noise artifact at the end of the audio.
|
||||
if codes_qt.shape[-1] >= 2:
|
||||
codes_qt = codes_qt[:, :-1]
|
||||
|
||||
if codes_qt.numel() == 0:
|
||||
code2wav_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=[],
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
additional_information=None,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# Code2Wav expects codebook-major flat: [Q * num_frames]
|
||||
codec_codes = codes_qt.reshape(-1).tolist()
|
||||
|
||||
code2wav_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=codec_codes,
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
additional_information=None,
|
||||
)
|
||||
)
|
||||
return code2wav_inputs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async-chunk streaming adapter (sliding window + holdback).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_last_step_row(pooling_output: OmniPayload) -> torch.Tensor | None:
|
||||
"""Return the talker's emit row for this AR step as a flat ``[Q]`` tensor.
|
||||
|
||||
The talker postprocess emits ``codes.audio`` as ``[1, Q]`` (most often) or
|
||||
``[N, Q]`` for multi-active stages. In the async-chunk path
|
||||
``pooling_output`` is already request-scoped so the typical shape is
|
||||
``[1, Q]``. Returns ``None`` if there is no row to append at this step.
|
||||
|
||||
Rows may carry BOC=1024 in trailing codebooks during the ramp-up phase
|
||||
(AR steps 0..Q-2). Those are KEPT verbatim — the de-delay window
|
||||
operation downstream slices them away by construction.
|
||||
"""
|
||||
codes = pooling_output.get("codes", {}) if isinstance(pooling_output, dict) else {}
|
||||
audio_codes = codes.get("audio") if isinstance(codes, dict) else None
|
||||
if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
|
||||
return None
|
||||
if audio_codes.ndim == 2:
|
||||
row = audio_codes[-1]
|
||||
elif audio_codes.ndim == 1:
|
||||
row = audio_codes
|
||||
else:
|
||||
raise ValueError(f"unexpected audio_codes shape for higgs_audio_v3 async_chunk: {tuple(audio_codes.shape)}")
|
||||
if row.numel() != _NUM_CODEBOOKS:
|
||||
raise ValueError(
|
||||
f"talker emit row has {row.numel()} codebooks; expected {_NUM_CODEBOOKS} for higgs_audio_v3 async_chunk."
|
||||
)
|
||||
return row.to(torch.long).reshape(-1)
|
||||
|
||||
|
||||
def talker2code2wav_async_chunk(
|
||||
transfer_manager: Any,
|
||||
pooling_output: OmniPayload | None,
|
||||
request: Any,
|
||||
is_finished: bool = False,
|
||||
) -> OmniPayloadStruct | None:
|
||||
"""Streaming adapter: per-step accumulator + sliding-window de-delay flush.
|
||||
|
||||
Algorithm:
|
||||
|
||||
* Append each step's raw delay-pattern row to a per-request buffer.
|
||||
* Track ``emitted_frames``: how many de-delayed frames' worth of audio
|
||||
Stage 1 has already emitted for this request.
|
||||
* With ``N`` accumulated rows we can recover ``N - Q + 1`` de-delayed
|
||||
frames (the de-delay drops Q-1 leading BOC pads).
|
||||
* Hold back the trailing ``H = codec_right_holdback_frames`` de-delayed
|
||||
frames so the codec has future context for the emit region's last
|
||||
samples. The emit boundary therefore lives at
|
||||
``target_emit = max(0, de_delayed_total - H)`` for non-final calls
|
||||
and ``de_delayed_total`` for the final flush.
|
||||
* When ``target_emit - emitted_frames >= chunk_size`` (or the request
|
||||
finishes), build a sliding window spanning
|
||||
``[emitted_frames - L, target_emit + H]`` de-delayed frames; that
|
||||
corresponds to AR rows
|
||||
``[emitted_frames - L, target_emit + H + Q - 1]``. Apply de-delay,
|
||||
replace out-of-range codes with 0, and emit codebook-major flat.
|
||||
* Stage 1 honors ``meta.left_context_size = L`` (trim front) and
|
||||
``meta.right_holdback_size = H`` (trim end). Net new audio per call
|
||||
is ``(target_emit - emitted_frames) * hop_length`` samples.
|
||||
"""
|
||||
request_id = request.external_req_id
|
||||
finished = bool(is_finished or request.is_finished())
|
||||
|
||||
# Per-request rolling state. ``code_prompt_token_ids`` is the
|
||||
# framework-provided buffer; ``higgs_v3_emitted_frames`` tracks how many
|
||||
# de-delayed frames we have already flushed for this request.
|
||||
emitted_frames = getattr(transfer_manager, "higgs_v3_emitted_frames", None)
|
||||
if emitted_frames is None:
|
||||
emitted_frames = {}
|
||||
transfer_manager.higgs_v3_emitted_frames = emitted_frames
|
||||
|
||||
if isinstance(pooling_output, dict):
|
||||
row = _extract_last_step_row(pooling_output)
|
||||
if row is not None:
|
||||
transfer_manager.code_prompt_token_ids[request_id].append(row.cpu().tolist())
|
||||
elif not finished:
|
||||
return None
|
||||
|
||||
connector = getattr(transfer_manager, "connector", None)
|
||||
raw_cfg = getattr(connector, "config", {}) or {}
|
||||
cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
|
||||
chunk_size = int(cfg.get("codec_chunk_frames", _DEFAULT_CODEC_CHUNK_FRAMES))
|
||||
left_context_size_config = int(cfg.get("codec_left_context_frames", _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES))
|
||||
right_holdback_size_config = int(cfg.get("codec_right_holdback_frames", _DEFAULT_CODEC_RIGHT_HOLDBACK_FRAMES))
|
||||
configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0)
|
||||
|
||||
if (
|
||||
chunk_size <= 0
|
||||
or left_context_size_config < 0
|
||||
or right_holdback_size_config < 0
|
||||
or configured_initial_chunk_size < 0
|
||||
):
|
||||
raise ValueError(
|
||||
f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
|
||||
f"codec_left_context_frames={left_context_size_config}, "
|
||||
f"codec_right_holdback_frames={right_holdback_size_config}, "
|
||||
f"initial_codec_chunk_frames={configured_initial_chunk_size}"
|
||||
)
|
||||
|
||||
n_rows = len(transfer_manager.code_prompt_token_ids[request_id])
|
||||
de_delayed_total = max(0, n_rows - (_NUM_CODEBOOKS - 1))
|
||||
emitted = int(emitted_frames.get(request_id, 0))
|
||||
|
||||
# Hold back the trailing H frames from the emit boundary unless the
|
||||
# request is finishing - the final chunk needs no future context.
|
||||
H = 0 if finished else right_holdback_size_config
|
||||
target_emit = max(0, de_delayed_total - H)
|
||||
pending = target_emit - emitted
|
||||
|
||||
# First-chunk fast path for TTFA.
|
||||
if emitted == 0 and 0 < configured_initial_chunk_size < chunk_size:
|
||||
target_chunk = configured_initial_chunk_size
|
||||
else:
|
||||
target_chunk = chunk_size
|
||||
|
||||
if not finished and pending < target_chunk:
|
||||
return None
|
||||
|
||||
if pending <= 0:
|
||||
if finished:
|
||||
emitted_frames.pop(request_id, None)
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
|
||||
meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
|
||||
)
|
||||
return None
|
||||
|
||||
actual_chunk = pending
|
||||
|
||||
desired_left_context = min(left_context_size_config, emitted)
|
||||
|
||||
# Window in AR-row space: span the de-delayed frames
|
||||
# [emitted - L, target_emit + H) which require AR rows
|
||||
# [emitted - L, target_emit + H + Q - 1).
|
||||
window_row_start = emitted - desired_left_context
|
||||
window_row_end_exclusive = target_emit + H + (_NUM_CODEBOOKS - 1)
|
||||
if window_row_end_exclusive > n_rows:
|
||||
# Defensive: pending >= 0 was checked above; this should be unreachable
|
||||
# for non-final calls. On the final call we cap to n_rows.
|
||||
if not finished:
|
||||
return None
|
||||
window_row_end_exclusive = n_rows
|
||||
|
||||
window_rows = transfer_manager.code_prompt_token_ids[request_id][window_row_start:window_row_end_exclusive]
|
||||
if len(window_rows) < _NUM_CODEBOOKS:
|
||||
if finished:
|
||||
emitted_frames.pop(request_id, None)
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
|
||||
meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
|
||||
)
|
||||
return None
|
||||
|
||||
codes_qt = torch.tensor(window_rows, dtype=torch.long).t().contiguous()
|
||||
try:
|
||||
de_delayed = _revert_delay_pattern(codes_qt)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"async_chunk: insufficient frames for delay pattern reversal "
|
||||
"(T=%d < Q=%d, request=%s). Returning empty chunk.",
|
||||
codes_qt.shape[1],
|
||||
codes_qt.shape[0],
|
||||
request_id,
|
||||
)
|
||||
emitted_frames.pop(request_id, None)
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
|
||||
meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
|
||||
)
|
||||
# Replace BOC=1024/EOC=1025 (and any negative pads) with 0; matches the
|
||||
# sync-path substitution. Clamp would turn 1025 into 1023 which is a
|
||||
# VALID codec code and decodes to audible artifacts.
|
||||
de_delayed = torch.where(
|
||||
(de_delayed >= _NUM_REAL_CODES) | (de_delayed < 0),
|
||||
torch.zeros_like(de_delayed),
|
||||
de_delayed,
|
||||
)
|
||||
|
||||
# On the FINAL chunk only: trim the trailing residual frame (the last
|
||||
# de-delayed frame still carries EOC-substituted codes from ramp-down
|
||||
# which decode to a brief noise artifact). Mirrors the sync-path trim.
|
||||
if finished and window_row_end_exclusive == n_rows and de_delayed.shape[-1] >= 2:
|
||||
de_delayed = de_delayed[:, :-1]
|
||||
actual_chunk = max(actual_chunk - 1, 0)
|
||||
|
||||
codec_codes = de_delayed.reshape(-1)
|
||||
left_context_emitted = desired_left_context
|
||||
right_holdback_emitted = H
|
||||
|
||||
emitted_frames[request_id] = emitted + actual_chunk
|
||||
if finished:
|
||||
emitted_frames.pop(request_id, None)
|
||||
|
||||
meta = MetaStruct(
|
||||
left_context_size=left_context_emitted,
|
||||
right_holdback_size=right_holdback_emitted,
|
||||
finished=torch.tensor(finished, dtype=torch.bool),
|
||||
)
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=codec_codes),
|
||||
meta=meta,
|
||||
)
|
||||
@@ -217,55 +217,3 @@ def _patch_scaled_mm_fp8_contiguous_activation():
|
||||
|
||||
|
||||
_patch_scaled_mm_fp8_contiguous_activation()
|
||||
|
||||
|
||||
def _patch_flashinfer_fp8_scaled_mm_output_shape():
|
||||
"""Restore the N-D output shape for the FlashInfer FP8 ScaledMM kernel.
|
||||
|
||||
``FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm`` returns the raw 2-D
|
||||
GEMM result and ignores ``output_shape``, unlike the CUTLASS / PyTorch
|
||||
ScaledMM kernels which reshape to it. A 3-D activation ``(B, S, D)`` thus
|
||||
collapses to ``(B*S, D)``, breaking diffusion DiTs that reshape the linear
|
||||
output by absolute dim (e.g. Wan2.2 ``qkv.unflatten(2, ...)``) with
|
||||
``IndexError: Dimension out of range``. It only bites >2-D inputs, so LLM
|
||||
(token-flattened, 2-D) paths are unaffected.
|
||||
|
||||
Carried here because the upstream fix may not have landed yet; this override
|
||||
becomes a harmless no-op once vLLM honors ``output_shape`` itself.
|
||||
"""
|
||||
try:
|
||||
from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
)
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
_original_apply_scaled_mm = FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm
|
||||
if getattr(_original_apply_scaled_mm, "_omni_output_shape_patched", False):
|
||||
return
|
||||
|
||||
def _patched_apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape): # noqa: N803
|
||||
out = _original_apply_scaled_mm(
|
||||
self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape
|
||||
)
|
||||
if tuple(out.shape) != tuple(output_shape):
|
||||
out = out.view(*output_shape)
|
||||
return out
|
||||
|
||||
_patched_apply_scaled_mm._omni_output_shape_patched = True
|
||||
FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm = _patched_apply_scaled_mm
|
||||
|
||||
|
||||
_patch_flashinfer_fp8_scaled_mm_output_shape()
|
||||
|
||||
|
||||
def _patch_fp8_use_quack_fused_bias():
|
||||
try:
|
||||
from vllm_omni.quantization.quack_fp8 import install_quack_fp8_patch
|
||||
|
||||
install_quack_fp8_patch()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
|
||||
_patch_fp8_use_quack_fused_bias()
|
||||
@@ -136,6 +136,11 @@ class OmniPlatform(Platform):
|
||||
"""
|
||||
return "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner"
|
||||
|
||||
@classmethod
|
||||
def init_diffusion_worker_vllm_config(cls, vllm_config: Any) -> None:
|
||||
"""Initialize platform-specific state for diffusion worker VllmConfig."""
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -55,6 +55,12 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
|
||||
def get_omni_generation_worker_cls(cls) -> str:
|
||||
return "vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker"
|
||||
|
||||
@classmethod
|
||||
def init_diffusion_worker_vllm_config(cls, vllm_config: Any) -> None:
|
||||
from vllm_ascend.ascend_config import init_ascend_config
|
||||
|
||||
init_ascend_config(vllm_config)
|
||||
|
||||
@classmethod
|
||||
def get_default_stage_config_path(cls) -> str:
|
||||
return "vllm_omni/platforms/npu/stage_configs"
|
||||
|
||||
@@ -33,7 +33,6 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
|
||||
# in _pad_for_sequence_parallelism during execute_model.
|
||||
# This is a workaround for vllm-ascend not passing vllm_config to enable_sp().
|
||||
enable_sp(self.vllm_config)
|
||||
self._maybe_enable_output_token_ids_for_model_sampler()
|
||||
self._init_talker_mtp()
|
||||
|
||||
@torch.inference_mode()
|
||||
|
||||
@@ -1,134 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""FP8 scaled-MM with a fused bias epilogue, backed by quack's CuteDSL GEMM."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_TRUTHY = {"1", "true", "yes", "on"}
|
||||
_gemm_interface = None
|
||||
|
||||
|
||||
def _is_blackwell() -> bool:
|
||||
try:
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
return torch.cuda.get_device_capability()[0] >= 10
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
|
||||
|
||||
def quack_enabled() -> bool:
|
||||
override = os.environ.get("VLLM_OMNI_USE_QUACK_FP8")
|
||||
if override is not None:
|
||||
return override.lower() in _TRUTHY
|
||||
return _is_blackwell()
|
||||
|
||||
|
||||
def _set_persistent_cache_dir() -> None:
|
||||
if os.environ.get("QUACK_CACHE_DIR"):
|
||||
return
|
||||
root = (
|
||||
os.environ.get("VLLM_CACHE_ROOT")
|
||||
or os.environ.get("XDG_CACHE_HOME")
|
||||
or os.path.join(os.path.expanduser("~"), ".cache")
|
||||
)
|
||||
os.environ["QUACK_CACHE_DIR"] = os.path.join(root, "vllm_omni", "quack")
|
||||
|
||||
|
||||
def _load_quack():
|
||||
global _gemm_interface
|
||||
if _gemm_interface is not None:
|
||||
return _gemm_interface or None
|
||||
try:
|
||||
_set_persistent_cache_dir()
|
||||
|
||||
import cutlass
|
||||
import cutlass.base_dsl
|
||||
import cutlass.base_dsl.arch as arch
|
||||
|
||||
if not hasattr(cutlass.base_dsl, "Arch"):
|
||||
cutlass.base_dsl.Arch = arch.Arch
|
||||
|
||||
import quack.gemm_interface as gemm_interface
|
||||
from quack.cute_dsl_utils import torch2cute_dtype_map
|
||||
|
||||
torch2cute_dtype_map.setdefault(torch.float8_e4m3fn, cutlass.Float8E4M3FN)
|
||||
torch2cute_dtype_map.setdefault(torch.float8_e5m2, cutlass.Float8E5M2)
|
||||
|
||||
_gemm_interface = gemm_interface
|
||||
logger.info("Quack FP8 fused-bias GEMM enabled (CuteDSL).")
|
||||
return gemm_interface
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("Quack FP8 unavailable, using FlashInfer: %s", exc)
|
||||
_gemm_interface = False
|
||||
return None
|
||||
|
||||
|
||||
def quack_scaled_fp8_mm(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
scale_a: torch.Tensor,
|
||||
scale_b: torch.Tensor,
|
||||
out_dtype: torch.dtype,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor | None:
|
||||
gemm = _load_quack()
|
||||
if gemm is None:
|
||||
return None
|
||||
out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=out_dtype)
|
||||
alpha = scale_a.reshape(1).float() * scale_b.reshape(1).float()
|
||||
gemm.gemm_tuned(a, b, out, bias=bias, alpha=alpha)
|
||||
return out
|
||||
|
||||
|
||||
def install_quack_fp8_patch() -> None:
|
||||
if not quack_enabled():
|
||||
return
|
||||
if _load_quack() is None:
|
||||
return
|
||||
try:
|
||||
from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
|
||||
FlashInferFP8ScaledMMLinearKernel,
|
||||
)
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
original = FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm
|
||||
if getattr(original, "_omni_quack_patched", False):
|
||||
return
|
||||
|
||||
def apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape): # noqa: N803
|
||||
try:
|
||||
out = quack_scaled_fp8_mm(A, B, As, Bs, out_dtype, bias)
|
||||
if out is not None:
|
||||
return out.view(*output_shape)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning_once("Quack FP8 GEMM failed (%s); using FlashInfer.", exc)
|
||||
return original(self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape)
|
||||
|
||||
apply_scaled_mm._omni_quack_patched = True
|
||||
FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm = apply_scaled_mm
|
||||
logger.info("Patched FlashInfer FP8 ScaledMM to use quack fused-bias GEMM.")
|
||||
|
||||
|
||||
def warmup_quack_fp8(
|
||||
shapes: list[tuple[int, int, int]],
|
||||
device: str = "cuda",
|
||||
out_dtype: torch.dtype = torch.bfloat16,
|
||||
) -> None:
|
||||
if _load_quack() is None:
|
||||
return
|
||||
scale = torch.ones(1, device=device, dtype=torch.float32)
|
||||
for m, k, n in shapes:
|
||||
a = torch.zeros(m, k, device=device, dtype=torch.float8_e4m3fn)
|
||||
b = torch.zeros(k, n, device=device, dtype=torch.float8_e4m3fn)
|
||||
quack_scaled_fp8_mm(a, b, scale, scale, out_dtype)
|
||||
if torch.cuda.is_available():
|
||||
torch.accelerator.synchronize()
|
||||
@@ -143,14 +143,9 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
@instrument(span_name="Loading (GPU)")
|
||||
def load_model(self, *args, **kwargs) -> None:
|
||||
super().load_model(*args, **kwargs)
|
||||
self._maybe_enable_output_token_ids_for_model_sampler()
|
||||
self._init_talker_mtp()
|
||||
self._prewarm_attention_capture_workspaces()
|
||||
|
||||
def _maybe_enable_output_token_ids_for_model_sampler(self) -> None:
|
||||
if getattr(self.model, "logitsprocs_need_output_token_ids", False):
|
||||
self.input_batch.logitsprocs_need_output_token_ids = True
|
||||
|
||||
def _init_talker_mtp(self) -> None:
|
||||
# TODO move this model specific logic to a separate class
|
||||
# TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
|
||||
|
||||
Reference in new issue
Block a user