Compare commits
68
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d98f94fa29 | ||
|
|
908e7d32d9 | ||
|
|
d6f1d3d61c | ||
|
|
c39ac9fffd | ||
|
|
580a0efe40 | ||
|
|
329c98ff9d | ||
|
|
36a5e9cfd2 | ||
|
|
0b04f2a0d6 | ||
|
|
5db89b41c4 | ||
|
|
e500c00545 | ||
|
|
8ddcb3621f | ||
|
|
2e6c83c89c | ||
|
|
d9063d7141 | ||
|
|
8c33793ffc | ||
|
|
4024a8ccd6 | ||
|
|
fb563fbd19 | ||
|
|
0fd2620d56 | ||
|
|
0d0f4e669e | ||
|
|
d94108693a | ||
|
|
d612233bef | ||
|
|
f6e8b77642 | ||
|
|
d8edf0c9d5 | ||
|
|
67bd02f767 | ||
|
|
0814d845dc | ||
|
|
bb973cb10c | ||
|
|
fb3abb38fb | ||
|
|
eb81914c64 | ||
|
|
1951dca786 | ||
|
|
d71cacc74c | ||
|
|
465f1464d2 | ||
|
|
f6894f4308 | ||
|
|
a2d1dafc65 | ||
|
|
6dc4638a7e | ||
|
|
d35cdd2dab | ||
|
|
0958e62ecc | ||
|
|
3901ad8165 | ||
|
|
b84c99c122 | ||
|
|
733f4f6a36 | ||
|
|
67d14b5ac4 | ||
|
|
64efa66869 | ||
|
|
6646a66d41 | ||
|
|
378359fc11 | ||
|
|
9c260b3774 | ||
|
|
9177b37e2f | ||
|
|
50bc51ffdf | ||
|
|
044240d8eb | ||
|
|
f4d922d8f6 | ||
|
|
eaa16e1941 | ||
|
|
8a6786731e | ||
|
|
fb0f17858f | ||
|
|
2447740308 | ||
|
|
a10eac0f8c | ||
|
+2 |
ec4ec0f604 | ||
|
|
ad07f98b73 | ||
|
|
a60242092a | ||
|
|
991b1ec752 | ||
|
|
6bdd85b384 | ||
|
|
36bd8029f3 | ||
|
|
f0d7fe18f2 | ||
|
|
ccb5c53ea5 | ||
|
|
8efd9d832d | ||
|
|
7ff11d8d0f | ||
|
|
fc32d9fe6f | ||
|
|
afbcc51672 | ||
|
|
1292c47804 | ||
|
|
e241005a71 | ||
|
|
950e66a44c | ||
|
|
a1e1ac0155 |
No files matched your search
+14
-4
@@ -26,23 +26,33 @@ steps:
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L2 Test
|
||||
# L2 Test — PR ready label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group)
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
if: __UPLOAD_READY_IF__
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
- |
|
||||
EXTRA=""
|
||||
if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
|
||||
EXTRA="--e2e"
|
||||
fi
|
||||
python3 .buildkite/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L3 Test
|
||||
# L3 Test — PR merge-test label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group)
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
if: __UPLOAD_MERGE_IF__
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
- |
|
||||
EXTRA=""
|
||||
if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
|
||||
EXTRA="--e2e"
|
||||
fi
|
||||
python3 .buildkite/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -23,7 +23,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -33,6 +33,6 @@ steps:
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -13,13 +13,19 @@ Test pipeline mode (e.g. test-merge.yml):
|
||||
recognize this uploader-only key).
|
||||
|
||||
Usage:
|
||||
python3 upload_pipeline.py [--upload] <pipeline.yml>
|
||||
python3 upload_pipeline.py [--upload] [--all | --e2e] <pipeline.yml>
|
||||
|
||||
# Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
|
||||
python3 upload_pipeline.py --upload .buildkite/pipeline.yml
|
||||
|
||||
# Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
|
||||
python3 upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
|
||||
# Full suite (rebase pipeline): keep all steps, still strip the uploader-only key:
|
||||
python3 upload_pipeline.py --upload --all .buildkite/test-merge.yml
|
||||
|
||||
# Nightly L2 E2E only: keep the E2E Test group, ignore source_file_dependencies:
|
||||
python3 upload_pipeline.py --upload --e2e .buildkite/test-ready.yml
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -44,6 +50,7 @@ LOG = "upload_pipeline"
|
||||
ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
DOC_SEP = "\n---\n"
|
||||
BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
|
||||
E2E_GROUP_MARKER = "E2E Test"
|
||||
|
||||
|
||||
def _log(message: str) -> None:
|
||||
@@ -319,18 +326,18 @@ def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
|
||||
nightly_only = (
|
||||
'(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
|
||||
)
|
||||
nightly_main = 'build.branch == "main" && build.env("NIGHTLY") == "1"'
|
||||
ready_pr = 'build.branch != "main" && build.pull_request.labels includes "ready"'
|
||||
merge_main = 'build.branch == "main" && build.env("NIGHTLY") != "1" && build.env("WEEKLY") != "1"'
|
||||
merge_pr = 'build.branch != "main" && build.pull_request.labels includes "merge-test"'
|
||||
if skip_ci:
|
||||
image_if = f"'{nightly_only}'"
|
||||
ready_if = "'false'"
|
||||
merge_if = "'false'"
|
||||
ready_if = f"'({nightly_main})'"
|
||||
merge_if = f"'({nightly_main})'"
|
||||
else:
|
||||
image_if = "'true'"
|
||||
ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\''
|
||||
merge_if = (
|
||||
'\'(build.branch == "main" && build.env("NIGHTLY") != "1" '
|
||||
'&& build.env("WEEKLY") != "1") || '
|
||||
'(build.branch != "main" && build.pull_request.labels includes "merge-test")\''
|
||||
)
|
||||
ready_if = f"'(({nightly_main}) || ({ready_pr}))'"
|
||||
merge_if = f"'(({nightly_main}) || (({merge_main}) || ({merge_pr})))'"
|
||||
|
||||
return (
|
||||
continuation.replace("__IMAGE_BUILD_IF__", image_if)
|
||||
@@ -393,6 +400,20 @@ def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
|
||||
return filtered
|
||||
|
||||
|
||||
def _is_e2e_group(step: dict[str, Any]) -> bool:
|
||||
group = step.get("group")
|
||||
return isinstance(group, str) and E2E_GROUP_MARKER in group
|
||||
|
||||
|
||||
def _select_e2e_group_steps(steps: list[Any]) -> list[Any]:
|
||||
selected = [step for step in steps if isinstance(step, dict) and _is_e2e_group(step)]
|
||||
if not selected:
|
||||
_log(f"no group matching {E2E_GROUP_MARKER!r} found")
|
||||
else:
|
||||
_log(f"keep {len(selected)} group(s) matching {E2E_GROUP_MARKER!r}")
|
||||
return selected
|
||||
|
||||
|
||||
def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
|
||||
"""Remove uploader-only keys while keeping all steps (no diff filtering)."""
|
||||
stripped: list[Any] = []
|
||||
@@ -419,11 +440,16 @@ def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
|
||||
def render_test_pipeline(
|
||||
doc: dict[str, Any],
|
||||
changed_files: list[str] | None,
|
||||
*,
|
||||
e2e_only: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
steps = doc.get("steps")
|
||||
if not isinstance(steps, list):
|
||||
return doc
|
||||
if changed_files is not None:
|
||||
if e2e_only:
|
||||
steps = _select_e2e_group_steps(steps)
|
||||
steps = _strip_uploader_metadata_from_steps(steps)
|
||||
elif changed_files is not None:
|
||||
steps = _filter_steps(steps, changed_files)
|
||||
else:
|
||||
steps = _strip_uploader_metadata_from_steps(steps)
|
||||
@@ -437,11 +463,22 @@ def resolve_pipeline_path(arg: str) -> Path:
|
||||
return ROOT / path
|
||||
|
||||
|
||||
def render_pipeline(path: Path) -> str:
|
||||
def render_pipeline(
|
||||
path: Path,
|
||||
*,
|
||||
force_all: bool = False,
|
||||
e2e_only: bool = False,
|
||||
) -> str:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
diff_range = resolve_diff_range()
|
||||
changed_files = resolve_changed_files()
|
||||
|
||||
# ``--all`` forces the keep-all-steps path (no diff-aware skipping) while still
|
||||
# stripping ``source_file_dependencies``. Used by the rebase pipeline so main builds
|
||||
# run the full e2e suite (see .buildkite/rebase-pipeline.yaml).
|
||||
if force_all or e2e_only:
|
||||
changed_files = None
|
||||
|
||||
if BOOTSTRAP_MARKER in text:
|
||||
rendered = render_bootstrap_pipeline(
|
||||
text,
|
||||
@@ -453,7 +490,7 @@ def render_pipeline(path: Path) -> str:
|
||||
if not isinstance(doc, dict):
|
||||
raise ValueError(f"invalid pipeline YAML: {path}")
|
||||
|
||||
doc = render_test_pipeline(doc, changed_files)
|
||||
doc = render_test_pipeline(doc, changed_files, e2e_only=e2e_only)
|
||||
|
||||
return yaml.safe_dump(doc, sort_keys=False)
|
||||
|
||||
@@ -480,6 +517,17 @@ def main() -> int:
|
||||
action="store_true",
|
||||
help="Pipe rendered YAML to buildkite-agent pipeline upload",
|
||||
)
|
||||
mode = parser.add_mutually_exclusive_group()
|
||||
mode.add_argument(
|
||||
"--all",
|
||||
action="store_true",
|
||||
help="Keep all steps (disable diff-aware skipping); still strips source_file_dependencies",
|
||||
)
|
||||
mode.add_argument(
|
||||
"--e2e",
|
||||
action="store_true",
|
||||
help="Keep only the E2E Test group; disable diff-aware skipping for those steps",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
path = resolve_pipeline_path(args.pipeline)
|
||||
@@ -487,7 +535,7 @@ def main() -> int:
|
||||
_log(f"missing pipeline file: {path}")
|
||||
return 1
|
||||
|
||||
rendered = render_pipeline(path)
|
||||
rendered = render_pipeline(path, force_all=args.all, e2e_only=args.e2e)
|
||||
if args.upload:
|
||||
upload_to_buildkite(rendered)
|
||||
else:
|
||||
|
||||
@@ -363,6 +363,42 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "TTS · IndexTTS2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_indextts2.py
|
||||
- tests/e2e/offline_inference/test_indextts2.py
|
||||
- tests/model_executor/models/indextts2/
|
||||
- tests/model_executor/stage_input_processors/test_indextts2.py
|
||||
- tests/assets/indextts2/
|
||||
- vllm_omni/model_executor/models/indextts2/
|
||||
- vllm_omni/model_executor/stage_input_processors/indextts2.py
|
||||
- vllm_omni/deploy/indextts2.yaml
|
||||
- vllm_omni/deploy/indextts2_low_latency.yaml
|
||||
- vllm_omni/entrypoints/openai/serving_speech.py
|
||||
- vllm_omni/entrypoints/openai/tts_adapters/indextts2.py
|
||||
- vllm_omni/config/pipeline_registry.py
|
||||
- vllm_omni/model_executor/models/registry.py
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
pytest -s -v tests/e2e/online_serving/test_indextts2.py tests/e2e/offline_inference/test_indextts2.py -m 'advanced_model and cuda' --run-level 'advanced_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
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
gpus: all
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
@@ -509,6 +545,34 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · SoulX-Singer Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_soulxsinger.py
|
||||
- vllm_omni/diffusion/models/soulx_singer/
|
||||
- vllm_omni/deploy/soulxsinger_svc.yaml
|
||||
- vllm_omni/deploy/soulxsinger_svs.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_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"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
|
||||
@@ -1279,7 +1279,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2V · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -1322,6 +1322,7 @@ steps:
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
|
||||
@@ -85,6 +85,66 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Reliability Test - hunyuan-image"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Reliability Test - voxcpm2"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
|
||||
agents:
|
||||
queue: "gpu_1_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"
|
||||
|
||||
# Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips).
|
||||
- group: "Reliability Test - Invalid parameters Test"
|
||||
depends_on: upload-weekly-pipeline
|
||||
|
||||
@@ -24,7 +24,7 @@ Check the model's HF repo for `model_index.json`. This determines your path:
|
||||
|
||||
| Scenario | How to identify | Migration path |
|
||||
|----------|----------------|----------------|
|
||||
| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
|
||||
| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip implementation, then validate model-specific examples, tests, and docs as needed |
|
||||
| **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
|
||||
| **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
|
||||
| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
|
||||
@@ -52,7 +52,7 @@ vllm_omni/diffusion/models/your_model_name/
|
||||
2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
|
||||
3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
|
||||
4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
|
||||
5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
|
||||
5. Add class attributes for acceleration features such as `_repeated_blocks` and `_layerwise_offload_blocks_attrs` (see [references/transformer-adaptation.md](references/transformer-adaptation.md) for examples).
|
||||
|
||||
### A4. Adapt pipeline
|
||||
|
||||
@@ -78,13 +78,34 @@ class YourPipeline(nn.Module):
|
||||
|
||||
Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
|
||||
|
||||
### A4.1 Add progress bar support (recommended)
|
||||
|
||||
For pipelines with a standard denoising loop, prefer the existing progress bar pattern instead of hand-rolled logging.
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||||
|
||||
class YourPipeline(nn.Module, ProgressBarMixin):
|
||||
def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
|
||||
# ... prepare timesteps / latents ...
|
||||
with self.progress_bar(total=len(timesteps)) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# predict noise / scheduler step
|
||||
latents = ...
|
||||
progress_bar.update()
|
||||
|
||||
return DiffusionOutput(output=output)
|
||||
```
|
||||
|
||||
For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py` and existing pipelines using `ProgressBarMixin`.
|
||||
|
||||
### A5. Register, test, docs → continue at Step 4 below.
|
||||
|
||||
---
|
||||
|
||||
## Path B: Custom/Private Repo Model
|
||||
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
|
||||
|
||||
### B1. Understand the reference repo
|
||||
|
||||
@@ -134,7 +155,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
|
||||
|
||||
### B4. Handle custom weight loading
|
||||
|
||||
Custom models have two patterns for weight loading:
|
||||
Custom models have two common patterns for weight loading:
|
||||
|
||||
**Pattern 1: Bypass standard loader** (DreamID-Omni style)
|
||||
|
||||
@@ -149,8 +170,10 @@ class CustomPipeline(nn.Module):
|
||||
self.vae = custom_init_vae(model, device=self.device)
|
||||
self.text_encoder = custom_init_text_encoder(model, device=self.device)
|
||||
self.transformer = CustomFusionModel(CONFIG)
|
||||
load_custom_checkpoint(self.transformer,
|
||||
checkpoint_path=os.path.join(model, "model.safetensors"))
|
||||
load_custom_checkpoint(
|
||||
self.transformer,
|
||||
checkpoint_path=os.path.join(model, "model.safetensors"),
|
||||
)
|
||||
# NO weights_sources defined — bypasses standard loader
|
||||
|
||||
def load_weights(self, weights):
|
||||
@@ -162,6 +185,8 @@ class CustomPipeline(nn.Module):
|
||||
When weights are in safetensors format but need name remapping:
|
||||
|
||||
```python
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
class CustomPipeline(nn.Module):
|
||||
def __init__(self, *, od_config, prefix=""):
|
||||
super().__init__()
|
||||
@@ -187,7 +212,9 @@ class CustomPipeline(nn.Module):
|
||||
# Remap original weight names to vllm-omni module names
|
||||
name = self._remap_weight_name(name)
|
||||
if name in params:
|
||||
default_weight_loader(params[name], tensor)
|
||||
param = params[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, tensor)
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
```
|
||||
@@ -281,12 +308,12 @@ For Omni or custom models, create:
|
||||
### Step 7: Update Documentation
|
||||
|
||||
Required updates:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
|
||||
1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
|
||||
2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
|
||||
3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
|
||||
4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
|
||||
5. `examples/offline_inference/xxx/README.md` — offline example docs
|
||||
6. `examples/online_serve/xxx/README.md` — online serve docs
|
||||
3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
|
||||
4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
|
||||
5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)
|
||||
6. `examples/online_serving/<name>/README.md` — online serving docs
|
||||
|
||||
### Step 8: Add E2E Tests
|
||||
|
||||
@@ -537,7 +564,7 @@ class YourTransformer(nn.Module):
|
||||
#### 10e. Update parallelism documentation
|
||||
|
||||
After adding parallelism support, update:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
|
||||
1. `docs/user_guide/diffusion/parallelism/overview.md` — add your model to the support overview/table
|
||||
2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
|
||||
|
||||
### Step 11: Add CPU Offload Support
|
||||
|
||||
@@ -569,4 +569,4 @@ Check which parallelism methods a model supports:
|
||||
| **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
|
||||
| **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
|
||||
|
||||
The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.
|
||||
The canonical per-model support overview/table is in `docs/user_guide/diffusion/parallelism/overview.md`.
|
||||
@@ -219,13 +219,14 @@ def load_weights(self, weights):
|
||||
|
||||
### Class-Level Attributes for Features
|
||||
|
||||
|
||||
```python
|
||||
class MyTransformer(nn.Module):
|
||||
# torch.compile: list block class names that repeat and can be compiled
|
||||
_repeated_blocks = ["MyTransformerBlock"]
|
||||
|
||||
# CPU offload: attribute name of the nn.ModuleList containing blocks
|
||||
_layerwise_offload_blocks_attr = "blocks"
|
||||
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||
|
||||
# LoRA: mapping of fused param names to original param names
|
||||
packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
|
||||
|
||||
@@ -63,6 +63,8 @@ If multiple rows apply (e.g., a diffusion model is also a new model), union the
|
||||
|
||||
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 ⚠.
|
||||
|
||||
**Also run the Code-Quality sweep on every PR**, regardless of type or mode: the five diff-scoped checks in [references/code-quality.md](references/code-quality.md) — kwargs fragility, broad-except swallow, `Any`/wrong type hints, hot-path `.clone()`/`deepcopy`, and event-loop blocking — plus the advisory conventions (log level, structured logging, synchronization, cleanup, dependencies, naming) in [checklists.md](references/checklists.md). These count only lines the PR *adds* — the pre-existing backlog across the repo is out of scope.
|
||||
|
||||
### Step 5: Print Report
|
||||
|
||||
```
|
||||
@@ -74,6 +76,7 @@ Pre-check report for <branch>
|
||||
Dimension Result
|
||||
───────────────── ──────
|
||||
PR title format ✓
|
||||
Code quality ⚠ 1 broad except, 2 Any hints
|
||||
PR desc integrity ✓
|
||||
Registry/config ✓
|
||||
Dead code ⚠ 2 warnings
|
||||
|
||||
@@ -35,6 +35,31 @@ PR title must follow the project convention documented in `docs/contributing/REA
|
||||
|
||||
---
|
||||
|
||||
## All PRs: Code Quality (run on every PR, quick & full)
|
||||
|
||||
Diff-scoped sweep for five fragility patterns that are pervasive in the existing codebase — **only lines the PR adds count**, not the pre-existing backlog. Detection commands, codebase examples, and full severity rules live in [code-quality.md](code-quality.md).
|
||||
|
||||
- [ ] **No new `**kwargs` string-lookup plumbing:** the diff adds no `**kwargs` + `kwargs["..."]` / `kwargs.get("...")` / `"..." in kwargs`. ⚠ any new instance; ✗ if a new string key is duplicated across ≥2 files without a shared constant, or unknown keys are silently dropped on a fail-fast path.
|
||||
- [ ] **No new broad exception swallow:** the diff adds no `except Exception:` / bare `except:`. ⚠ any new instance; ✗ for `except: pass` / `except Exception: return None` / `continue` on a fail-fast path (init, config validation, weight loading, request handling).
|
||||
- [ ] **No new `Any` / wrong type hints:** the diff adds no `: Any` / `-> Any` / untyped production signature, and no new `SimpleNamespace` in tests faking an object that already has a real typed stub. ⚠ any new instance; ✗ for a *wrong* (actively misleading) annotation or a `SimpleNamespace` test fake of a typed object.
|
||||
- [ ] **No new hot-path copy:** the diff adds no `.clone()` / `.copy_()` / `copy.deepcopy(...)` inside a per-step loop or per-request path. ⚠ any new instance; ✗ for `.clone()` of a full latent/activations tensor inside an AR or diffusion step loop without a comment, or `deepcopy` of a scheduler / sampling-config on the request path.
|
||||
- [ ] **No new event-loop blocking:** the diff adds no `time.sleep` / blocking HTTP (`requests`, `urllib`) / `await` inside a held lock in `async def` code. ⚠ any new instance; ✗ for blocking sleep or HTTP on the serving path (`engine/`, `entrypoints/`, async `connectors/`), or a lock held across an `await`. Blocking inside a dedicated worker thread is fine.
|
||||
|
||||
Roll the ⚠/✗ counts into the report as a single **Code quality** dimension row.
|
||||
|
||||
### Conventions (eyeball — no grep)
|
||||
|
||||
Judgment calls that don't grep cleanly — apply them while reading the diff, not as a mechanical check:
|
||||
|
||||
- **Right log level** — `info` is for end-user-visible events; routine internal events are `debug`; hot paths are `trace` or removed. Logging is not free.
|
||||
- **Structured fields over f-strings** — pass dynamic values as structured fields (`logger.info("msg", extra={...})`), not interpolated strings, so they stay greppable and parseable.
|
||||
- **No reflexive synchronization** — don't add `threading.Lock` / `asyncio.Lock` / `deepcopy` / re-wrap shared state unless there is real concurrent mutation; the GIL covers most Python code. Owners decide their own synchronization — don't pre-lock in a constructor.
|
||||
- **Context managers for cleanup** — use `with` / `contextlib` for files, locks, CUDA streams, and profilers; not raw `open()` + manual `.close()` or ad-hoc `finally`.
|
||||
- **No new dependencies without justification** — prefer the stdlib and existing vllm-omni deps; a new top-level dep needs a reason (check `pyproject.toml` / `requirements*.txt`).
|
||||
- **Naming / interface discipline** — don't rename or change interfaces just for taste in this PR, but do fix names that actively mislead (e.g. `serve` implying a long-running server, `Instance` where multiple instance kinds coexist).
|
||||
|
||||
---
|
||||
|
||||
## Bug Fix PRs
|
||||
|
||||
### Quick
|
||||
@@ -49,7 +74,7 @@ PR title must follow the project convention documented in `docs/contributing/REA
|
||||
- [ ] **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
|
||||
- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added. See the broad-except pattern in [code-quality.md](code-quality.md#2-broad-exception-swallow) for the full severity rule (⚠ any new broad catch; ✗ for a swallow on a fail-fast path).
|
||||
- [ ] **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
|
||||
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
# Code-Quality Patterns
|
||||
|
||||
Five fragility patterns to sweep on **every** PR, in both quick and full mode. These are pervasive in the existing codebase (hundreds of pre-existing sites) — so the checks are **diff-scoped**: only lines the PR *adds* count. Whole-repo greps light up every PR with the backlog and are useless.
|
||||
|
||||
All commands assume:
|
||||
```bash
|
||||
BASE="$(git merge-base HEAD origin/main)"
|
||||
```
|
||||
|
||||
**Severity:** ⚠ for any *new* instance introduced by the diff; ✗ for the acute sub-cases defined per pattern. Acute ✗ findings belong in the report's blocking count.
|
||||
|
||||
---
|
||||
|
||||
## 1. `**kwargs` string-lookup passthrough
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
def forward(self, **kwargs):
|
||||
# BLOCKER-prone: raw dict + string keys duplicated across files
|
||||
if "runtime_additional_information" in kwargs and \
|
||||
"model_intermediate_buffer" not in kwargs:
|
||||
...
|
||||
```
|
||||
Real instances of this exact guard are copy-pasted across `cosyvoice3.py`, `indextts2_talker.py` (×2), `indextts2_s2mel_decoder.py`, `qwen3_tts_talker.py`, `qwen3_omni.py`. A typo in any string key silently no-ops; there is no type check, no IDE completion, and unknown keys are silently dropped.
|
||||
|
||||
### Why dangerous
|
||||
Raw `dict` + `**kwargs` + lists of param-name strings is **fragile by construction**: the compiler can't see a missing or misspelled key, so failures are silent. Each duplicated copy is a new place for the strings to drift.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `**kwargs` signature whose body does `kwargs["..."]` / `kwargs.get("...")` / `"..." in kwargs`.
|
||||
- ✗ when the diff introduces a **new** string key that is duplicated across ≥2 files without a shared module constant, or when `**kwargs` plumbing **silently drops unknown keys** on a fail-fast path (init, config validation, weight loading).
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E '\*\*kwargs|kwargs(\.get|\[)|"[a-z_]+"\s+in\s+kwargs'
|
||||
```
|
||||
|
||||
### Fix
|
||||
Prefer explicit typed params. The repo already has the right primitive: `msgspec.Struct` with `forbid_unknown_fields=True` rejects unknown keys instead of silently dropping them — see `_StructBase` in `vllm_omni/data_entry_keys.py:105` and `MoriPullRequest` in `vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py:81`. A `@dataclass` or `TypedDict` works too.
|
||||
|
||||
If `**kwargs` is unavoidable (vLLM base-class compat), centralize the key strings as module-level constants behind **one** typed accessor, instead of repeating the literal across files.
|
||||
|
||||
---
|
||||
|
||||
## 2. Broad exception swallow
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
# BLOCKER-prone: catches everything, returns None on any failure
|
||||
try:
|
||||
evt = build_event(...)
|
||||
return evt
|
||||
except Exception:
|
||||
return None
|
||||
```
|
||||
Real clusters: `vllm_omni/metrics/stats.py:240,274,367,450` and `vllm_omni/metrics/{modality,utils}.py` — all `except Exception: return None`.
|
||||
|
||||
### Why dangerous
|
||||
A generic `except Exception:` is **too flexible**. If we catch an exception type we don't expect, it is almost always an unhandled edge case that needs to be *fixed* — not swallowed into `None`/`pass`/`continue`. Broad catches turn fixable bugs into silent wrong behavior and make debugging impossible.
|
||||
|
||||
The repo backs this with ruff `BLE001` (broad-exception catch); it is currently suppressed via `# noqa: BLE001` in `vllm_omni/patch.py:413`. Prefer fixing over suppressing.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `except Exception:` / `except BaseException:` / bare `except:`.
|
||||
- ✗ for `except:` / `except Exception: pass|return None|continue` on a **fail-fast path** (init, config validation, weight loading, request handling, connector setup). This extends the Bug-Fix "No silent failure risk" ✗ in [checklists.md](checklists.md).
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
# Any new broad catch (⚠ baseline):
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E 'except\s*(Exception|BaseException)?\s*:'
|
||||
# Note: this matches only the immediate-colon swallow forms (`except :`,
|
||||
# `except Exception:`) — it deliberately excludes specific catches like
|
||||
# `except ValueError:` and the `except Exception as e:` variant. Eyeball
|
||||
# the latter; `as e` usually means it is logged, but still verify.
|
||||
```
|
||||
|
||||
### Fix
|
||||
Catch the **specific** types you actually expect (`ValueError`, `KeyError`, `AttributeError`, `OSError`, `torch.cuda.OutOfMemoryError`, `msgspec.ValidationError`, …). At a genuine top-level / best-effort boundary (metrics, signal handlers), at minimum **log** the exception; never swallow into a bare `pass`/`return None` on a path that should fail loudly.
|
||||
|
||||
---
|
||||
|
||||
## 3. `Any` / wrong type annotations
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
# ⚠ Any param + Any return — no contract, no narrowing
|
||||
def _extract_mm_output(engine_outputs: Any) -> dict[str, Any]:
|
||||
...
|
||||
```
|
||||
|
||||
### Why dangerous (root cause)
|
||||
`Any` disables the type checker for that position, so typos and wrong shapes pass silently. The repo has 636 `: Any` params, 244 `-> Any` returns, and **1,173 `SimpleNamespace` in `tests/`**. The `Any` leak is largely test-driven: vibe-coded unit tests fake objects with `SimpleNamespace`, which forces the production code consuming them to be typed `Any`, and the `Any` then propagates outward into public signatures.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `: Any` / `-> Any` / fully-untyped signature in production code.
|
||||
- ✗ for a **wrong** annotation (actively misleading — e.g. `-> bool` that returns `Optional[bool]`), or a *new* `SimpleNamespace` in a test that mimics an object which already has a real typed stub / `@dataclass` / `TypedDict` / `Protocol`.
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
# Any leaking into new signatures:
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E ':\s*Any\b|->\s*Any\b'
|
||||
# SimpleNamespace newly added in tests (the Any-leak source):
|
||||
git diff ${BASE}...HEAD -- "tests/" | grep "^+[^+]" | grep "SimpleNamespace"
|
||||
```
|
||||
|
||||
### Fix
|
||||
Replace `Any` with the concrete type, a `Protocol`, or a `Union`. If a type is genuinely dynamic, prefer `object` + `isinstance` narrowing over `Any` — `object` keeps the checker engaged. In tests, replace `SimpleNamespace` with the real class, a small `@dataclass`/`TypedDict`, or a `Protocol` + fake, so the system under test stays fully typed.
|
||||
|
||||
---
|
||||
|
||||
## 4. Unnecessary copy — `.clone()` / `copy.deepcopy` in hot paths
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
# Per-step in a diffusion/AR loop: clones the full latent every iteration
|
||||
for t in timesteps:
|
||||
trajectory_latents.append(x_t.clone()) # O(N) GPU alloc + memcpy each step
|
||||
```
|
||||
Real clusters: `vllm_omni/diffusion/models/bagel/bagel_transformer.py` calls `x_t.clone()` ~9× across a trajectory loop (lines 1820–2186); the cache hooks `vllm_omni/diffusion/cache/teacache/hook.py:155` and `magcache/hook.py:182` clone `hidden_states` per step. Separately, `copy.deepcopy(self.scheduler)` / `copy.deepcopy(req.sampling_params)` is copy-pasted across diffusion pipelines — `pipeline_qwen_image.py:786`, `pipeline_hunyuan_image3.py:1904`, `pipeline_ltx2*.py` (×3), `pipeline_dreamzero.py:1085`, `diffusion_model_runner.py:398` — a full object-graph deepcopy on the request path. (Contrast `dreamzero/causal_wan_model.py:203`, which carries a `# No .clone() needed.` comment — the aware pattern.)
|
||||
|
||||
### Why dangerous
|
||||
A `.clone()` of a GPU tensor or a `copy.deepcopy` of a config object inside a per-step / per-request path is a hidden allocator: it reads as one Python line but triggers a GPU memcpy or a deep object-graph walk every iteration. Across 50 diffusion steps or hundreds of AR steps it dominates latency and VRAM bandwidth. `deepcopy` additionally ignores any custom device / `__copy__` logic and can drag tensors through CPU.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `.clone()` / `.copy_()` / `copy.deepcopy(...)` inside a per-step loop or per-request path.
|
||||
- ✗ for `.clone()` of a full latent/activations tensor inside an AR or diffusion step loop without an explanatory comment, or `copy.deepcopy` of a scheduler / sampling-config on the request hot path.
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E '\.clone\(\)|\.copy_\(|copy\.deepcopy|deepcopy\('
|
||||
```
|
||||
|
||||
### Fix
|
||||
Tensors: prefer a view (`x[..., :k]`, `torch.narrow`, `as_strided`) or write into a pre-allocated buffer with `dst.copy_(src)` instead of `.clone()`; if the caller no longer needs the original, `move`/`del` it rather than copy. Schedulers / configs: build a fresh lightweight instance per request, or split out the mutable per-request bits so the shared immutable part isn't deepcopied. If a copy is genuinely required (e.g. the source is mutated in place later), keep it and add a one-line comment saying why.
|
||||
|
||||
---
|
||||
|
||||
## 5. Blocking the asyncio event loop / lock held across `await`
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
# BLOCKER-prone: blocking sleep / blocking HTTP / lock held across an await,
|
||||
# all inside an async (event-loop) function
|
||||
async def handle(self, req):
|
||||
with self._lock:
|
||||
resp = await self._client.call(req) # lock held across await -> serializes all requests
|
||||
time.sleep(0.1) # stalls the single loop for every concurrent request
|
||||
```
|
||||
|
||||
### Why dangerous
|
||||
vLLM-Omni serves over a single asyncio event loop. Any synchronous blocking work on an `async def` path — `time.sleep`, blocking HTTP (`requests`/`urllib`), heavy CPU work, or a lock held across an `await` — stalls the loop for *every* concurrent request, not just the blocking one. A lock held across `await` silently serializes the whole pipeline; blocking I/O tanks throughput and inflates tail latency.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `time.sleep(`, blocking HTTP call (`requests.get/post`, `urllib.request`, `urlopen`), or `await` inside a `with Lock` / `async with` critical section, in `async def` code.
|
||||
- ✗ for `time.sleep` or blocking HTTP on the serving path (`engine/`, `entrypoints/`, async `connectors/`), or a lock acquired and held across an `await` on the request path.
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
# New blocking primitives introduced by the diff — then eyeball whether each
|
||||
# sits in an async function on the loop vs a dedicated worker thread:
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E 'time\.sleep|requests\.(get|post|put|delete)|urllib\.request|urlopen'
|
||||
# Lock acquired with an await potentially inside the critical section (manual eyeball):
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" | grep -E 'async with|\.acquire\(\)'
|
||||
```
|
||||
Note: blocking inside a **dedicated worker thread is correct** — the repo already does this (`diffusion/model_loader/hub_prefetch.py`, `distributed/omni_connectors/kv_transfer_manager.py`, the `mooncake`/`mori` connectors poll with `time.sleep` on their own threads). The check targets the *event-loop* path only.
|
||||
|
||||
### Fix
|
||||
Move blocking work off the loop with `asyncio.to_thread(...)` / `run_in_executor`, or use a native async client (`httpx.AsyncClient` / `aiohttp`). Never `await` while holding a lock — drop the lock before the `await` and re-acquire after, or restructure so the critical section contains no `await`.
|
||||
|
||||
---
|
||||
|
||||
## Running the sweep
|
||||
|
||||
Run all five detections against `${BASE}...HEAD`, then for each hit decide ⚠ vs ✗ using the severity rules above. Roll the results into the report as a single **Code quality** dimension row (count of ⚠ and ✗), alongside the type-specific checklist from [checklists.md](checklists.md).
|
||||
@@ -19,6 +19,7 @@ body:
|
||||
|
||||
Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
|
||||
- type: textarea
|
||||
id: current-environment
|
||||
attributes:
|
||||
label: Your current environment
|
||||
description: |
|
||||
@@ -42,6 +43,7 @@ body:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: code-version
|
||||
attributes:
|
||||
label: Your code version
|
||||
description: |
|
||||
@@ -66,6 +68,7 @@ body:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: bug-description
|
||||
attributes:
|
||||
label: 🐛 Describe the bug
|
||||
description: |
|
||||
|
||||
@@ -241,6 +241,7 @@ datasets/
|
||||
*.json
|
||||
!tests/dfx/perf/tests/*.json
|
||||
!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
|
||||
!tests/assets/soulxsinger/*.json
|
||||
*.jsonl
|
||||
*.parquet
|
||||
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
# LVSA Long-Video Showcase
|
||||
|
||||
**Training-free block-sparse attention for long-video diffusion**, running on
|
||||
vLLM-Omni via the public plugin entry points (no core changes). LVSA speeds up
|
||||
extended-horizon generation and prevents the freeze/loop failure mode dense
|
||||
attention exhibits beyond a model's training length — with no fine-tuning and no
|
||||
weight changes.
|
||||
|
||||
Upstream: **[LongVideoSparseAttention](https://github.com/JiusiServe/LongVideoSparseAttention)**
|
||||
(algorithm, all model adapters, full docs).
|
||||
|
||||
## Results
|
||||
|
||||
Wan 2.1 1.3B, single A100 80 GB, mean over 5 prompts (Dense vs LVSA-FlashInfer):
|
||||
|
||||
| Horizon | 1× | 2× | 3× | 4× | 5× | 6× |
|
||||
|---|---|---|---|---|---|---|
|
||||
| **Speedup** | 0.9× | 1.4× | 1.9× | 2.5× | 3.0× | **3.5×** |
|
||||
| **VQeval composite** (Dense→LVSA) | 63→63 | 59→63 | 61→63 | 61→64 | 59→64 | 59→63 |
|
||||
|
||||
The speedup **grows with length** (LVSA is the dense regime at 1× by design, so
|
||||
the win starts at 2×), and quality *improves* at extension — LVSA's rotating
|
||||
keyframes prevent the looping/static output dense produces beyond its horizon
|
||||
(VQeval loop-quality +27 to +34 at ≥3×). At long horizons dense **OOMs on 80 GB**
|
||||
(HunyuanVideo ≥2×, Wan2.1-14B and Cosmos at 6×) where LVSA still generates.
|
||||
|
||||
→ Full sweep (5 models × 6 horizons, VQeval + VBench-Long) and the
|
||||
quality-metric methodology:
|
||||
**[benchmarks/](https://github.com/JiusiServe/LongVideoSparseAttention/tree/main/benchmarks)**.
|
||||
|
||||
## Quickstart
|
||||
|
||||
```bash
|
||||
# 1. vLLM-Omni 0.22.0 (stable) + matching vLLM
|
||||
pip install "vllm==0.22.0"
|
||||
pip install --no-build-isolation \
|
||||
"vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@v0.22.0"
|
||||
|
||||
# 2. LVSA core + plugin
|
||||
git clone https://github.com/JiusiServe/LongVideoSparseAttention.git
|
||||
( cd LongVideoSparseAttention && pip install -e . && pip install -e lvsa-vllm-omni/ )
|
||||
|
||||
# 3. this showcase's client dep
|
||||
pip install requests
|
||||
```
|
||||
|
||||
FlashInfer (used by the default `flashinfer` backend) **ships with vLLM** — step 1
|
||||
already pulls the pinned `flashinfer-python` + `flashinfer-cubin`, so the fused
|
||||
backend works with no extra install. Two optional extras:
|
||||
|
||||
```bash
|
||||
# (optional) faster startup — prebuilt JIT cache, matched to your CUDA + flashinfer
|
||||
CU=cu128 # match your CUDA: cu121 / cu124 / cu128 / …
|
||||
FIVER=$(python -c "import importlib.metadata as m; print(m.version('flashinfer-python'))")
|
||||
pip install --extra-index-url "https://flashinfer.ai/whl/${CU}" \
|
||||
"flashinfer-jit-cache==${FIVER}+${CU}"
|
||||
|
||||
# (fallback) no FlashInfer in your environment? run the SDPA backend instead:
|
||||
export LVSA_BACKEND=sdpa
|
||||
```
|
||||
|
||||
```bash
|
||||
# Serve an LVSA-enabled endpoint (Wan 2.1 1.3B at 2× horizon):
|
||||
MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers MODEL_FAMILY=wan FRAMES=161 \
|
||||
bash serve_lvsa.sh
|
||||
|
||||
# In another shell — generate:
|
||||
python generate.py --frames 161 --height 480 --width 832 \
|
||||
--prompt "A dog running through a sunlit forest." --out wan_lvsa.mp4
|
||||
```
|
||||
|
||||
`serve_lvsa.sh` sets the LVSA env vars for the chosen `MODEL_FAMILY`
|
||||
(`wan` / `hunyuan`). Wan and HunyuanVideo run through the LVSA **attention
|
||||
backend** by default — it stays sparse even under sequence-parallel, unlike the
|
||||
monkey-patch hooks.
|
||||
|
||||
## Configuration
|
||||
|
||||
The one knob you must get right is the model's training horizon
|
||||
(`LVSA_REFERENCE_LATENT_FRAMES`: Wan = 21, HunyuanVideo = 33); `serve_lvsa.sh`
|
||||
sets it per family. For the full environment-variable reference, the multi-GPU
|
||||
support matrix, and tuning (`sparsity_scale`, rotating keyframes, non-standard
|
||||
resolutions), see the upstream docs:
|
||||
|
||||
- **[lvsa-vllm-omni plugin reference](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/lvsa-vllm-omni/README.md)** — every `LVSA_*` env var
|
||||
- **[docs/parallelism.md](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/docs/parallelism.md)** — TP / Ulysses / HSDP support per model
|
||||
- **[docs/tuning.md](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/docs/tuning.md)** — picking knobs for your model and horizon
|
||||
@@ -0,0 +1,93 @@
|
||||
#!/usr/bin/env python
|
||||
"""Generate a long video through an LVSA-enabled vLLM-Omni server.
|
||||
|
||||
A thin client over the Omni video API — start the server first with
|
||||
``serve_lvsa.sh``. The sparse attention runs server-side; this just submits a
|
||||
job and saves the result.
|
||||
|
||||
Uses the **asynchronous** endpoint (``POST /v1/videos`` → poll
|
||||
``GET /v1/videos/{id}`` → download ``GET /v1/videos/{id}/content``) rather than
|
||||
``POST /v1/videos/sync``. The sync endpoint is capped by a fixed server-side
|
||||
timeout (``VIDEO_SYNC_TIMEOUT_S``, 600 s) and is documented as
|
||||
"for benchmark and testing scenarios" — long videos (LVSA's use case) routinely
|
||||
exceed it. The async job has no such cap.
|
||||
|
||||
python generate.py --frames 481 --height 480 --width 832 \
|
||||
--prompt "A golden retriever running through a sunlit forest." \
|
||||
--out dog_forest_lvsa.mp4
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--server", default=os.environ.get("OMNI_SERVER", "http://localhost:8000"))
|
||||
ap.add_argument("--prompt", default="A dog running through a sunlit forest.")
|
||||
ap.add_argument("--frames", type=int, default=161) # Wan 2x horizon
|
||||
ap.add_argument("--height", type=int, default=480)
|
||||
ap.add_argument("--width", type=int, default=832)
|
||||
ap.add_argument("--fps", type=int, default=16)
|
||||
ap.add_argument("--steps", type=int, default=50)
|
||||
ap.add_argument("--guidance", type=float, default=5.0)
|
||||
ap.add_argument("--out", default="lvsa_long_video.mp4")
|
||||
ap.add_argument("--poll-interval", type=float, default=5.0, help="Seconds between status polls.")
|
||||
ap.add_argument(
|
||||
"--max-wait", type=float, default=7200.0, help="Give up if the job hasn't finished after this many seconds."
|
||||
)
|
||||
args = ap.parse_args()
|
||||
|
||||
base = args.server.rstrip("/")
|
||||
form = {
|
||||
"prompt": args.prompt,
|
||||
"size": f"{args.width}x{args.height}",
|
||||
"num_frames": str(args.frames),
|
||||
"fps": str(args.fps),
|
||||
"num_inference_steps": str(args.steps),
|
||||
"guidance_scale": str(args.guidance),
|
||||
}
|
||||
|
||||
# ── 1. submit the async job ──────────────────────────────────────────────
|
||||
print(f"[lvsa-app] POST {base}/v1/videos {args.width}x{args.height} {args.frames}f steps={args.steps}")
|
||||
resp = requests.post(f"{base}/v1/videos", data=form, timeout=120)
|
||||
resp.raise_for_status()
|
||||
job = resp.json()
|
||||
vid = job["id"]
|
||||
print(f"[lvsa-app] job {vid} ({job.get('status', '?')})")
|
||||
|
||||
# ── 2. poll until completed / failed ─────────────────────────────────────
|
||||
deadline = time.time() + args.max_wait
|
||||
last = None
|
||||
while True:
|
||||
if time.time() > deadline:
|
||||
raise SystemExit(f"[lvsa-app] timed out after {args.max_wait:.0f}s (job {vid} still running)")
|
||||
time.sleep(args.poll_interval)
|
||||
r = requests.get(f"{base}/v1/videos/{vid}", timeout=60)
|
||||
if r.status_code != 200:
|
||||
# retrieve_video returns a non-200 JSON when the job has FAILED
|
||||
raise SystemExit(f"[lvsa-app] job failed: {r.text}")
|
||||
st = r.json()
|
||||
status = st.get("status")
|
||||
progress = st.get("progress", 0)
|
||||
if (status, progress) != last:
|
||||
print(f"[lvsa-app] {status} {progress}%")
|
||||
last = (status, progress)
|
||||
if status == "completed":
|
||||
break
|
||||
if status == "failed":
|
||||
raise SystemExit(f"[lvsa-app] job failed: {st.get('error')}")
|
||||
|
||||
# ── 3. download the rendered mp4 ─────────────────────────────────────────
|
||||
c = requests.get(f"{base}/v1/videos/{vid}/content", timeout=600)
|
||||
c.raise_for_status()
|
||||
with open(args.out, "wb") as f:
|
||||
f.write(c.content)
|
||||
print(f"[lvsa-app] saved {args.out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env bash
|
||||
# Start an LVSA-enabled vLLM-Omni server for Wan or HunyuanVideo.
|
||||
#
|
||||
# LVSA is selected as the attention backend via --diffusion-attention-config (see
|
||||
# below); the `LVSA_*` env vars only tune it. The `lvsa-vllm-omni` plugin registers
|
||||
# itself through vLLM-Omni's `general_plugins` entry points — nothing in vLLM-Omni
|
||||
# core is modified. Watch the log for "[LVSA] Geometry detected" to confirm sparse
|
||||
# attention is active for your request geometry.
|
||||
#
|
||||
# MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers MODEL_FAMILY=wan FRAMES=161 bash serve_lvsa.sh
|
||||
# MODEL=/path/to/HunyuanVideo-1.5-... MODEL_FAMILY=hunyuan FRAMES=193 bash serve_lvsa.sh
|
||||
set -euo pipefail
|
||||
|
||||
MODEL=${MODEL:?set MODEL=/path/to/Wan-or-HunyuanVideo-Diffusers}
|
||||
MODEL_FAMILY=${MODEL_FAMILY:-wan} # wan | hunyuan
|
||||
PORT=${PORT:-8000}
|
||||
FRAMES=${FRAMES:-161} # target horizon in frames
|
||||
VAE_T=${VAE_T:-4} # VAE temporal compression (Wan/Hunyuan = 4)
|
||||
BACKEND=${LVSA_BACKEND:-flashinfer} # flashinfer (fused, fastest) | sdpa
|
||||
|
||||
T_LAT=$(( (FRAMES - 1) / VAE_T + 1 ))
|
||||
|
||||
# Per-family training reference horizon (Wan train ref = 21, Hunyuan = 33).
|
||||
case "$MODEL_FAMILY" in
|
||||
wan) REFERENCE=${REFERENCE:-21} ;;
|
||||
hunyuan) REFERENCE=${REFERENCE:-33} ;;
|
||||
*) echo "MODEL_FAMILY must be 'wan' or 'hunyuan'"; exit 1 ;;
|
||||
esac
|
||||
export LVSA_BACKEND="$BACKEND"
|
||||
export LVSA_TOTAL_LATENT_FRAMES="$T_LAT"
|
||||
export LVSA_REFERENCE_LATENT_FRAMES="$REFERENCE"
|
||||
export LVSA_ROTATE_KEYFRAMES=${LVSA_ROTATE_KEYFRAMES:-1}
|
||||
|
||||
echo "[lvsa-app] serving $MODEL ($MODEL_FAMILY) on :$PORT"
|
||||
echo "[lvsa-app] LVSA backend=$BACKEND T_lat=$T_LAT (frames=$FRAMES) reference=$REFERENCE"
|
||||
echo "[lvsa-app] confirm activation by grep '\\[LVSA\\] Geometry detected' in this log"
|
||||
|
||||
# vLLM-Omni 0.22 selects the LVSA attention backend per-role via
|
||||
# --diffusion-attention-config. The backend is what drives sparse attention
|
||||
# here; it engages under tensor-parallel too, unlike the LVSA_*_HOOK monkey-patch
|
||||
# path (which falls back to dense under sequence-parallel), so we deliberately do
|
||||
# NOT set a hook.
|
||||
exec vllm serve "$MODEL" --omni --host 0.0.0.0 --port "$PORT" \
|
||||
--diffusion-attention-config '{"per_role": {"self": {"backend": "LVSA"}}}' \
|
||||
--init-timeout 1800 ${EXTRA_ARGS:-}
|
||||
@@ -1031,6 +1031,14 @@ def _save_generated_outputs(
|
||||
if not out.success or not out.response_body:
|
||||
continue
|
||||
|
||||
if isinstance(out.response_body, bytes):
|
||||
fname = f"req_{idx:04d}.mp4"
|
||||
fpath = os.path.join(save_dir, fname)
|
||||
with open(fpath, "wb") as f:
|
||||
f.write(out.response_body)
|
||||
saved += 1
|
||||
continue
|
||||
|
||||
media_urls: list[str] = []
|
||||
|
||||
# Chat-completions style: choices[*].message.content[*].image_url.url
|
||||
|
||||
@@ -80,6 +80,17 @@ models:
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
# MOSS-TTS-Local-Transformer-v1.5: Qwen3 backbone + 1-layer GPT2-style local
|
||||
# depth transformer (n_vq=12), paired with the 48 kHz stereo
|
||||
# MOSS-Audio-Tokenizer-v2 codec. Same voice-clone request shape as the
|
||||
# other MOSS-TTS variants.
|
||||
OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
# MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
|
||||
# has single-speaker rows; bench piggybacks on it (both speakers share the
|
||||
# same ref_audio). For real dialogue evaluation use the ttsd dataset.
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.18.0rc1
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=nightly-main
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=nightly-main-a3
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
|
||||
@@ -101,7 +101,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=xpu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
@@ -18,7 +18,6 @@ nav:
|
||||
- Offline Inference:
|
||||
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
|
||||
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
|
||||
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
|
||||
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
|
||||
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 262 KiB After Width: | Height: | Size: 410 KiB |
@@ -155,7 +155,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
/tests/dfx/stability/tests/test_wan22.json<br>
|
||||
<strong>Reliability:</strong><br>
|
||||
tests/dfx/reliability/test_reliability_{model_key}.py<br>
|
||||
(e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>)
|
||||
(e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>, <code>test_reliability_hunyuan_image.py</code>, <code>test_reliability_voxcpm2.py</code>)
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
|
||||
@@ -741,14 +741,16 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
### 4.2 Testing Content and Scope
|
||||
|
||||
- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni or `multiprocessing.spawn` for Wan2.2 video workers), **large multimodal chat** or **`/v1/videos`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
|
||||
- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni, `multiprocessing.spawn` for Wan2.2 video workers, or `vLLM-Omni::` for HunyuanImage DiT workers), **large multimodal chat**, **`/v1/videos`**, **`/v1/images/generations`**, or **`/v1/audio/speech`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
|
||||
|
||||
### 4.3 Test Directory and Execution Files
|
||||
|
||||
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
|
||||
- ***Reliability Test Suite*** (`tests/dfx/reliability/`):
|
||||
- `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
|
||||
- `test_reliability_wan22.py` — Wan2.2 T2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
|
||||
- `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
|
||||
- `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
|
||||
- `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
|
||||
- `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
|
||||
- `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
|
||||
- `README.md` — Short local run commands for this directory.
|
||||
@@ -846,22 +848,33 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in
|
||||
| `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). |
|
||||
| `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). |
|
||||
| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
|
||||
| `test_reliability_wan22.py` | Wan2.2 T2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
|
||||
| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
|
||||
| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
|
||||
| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
|
||||
| `README.md` | Minimal run / collect examples. |
|
||||
|
||||
#### Parametrization and markers
|
||||
|
||||
- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni paths often require **2** cards; Wan2.2 video paths **1** card).
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
|
||||
- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
|
||||
|
||||
#### CI trigger
|
||||
|
||||
Weekly Buildkite (`.buildkite/test-weekly.yml`) runs, for example:
|
||||
Weekly Buildkite (`.buildkite/test-weekly.yml`) runs one step per model suite (trigger: `WEEKLY=1` or PR label `weekly-test`), for example:
|
||||
|
||||
| Buildkite step | Test file | CI hardware |
|
||||
| -------------- | --------- | ----------- |
|
||||
| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
|
||||
| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
|
||||
| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
|
||||
| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
|
||||
|
||||
```bash
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
|
||||
```
|
||||
|
||||
#### Local commands
|
||||
@@ -870,6 +883,8 @@ pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
pytest --collect-only tests/dfx/reliability
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
```
|
||||
|
||||
#### Adding a new model suite
|
||||
@@ -882,7 +897,7 @@ pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
</details>
|
||||
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow` and/or `pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow` (add `test_reliability_<suite>.py` for new models)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/test-weekly.yml` for new models)
|
||||
|
||||
## Summary
|
||||
|
||||
|
||||
@@ -96,15 +96,21 @@ Our test scripts use the pytest framework. First, please use `git clone https://
|
||||
# Stability: Wan2.2 (v1/videos diffusion benchmark loop)
|
||||
pytest -s -v dfx/stability/scripts/test_stability_wan22.py -m slow
|
||||
|
||||
# Reliability: Qwen3-Omni
|
||||
# Reliability: Qwen3-Omni (H100 × 2)
|
||||
pytest -s -v dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
|
||||
# Reliability: Wan2.2
|
||||
# Reliability: Wan2.2 (H100)
|
||||
pytest -s -v dfx/reliability/test_reliability_wan22.py -m slow
|
||||
|
||||
# Reliability: HunyuanImage DiT (H100 × 4)
|
||||
pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
|
||||
# Reliability: VoxCPM2 (L4)
|
||||
pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
|
||||
```
|
||||
|
||||
The latest L5 commands for CI can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
|
||||
The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-weekly.yml).
|
||||
|
||||
You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
|
||||
|
||||
|
||||
@@ -144,6 +144,7 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
|
||||
"vllm_omni.entrypoints.async_diffusion",
|
||||
"vllm_omni.entrypoints.openai",
|
||||
"vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
|
||||
"vllm_omni.experimental", # optional serving deps not installed in docs build
|
||||
]
|
||||
if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
|
||||
continue
|
||||
|
||||
@@ -36,9 +36,10 @@ 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22S2VPipeline` | 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
@@ -63,6 +64,8 @@ th {
|
||||
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
|
||||
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
|
||||
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
@@ -74,6 +77,7 @@ th {
|
||||
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
|
||||
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
|
||||
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
|
||||
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
|
||||
| `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` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -32,6 +32,20 @@ curl -X POST http://localhost:8000/v1/images/generations \
|
||||
}' | jq -r '.data[0].b64_json' | base64 -d > dragon.png
|
||||
```
|
||||
|
||||
**Using curl save to file:**
|
||||
|
||||
```bash
|
||||
curl -o dragon.png -X POST http://localhost:8000/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
|
||||
"size": "1024x1024",
|
||||
"seed": 42,
|
||||
"response_format":"file"
|
||||
}'
|
||||
```
|
||||
|
||||
|
||||
**Using Python:**
|
||||
|
||||
```python
|
||||
@@ -57,6 +71,36 @@ img = Image.open(io.BytesIO(img_bytes))
|
||||
img.save("cat.png")
|
||||
```
|
||||
|
||||
**Using Python save to file:**
|
||||
|
||||
```python
|
||||
import requests
|
||||
import base64
|
||||
from PIL import Image
|
||||
import io
|
||||
import re
|
||||
|
||||
response = requests.post(
|
||||
"http://localhost:8000/v1/images/generations",
|
||||
json={
|
||||
"prompt": "a black and white cat wearing a princess tiara",
|
||||
"size": "1024x1024",
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42,
|
||||
"response_format":"file"
|
||||
}
|
||||
)
|
||||
|
||||
# save to file
|
||||
content_disposition = response.headers.get("Content-Disposition", "")
|
||||
match = re.search(r'filename="?(.+)"?', content_disposition)
|
||||
filename = match.group(1) if match else "save.png"
|
||||
with open(filename, "wb") as f:
|
||||
for chunk in response.iter_content(8192):
|
||||
f.write(chunk)
|
||||
print("saved:", filename)
|
||||
```
|
||||
|
||||
**Using OpenAI SDK:**
|
||||
|
||||
```python
|
||||
|
||||
@@ -233,9 +233,9 @@ Server -> Client:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence |
|
||||
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
|
||||
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
|
||||
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence |
|
||||
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
|
||||
| `{"type": "session.done", "total_sentences": N}` | Session complete |
|
||||
| `{"type": "error", "message": "..."}` | Non-fatal error |
|
||||
|
||||
@@ -245,8 +245,7 @@ All REST API parameters are supported, plus:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket |
|
||||
| `split_granularity` | string | "sentence" | Text splitting granularity |
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |
|
||||
|
||||
|
||||
```bash
|
||||
|
||||
@@ -160,7 +160,12 @@ 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.
|
||||
decoded reference frames are used as clean conditioning. Cosmos3 transfer mode
|
||||
also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints plus
|
||||
transfer options such as `control_path`, `control_guidance`,
|
||||
`control_guidance_interval`, `num_video_frames_per_chunk`,
|
||||
`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
|
||||
Cosmos3 recipe for complete examples.
|
||||
|
||||
### Speech-to-Video
|
||||
|
||||
|
||||
@@ -197,6 +197,7 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
|
||||
|
||||
**Notes:**
|
||||
|
||||
@@ -146,7 +146,7 @@ The following tables show which models support each feature:
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
|
||||
@@ -1,203 +0,0 @@
|
||||
# Helios Video Generation
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/helios>.
|
||||
|
||||
|
||||
Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
|
||||
|
||||
| Variant | Description | Key Features |
|
||||
|---------|-------------|--------------|
|
||||
| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
|
||||
| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
|
||||
| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Run Examples
|
||||
|
||||
Get into the example folder:
|
||||
```bash
|
||||
cd examples/offline_inference/helios
|
||||
```
|
||||
|
||||
### Text-to-Video (T2V)
|
||||
|
||||
**Helios-Base** (Stage 1 only):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_t2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_t2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled** (Stage 2 pyramid + DMD):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_t2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_i2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_i2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_i2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Video-to-Video (V2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_v2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_v2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_v2v_distilled.mp4
|
||||
```
|
||||
|
||||
## Common Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
|
||||
| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
|
||||
| `--prompt` | — | Text prompt describing the video |
|
||||
| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
|
||||
| `--image-path` | — | Input image (required for `i2v`) |
|
||||
| `--video-path` | — | Input video (required for `v2v`) |
|
||||
| `--height` | `384` | Video height in pixels |
|
||||
| `--width` | `640` | Video width in pixels |
|
||||
| `--num-frames` | `99` | Number of output frames |
|
||||
| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
|
||||
| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--fps` | `16` | Output video frame rate |
|
||||
| `--output` | `helios_output.mp4` | Output file path |
|
||||
|
||||
### Stage 2 / Pyramid Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
|
||||
| `--pyramid-num-stages` | `3` | Number of pyramid stages |
|
||||
| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
|
||||
| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
|
||||
|
||||
### CFG-Zero* Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
|
||||
| `--use-zero-init` | off | Zero init for first steps |
|
||||
| `--zero-steps` | `1` | Number of zero-init steps |
|
||||
|
||||
### Memory & Parallelism
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--vae-use-slicing` | off | Enable VAE slicing |
|
||||
| `--vae-use-tiling` | off | Enable VAE tiling |
|
||||
| `--enforce-eager` | off | Disable torch.compile |
|
||||
| `--enable-cpu-offload` | off | CPU offloading |
|
||||
| `--enable-layerwise-offload` | off | Layerwise offloading |
|
||||
| `--ulysses-degree` | `1` | Ulysses SP degree |
|
||||
| `--ring-degree` | `1` | Ring SP degree |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
|
||||
| `--tensor-parallel-size` | `1` | Tensor parallelism size |
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/helios/end2end.py"
|
||||
``````
|
||||
@@ -0,0 +1,65 @@
|
||||
# AURA Omni Pipeline
|
||||
|
||||
`aura_omni` wires ASR, AURA, and Qwen3-TTS into one vLLM-Omni pipeline:
|
||||
|
||||
```text
|
||||
ASR -> AURA -> Qwen3-TTS Talker -> Code2Wav
|
||||
```
|
||||
|
||||
Qwen3-TTS remains two engine stages so the pipeline reuses the existing native
|
||||
Talker and Code2Wav implementation.
|
||||
|
||||
```bash
|
||||
vllm serve aurateam/AURA \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/aura_omni.yaml \
|
||||
--served-model-name aurateam/AURA \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
Configure local checkpoints by editing per-stage `model` values in
|
||||
`vllm_omni/deploy/aura_omni.yaml`. The deploy file sets
|
||||
`pipeline: aura_omni`, so the four-stage topology is used even if the
|
||||
command-line `--model` points at one of the component checkpoints.
|
||||
|
||||
Send requests with `"model": "aurateam/AURA"`. The ASR, AURA, and Qwen3-TTS
|
||||
checkpoint paths are internal stage models from the deploy YAML, not the
|
||||
OpenAI-facing served model name.
|
||||
|
||||
The AURA stage can emit `<|silent|>`. Silent outputs are treated as a gate:
|
||||
they produce no Qwen3-TTS Talker input, so no audio is synthesized for that
|
||||
turn.
|
||||
|
||||
## GPU Utilization Recommendation
|
||||
|
||||
`gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` controls how much
|
||||
VRAM each stage can reserve. Start with this split for a single GPU:
|
||||
|
||||
- Stage 0 (ASR): `0.10`
|
||||
- Stage 1 (AURA): `0.40`
|
||||
- Stage 2 (Qwen3-TTS Talker): `0.20`
|
||||
- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
|
||||
|
||||
## TTS Modes
|
||||
|
||||
`aura_omni` can pass AURA text to Qwen3-TTS in two task modes:
|
||||
|
||||
- `Base`: voice clone from `tts_ref_audio` with ICL enabled in the AURA
|
||||
pipeline. Provide both `tts_ref_audio` and `tts_ref_text`. Set
|
||||
`tts_x_vector_only_mode=true` to disable ICL and use speaker embedding only.
|
||||
- `CustomVoice`: predefined speaker mode. Use a Qwen3-TTS CustomVoice
|
||||
checkpoint for stages 2 and 3 in `aura_omni.yaml`, then pass
|
||||
`tts_task_type=CustomVoice` and `tts_speaker`.
|
||||
|
||||
By default, AURA responses are passed to Qwen3-TTS as text. Set
|
||||
`additional_information.tts_pass_token_ids=true` to pass AURA-generated
|
||||
assistant token ids directly instead. Even when token passthrough is disabled,
|
||||
the stage processor uses AURA token ids when available to estimate the Talker
|
||||
prefill length, so scheduling does not rely on raw character length.
|
||||
|
||||
The example client exposes this as:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--tts-pass-token-ids
|
||||
```
|
||||
@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
|
||||
|
||||
## Source of Conclusions
|
||||
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on `Qwen3-Omni` and `Wan2.2`.
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/test-weekly.yml`.
|
||||
|
||||
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
|
||||
|
||||
|
||||
@@ -85,7 +85,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | |
|
||||
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
|
||||
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
|
||||
| Cosmos3 | `nvidia/Cosmos3-Nano` | Yes | Not validated | All layers | None | |
|
||||
| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
# AURA Omni
|
||||
|
||||
This example runs the native AURA Omni pipeline offline:
|
||||
|
||||
```text
|
||||
Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav
|
||||
```
|
||||
|
||||
The first stage consumes speech and produces a transcript. The AURA stage then combines the transcript with the video frames from the original request and returns text or `<|silent|>`. Non-silent responses are passed to Qwen3-TTS as AURA-generated token ids.
|
||||
|
||||
## Run
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/aura_omni
|
||||
bash run_single_prompt.sh
|
||||
```
|
||||
|
||||
Use local media:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--audio-path /path/to/input.wav \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--modalities text,audio
|
||||
```
|
||||
|
||||
Base voice clone mode is the default. It uses the AURA reference audio unless
|
||||
you override it:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--tts-task-type Base \
|
||||
--tts-ref-audio /data/yrr/rein_test/shuhan.mp3 \
|
||||
--tts-x-vector-only-mode
|
||||
```
|
||||
|
||||
CustomVoice mode requires stages 2 and 3 in `vllm_omni/deploy/aura_omni.yaml`
|
||||
to point at a Qwen3-TTS CustomVoice checkpoint:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--tts-task-type CustomVoice \
|
||||
--tts-speaker Vivian
|
||||
```
|
||||
|
||||
For local checkpoints, edit the stage `model` entries in
|
||||
`vllm_omni/deploy/aura_omni.yaml` or pass a copied deploy config with
|
||||
`--deploy-config`.
|
||||
|
||||
## GPU Utilization Recommendation
|
||||
|
||||
Set `gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` by stage.
|
||||
Suggested starting point:
|
||||
|
||||
- Stage 0 (ASR): `0.10`
|
||||
- Stage 1 (AURA): `0.4`
|
||||
- Stage 2 (Qwen3-TTS Talker): `0.20`
|
||||
- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
|
||||
|
||||
Generated text and audio are written to `--output-dir`
|
||||
(default: `output_aura_omni`).
|
||||
@@ -0,0 +1,236 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline AURA Omni example.
|
||||
|
||||
The native AURA pipeline is:
|
||||
|
||||
Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav
|
||||
|
||||
Stage 0 consumes audio and emits a transcript. The stage input processor then
|
||||
combines that transcript with the original visual payload for AURA.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import vllm
|
||||
from vllm import SamplingParams
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
default_qwen3_tts_ref_audio_path,
|
||||
)
|
||||
|
||||
SEED = 42
|
||||
DEFAULT_MODEL = "aurateam/AURA"
|
||||
DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/aura_omni.yaml"
|
||||
|
||||
|
||||
class QueryResult(NamedTuple):
|
||||
inputs: dict
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
|
||||
|
||||
def _load_audio(audio_path: str | None, sampling_rate: int) -> tuple[np.ndarray, int]:
|
||||
if audio_path:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
|
||||
return audio_signal.astype(np.float32), int(sr)
|
||||
audio_signal, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
return audio_signal.astype(np.float32), int(sr)
|
||||
|
||||
|
||||
def _load_video(video_path: str | None, num_frames: int) -> list[np.ndarray]:
|
||||
if video_path:
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"Video file not found: {video_path}")
|
||||
return video_to_ndarrays(video_path, num_frames=num_frames)
|
||||
return VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
|
||||
|
||||
|
||||
def build_query(args) -> QueryResult:
|
||||
prompt = (
|
||||
"<|im_start|>user\n"
|
||||
"<|audio_start|><|audio_pad|><|audio_end|>"
|
||||
"Please transcribe this speech and use the video context to decide whether a reply is needed.<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
audio_data = _load_audio(args.audio_path, args.sampling_rate)
|
||||
video_frames = _load_video(args.video_path, args.num_frames)
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"audio": audio_data,
|
||||
"video": video_frames,
|
||||
},
|
||||
"additional_information": {
|
||||
"aura_system_prompt": args.aura_system_prompt,
|
||||
"tts_task_type": args.tts_task_type,
|
||||
"tts_language": args.tts_language,
|
||||
"tts_speaker": args.tts_speaker,
|
||||
"tts_instruct": args.tts_instruct,
|
||||
"tts_ref_audio": args.tts_ref_audio,
|
||||
"tts_ref_text": args.tts_ref_text,
|
||||
"tts_x_vector_only_mode": args.tts_x_vector_only_mode,
|
||||
"tts_pass_token_ids": args.tts_pass_token_ids,
|
||||
},
|
||||
},
|
||||
limit_mm_per_prompt={"audio": 1, "video": 1},
|
||||
)
|
||||
|
||||
|
||||
def build_sampling_params() -> list[SamplingParams]:
|
||||
asr_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=256,
|
||||
detokenize=True,
|
||||
seed=SEED,
|
||||
)
|
||||
aura_sampling = SamplingParams(
|
||||
temperature=0.5,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=256,
|
||||
detokenize=True,
|
||||
repetition_penalty=1.0,
|
||||
seed=SEED,
|
||||
)
|
||||
talker_sampling = SamplingParams(
|
||||
temperature=0.9,
|
||||
top_k=50,
|
||||
max_tokens=4096,
|
||||
detokenize=False,
|
||||
repetition_penalty=1.05,
|
||||
stop_token_ids=[2150],
|
||||
seed=SEED,
|
||||
)
|
||||
code2wav_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=65536,
|
||||
detokenize=True,
|
||||
repetition_penalty=1.0,
|
||||
seed=SEED,
|
||||
)
|
||||
return [asr_sampling, aura_sampling, talker_sampling, code2wav_sampling]
|
||||
|
||||
|
||||
def save_stage_output(stage_outputs, output_dir: str) -> None:
|
||||
output = stage_outputs.request_output
|
||||
request_id = output.request_id
|
||||
if stage_outputs.final_output_type == "text":
|
||||
text_output = output.outputs[0].text
|
||||
out_txt = os.path.join(output_dir, f"{request_id}.txt")
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
f.write(str(text_output).strip() + "\n")
|
||||
print(f"Request ID: {request_id}, text saved to {out_txt}")
|
||||
print(text_output)
|
||||
return
|
||||
|
||||
if stage_outputs.final_output_type == "audio":
|
||||
mm = output.outputs[0].multimodal_output
|
||||
if not mm or "audio" not in mm:
|
||||
print(f"Request ID: {request_id}, no audio output")
|
||||
return
|
||||
audio = mm["audio"]
|
||||
if isinstance(audio, list):
|
||||
import torch
|
||||
|
||||
audio = torch.cat([(item if hasattr(item, "detach") else torch.tensor(item)).flatten() for item in audio])
|
||||
sample_rate = int(mm.get("sr", 24000))
|
||||
audio_numpy = audio.float().detach().cpu().numpy().flatten()
|
||||
out_wav = os.path.join(output_dir, f"{request_id}.wav")
|
||||
sf.write(out_wav, audio_numpy, samplerate=sample_rate, format="WAV")
|
||||
print(f"Request ID: {request_id}, audio saved to {out_wav}")
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
query = build_query(args)
|
||||
prompts = [query.inputs for _ in range(args.num_prompts)]
|
||||
if args.modalities:
|
||||
modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
|
||||
for prompt in prompts:
|
||||
prompt["modalities"] = modalities
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
omni_kwargs["model"] = args.model
|
||||
omni = Omni(**omni_kwargs)
|
||||
sampling_params_list = build_sampling_params()[: omni.num_stages]
|
||||
|
||||
try:
|
||||
for stage_outputs in omni.generate(prompts, sampling_params_list):
|
||||
save_stage_output(stage_outputs, args.output_dir)
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="AURA Omni offline inference example")
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL, help="Top-level model name/path.")
|
||||
parser.add_argument("--deploy-config", default=DEFAULT_DEPLOY_CONFIG, help="AURA Omni deploy config.")
|
||||
parser.add_argument("--audio-path", "-a", default=None, help="Local audio file. Uses a built-in asset by default.")
|
||||
parser.add_argument("--video-path", "-v", default=None, help="Local video file. Uses a built-in asset by default.")
|
||||
parser.add_argument("--num-frames", type=int, default=16, help="Number of video frames to sample.")
|
||||
parser.add_argument("--sampling-rate", type=int, default=16000, help="Audio sampling rate for ASR input.")
|
||||
parser.add_argument("--num-prompts", type=int, default=1, help="Number of repeated prompts.")
|
||||
parser.add_argument("--modalities", default="text,audio", help="Comma-separated final modalities.")
|
||||
parser.add_argument("--output-dir", default="output_aura_omni", help="Directory for generated files.")
|
||||
parser.add_argument(
|
||||
"--aura-system-prompt",
|
||||
default=(
|
||||
"You are receiving a live video stream where the final frame is the present moment. "
|
||||
"Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in Chinese."
|
||||
),
|
||||
help="System prompt for the AURA stage.",
|
||||
)
|
||||
parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"], help="Qwen3-TTS task.")
|
||||
parser.add_argument("--tts-language", default="Chinese", help="Qwen3-TTS language.")
|
||||
parser.add_argument("--tts-speaker", default="Vivian", help="Qwen3-TTS speaker.")
|
||||
parser.add_argument("--tts-instruct", default="", help="Optional Qwen3-TTS style instruction.")
|
||||
parser.add_argument(
|
||||
"--tts-ref-audio",
|
||||
default=default_qwen3_tts_ref_audio_path(),
|
||||
help="Base-mode reference audio path/URL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-ref-text",
|
||||
default=DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
help="Base-mode reference audio transcript.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-x-vector-only-mode",
|
||||
action="store_true",
|
||||
help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-pass-token-ids",
|
||||
action="store_true",
|
||||
help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.",
|
||||
)
|
||||
parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
python end2end.py \
|
||||
--model aurateam/AURA \
|
||||
--deploy-config vllm_omni/deploy/aura_omni.yaml \
|
||||
--modalities text,audio \
|
||||
"$@"
|
||||
@@ -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,193 +0,0 @@
|
||||
# Helios Video Generation
|
||||
|
||||
Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
|
||||
|
||||
| Variant | Description | Key Features |
|
||||
|---------|-------------|--------------|
|
||||
| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
|
||||
| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
|
||||
| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Run Examples
|
||||
|
||||
Get into the example folder:
|
||||
```bash
|
||||
cd examples/offline_inference/helios
|
||||
```
|
||||
|
||||
### Text-to-Video (T2V)
|
||||
|
||||
**Helios-Base** (Stage 1 only):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_t2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_t2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled** (Stage 2 pyramid + DMD):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_t2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_i2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_i2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_i2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Video-to-Video (V2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_v2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_v2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_v2v_distilled.mp4
|
||||
```
|
||||
|
||||
## Common Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
|
||||
| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
|
||||
| `--prompt` | — | Text prompt describing the video |
|
||||
| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
|
||||
| `--image-path` | — | Input image (required for `i2v`) |
|
||||
| `--video-path` | — | Input video (required for `v2v`) |
|
||||
| `--height` | `384` | Video height in pixels |
|
||||
| `--width` | `640` | Video width in pixels |
|
||||
| `--num-frames` | `99` | Number of output frames |
|
||||
| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
|
||||
| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--fps` | `16` | Output video frame rate |
|
||||
| `--output` | `helios_output.mp4` | Output file path |
|
||||
|
||||
### Stage 2 / Pyramid Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
|
||||
| `--pyramid-num-stages` | `3` | Number of pyramid stages |
|
||||
| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
|
||||
| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
|
||||
|
||||
### CFG-Zero* Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
|
||||
| `--use-zero-init` | off | Zero init for first steps |
|
||||
| `--zero-steps` | `1` | Number of zero-init steps |
|
||||
|
||||
### Memory & Parallelism
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--vae-use-slicing` | off | Enable VAE slicing |
|
||||
| `--vae-use-tiling` | off | Enable VAE tiling |
|
||||
| `--enforce-eager` | off | Disable torch.compile |
|
||||
| `--enable-cpu-offload` | off | CPU offloading |
|
||||
| `--enable-layerwise-offload` | off | Layerwise offloading |
|
||||
| `--ulysses-degree` | `1` | Ulysses SP degree |
|
||||
| `--ring-degree` | `1` | Ring SP degree |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
|
||||
| `--tensor-parallel-size` | `1` | Tensor parallelism size |
|
||||
@@ -1,355 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Helios video generation example supporting T2V, I2V, and V2V.
|
||||
|
||||
Usage (T2V, Helios-Base, Stage 1 only):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--height 384 --width 640 --num-frames 99 \
|
||||
--num-inference-steps 50 --guidance-scale 5.0
|
||||
|
||||
Usage (I2V, Helios-Base):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "Description of desired animation." \
|
||||
--guidance-scale 5.0
|
||||
|
||||
Usage (V2V, Helios-Base):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "Description of desired transformation." \
|
||||
--guidance-scale 5.0
|
||||
|
||||
Usage (Helios-Mid, Stage 2 + CFG-Zero*):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Mid --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1
|
||||
|
||||
Usage (Helios-Distilled, Stage 2 pyramid + DMD):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Distilled --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor:
|
||||
"""Load an image and return a (1, C, H, W) tensor normalized to [-1, 1]."""
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
|
||||
image = Image.open(image_path).convert("RGB").resize((width, height))
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
return transform(image).unsqueeze(0)
|
||||
|
||||
|
||||
def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor:
|
||||
"""Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1]."""
|
||||
from torchvision import transforms
|
||||
from torchvision.io import read_video
|
||||
|
||||
video_frames, _, _ = read_video(video_path, output_format="TCHW")
|
||||
video_frames = video_frames.float() / 255.0
|
||||
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize((height, width), antialias=True),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
video_frames = torch.stack([transform(f) for f in video_frames])
|
||||
return video_frames.permute(1, 0, 2, 3).unsqueeze(0)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="BestWishYsh/Helios-Base",
|
||||
help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample-type",
|
||||
choices=["t2v", "i2v", "v2v"],
|
||||
default="t2v",
|
||||
help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).",
|
||||
)
|
||||
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
||||
parser.add_argument(
|
||||
"--negative-prompt",
|
||||
default=(
|
||||
"Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "
|
||||
"images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, "
|
||||
"incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "
|
||||
"misshapen limbs, fused fingers, still picture, messy background, three legs, many people "
|
||||
"in the background, walking backwards"
|
||||
),
|
||||
help="Negative prompt.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
|
||||
parser.add_argument("--height", type=int, default=384, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=640, help="Video width.")
|
||||
parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).")
|
||||
parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
|
||||
|
||||
# I2V / V2V
|
||||
parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.")
|
||||
parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.")
|
||||
|
||||
# Stage 2 (pyramid multi-stage denoising)
|
||||
parser.add_argument(
|
||||
"--is-enable-stage2",
|
||||
action="store_true",
|
||||
help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pyramid-num-stages",
|
||||
type=int,
|
||||
default=3,
|
||||
help="Number of pyramid stages for Stage 2.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pyramid-num-inference-steps-list",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[10, 10, 10],
|
||||
help="Inference steps per pyramid stage.",
|
||||
)
|
||||
|
||||
# DMD
|
||||
parser.add_argument(
|
||||
"--is-amplify-first-chunk",
|
||||
action="store_true",
|
||||
help="Enable DMD amplification for the first chunk (Helios-Distilled).",
|
||||
)
|
||||
|
||||
# CFG Zero Star
|
||||
parser.add_argument(
|
||||
"--use-cfg-zero-star",
|
||||
action="store_true",
|
||||
help="Enable CFG Zero Star guidance (recommended for Helios-Mid).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-zero-init",
|
||||
action="store_true",
|
||||
help="Use zero initialization for the first denoising steps with CFG-Zero*.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--zero-steps",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of initial denoising steps using zero prediction (default: 1).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Quantization method (e.g. 'fp8').",
|
||||
)
|
||||
|
||||
# Memory & parallelism
|
||||
parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.")
|
||||
parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.")
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
|
||||
parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.")
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
enable_layerwise_offload=args.enable_layerwise_offload,
|
||||
vae_use_slicing=args.vae_use_slicing,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
quantization=args.quantization,
|
||||
)
|
||||
|
||||
# Validate I2V / V2V arguments
|
||||
if args.sample_type == "i2v" and not args.image_path:
|
||||
raise ValueError("--image-path is required for I2V mode.")
|
||||
if args.sample_type == "v2v" and not args.video_path:
|
||||
raise ValueError("--video-path is required for V2V mode.")
|
||||
|
||||
# Build extra_args for Helios-specific parameters
|
||||
extra_args = {}
|
||||
if args.is_enable_stage2:
|
||||
extra_args["is_enable_stage2"] = True
|
||||
extra_args["pyramid_num_stages"] = args.pyramid_num_stages
|
||||
extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list
|
||||
if args.is_amplify_first_chunk:
|
||||
extra_args["is_amplify_first_chunk"] = True
|
||||
if args.use_cfg_zero_star:
|
||||
extra_args["use_cfg_zero_star"] = True
|
||||
if args.use_zero_init:
|
||||
extra_args["use_zero_init"] = True
|
||||
extra_args["zero_steps"] = args.zero_steps
|
||||
|
||||
if args.sample_type == "i2v":
|
||||
image_tensor = load_image_as_tensor(args.image_path, args.height, args.width)
|
||||
extra_args["image"] = image_tensor
|
||||
elif args.sample_type == "v2v":
|
||||
video_tensor = load_video_as_tensor(args.video_path, args.height, args.width)
|
||||
extra_args["video"] = video_tensor
|
||||
|
||||
# Print generation configuration
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Helios Generation Configuration:")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Sample type: {args.sample_type.upper()}")
|
||||
if args.sample_type == "i2v":
|
||||
print(f" Image: {args.image_path}")
|
||||
elif args.sample_type == "v2v":
|
||||
print(f" Video: {args.video_path}")
|
||||
print(f" Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}")
|
||||
print(f" Video size: {args.width}x{args.height}, {args.num_frames} frames")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Guidance scale: {args.guidance_scale}")
|
||||
if args.is_enable_stage2:
|
||||
print(f" Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})")
|
||||
if args.is_amplify_first_chunk:
|
||||
print(" DMD amplify first chunk: enabled")
|
||||
if args.use_cfg_zero_star:
|
||||
print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
|
||||
else:
|
||||
if args.use_cfg_zero_star:
|
||||
print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
|
||||
print(" Stage 2: disabled (Stage 1 only)")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
frames = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
guidance_scale=args.guidance_scale,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_frames=args.num_frames,
|
||||
extra_args=extra_args,
|
||||
),
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
|
||||
|
||||
# Extract video frames from OmniRequestOutput
|
||||
if isinstance(frames, list) and len(frames) > 0:
|
||||
first_item = frames[0]
|
||||
|
||||
if hasattr(first_item, "final_output_type"):
|
||||
if first_item.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
|
||||
if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
|
||||
inner_output = first_item.request_output
|
||||
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
|
||||
frames = inner_output.images[0] if inner_output.images else None
|
||||
if frames is None:
|
||||
raise ValueError("No video frames found in output.")
|
||||
elif hasattr(first_item, "images") and first_item.images:
|
||||
frames = first_item.images
|
||||
else:
|
||||
raise ValueError("No video frames found in OmniRequestOutput.")
|
||||
|
||||
# Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C)
|
||||
if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray):
|
||||
frames = frames[0]
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
from diffusers.utils import export_to_video
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for export_to_video.")
|
||||
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video_tensor = frames.detach().cpu()
|
||||
if video_tensor.dim() == 5:
|
||||
if video_tensor.shape[1] in (3, 4):
|
||||
video_tensor = video_tensor[0].permute(1, 2, 3, 0)
|
||||
else:
|
||||
video_tensor = video_tensor[0]
|
||||
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
|
||||
video_tensor = video_tensor.permute(1, 2, 3, 0)
|
||||
if video_tensor.is_floating_point():
|
||||
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
video_array = video_tensor.float().numpy()
|
||||
elif isinstance(frames, np.ndarray):
|
||||
video_array = frames
|
||||
else:
|
||||
video_array = frames
|
||||
|
||||
if isinstance(video_array, np.ndarray) and video_array.ndim == 5:
|
||||
video_array = video_array[0]
|
||||
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -74,6 +74,12 @@ def parse_args():
|
||||
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
|
||||
parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
|
||||
parser.add_argument(
|
||||
"--stream",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable streaming CoT display; print AR text token-by-token in real time.",
|
||||
)
|
||||
parser.add_argument("--log-stats", action="store_true", default=False)
|
||||
parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
@@ -222,6 +228,8 @@ def main():
|
||||
prompt_dict["modalities"] = ["text"]
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
|
||||
params_list = list(omni.default_sampling_params_list)
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -253,6 +261,10 @@ def main():
|
||||
sp.width = args.width
|
||||
elif hasattr(sp, "stop_token_ids"):
|
||||
sp.stop_token_ids = ar_stop_token_ids
|
||||
# When --stream is set, request DELTA output from the AR stage
|
||||
# so we can display CoT text token-by-token in real time.
|
||||
if args.stream and hasattr(sp, "output_kind"):
|
||||
sp.output_kind = RequestOutputKind.DELTA
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print("HunyuanImage-3.0 Generation Configuration:")
|
||||
@@ -281,22 +293,32 @@ def main():
|
||||
print(f" Prompts: {prompts}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
# When --stream is set, print AR CoT text token-by-token in real time.
|
||||
# Otherwise, collect and print the full AR text once when stage 0 finishes.
|
||||
omni_outputs = omni.generate(
|
||||
prompts=formatted_prompts,
|
||||
sampling_params_list=params_list,
|
||||
py_generator=True,
|
||||
use_tqdm=False,
|
||||
)
|
||||
img_idx = 0
|
||||
for req_output in omni_outputs:
|
||||
ro = getattr(req_output, "request_output", None)
|
||||
txt = ""
|
||||
if ro and getattr(ro, "outputs", None):
|
||||
txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
|
||||
if not txt:
|
||||
ar_text = getattr(req_output, "custom_output", {}).get("ar_generated_text")
|
||||
if isinstance(ar_text, list):
|
||||
txt = "\n".join(text for text in ar_text if text)
|
||||
else:
|
||||
txt = ar_text or ""
|
||||
if txt:
|
||||
print(f"[Output] Text:\n{txt}")
|
||||
stage_id = getattr(req_output, "stage_id", None)
|
||||
|
||||
# AR stage text — each CompletionOutput.text is already a delta when
|
||||
# output_kind=DELTA, so we can print it directly (matching the pattern
|
||||
# in serving_chat.py).
|
||||
if stage_id == 0 and ro and getattr(ro, "outputs", None):
|
||||
for o in ro.outputs:
|
||||
text = getattr(o, "text", "") or ""
|
||||
if text:
|
||||
print(text, end="", flush=True)
|
||||
# Non-streaming: one shot with full text — emit a trailing newline.
|
||||
if not args.stream:
|
||||
print(flush=True)
|
||||
|
||||
# Collect images from diffusion stage
|
||||
images = getattr(req_output, "images", None)
|
||||
if not images and ro and hasattr(ro, "images"):
|
||||
images = ro.images
|
||||
@@ -304,7 +326,7 @@ def main():
|
||||
for j, img in enumerate(images):
|
||||
save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
|
||||
img.save(save_path)
|
||||
print(f"[Output] Saved image to {save_path}")
|
||||
print(f"\n[Output] Saved image to {save_path}")
|
||||
img_idx += 1
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Image-To-Video
|
||||
|
||||
This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API.
|
||||
This shared example generates videos from images with VACE, Wan2.2, LTX2,
|
||||
HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -48,19 +49,38 @@ python image_to_video.py \
|
||||
--output i2v_output.mp4
|
||||
```
|
||||
|
||||
### Cosmos3
|
||||
|
||||
```bash
|
||||
# Cosmos3 bundles example frames under assets/ (any RGB image works too):
|
||||
python image_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
|
||||
--prompt "The scene comes to life with smooth, natural motion." \
|
||||
--negative-prompt "blurry, distorted, low quality" \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
|
||||
--output cosmos3_i2v.mp4
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
|
||||
- `--image`: Path to input image (required).
|
||||
- `--image`: Path to the first-frame or source image.
|
||||
- `--last-image`: Optional last-frame condition for models such as VACE.
|
||||
- `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
|
||||
- `--reference-image`: Optional reference image; repeat it to provide multiple references.
|
||||
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3.
|
||||
- `--prompt`: Text description of desired motion/animation.
|
||||
- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
|
||||
- `--num-frames`: Number of frames (default 81).
|
||||
- `--num-frames`: Number of frames (default: model-specific — Wan 81, LTX2 121, Cosmos3 189).
|
||||
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
|
||||
- `--negative-prompt`: Optional list of artifacts to suppress.
|
||||
- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
|
||||
- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
|
||||
- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0).
|
||||
- `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
|
||||
- `--num-inference-steps`: Number of denoising steps (default 50).
|
||||
- `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
|
||||
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
|
||||
- `--output`: Path to save the generated video.
|
||||
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
|
||||
@@ -76,5 +96,98 @@ Key arguments:
|
||||
|
||||
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
|
||||
## Wan2.1 VACE Conditional Tasks
|
||||
|
||||
The shared script selects the VACE conditioning structure from the media inputs.
|
||||
No explicit mode parameter is required: the script constructs the source video,
|
||||
mask, or reference images consumed by the VACE pipeline.
|
||||
Download the same Hugging Face assets used by the original VACE example:
|
||||
|
||||
```bash
|
||||
wget -O astronaut.jpg https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
|
||||
wget -O vace_first_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png
|
||||
wget -O vace_last_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_i2v_output.mp4
|
||||
```
|
||||
|
||||
### Video-to-Last-Frame (V2LF)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--last-image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_v2lf_output.mp4
|
||||
```
|
||||
|
||||
### First-Last-Frame-to-Video (FLF2V)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image vace_first_frame.png \
|
||||
--last-image vace_last_frame.png \
|
||||
--prompt "CG animation style, a small blue bird takes off from a branch and lands on another branch" \
|
||||
--seed 42 --height 512 --width 512 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_flf2v_output.mp4
|
||||
```
|
||||
|
||||
### Inpainting
|
||||
|
||||
Create a mask matching the original VACE example: a 160-pixel-wide white
|
||||
vertical stripe marks the region to regenerate.
|
||||
|
||||
```bash
|
||||
python - <<'PY'
|
||||
from PIL import Image
|
||||
|
||||
mask = Image.new("L", (832, 480), 0)
|
||||
mask.paste(255, (336, 0, 496, 480))
|
||||
mask.save("vace_center_mask.png")
|
||||
PY
|
||||
```
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image astronaut.jpg \
|
||||
--mask-image vace_center_mask.png \
|
||||
--prompt "Shrek, the ogre, walks out of a building in a happy mood" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_inpaint_output.mp4
|
||||
```
|
||||
|
||||
### Reference-to-Video (R2V)
|
||||
|
||||
Repeat `--reference-image` to provide more than one reference image.
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--reference-image astronaut.jpg \
|
||||
--prompt "Camera slowly zooms out from the character walking in a garden" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_r2v_output.mp4
|
||||
```
|
||||
|
||||
The VACE T2V command is documented in the shared
|
||||
[`text_to_video.py`](../text_to_video/text_to_video.md#wan21-vace-t2v) example.
|
||||
|
||||
For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA
|
||||
assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly).
|
||||
@@ -2,13 +2,14 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5.
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE.
|
||||
|
||||
Supports:
|
||||
- Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
|
||||
- Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
|
||||
- LTX2 image-to-video pipeline
|
||||
- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning
|
||||
- Wan2.1 VACE: first/last-frame, inpainting, and reference conditioning
|
||||
|
||||
Usage:
|
||||
# Wan I2V-A14B (MoE)
|
||||
@@ -30,9 +31,16 @@ Usage:
|
||||
python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
|
||||
--image input.jpg --prompt "A cat playing with yarn" \
|
||||
--flow-shift 5.0 --guidance-scale 6.0
|
||||
|
||||
# Cosmos3 I2V (image conditioning)
|
||||
python image_to_video.py --model nvidia/Cosmos3-Nano \
|
||||
--image input.jpg --prompt "The scene comes to life with smooth, natural motion." \
|
||||
--num-frames 189 --num-inference-steps 35 --guidance-scale 6.0 --fps 24 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}'
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -43,39 +51,59 @@ import PIL.Image
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import (
|
||||
build_image_to_video_prompt,
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
|
||||
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a video from one or more images (Wan2.2, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).",
|
||||
help="Diffusers I2V model ID or local path (Wan2.2, HunyuanVideo-1.5, or Wan2.1 VACE).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-class-name",
|
||||
default=None,
|
||||
help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
|
||||
)
|
||||
parser.add_argument("--image", required=True, help="Path to input image.")
|
||||
parser.add_argument("--image", help="Path to the first-frame or source image.")
|
||||
parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).")
|
||||
parser.add_argument("--mask-image", help="Path to an inpainting mask (used by models such as VACE).")
|
||||
parser.add_argument(
|
||||
"--reference-image",
|
||||
action="append",
|
||||
default=None,
|
||||
help="Path to a reference image. Repeat to provide multiple references.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
"--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)."
|
||||
)
|
||||
@@ -83,8 +111,10 @@ def parse_args() -> argparse.Namespace:
|
||||
"--height", type=int, default=None, help="Video height (auto-calculated from image if not set)."
|
||||
)
|
||||
parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).")
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
|
||||
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
"--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
|
||||
)
|
||||
parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
|
||||
parser.add_argument(
|
||||
"--frame-rate",
|
||||
@@ -93,7 +123,10 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
|
||||
"--flow-shift",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Scheduler flow_shift. Default: model-specific (Wan 5.0, Cosmos3 10.0).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample-solver",
|
||||
@@ -244,6 +277,19 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=functools.partial(parse_json_object, flag_name="--extra-body"),
|
||||
default=None,
|
||||
help=(
|
||||
"Model-specific generation params as a JSON object. Keys are filtered "
|
||||
"against the model's declared extra_body_params (see vllm_omni/model_extras), "
|
||||
"so unknown keys for the chosen model are silently dropped. "
|
||||
'Cosmos3 V2V example: \'{"condition_frame_indexes_vision": [0, 1], '
|
||||
'"condition_video_keep": "first", "flow_shift": 10.0, '
|
||||
'"max_sequence_length": 4096, "guardrails": false}\'.'
|
||||
),
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -267,31 +313,63 @@ def main():
|
||||
model_name = str(args.model).lower() if args.model is not None else ""
|
||||
model_class_name = args.model_class_name
|
||||
is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
|
||||
is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
|
||||
if model_class_name is None and is_ltx2:
|
||||
model_class_name = "LTX2ImageToVideoPipeline"
|
||||
|
||||
# Load input image
|
||||
image = PIL.Image.open(args.image).convert("RGB")
|
||||
image = PIL.Image.open(args.image).convert("RGB") if args.image else None
|
||||
last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
|
||||
mask_image = PIL.Image.open(args.mask_image).convert("L") if args.mask_image else None
|
||||
reference_images = (
|
||||
[PIL.Image.open(path).convert("RGB") for path in args.reference_image] if args.reference_image else None
|
||||
)
|
||||
dimension_image = image or last_image or (reference_images[0] if reference_images else None)
|
||||
if dimension_image is None:
|
||||
raise ValueError("Provide --image, --last-image, or at least one --reference-image.")
|
||||
|
||||
fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16)
|
||||
# Per-model generation defaults, applied only when the matching flag is omitted.
|
||||
# Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
|
||||
if is_cosmos:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
|
||||
24,
|
||||
6.0,
|
||||
189,
|
||||
35,
|
||||
10.0,
|
||||
1280 * 720,
|
||||
16,
|
||||
)
|
||||
elif is_ltx2:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 24, 4.0, 121, 40, 5.0, 512 * 768, 32
|
||||
else: # Wan2.2 / HunyuanVideo-1.5
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16
|
||||
|
||||
fps = args.fps if args.fps is not None else d_fps
|
||||
frame_rate = args.frame_rate if args.frame_rate is not None else float(fps)
|
||||
guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0)
|
||||
num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81)
|
||||
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50)
|
||||
guidance_scale = args.guidance_scale if args.guidance_scale is not None else d_guidance
|
||||
num_frames = args.num_frames if args.num_frames is not None else d_num_frames
|
||||
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else d_steps
|
||||
flow_shift = args.flow_shift if args.flow_shift is not None else d_flow_shift
|
||||
|
||||
# Calculate dimensions if not provided
|
||||
# Calculate dimensions if not provided (model-aware max area).
|
||||
height = args.height
|
||||
width = args.width
|
||||
if height is None or width is None:
|
||||
# Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2
|
||||
max_area = 512 * 768 if is_ltx2 else 480 * 832
|
||||
mod_value = 32 if is_ltx2 else 16
|
||||
calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value)
|
||||
calc_height, calc_width = calculate_dimensions(dimension_image, max_area=d_max_area, mod_value=d_mod)
|
||||
height = height or calc_height
|
||||
width = width or calc_width
|
||||
|
||||
# Resize image to target dimensions
|
||||
image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
media_inputs: dict[str, Any] = {}
|
||||
if image is not None:
|
||||
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if last_image is not None:
|
||||
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if mask_image is not None:
|
||||
media_inputs["mask"] = mask_image.resize((width, height), PIL.Image.Resampling.NEAREST)
|
||||
if reference_images is not None:
|
||||
media_inputs["reference_images"] = [
|
||||
reference.resize((width, height), PIL.Image.Resampling.LANCZOS) for reference in reference_images
|
||||
]
|
||||
|
||||
# Configure cache based on backend type
|
||||
cache_config = None
|
||||
@@ -342,7 +420,7 @@ def main():
|
||||
vae_use_slicing=args.vae_use_slicing,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
flow_shift=args.flow_shift,
|
||||
flow_shift=flow_shift,
|
||||
diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
|
||||
diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
|
||||
diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
|
||||
@@ -357,7 +435,13 @@ def main():
|
||||
)
|
||||
if args.quantization is not None:
|
||||
omni_kwargs["quantization"] = args.quantization
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni) or model_class_name
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -367,8 +451,8 @@ def main():
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Generation Configuration:")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Frames: {args.num_frames}")
|
||||
print(f" Inference steps: {num_inference_steps}")
|
||||
print(f" Frames: {num_frames}")
|
||||
print(f" Solver: {args.sample_solver}")
|
||||
print(f" diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}")
|
||||
print(f" diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}")
|
||||
@@ -381,30 +465,43 @@ def main():
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
prompt_dict = build_image_to_video_prompt(
|
||||
model_class_name=model_class_name,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
media_inputs=media_inputs,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
)
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
generator=generator,
|
||||
guidance_scale=guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
num_inference_steps=num_inference_steps,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
extra_args={
|
||||
"sample_solver": args.sample_solver,
|
||||
"flow_shift": flow_shift,
|
||||
},
|
||||
)
|
||||
|
||||
# Route model-specific knobs through extra_body, filtered against the model's
|
||||
# declared extra_body_params. Models without a declaration only forward explicit
|
||||
# --extra-body JSON (e.g. Cosmos3 V2V's condition_frame_indexes_vision).
|
||||
extra_body = dict(args.extra_body or {})
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
|
||||
elif extra_body:
|
||||
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
# omni.generate() returns Generator[OmniRequestOutput, None, None]
|
||||
frames = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"multi_modal_data": {"image": image},
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
generator=generator,
|
||||
guidance_scale=guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
num_inference_steps=num_inference_steps,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
extra_args={
|
||||
"sample_solver": args.sample_solver,
|
||||
"flow_shift": args.flow_shift,
|
||||
},
|
||||
),
|
||||
)
|
||||
frames = omni.generate(prompt_dict, sampling_params)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
|
||||
@@ -1,72 +0,0 @@
|
||||
# MagiHuman Generation
|
||||
|
||||
MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt.
|
||||
|
||||
Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks.
|
||||
|
||||
## Setup
|
||||
|
||||
### Install MagiCompiler (recommended)
|
||||
|
||||
MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour.
|
||||
|
||||
```bash
|
||||
# Clone the repo
|
||||
git clone https://github.com/SandAI-org/MagiCompiler.git
|
||||
cd MagiCompiler
|
||||
|
||||
# System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
|
||||
sudo apt update && sudo apt install -y graphviz
|
||||
|
||||
# Python dependencies
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Install MagiCompiler
|
||||
pip install . # end users (recommended)
|
||||
# pip install -e . # developers (editable install)
|
||||
```
|
||||
|
||||
### Hardware requirements
|
||||
|
||||
Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference).
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory.
|
||||
|
||||
## Run Examples
|
||||
|
||||
Get into the example folder:
|
||||
```bash
|
||||
cd examples/offline_inference/magi_human
|
||||
```
|
||||
|
||||
### End-to-End Generation (Text to Video+Audio)
|
||||
|
||||
Generate a video with synchronized speech natively generated by the model.
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \
|
||||
--prompt "A young woman with long, wavy golden blonde hair..." \
|
||||
--tensor-parallel-size 4 \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
## Common Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--model` | *(Required)* | Local model path or HuggingFace ID |
|
||||
| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text |
|
||||
| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) |
|
||||
| `--height` | `256` | Initial resolution height |
|
||||
| `--width` | `448` | Initial resolution width |
|
||||
| `--num-inference-steps` | `8` | Denoising steps |
|
||||
| `--seed` | `52` | Random seed |
|
||||
| `--output` | `output_magihuman.mp4` | Output video with audio path |
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/magi_human/end2end.py"
|
||||
``````
|
||||
@@ -1,122 +0,0 @@
|
||||
import argparse
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.")
|
||||
parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.")
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
default="",
|
||||
help="Text prompt containing visual description, dialogue, and background sound.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file."
|
||||
)
|
||||
parser.add_argument("--height", type=int, default=256, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=448, help="Video width.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
|
||||
parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...")
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
init_timeout=1200,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
devices=list(range(args.tensor_parallel_size)),
|
||||
)
|
||||
|
||||
prompt = args.prompt
|
||||
if not prompt:
|
||||
prompt = (
|
||||
"A young woman with long, wavy golden blonde hair and bright blue eyes, "
|
||||
"wearing a fitted ivory silk blouse with a delicate lace collar, sits "
|
||||
"stationary in front of a softly lit, blurred warm-toned interior. Her "
|
||||
"overall disposition is warm, composed, and gently confident. The camera "
|
||||
"holds a static medium close-up, framing her from the shoulders up, "
|
||||
"with shallow depth of field keeping her face in sharp focus. Soft "
|
||||
"directional key light falls from the upper left, casting a gentle "
|
||||
"highlight along her cheekbone and nose bridge. She draws a quiet breath, "
|
||||
"the levator labii superiors relaxing as her lips part. She speaks in "
|
||||
"clear, warm, unhurried American English: "
|
||||
"\"The most beautiful things in life aren't things at all — "
|
||||
"they're moments, feelings, and the people who make you feel truly alive.\" "
|
||||
"Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
|
||||
"shapes each vowel with precision. A faint, genuine smile engages the "
|
||||
"zygomaticus major, lifting her lip corners fractionally. Her brows rest "
|
||||
"in a soft, neutral arch throughout. She maintains steady, forward-facing "
|
||||
"eye contact. Head position remains level; no torso displacement occurs.\n\n"
|
||||
"Dialogue:\n"
|
||||
"<Young blonde woman, American English>: "
|
||||
"\"The most beautiful things in life aren't things at all — "
|
||||
"they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
|
||||
"Background Sound:\n"
|
||||
"<Soft, warm indoor ambience with a faint distant piano melody>"
|
||||
)
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"seconds": 5,
|
||||
"sr_height": 1080,
|
||||
"sr_width": 1920,
|
||||
"sr_num_inference_steps": 5,
|
||||
},
|
||||
)
|
||||
|
||||
print(f"Generating with prompt: {prompt[:80]}...")
|
||||
outputs = omni.generate(
|
||||
prompts=[prompt],
|
||||
sampling_params_list=[sampling_params],
|
||||
)
|
||||
|
||||
print(f"Generation complete. Output type: {type(outputs)}")
|
||||
if outputs:
|
||||
first = outputs[0]
|
||||
|
||||
if hasattr(first, "images") and first.images:
|
||||
video_frames = first.images[0]
|
||||
print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
|
||||
|
||||
audio_waveform = None
|
||||
mm = first.multimodal_output or {}
|
||||
if mm:
|
||||
audio_waveform = mm.get("audio")
|
||||
if audio_waveform is not None:
|
||||
print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
|
||||
|
||||
output_fps = float(mm.get("fps", 25))
|
||||
output_sr = int(mm.get("audio_sample_rate", 24000))
|
||||
print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output")
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
video_frames,
|
||||
audio_waveform,
|
||||
fps=output_fps,
|
||||
audio_sample_rate=output_sr,
|
||||
)
|
||||
with open(args.output, "wb") as f:
|
||||
f.write(video_bytes)
|
||||
print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}")
|
||||
print("SUCCESS: MagiHuman pipeline generation completed.")
|
||||
else:
|
||||
print("WARNING: No outputs returned.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -38,16 +39,20 @@ def is_nextstep_model(model_name: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
|
||||
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
|
||||
parser.add_argument(
|
||||
@@ -269,7 +274,7 @@ def parse_args() -> argparse.Namespace:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=parse_profiler_config,
|
||||
type=functools.partial(parse_json_object, flag_name="--extra-body"),
|
||||
default=None,
|
||||
help=(
|
||||
"Model-specific generation params as a JSON object, e.g. "
|
||||
@@ -425,6 +430,10 @@ def main():
|
||||
if use_nextstep:
|
||||
# NextStep-1.1 requires explicit pipeline class
|
||||
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni)
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
|
||||
|
||||
## Common Quick Start
|
||||
@@ -369,6 +370,20 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
```
|
||||
Streaming requires `async_chunk: true` in the stage config.
|
||||
|
||||
### Word Timestamps
|
||||
Generate a WAV offline and a JSON sidecar with word-level timestamps from
|
||||
`Qwen/Qwen3-ForcedAligner-0.6B`:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \
|
||||
--text "Hello world." \
|
||||
--output-dir /tmp/qwen3_tts_timestamps
|
||||
```
|
||||
The script writes `qwen3_tts_word_timestamps.wav` and
|
||||
`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit,
|
||||
set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT.
|
||||
|
||||
### Batched decoding
|
||||
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):
|
||||
```bash
|
||||
@@ -422,6 +437,50 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2
|
||||
|
||||
2-stage TTS pipeline (GPT AR talker + S2Mel CFM DiT + BigVGAN vocoder) at 22.05 kHz. Every request requires reference audio for zero-shot voice cloning. Supports emotion conditioning via audio, text, or 8-dim vector.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
```
|
||||
|
||||
### Emotion control
|
||||
```bash
|
||||
# Emotion from reference audio
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
|
||||
# Emotion from 8-dim vector (happy angry sad afraid disgusted melancholy surprised calm)
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-vector 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
|
||||
|
||||
# Emotion from text description
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-text "happy and excited"
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `--ref-audio` is **required** — IndexTTS-2 does not support text-only synthesis.
|
||||
- Stage 0 (AR Talker): GPT-2 generates mel codes from text + reference audio.
|
||||
- Stage 1 (S2Mel + BigVGAN): CFM DiT converts mel codes to waveform at 22.05 kHz.
|
||||
- Deploy config: `vllm_omni/deploy/indextts2.yaml`. Stage 1 runs with `enforce_eager: true` (DiT has dynamic shapes).
|
||||
|
||||
---
|
||||
|
||||
## Voxtral TTS
|
||||
|
||||
Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.
|
||||
@@ -462,3 +521,100 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
|
||||
- `--num-prompts N` replicates the prompt for performance measurement.
|
||||
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
|
||||
- Run `--help` for the full argument surface.
|
||||
|
||||
---
|
||||
|
||||
## SoulX-Singer
|
||||
|
||||
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it.
|
||||
|
||||
```bash
|
||||
# 1. DiT weights
|
||||
export BASE=path/to/SoulX-Singer
|
||||
export PREPROCESS=path/to/SoulX-Singer-Preprocess
|
||||
export SVC_DIR=path/to/SoulX-Singer-svc
|
||||
|
||||
huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE"
|
||||
|
||||
# 2. Preprocess weights (required)
|
||||
huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS"
|
||||
export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS"
|
||||
|
||||
# 3. SVS / SVC view directories
|
||||
mkdir -p "$SVC_DIR"
|
||||
cp $BASE/{config.yaml,README.md,assets} $SVC_DIR
|
||||
mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt
|
||||
|
||||
cat > "$BASE/config.json" <<'EOF'
|
||||
{
|
||||
"model_type": "soulxsinger",
|
||||
"architectures": ["SoulXSingerPipeline"],
|
||||
"max_num_seqs": 1
|
||||
}
|
||||
EOF
|
||||
|
||||
cat > "$SVC_DIR/config.json" <<'EOF'
|
||||
{
|
||||
"model_type": "soulxsinger",
|
||||
"architectures": ["SoulXSingerSVCPipeline"],
|
||||
"max_num_seqs": 1
|
||||
}
|
||||
EOF
|
||||
```
|
||||
|
||||
`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches.
|
||||
|
||||
**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs:
|
||||
|
||||
```bash
|
||||
pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC
|
||||
```
|
||||
|
||||
Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`:
|
||||
|
||||
```bash
|
||||
# install optional dependencies:
|
||||
pip install -e ".[soulx-svs]"
|
||||
```
|
||||
|
||||
English SVS adds NeMo ASR and NLTK data; pass `--language English`:
|
||||
|
||||
```bash
|
||||
pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2
|
||||
python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')"
|
||||
```
|
||||
|
||||
**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
|
||||
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
|
||||
--model "$BASE" \
|
||||
--preprocess-weights-dir "$PREPROCESS" \
|
||||
--control score \
|
||||
--num-inference-steps 32 \
|
||||
-o output.wav
|
||||
|
||||
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
|
||||
--model "$SVC_DIR" \
|
||||
--preprocess-weights-dir "$PREPROCESS" \
|
||||
--svc \
|
||||
--num-inference-steps 32 \
|
||||
-o output_svc.wav
|
||||
```
|
||||
|
||||
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 24 kHz mono WAV; batch only.
|
||||
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
|
||||
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
|
||||
|
||||
---
|
||||
@@ -0,0 +1,187 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference example for IndexTTS2 via vLLM-Omni.
|
||||
|
||||
Two-stage pipeline: GPT AR (Stage 0) → S2Mel + BigVGAN (Stage 1).
|
||||
Output is 22050 Hz mono WAV.
|
||||
|
||||
Usage:
|
||||
python end2end.py \
|
||||
--model /path/to/IndexTeam/IndexTTS-2 \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/ref.wav
|
||||
|
||||
# With emotion audio:
|
||||
python end2end.py \
|
||||
--model /path/to/IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
from vllm_omni import Omni # noqa: E402
|
||||
from vllm_omni.model_executor.models.indextts2.prompt_utils import ( # noqa: E402
|
||||
build_indextts2_prefill_prompt_ids,
|
||||
)
|
||||
|
||||
|
||||
def build_request(
|
||||
model: str,
|
||||
text: str,
|
||||
ref_audio_path: str | None = None,
|
||||
emo_audio_path: str | None = None,
|
||||
emo_text: str | None = None,
|
||||
emo_vector: list[float] | None = None,
|
||||
emo_alpha: float | None = None,
|
||||
use_emo_text: bool = False,
|
||||
use_random: bool = False,
|
||||
) -> dict:
|
||||
additional: dict = {"text": [text]}
|
||||
if ref_audio_path:
|
||||
additional["voice"] = [str(ref_audio_path)]
|
||||
if emo_audio_path:
|
||||
additional["emo_audio"] = [str(emo_audio_path)]
|
||||
if emo_text:
|
||||
additional["emo_text"] = [emo_text]
|
||||
if emo_vector is not None:
|
||||
additional["emo_vector"] = [emo_vector]
|
||||
if emo_alpha is not None:
|
||||
additional["emo_alpha"] = [emo_alpha]
|
||||
if use_emo_text:
|
||||
additional["use_emo_text"] = [True]
|
||||
if use_random:
|
||||
additional["use_random"] = [True]
|
||||
return {
|
||||
"prompt_token_ids": build_indextts2_prefill_prompt_ids(model, text),
|
||||
"additional_information": additional,
|
||||
}
|
||||
|
||||
|
||||
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 22050) -> None:
|
||||
audio_np = waveform.float().numpy()
|
||||
sf.write(path, audio_np, sample_rate)
|
||||
print(f" Saved {path} ({audio_np.shape}, {sample_rate} Hz)")
|
||||
|
||||
|
||||
def extract_audio(mm: dict) -> tuple[torch.Tensor | None, int]:
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if isinstance(audio, list):
|
||||
chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
|
||||
audio = torch.cat(chunks, dim=0) if chunks else None
|
||||
|
||||
sr_val = mm.get("sr")
|
||||
if isinstance(sr_val, list):
|
||||
sr_val = sr_val[-1] if sr_val else None
|
||||
if hasattr(sr_val, "item"):
|
||||
sample_rate = int(sr_val.item())
|
||||
else:
|
||||
sample_rate = int(sr_val) if sr_val is not None else 22050
|
||||
|
||||
return audio if isinstance(audio, torch.Tensor) else None, sample_rate
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
deploy_config=args.deploy_config,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
# Stage 0: GPT AR (text → mel codes)
|
||||
gpt_sampling = SamplingParams(
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
top_k=30,
|
||||
max_tokens=1500,
|
||||
repetition_penalty=10.0,
|
||||
stop_token_ids=[8193],
|
||||
seed=args.seed if args.seed is not None else 42,
|
||||
detokenize=False,
|
||||
)
|
||||
# Stage 1: S2Mel + BigVGAN (non-AR, params mostly ignored)
|
||||
s2mel_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=65536,
|
||||
detokenize=True,
|
||||
)
|
||||
sampling_params = [gpt_sampling, s2mel_sampling]
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Synthesizing: {args.text!r}")
|
||||
if args.ref_audio:
|
||||
print(f" ref_audio: {args.ref_audio}")
|
||||
inputs = build_request(
|
||||
model=args.model,
|
||||
text=args.text,
|
||||
ref_audio_path=args.ref_audio,
|
||||
emo_audio_path=args.emo_audio,
|
||||
emo_text=args.emo_text,
|
||||
emo_vector=args.emo_vector,
|
||||
emo_alpha=args.emo_alpha,
|
||||
use_emo_text=args.use_emo_text,
|
||||
use_random=args.use_random,
|
||||
)
|
||||
|
||||
for i, omni_out in enumerate(omni.generate(inputs, sampling_params_list=sampling_params)):
|
||||
mm = omni_out.multimodal_output
|
||||
if not mm:
|
||||
print(f" [req {i}] No multimodal output.")
|
||||
continue
|
||||
audio, sr = extract_audio(mm)
|
||||
if audio is None:
|
||||
print(f" [req {i}] No waveform in multimodal_output.")
|
||||
continue
|
||||
out_path = str(output_dir / f"output_{i}.wav")
|
||||
save_audio(audio.cpu(), out_path, sr)
|
||||
|
||||
print("Done.")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="IndexTTS2 offline inference")
|
||||
parser.add_argument("--model", required=True, help="HF model path for IndexTTS2.")
|
||||
parser.add_argument("--text", default="你好,这是IndexTTS2语音合成测试。")
|
||||
parser.add_argument("--ref-audio", required=True, help="Reference audio for voice cloning.")
|
||||
parser.add_argument("--emo-audio", default=None, help="Emotion reference audio.")
|
||||
parser.add_argument("--emo-text", default=None, help="Emotion description text.")
|
||||
parser.add_argument(
|
||||
"--emo-vector",
|
||||
type=float,
|
||||
nargs=8,
|
||||
default=None,
|
||||
help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm.",
|
||||
)
|
||||
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1].")
|
||||
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text.")
|
||||
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes.")
|
||||
parser.add_argument("--seed", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default=os.path.join(
|
||||
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
|
||||
"indextts2_output",
|
||||
),
|
||||
)
|
||||
parser.add_argument("--deploy-config", default=None)
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -1,11 +1,12 @@
|
||||
"""Offline inference example for MOSS-TTS variants.
|
||||
|
||||
Covers all five models:
|
||||
- OpenMOSS-Team/MOSS-TTS (8B, general TTS)
|
||||
- OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS)
|
||||
- OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis)
|
||||
- OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming)
|
||||
- OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
|
||||
Covers six models:
|
||||
- OpenMOSS-Team/MOSS-TTS (8B, general TTS)
|
||||
- OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS)
|
||||
- OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis)
|
||||
- OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming)
|
||||
- OpenMOSS-Team/MOSS-VoiceGenerator (1.7B, voice design)
|
||||
- OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 (n_vq=12, 48 kHz stereo)
|
||||
|
||||
Usage
|
||||
-----
|
||||
@@ -210,12 +211,18 @@ def run_tts(args: argparse.Namespace) -> None:
|
||||
"OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
|
||||
"OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
|
||||
"OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
|
||||
"OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5": "moss_tts_local",
|
||||
}
|
||||
deploy_config = deploy_map.get(args.model, "moss_tts")
|
||||
|
||||
is_sound_effect = "SoundEffect" in args.model
|
||||
is_voice_gen = "VoiceGenerator" in args.model
|
||||
is_realtime = "Realtime" in args.model
|
||||
# Local-v1.5 ships its own AutoProcessor (processor_config.json +
|
||||
# processing_moss_tts.py), so it reuses the same _build_unified_codes
|
||||
# path as the Delay variants below — just a different n_vq/codec.
|
||||
is_local = "Local" in args.model
|
||||
output_sr = 48000 if is_local else 24000
|
||||
|
||||
ref_audio = None
|
||||
builder_kwargs: dict = {}
|
||||
@@ -314,27 +321,33 @@ def run_tts(args: argparse.Namespace) -> None:
|
||||
wav = mm.get("model_outputs")
|
||||
if wav is None:
|
||||
continue
|
||||
# Local-v1.5's codec emits stereo chunks shaped (C, T); keep the
|
||||
# channel axis intact (don't flatten) so concatenation below stays
|
||||
# along the time axis. Mono variants stay (T,) as before.
|
||||
if isinstance(wav, list):
|
||||
non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
|
||||
non_empty = [w for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
|
||||
if non_empty and t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
chunks.extend(non_empty)
|
||||
elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
|
||||
if t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
chunks.append(wav.reshape(-1))
|
||||
chunks.append(wav)
|
||||
|
||||
t_gen_end = _time.monotonic()
|
||||
if not chunks:
|
||||
print("No audio in output — check model loading logs.")
|
||||
return
|
||||
|
||||
wav_t = torch.cat(chunks).float().cpu()
|
||||
sf.write(args.output, wav_t.numpy(), 24000)
|
||||
print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
|
||||
cat_dim = -1 if chunks[0].dim() == 2 else 0
|
||||
wav_t = torch.cat(chunks, dim=cat_dim).float().cpu()
|
||||
n_samples = wav_t.shape[-1]
|
||||
# soundfile wants (T,) mono or (T, C) multi-channel.
|
||||
sf.write(args.output, (wav_t.transpose(0, 1) if wav_t.dim() == 2 else wav_t).numpy(), output_sr)
|
||||
print(f"Saved {n_samples / output_sr:.2f}s of audio to {args.output}")
|
||||
ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
|
||||
gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
|
||||
print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
|
||||
print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={n_samples / output_sr:.2f}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Offline Qwen3-TTS example with word-level timestamps.
|
||||
|
||||
This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with
|
||||
the same shared forced aligner used by the streaming server path, then writes
|
||||
the WAV and JSON sidecar:
|
||||
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
|
||||
|
||||
On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0``
|
||||
to avoid FlashInfer sampler JIT during warmup.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
# The in-process offline aligner consumes token_classify outputs. Keeping V1
|
||||
# multiprocessing off avoids msgspec serialization issues for token-wise logits.
|
||||
os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from end2end import _estimate_prompt_len
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
|
||||
|
||||
|
||||
def _default_stage_config() -> str:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
|
||||
|
||||
|
||||
def _build_custom_voice_input(args: Any) -> dict[str, Any]:
|
||||
additional_information = {
|
||||
"task_type": ["CustomVoice"],
|
||||
"text": [args.text],
|
||||
"language": [args.language],
|
||||
"speaker": [args.speaker],
|
||||
"instruct": [args.instructions],
|
||||
"max_new_tokens": [args.max_new_tokens],
|
||||
}
|
||||
return {
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model),
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
|
||||
def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]:
|
||||
audio_data = mm["audio"]
|
||||
sr_raw = mm["sr"]
|
||||
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
|
||||
sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val)
|
||||
audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
|
||||
return audio_tensor.float().detach().cpu().flatten(), sample_rate
|
||||
|
||||
|
||||
def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes:
|
||||
samples = audio.numpy()
|
||||
pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype("<i2")
|
||||
return pcm.tobytes()
|
||||
|
||||
|
||||
async def _run_alignment(args: Any, audio: torch.Tensor, sample_rate: int) -> list[dict[str, Any]] | None:
|
||||
aligner_config = build_forced_aligner_config(args)
|
||||
if aligner_config is None:
|
||||
raise ValueError("--forced-aligner or --forced-aligner-config is required")
|
||||
timestamps = await align(
|
||||
audio=_float_audio_to_pcm16_bytes(audio),
|
||||
text=args.text,
|
||||
sample_rate=sample_rate,
|
||||
config=aligner_config,
|
||||
language=args.language,
|
||||
)
|
||||
if timestamps is None:
|
||||
return None
|
||||
return [
|
||||
{
|
||||
"word": item.word,
|
||||
"start_ms": item.start_ms,
|
||||
"end_ms": item.end_ms,
|
||||
}
|
||||
for item in timestamps
|
||||
]
|
||||
|
||||
|
||||
def main(args: Any) -> None:
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
for key in (
|
||||
"forced_aligner",
|
||||
"text",
|
||||
"language",
|
||||
"speaker",
|
||||
"instructions",
|
||||
"max_new_tokens",
|
||||
"output_dir",
|
||||
):
|
||||
omni_kwargs.pop(key, None)
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
|
||||
omni_kwargs["log_stats"] = args.log_stats
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
prompt = _build_custom_voice_input(args)
|
||||
final_output = None
|
||||
for stage_outputs in omni.generate([prompt]):
|
||||
final_output = stage_outputs.request_output
|
||||
if final_output is None:
|
||||
raise RuntimeError("Qwen3-TTS did not produce an output.")
|
||||
|
||||
mm = final_output.outputs[0].multimodal_output
|
||||
audio, sample_rate = _audio_tensor_and_sample_rate(mm)
|
||||
|
||||
wav_path = output_dir / "qwen3_tts_word_timestamps.wav"
|
||||
# Align the in-memory tensor, not the saved WAV. The WAV is only a demo
|
||||
# artifact so users can listen to the same audio referenced by the sidecar.
|
||||
timestamps = asyncio.run(_run_alignment(args, audio, sample_rate))
|
||||
sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV")
|
||||
|
||||
sidecar = {
|
||||
"text": args.text,
|
||||
"sample_rate": sample_rate,
|
||||
"audio_path": str(wav_path),
|
||||
"timestamps": timestamps,
|
||||
}
|
||||
json_path = output_dir / "qwen3_tts_word_timestamps.json"
|
||||
json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
print(f"Saved audio: {wav_path}")
|
||||
print(f"Saved timestamps: {json_path}")
|
||||
print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
def parse_args() -> Any:
|
||||
parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example")
|
||||
parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name")
|
||||
parser.add_argument(
|
||||
"--forced-aligner",
|
||||
default=None,
|
||||
help="Qwen3 forced aligner model path/name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--forced-aligner-config",
|
||||
default=None,
|
||||
help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
|
||||
)
|
||||
parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
|
||||
parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
|
||||
parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
|
||||
parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
|
||||
parser.add_argument("--instructions", default="", help="Optional speaking style instruction")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
|
||||
parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings).
|
||||
|
||||
Usage:
|
||||
python benchmark.py --model /path/to/SoulX-Singer --svs \\
|
||||
--prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\
|
||||
--target-audio ../SoulX-Singer/example/audio/music.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o benchmark.wav
|
||||
|
||||
python benchmark.py ... --enable-diffusion-pipeline-profiler
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import statistics
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
from end2end import (
|
||||
SVC_DEPLOY_CONFIG,
|
||||
SVS_DEPLOY_CONFIG,
|
||||
add_inference_args,
|
||||
build_sampling,
|
||||
extract_audio,
|
||||
resolve_preprocess_weights_dir,
|
||||
)
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
|
||||
def _audio_duration_sec(outputs) -> float:
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
return audio_np.size / sr
|
||||
|
||||
|
||||
def _stage_durations_from_output(outputs) -> dict[str, float]:
|
||||
omni_out = outputs[0]
|
||||
durations = getattr(omni_out, "stage_durations", None) or {}
|
||||
if durations:
|
||||
return dict(durations)
|
||||
ro = getattr(omni_out, "request_output", None)
|
||||
if ro is not None:
|
||||
return dict(getattr(ro, "stage_durations", None) or {})
|
||||
return {}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark")
|
||||
add_inference_args(parser)
|
||||
parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)")
|
||||
parser.add_argument("--runs", type=int, default=3, help="Measured runs")
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable low-overhead stage timing in DiffusionOutput.stage_durations",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable torch.compile on diff_estimator (eager baseline for A/B)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["bfloat16", "bf16", "float16", "fp16", "half"],
|
||||
help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
"-o",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Save WAV from the last measured run (after warmup)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-each-run",
|
||||
action="store_true",
|
||||
help="With --output, also write measured runs as <stem>_run1.wav, _run2.wav, ...",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
|
||||
deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG)
|
||||
mode = "SVS" if args.svs else "SVC"
|
||||
|
||||
omni_kwargs: dict = {
|
||||
"model": args.model,
|
||||
"deploy_config": deploy_config,
|
||||
"async_chunk": False, # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed)
|
||||
}
|
||||
if args.enable_diffusion_pipeline_profiler:
|
||||
omni_kwargs["enable_diffusion_pipeline_profiler"] = True
|
||||
if args.enforce_eager:
|
||||
omni_kwargs["enforce_eager"] = True
|
||||
if args.dtype is not None:
|
||||
omni_kwargs["dtype"] = args.dtype
|
||||
|
||||
compile_mode = "eager" if args.enforce_eager else "torch.compile"
|
||||
dtype_label = args.dtype or "from deploy YAML"
|
||||
print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]")
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
kind = "svs" if args.svs else "svc"
|
||||
sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind)
|
||||
prompt = {"prompt_token_ids": [0]}
|
||||
|
||||
latencies_ms: list[float] = []
|
||||
rtfs: list[float] = []
|
||||
last_stages: dict[str, float] = {}
|
||||
last_measured_outputs = None
|
||||
|
||||
total_runs = args.warmup + args.runs
|
||||
for run_idx in range(total_runs):
|
||||
is_warmup = run_idx < args.warmup
|
||||
label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}"
|
||||
t0 = time.perf_counter()
|
||||
outputs = list(omni.generate([prompt], sampling))
|
||||
elapsed_ms = (time.perf_counter() - t0) * 1000.0
|
||||
audio_sec = _audio_duration_sec(outputs)
|
||||
rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf")
|
||||
print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}")
|
||||
if not is_warmup:
|
||||
latencies_ms.append(elapsed_ms)
|
||||
rtfs.append(rtf)
|
||||
last_stages = _stage_durations_from_output(outputs)
|
||||
last_measured_outputs = outputs
|
||||
if args.output and args.save_each_run:
|
||||
measured_idx = run_idx - args.warmup + 1
|
||||
out_path = Path(args.output)
|
||||
run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}")
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
sf.write(str(run_path), audio_np, sr)
|
||||
print(f" saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
|
||||
|
||||
omni.close()
|
||||
|
||||
if args.output and last_measured_outputs is not None:
|
||||
audio_np, sr = extract_audio(last_measured_outputs)
|
||||
out_path = Path(args.output)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
sf.write(str(out_path), audio_np, sr)
|
||||
print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
|
||||
|
||||
if latencies_ms:
|
||||
print("\n=== Summary (measured runs) ===")
|
||||
print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}")
|
||||
print(f"RTF: mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}")
|
||||
if last_stages:
|
||||
print("\nStage durations (last measured run):")
|
||||
for name, value in sorted(last_stages.items()):
|
||||
if name.endswith("_ms"):
|
||||
print(f" {name}: {value:.1f} ms")
|
||||
else:
|
||||
print(f" {name}: {value * 1000.0:.1f} ms")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline).
|
||||
|
||||
Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference.
|
||||
|
||||
Usage:
|
||||
# SVS (default)
|
||||
python end2end.py --model /path/to/SoulX-Singer \\
|
||||
--prompt-audio path/to/prompt/audio.mp3 \\
|
||||
--target-audio path/to/target/audio.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o output.wav
|
||||
|
||||
# SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view)
|
||||
python end2end.py --model /path/to/SoulX-Singer-svc \\
|
||||
--prompt-audio path/to/prompt/audio.mp3 \\
|
||||
--target-audio path/to/target/audio.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o output_svc.wav
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml"
|
||||
SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml"
|
||||
DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3"
|
||||
DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3"
|
||||
_SAMPLE_RATE = 24000
|
||||
|
||||
|
||||
def _require_paths(paths: dict[str, str | None]) -> None:
|
||||
missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)]
|
||||
if missing:
|
||||
raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing))
|
||||
|
||||
|
||||
def resolve_preprocess_weights_dir(explicit: str | None) -> str:
|
||||
if explicit:
|
||||
path = Path(explicit).expanduser().resolve()
|
||||
if not (path / "rmvpe" / "rmvpe.pt").is_file():
|
||||
raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}")
|
||||
return str(path)
|
||||
env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR")
|
||||
if env:
|
||||
path = Path(env).expanduser().resolve()
|
||||
if (path / "rmvpe" / "rmvpe.pt").is_file():
|
||||
return str(path)
|
||||
raise FileNotFoundError(
|
||||
"SoulX preprocess weights not found. Pass --preprocess-weights-dir or set "
|
||||
"SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt "
|
||||
"(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)."
|
||||
)
|
||||
|
||||
|
||||
def build_sampling(
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
preprocess_weights_dir: str,
|
||||
kind: str,
|
||||
) -> OmniDiffusionSamplingParams:
|
||||
extra: dict = {
|
||||
"prompt_audio": os.path.abspath(args.prompt_audio),
|
||||
"target_audio": os.path.abspath(args.target_audio),
|
||||
"preprocess_weights_dir": preprocess_weights_dir,
|
||||
"vocal_sep": args.vocal_sep,
|
||||
"auto_shift": args.auto_shift,
|
||||
"pitch_shift": args.pitch_shift,
|
||||
}
|
||||
if kind == "svs":
|
||||
extra["language"] = args.language
|
||||
extra["control"] = args.control
|
||||
if args.prompt_metadata_path:
|
||||
extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path)
|
||||
if args.target_metadata_path:
|
||||
extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path)
|
||||
if args.prompt_metadata_path or args.target_metadata_path:
|
||||
extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio)
|
||||
|
||||
extra = validate_soulx_extra_args(kind, extra)
|
||||
|
||||
return OmniDiffusionSamplingParams(
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
extra_args=extra,
|
||||
)
|
||||
|
||||
|
||||
def extract_audio(outputs) -> tuple[np.ndarray, int]:
|
||||
out = outputs[0]
|
||||
mm = getattr(out, "multimodal_output", None)
|
||||
if mm is None:
|
||||
ro = getattr(out, "request_output", None)
|
||||
if ro is not None:
|
||||
mm = getattr(ro, "multimodal_output", None)
|
||||
if mm is None and ro.outputs:
|
||||
mm = getattr(ro.outputs[0], "multimodal_output", None)
|
||||
if not mm or "audio" not in mm:
|
||||
raise RuntimeError("No audio in multimodal_output")
|
||||
|
||||
audio_val = mm["audio"]
|
||||
if isinstance(audio_val, list):
|
||||
chunks = [
|
||||
np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1)
|
||||
for chunk in audio_val
|
||||
if chunk is not None
|
||||
]
|
||||
audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32)
|
||||
elif hasattr(audio_val, "cpu"):
|
||||
audio_np = audio_val.detach().cpu().numpy().squeeze()
|
||||
else:
|
||||
audio_np = np.asarray(audio_val, dtype=np.float32).squeeze()
|
||||
|
||||
sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE))
|
||||
if isinstance(sr_val, list) and sr_val:
|
||||
sr_val = sr_val[-1]
|
||||
if hasattr(sr_val, "item"):
|
||||
sr_val = sr_val.item()
|
||||
return audio_np.astype(np.float32, copy=False), int(sr_val)
|
||||
|
||||
|
||||
def add_inference_args(parser: argparse.ArgumentParser) -> None:
|
||||
mode = parser.add_mutually_exclusive_group()
|
||||
mode.add_argument(
|
||||
"--svc",
|
||||
action="store_true",
|
||||
help="Assert SVC mode (must match config.json architectures)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-audio",
|
||||
type=str,
|
||||
default=str(DEFAULT_PROMPT_AUDIO),
|
||||
help="Prompt / reference audio (mp3 or wav)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-audio",
|
||||
type=str,
|
||||
default=str(DEFAULT_TARGET_AUDIO),
|
||||
help="Target content audio (mp3 or wav)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-metadata-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-metadata-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-wav-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preprocess-weights-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--language",
|
||||
type=str,
|
||||
default="Mandarin",
|
||||
help="SVS lyrics language (Mandarin, Cantonese, English)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--control",
|
||||
type=str,
|
||||
default="melody",
|
||||
choices=["score", "melody"],
|
||||
help="SVS control mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vocal-sep",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=False,
|
||||
help="Run vocal separation during preprocess",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auto-shift",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Auto pitch shift to align target pitch with prompt",
|
||||
)
|
||||
parser.add_argument("--pitch-shift", type=int, default=0)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=32)
|
||||
parser.add_argument("--guidance-scale", type=float, default=3.0)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Optional CFM RNG seed. Omit for non-deterministic sampling.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)")
|
||||
add_inference_args(parser)
|
||||
parser.add_argument("--output", "-o", type=str, default="output.wav")
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
kind = resolve_soulx_kind(args.model)
|
||||
|
||||
if not (args.svc ^ (kind == "svs")):
|
||||
raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}")
|
||||
|
||||
required = {
|
||||
"prompt_audio": args.prompt_audio,
|
||||
"target_audio": args.target_audio,
|
||||
}
|
||||
if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path):
|
||||
raise ValueError(
|
||||
"SVS precomputed metadata requires both --prompt-metadata-path and "
|
||||
"--target-metadata-path (or neither, for online preprocess)."
|
||||
)
|
||||
if args.prompt_metadata_path:
|
||||
required["prompt_metadata_path"] = args.prompt_metadata_path
|
||||
required["target_metadata_path"] = args.target_metadata_path
|
||||
if args.prompt_wav_path:
|
||||
required["prompt_wav_path"] = args.prompt_wav_path
|
||||
_require_paths(required)
|
||||
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
|
||||
deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG)
|
||||
mode = kind.upper()
|
||||
|
||||
print(f"Loading SoulX-Singer {mode} from {args.model}")
|
||||
print(f" deploy: {deploy_config}")
|
||||
print(f" prompt: {args.prompt_audio}")
|
||||
print(f" target: {args.target_audio}")
|
||||
|
||||
omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only
|
||||
sampling = build_sampling(
|
||||
args,
|
||||
preprocess_weights_dir=preprocess_dir,
|
||||
kind=kind,
|
||||
)
|
||||
prompt = {"prompt_token_ids": [0], "multi_modal_data": {}}
|
||||
|
||||
print("Running generate (inline preprocess + DiT)...")
|
||||
outputs = list(omni.generate([prompt], sampling))
|
||||
if not outputs:
|
||||
raise RuntimeError("No output from omni.generate")
|
||||
|
||||
err = getattr(outputs[0], "error", None)
|
||||
if err is not None:
|
||||
raise RuntimeError(f"Pipeline error: {err}")
|
||||
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
sf.write(args.output, audio_np, sr)
|
||||
print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)")
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -6,9 +6,12 @@ A unified script for text-to-video generation. Supports multiple models with mod
|
||||
|
||||
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|
||||
|---|---|---|---|---|---|
|
||||
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
|
||||
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
|
||||
| `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -29,6 +32,28 @@ python text_to_video.py \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
### Wan2.1 VACE (T2V)
|
||||
|
||||
VACE text-to-video uses this shared entrypoint. Conditional VACE tasks use
|
||||
the shared [`image_to_video.py`](../image_to_video/README.md#wan21-vace-conditional-tasks)
|
||||
entrypoint, which constructs the pipeline-native conditioning data from the
|
||||
provided media inputs. No explicit mode parameter is required.
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--prompt "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves." \
|
||||
--seed 0 \
|
||||
--height 480 \
|
||||
--width 832 \
|
||||
--num-frames 81 \
|
||||
--num-inference-steps 30 \
|
||||
--guidance-scale 5.0 \
|
||||
--flow-shift 5.0 \
|
||||
--vae-use-tiling \
|
||||
--output vace_t2v_output.mp4
|
||||
```
|
||||
|
||||
LTX2 example:
|
||||
|
||||
```bash
|
||||
@@ -100,6 +125,63 @@ python text_to_video.py \
|
||||
--output quick_test.mp4
|
||||
```
|
||||
|
||||
### Cosmos3
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--prompt "A robot arm is cleaning a plate in the kitchen." \
|
||||
--negative-prompt "blurry, distorted, low quality, jittery, deformed" \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
|
||||
"use_resolution_template": false, "use_duration_template": false}' \
|
||||
--output cosmos3_t2v.mp4
|
||||
```
|
||||
|
||||
### Helios (T2V)
|
||||
|
||||
Helios ships three variants. Model-specific knobs (declared in
|
||||
`vllm_omni/model_extras/helios.py`) are passed via the generic `--extra-body`
|
||||
JSON flag rather than bespoke per-model flags.
|
||||
|
||||
**Helios-Base** (Stage 1 only):
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_t2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 pyramid + CFG-Zero*):
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [20, 20, 20], "use_cfg_zero_star": true, "use_zero_init": true, "zero_steps": 1}' \
|
||||
--output helios_t2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled** (Stage 2 pyramid + DMD, few-step):
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}' \
|
||||
--output helios_t2v_distilled.mp4
|
||||
```
|
||||
|
||||
> Helios image-to-video (I2V) and video-to-video (V2V) require image/video
|
||||
> conditioning tensors that cannot be passed through the JSON `--extra-body`
|
||||
> flag; they are out of scope for this text-to-video example.
|
||||
|
||||
## Key Arguments
|
||||
|
||||
### Common
|
||||
@@ -112,6 +194,7 @@ python text_to_video.py \
|
||||
- `--num-inference-steps`: sampling steps. Default depends on model.
|
||||
- `--fps`: frames per second for the saved MP4.
|
||||
- `--output`: path to save the generated video.
|
||||
- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. See the Helios recipes above.
|
||||
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
|
||||
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
|
||||
@@ -122,6 +205,7 @@ python text_to_video.py \
|
||||
- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
|
||||
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
|
||||
- `--flow-shift`: scheduler flow_shift parameter.
|
||||
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
|
||||
|
||||
### Wan2.2-specific
|
||||
|
||||
|
||||
@@ -11,12 +11,23 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
_MODEL_PRESETS = {
|
||||
"vace": {
|
||||
"height": 480,
|
||||
"width": 832,
|
||||
"num_frames": 81,
|
||||
"num_inference_steps": 30,
|
||||
"guidance_scale": 5.0,
|
||||
"fps": 16,
|
||||
"output": "vace_t2v_output.mp4",
|
||||
},
|
||||
"wan": {
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
@@ -35,13 +46,38 @@ _MODEL_PRESETS = {
|
||||
"fps": 24,
|
||||
"output": "hunyuan_video_15_output.mp4",
|
||||
},
|
||||
"cosmos": {
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 6.0,
|
||||
"fps": 24,
|
||||
"flow_shift": 10.0,
|
||||
"output": "cosmos3_t2v_output.mp4",
|
||||
},
|
||||
"helios": {
|
||||
"height": 384,
|
||||
"width": 640,
|
||||
"num_frames": 99,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 5.0,
|
||||
"fps": 16,
|
||||
"output": "helios_output.mp4",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _detect_preset(model: str) -> dict:
|
||||
model_lower = model.lower()
|
||||
if "vace" in model_lower:
|
||||
return _MODEL_PRESETS["vace"]
|
||||
if "cosmos" in model_lower:
|
||||
return _MODEL_PRESETS["cosmos"]
|
||||
if "hunyuan" in model_lower:
|
||||
return _MODEL_PRESETS["hunyuan"]
|
||||
if "helios" in model_lower:
|
||||
return _MODEL_PRESETS["helios"]
|
||||
return _MODEL_PRESETS["wan"]
|
||||
|
||||
|
||||
@@ -55,10 +91,26 @@ def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
return config
|
||||
|
||||
|
||||
def parse_extra_body(value: str) -> dict[str, Any]:
|
||||
"""Parse a JSON object of model-specific extra_body params.
|
||||
|
||||
Pipeline-declared knobs (see vllm_omni/model_extras/) are merged into
|
||||
OmniDiffusionSamplingParams.extra_args, so a single generic example can
|
||||
drive model-specific behaviour without bespoke per-model flags.
|
||||
"""
|
||||
try:
|
||||
body = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e
|
||||
if not isinstance(body, dict):
|
||||
raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
|
||||
return body
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a video from a text prompt. "
|
||||
"Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models."
|
||||
"Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
@@ -74,6 +126,18 @@ def parse_args() -> argparse.Namespace:
|
||||
)
|
||||
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=parse_extra_body,
|
||||
default=None,
|
||||
help="JSON dict of model-specific extra_body params (declared in vllm_omni/model_extras/), "
|
||||
"merged into sampling extra_args. Unknown keys for the chosen model are dropped. "
|
||||
'Cosmos3 example: \'{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, '
|
||||
'"use_resolution_template": false, "use_duration_template": false}\' '
|
||||
'(add "generate_sound": true and optional "sound_duration" for synchronized audio). '
|
||||
'Helios-Distilled example: \'{"is_enable_stage2": true, '
|
||||
'"pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.',
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
@@ -322,7 +386,14 @@ def main():
|
||||
omni_kwargs["cache_config"] = cache_config
|
||||
omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
|
||||
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni) or model_class_name
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -358,10 +429,20 @@ def main():
|
||||
if args.guidance_scale_high is not None:
|
||||
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(**sampling_kwargs)
|
||||
# Route model-specific knobs through extra_body, filtered against the model's
|
||||
# declared extra_body_params. Models without a declaration only forward explicit
|
||||
# --extra-body JSON (preserving the generic flags' legacy behavior).
|
||||
extra_body = dict(args.extra_body or {})
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
|
||||
elif extra_body:
|
||||
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
frames = omni.generate(
|
||||
prompt_dict,
|
||||
OmniDiffusionSamplingParams(**sampling_kwargs),
|
||||
sampling_params,
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
@@ -1,148 +0,0 @@
|
||||
# VACE Video Generation
|
||||
|
||||
Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline.
|
||||
|
||||
- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Key Arguments](#key-arguments)
|
||||
- [More CLI Examples](#more-cli-examples)
|
||||
|
||||
## Overview
|
||||
|
||||
[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation.
|
||||
|
||||
### Supported Models
|
||||
|
||||
| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace |
|
||||
|-------|-------------|-------------------|---------------------|-------------|
|
||||
| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
|
||||
| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
|
||||
|
||||
!!! info
|
||||
*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features.
|
||||
|
||||
Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
Text-to-video generation:
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers")
|
||||
prompt = "A sleek robot stands in a vast warehouse filled with boxes"
|
||||
outputs = omni.generate(
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=81,
|
||||
num_inference_steps=30,
|
||||
guidance_scale=5.0,
|
||||
),
|
||||
)
|
||||
video = outputs[0].images
|
||||
|
||||
from diffusers.utils import export_to_video
|
||||
export_to_video(list(video[0]), "t2v_output.mp4", fps=16)
|
||||
omni.close()
|
||||
```
|
||||
|
||||
### Local CLI Usage
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode t2v \
|
||||
--prompt "A sleek robot stands in a vast warehouse filled with boxes" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output t2v_output.mp4
|
||||
```
|
||||
|
||||
## Key Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` |
|
||||
| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path |
|
||||
| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video |
|
||||
| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance |
|
||||
| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) |
|
||||
| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) |
|
||||
| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) |
|
||||
| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) |
|
||||
| `--num-frames` | int | `81` | Number of video frames to generate |
|
||||
| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) |
|
||||
| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale |
|
||||
| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter |
|
||||
| `--seed` | int | `42` | Random seed for deterministic sampling |
|
||||
| `--fps` | int | `16` | Frames per second for the saved MP4 |
|
||||
| `--output` | str | `"vace_output.mp4"` | Path to save the generated video |
|
||||
| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization |
|
||||
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
|
||||
| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
|
||||
| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
|
||||
| `--tensor-parallel-size` | int | — | Tensor parallel size |
|
||||
| `--enforce-eager` | flag | off | Disable torch.compile |
|
||||
|
||||
> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`).
|
||||
|
||||
## More CLI Examples
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
First frame is kept, remaining frames are generated:
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode i2v \
|
||||
--image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked egg on the moon" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output i2v_output.mp4
|
||||
```
|
||||
|
||||
### First-Last-Frame Interpolation (FLF2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode flf2v \
|
||||
--image first_frame.jpg --last-image last_frame.jpg \
|
||||
--prompt "A bird takes off from a branch and lands on another" \
|
||||
--height 512 --width 512 --num-frames 81 \
|
||||
--output flf2v_output.mp4
|
||||
```
|
||||
|
||||
### Inpainting
|
||||
|
||||
Center vertical stripe is masked and regenerated:
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode inpaint \
|
||||
--image scene.jpg \
|
||||
--prompt "Shrek walks out of a building" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output inpaint_output.mp4
|
||||
```
|
||||
|
||||
### Reference Image-guided (R2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode r2v \
|
||||
--image reference.jpg \
|
||||
--prompt "Camera slowly zooms out from the character" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output r2v_output.mp4
|
||||
```
|
||||
@@ -1,209 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""VACE video generation example.
|
||||
|
||||
VACE (Video All-in-one Creation Engine) supports multiple video tasks:
|
||||
- T2V: Text-to-Video
|
||||
- I2V: Image-to-Video (first frame conditioning)
|
||||
- V2LF: Video-to-Last-Frame
|
||||
- FLF2V: First-Last-Frame interpolation
|
||||
- Inpainting: Masked region generation
|
||||
- R2V: Reference image-guided generation
|
||||
|
||||
Usage examples:
|
||||
# T2V (text-to-video)
|
||||
python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
|
||||
|
||||
# I2V (image-to-video, first frame kept)
|
||||
python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
|
||||
|
||||
# FLF2V (first-last frame interpolation)
|
||||
python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
|
||||
|
||||
# R2V (reference image guided)
|
||||
python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
|
||||
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
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="VACE video generation.")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Wan-AI/Wan2.1-VACE-14B-diffusers",
|
||||
help="VACE model ID or local path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
default="t2v",
|
||||
choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
|
||||
help="Generation mode.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
|
||||
parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
|
||||
parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
|
||||
parser.add_argument("--height", type=int, default=480, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=832, help="Video width.")
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
|
||||
parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
|
||||
parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
|
||||
parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_prompts(args):
|
||||
"""Build prompt dict with multi_modal_data based on mode."""
|
||||
h, w, nf = args.height, args.width, args.num_frames
|
||||
|
||||
gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
|
||||
mask_black = PIL.Image.new("L", (w, h), 0)
|
||||
mask_white = PIL.Image.new("L", (w, h), 255)
|
||||
|
||||
prompt_data = {
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
}
|
||||
|
||||
if args.mode == "t2v":
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "r2v":
|
||||
assert args.image, "--image required for R2V mode"
|
||||
ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "i2v":
|
||||
assert args.image, "--image required for I2V mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [img] + [gray] * (nf - 1),
|
||||
"mask": [mask_black] + [mask_white] * (nf - 1),
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "v2lf":
|
||||
assert args.image, "--image required for V2LF mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [gray] * (nf - 1) + [img],
|
||||
"mask": [mask_white] * (nf - 1) + [mask_black],
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "flf2v":
|
||||
assert args.image and args.last_image, "--image and --last-image required for FLF2V"
|
||||
first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [first] + [gray] * (nf - 2) + [last],
|
||||
"mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "inpaint":
|
||||
assert args.image, "--image required for inpaint mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
d = 80
|
||||
frames, masks = [], []
|
||||
for _ in range(nf):
|
||||
base = np.array(img).copy()
|
||||
mask = PIL.Image.new("L", (w, h), 0)
|
||||
stripe = PIL.Image.new("L", (2 * d, h), 255)
|
||||
mask.paste(stripe, (w // 2 - d, 0))
|
||||
base[np.array(mask) > 128] = 128
|
||||
frames.append(PIL.Image.fromarray(base))
|
||||
masks.append(mask)
|
||||
prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
|
||||
return prompt_data
|
||||
|
||||
raise ValueError(f"Unknown mode: {args.mode}")
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
flow_shift=args.flow_shift,
|
||||
enforce_eager=args.enforce_eager,
|
||||
parallel_config=parallel_config,
|
||||
)
|
||||
|
||||
prompt_data = build_prompts(args)
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"VACE {args.mode.upper()} Generation")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
prompt_data,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
generator=generator,
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
video = outputs[0].images
|
||||
if isinstance(video, list):
|
||||
video = video[0]
|
||||
if isinstance(video, torch.Tensor):
|
||||
video = video.cpu().numpy()
|
||||
if video.ndim == 5:
|
||||
video = video[0]
|
||||
print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
if np.issubdtype(video.dtype, np.integer):
|
||||
video = video.astype(np.float32) / 255.0
|
||||
export_to_video(list(video), str(output_path), fps=args.fps)
|
||||
print(f"Saved to {output_path}")
|
||||
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -87,6 +87,63 @@ python x_to_video_audio.py \
|
||||
```
|
||||
|
||||
|
||||
## MagiHuman (`--model-type magi-human`)
|
||||
|
||||
MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
|
||||
T5Gemma text encoder. A detailed text prompt is the only required input; an optional
|
||||
image and/or audio file may be supplied for conditioning. Natively supports Tensor
|
||||
Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
|
||||
the MoE weights and the text encoder.
|
||||
|
||||
> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
|
||||
> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
|
||||
|
||||
Text-only generation:
|
||||
|
||||
```bash
|
||||
python x_to_video_audio.py \
|
||||
--model-type magi-human \
|
||||
--model /path/to/daVinci-MagiHuman \
|
||||
--prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \
|
||||
--tensor-parallel-size 4 \
|
||||
--height 256 --width 448 \
|
||||
--num-inference-steps 8 \
|
||||
--seed 52 \
|
||||
--extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
With optional image and audio conditioning:
|
||||
|
||||
```bash
|
||||
python x_to_video_audio.py \
|
||||
--model-type magi-human \
|
||||
--model /path/to/daVinci-MagiHuman \
|
||||
--prompt "A young woman..." \
|
||||
--tensor-parallel-size 4 \
|
||||
--height 256 --width 448 \
|
||||
--num-inference-steps 8 \
|
||||
--seed 52 \
|
||||
--extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
|
||||
in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
|
||||
|
||||
- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
|
||||
because the audio length then determines the number of frames).
|
||||
- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
|
||||
the audio drives the frame count and is encoded as a latent condition into the DiT
|
||||
(`is_a2v` mode). Omit for pure text-to-video+audio generation.
|
||||
- `image_path`: path to an image file for visual conditioning. Applied at both the
|
||||
base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
|
||||
- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
|
||||
when these are omitted.
|
||||
- `sr_num_inference_steps`: denoising steps for the SR stage.
|
||||
|
||||
## DreamID-Omni arguments
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
@@ -20,14 +21,26 @@ AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).")
|
||||
parser.add_argument("--model", required=True, help="DreamID ckpt root directory.")
|
||||
parser.add_argument("--model-type", default="dreamid-omni", help="Model type.")
|
||||
parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
|
||||
parser.add_argument("--model", required=True, help="Model ckpt root directory.")
|
||||
parser.add_argument(
|
||||
"--model-type",
|
||||
default="dreamid-omni",
|
||||
choices=["dreamid-omni", "magi-human"],
|
||||
help="Model type.",
|
||||
)
|
||||
parser.add_argument("--prompt", default=None, help="Text prompt.")
|
||||
|
||||
parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
|
||||
parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
|
||||
parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=str,
|
||||
default=None,
|
||||
help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
|
||||
'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
|
||||
)
|
||||
|
||||
parser.add_argument("--height", type=int, default=704, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=1280, help="Video width.")
|
||||
@@ -142,8 +155,6 @@ def main() -> None:
|
||||
|
||||
text_prompt = args.prompt
|
||||
if args.prompt_file:
|
||||
import json
|
||||
|
||||
with open(args.prompt_file) as f:
|
||||
text_prompt = json.load(f)
|
||||
text_prompt = re.sub(
|
||||
@@ -155,25 +166,36 @@ def main() -> None:
|
||||
text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
|
||||
text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
|
||||
|
||||
image, audio = load_image_and_audio(args.image_path, args.audio_path)
|
||||
if args.model_type == "magi-human":
|
||||
prompt = text_prompt
|
||||
extra_args = json.loads(args.extra_body) if args.extra_body else {}
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args=extra_args,
|
||||
)
|
||||
else:
|
||||
image, audio = load_image_and_audio(args.image_path, args.audio_path)
|
||||
|
||||
prompt = {
|
||||
"prompt": text_prompt,
|
||||
"video_negative_prompt": args.video_negative_prompt,
|
||||
"audio_negative_prompt": args.audio_negative_prompt,
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
}
|
||||
prompt = {
|
||||
"prompt": text_prompt,
|
||||
"video_negative_prompt": args.video_negative_prompt,
|
||||
"audio_negative_prompt": args.audio_negative_prompt,
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
}
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"solver_name": args.solver_name,
|
||||
"shift": args.shift,
|
||||
},
|
||||
)
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"solver_name": args.solver_name,
|
||||
"shift": args.shift,
|
||||
},
|
||||
)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
@@ -215,22 +237,28 @@ def main() -> None:
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
if not outputs:
|
||||
raise RuntimeError("No output returned from DreamID-Omni.")
|
||||
raise RuntimeError("No output returned from the model.")
|
||||
result = outputs[0]
|
||||
if not result.images:
|
||||
raise RuntimeError("No video frames found in DreamID-Omni output.")
|
||||
raise RuntimeError("No video frames found in model output.")
|
||||
generated_video = result.images[0]
|
||||
mm = result.multimodal_output or {}
|
||||
generated_audio = mm.get("audio")
|
||||
fps = int(mm.get("fps", 24))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 16000))
|
||||
|
||||
# DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
|
||||
# mux_video_audio_bytes expects (F, H, W, C) uint8.
|
||||
if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
|
||||
raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
|
||||
frames = generated_video.transpose(1, 2, 3, 0)
|
||||
frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||||
if args.model_type == "magi-human":
|
||||
# MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
|
||||
fps = float(mm.get("fps", 25))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 24000))
|
||||
frames = generated_video
|
||||
else:
|
||||
fps = float(mm.get("fps", 24))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 16000))
|
||||
# DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
|
||||
# mux_video_audio_bytes expects (F, H, W, C) uint8.
|
||||
if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
|
||||
raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
|
||||
frames = generated_video.transpose(1, 2, 3, 0)
|
||||
frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||||
|
||||
audio_np = None
|
||||
if generated_audio is not None:
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
# AURA Omni Native Pipeline
|
||||
|
||||
`aura_omni` serves AURA as a native multi-stage vLLM-Omni pipeline:
|
||||
|
||||
```text
|
||||
Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav
|
||||
```
|
||||
|
||||
The pipeline has three semantic modules, but four engine stages because the
|
||||
existing Qwen3-TTS implementation is natively split into Talker and Code2Wav.
|
||||
|
||||
Start the server with the deploy profile:
|
||||
|
||||
```bash
|
||||
vllm serve aurateam/AURA \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--deploy-config vllm_omni/deploy/aura_omni.yaml \
|
||||
--served-model-name aurateam/AURA \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
The deploy file sets per-stage model repos:
|
||||
|
||||
- Stage 0 ASR: `Qwen/Qwen3-ASR-1.7B`
|
||||
- Stage 1 AURA: `aurateam/AURA`
|
||||
- Stage 2/3 TTS: `Qwen/Qwen3-TTS-12Hz-1.7B-Base`
|
||||
|
||||
For local weights, edit the `model` value on each stage in
|
||||
`vllm_omni/deploy/aura_omni.yaml`. The deploy profile includes
|
||||
`pipeline: aura_omni`, so the server uses this four-stage topology even when
|
||||
the command-line model path points at one component checkpoint.
|
||||
|
||||
Expected request shape:
|
||||
|
||||
- Send microphone audio as the Stage 0 multimodal audio input.
|
||||
- Include video frames in the original request `multi_modal_data`; the
|
||||
`asr2aura` processor carries them forward to AURA.
|
||||
- Optional `additional_information` keys:
|
||||
- `aura_system_prompt`
|
||||
- `tts_task_type`
|
||||
- `tts_language`
|
||||
- `tts_speaker`
|
||||
- `tts_instruct`
|
||||
- `tts_ref_audio`
|
||||
- `tts_ref_text`
|
||||
- `tts_x_vector_only_mode`
|
||||
- `tts_pass_token_ids`
|
||||
|
||||
If AURA emits `<|silent|>`, the `aura2tts` processor returns no TTS request, so
|
||||
the TTS stages are skipped for that turn.
|
||||
|
||||
## GPU Utilization Recommendation
|
||||
|
||||
Tune `gpu_memory_utilization` per stage in `vllm_omni/deploy/aura_omni.yaml`.
|
||||
Recommended baseline on one GPU for H200
|
||||
|
||||
- Stage 0 (ASR): `0.10`
|
||||
- Stage 1 (AURA): `0.4`
|
||||
- Stage 2 (Qwen3-TTS Talker): `0.20`
|
||||
- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
|
||||
|
||||
## Python Client
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--host localhost \
|
||||
--port 8091 \
|
||||
--model aurateam/AURA \
|
||||
--modalities text,audio
|
||||
```
|
||||
|
||||
Use local media:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--audio-path /path/to/input.wav \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--output-dir output_aura_omni_online
|
||||
```
|
||||
|
||||
Base voice clone mode (default, recommended as x-vector while debugging ICL):
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--tts-task-type Base \
|
||||
--tts-ref-audio vllm-omni/tests/assets/qwen3_tts/clone_2.wav \
|
||||
--tts-ref-text "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
```
|
||||
|
||||
Enable AURA token-id passthrough explicitly:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--tts-pass-token-ids
|
||||
```
|
||||
|
||||
CustomVoice mode requires stages 2 and 3 in `aura_omni.yaml` to point at a
|
||||
Qwen3-TTS CustomVoice checkpoint:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--tts-task-type CustomVoice \
|
||||
--tts-speaker Vivian
|
||||
```
|
||||
|
||||
By default, AURA responses are passed to Qwen3-TTS as text. Set
|
||||
`tts_pass_token_ids=true` to pass AURA-generated assistant token ids directly
|
||||
to Qwen3-TTS instead. The processor still uses AURA token ids, when available,
|
||||
to estimate the Talker prompt length in the default text path.
|
||||
|
||||
## Curl
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/aura_omni
|
||||
bash run_curl_multimodal_generation.sh
|
||||
```
|
||||
|
||||
Set `PORT`, `MODEL`, or `OUTPUT_DIR` to override defaults:
|
||||
|
||||
```bash
|
||||
PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh
|
||||
TTS_PASS_TOKEN_IDS=true PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh
|
||||
```
|
||||
|
||||
## Gradio
|
||||
|
||||
Launch the server and Gradio UI together:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/aura_omni
|
||||
bash run_gradio_demo.sh
|
||||
```
|
||||
|
||||
If the server is already running:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/gradio_demo.py \
|
||||
--model aurateam/AURA \
|
||||
--api-base http://localhost:8091/v1
|
||||
```
|
||||
|
||||
## Offline
|
||||
|
||||
For offline inference, see
|
||||
[`examples/offline_inference/aura_omni`](../../offline_inference/aura_omni/).
|
||||
@@ -0,0 +1,305 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Gradio demo for AURA Omni online serving."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import gradio as gr
|
||||
except ImportError:
|
||||
raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from openai import OpenAI
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
default_qwen3_tts_ref_audio_path,
|
||||
)
|
||||
|
||||
SEED = 42
|
||||
DEFAULT_TTS_REF_AUDIO = default_qwen3_tts_ref_audio_path()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Gradio demo for AURA Omni online inference.")
|
||||
parser.add_argument("--model", default="aurateam/AURA", help="Model name served by vLLM.")
|
||||
parser.add_argument("--api-base", default="http://localhost:8091/v1", help="OpenAI-compatible API base.")
|
||||
parser.add_argument("--ip", default="127.0.0.1", help="Gradio host.")
|
||||
parser.add_argument("--port", type=int, default=7862, help="Gradio port.")
|
||||
parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _audio_to_data_url(audio_file: Any | None) -> str | None:
|
||||
if audio_file is None:
|
||||
return None
|
||||
|
||||
def _path_to_data_url(path_str: str) -> str | None:
|
||||
if not path_str:
|
||||
return None
|
||||
path = Path(path_str)
|
||||
if not path.exists():
|
||||
return None
|
||||
with open(path, "rb") as f:
|
||||
return f"data:audio/wav;base64,{base64.b64encode(f.read()).decode('utf-8')}"
|
||||
|
||||
if isinstance(audio_file, str):
|
||||
return _path_to_data_url(audio_file)
|
||||
|
||||
# Gradio >=5 may return dict payloads with local temp path.
|
||||
if isinstance(audio_file, dict):
|
||||
path_val = audio_file.get("path") or audio_file.get("name")
|
||||
if isinstance(path_val, str):
|
||||
return _path_to_data_url(path_val)
|
||||
return None
|
||||
|
||||
if isinstance(audio_file, tuple):
|
||||
# Common shape: (sample_rate, np.ndarray)
|
||||
if len(audio_file) == 2 and isinstance(audio_file[0], (int, float)):
|
||||
sample_rate, audio_np = audio_file
|
||||
audio_np = np.asarray(audio_np)
|
||||
if audio_np.ndim > 1:
|
||||
audio_np = audio_np[:, 0]
|
||||
if audio_np.dtype != np.int16:
|
||||
audio_np = np.clip(audio_np.astype(np.float32), -1.0, 1.0)
|
||||
buf = io.BytesIO()
|
||||
sf.write(buf, audio_np, int(sample_rate), format="WAV")
|
||||
return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
|
||||
# Compatibility shape from some gradio builds: (filepath, (sr, np.ndarray))
|
||||
if len(audio_file) >= 1 and isinstance(audio_file[0], str):
|
||||
path_url = _path_to_data_url(audio_file[0])
|
||||
if path_url is not None:
|
||||
return path_url
|
||||
return None
|
||||
|
||||
|
||||
def _video_to_data_url(video_file: str | None) -> str | None:
|
||||
if not video_file:
|
||||
return None
|
||||
path = Path(video_file)
|
||||
mime = {
|
||||
".mp4": "video/mp4",
|
||||
".webm": "video/webm",
|
||||
".mov": "video/quicktime",
|
||||
".avi": "video/x-msvideo",
|
||||
".mkv": "video/x-matroska",
|
||||
}.get(path.suffix.lower(), "video/mp4")
|
||||
with open(path, "rb") as f:
|
||||
return f"data:{mime};base64,{base64.b64encode(f.read()).decode('utf-8')}"
|
||||
|
||||
|
||||
def _sampling_params_list() -> list[dict[str, Any]]:
|
||||
return [
|
||||
{"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED},
|
||||
{
|
||||
"temperature": 0.5,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 256,
|
||||
"seed": SEED,
|
||||
"repetition_penalty": 1.0,
|
||||
},
|
||||
{
|
||||
"temperature": 0.9,
|
||||
"top_k": 50,
|
||||
"max_tokens": 4096,
|
||||
"seed": SEED,
|
||||
"detokenize": False,
|
||||
"repetition_penalty": 1.05,
|
||||
"stop_token_ids": [2150],
|
||||
},
|
||||
{
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 65536,
|
||||
"seed": SEED,
|
||||
"repetition_penalty": 1.0,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _decode_audio_choice(audio_data: str | None) -> tuple[int, np.ndarray] | None:
|
||||
if not audio_data:
|
||||
return None
|
||||
wav_bytes = base64.b64decode(audio_data)
|
||||
audio_np, sample_rate = sf.read(io.BytesIO(wav_bytes))
|
||||
if audio_np.ndim > 1:
|
||||
audio_np = audio_np[:, 0]
|
||||
return int(sample_rate), audio_np.astype(np.float32)
|
||||
|
||||
|
||||
def build_interface(client: OpenAI, model: str):
|
||||
def run(
|
||||
audio_file,
|
||||
video_file,
|
||||
prompt: str,
|
||||
aura_system_prompt: str,
|
||||
tts_task_type: str,
|
||||
tts_language: str,
|
||||
tts_speaker: str,
|
||||
tts_instruct: str,
|
||||
tts_ref_audio: str,
|
||||
tts_ref_text: str,
|
||||
tts_x_vector_only_mode: bool,
|
||||
tts_pass_token_ids: bool,
|
||||
):
|
||||
audio_url = _audio_to_data_url(audio_file)
|
||||
video_url = _video_to_data_url(video_file)
|
||||
if not audio_url:
|
||||
return "Please provide audio input for the ASR stage.", None
|
||||
if tts_task_type == "Base" and (not (tts_ref_audio or "").strip() or not (tts_ref_text or "").strip()):
|
||||
return "Base TTS requires both reference audio and reference transcript.", None
|
||||
content: list[dict[str, Any]] = [{"type": "audio_url", "audio_url": {"url": audio_url}}]
|
||||
if video_url:
|
||||
content.append({"type": "video_url", "video_url": {"url": video_url}})
|
||||
content.append(
|
||||
{"type": "text", "text": prompt or "Use the audio and video together to decide whether a reply is needed."}
|
||||
)
|
||||
|
||||
additional_information: dict[str, Any] = {
|
||||
"aura_system_prompt": aura_system_prompt,
|
||||
"tts_task_type": tts_task_type,
|
||||
"tts_instruct": tts_instruct,
|
||||
"tts_pass_token_ids": bool(tts_pass_token_ids),
|
||||
}
|
||||
if tts_task_type == "CustomVoice":
|
||||
additional_information.update(
|
||||
{
|
||||
"tts_language": tts_language or "English",
|
||||
"tts_speaker": tts_speaker or "Vivian",
|
||||
}
|
||||
)
|
||||
else:
|
||||
additional_information.update(
|
||||
{
|
||||
"tts_ref_audio": tts_ref_audio.strip() if tts_ref_audio else None,
|
||||
"tts_ref_text": tts_ref_text.strip() if tts_ref_text else None,
|
||||
"tts_x_vector_only_mode": bool(tts_x_vector_only_mode),
|
||||
}
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": content}],
|
||||
modalities=["text", "audio"],
|
||||
extra_body={
|
||||
"sampling_params_list": _sampling_params_list(),
|
||||
"additional_information": additional_information,
|
||||
},
|
||||
timeout=600.0,
|
||||
)
|
||||
|
||||
text_parts: list[str] = []
|
||||
audio_output = None
|
||||
for choice in response.choices:
|
||||
message = choice.message
|
||||
if message.content:
|
||||
text_parts.append(str(message.content))
|
||||
if getattr(message, "audio", None):
|
||||
audio_output = _decode_audio_choice(message.audio.data)
|
||||
return "\n\n".join(text_parts) or "No text output.", audio_output
|
||||
|
||||
default_system = (
|
||||
"You are receiving a live video stream where the final frame is the present moment. "
|
||||
"Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English."
|
||||
)
|
||||
|
||||
def _toggle_tts_advanced(task_type: str):
|
||||
is_base = task_type == "Base"
|
||||
return (
|
||||
gr.update(visible=not is_base),
|
||||
gr.update(visible=is_base),
|
||||
)
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
gr.Markdown("# AURA Omni")
|
||||
gr.Markdown("ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav")
|
||||
with gr.Row():
|
||||
audio_input = gr.Audio(label="Audio input", sources=["upload", "microphone"], type="numpy")
|
||||
video_input = gr.Video(label="Video input (optional)", sources=["upload"])
|
||||
prompt = gr.Textbox(
|
||||
label="Instruction",
|
||||
value="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.",
|
||||
lines=2,
|
||||
)
|
||||
with gr.Accordion("Advanced", open=False):
|
||||
aura_system_prompt = gr.Textbox(label="AURA system prompt", value=default_system, lines=4)
|
||||
tts_task_type = gr.Radio(choices=["Base", "CustomVoice"], value="Base", label="TTS task type")
|
||||
tts_instruct = gr.Textbox(label="TTS instruction", value="")
|
||||
tts_pass_token_ids = gr.Checkbox(
|
||||
label="Pass AURA token ids directly to TTS",
|
||||
value=False,
|
||||
)
|
||||
|
||||
with gr.Group(visible=False) as customvoice_config_group:
|
||||
tts_language = gr.Dropdown(
|
||||
choices=["Chinese", "English", "Japanese", "Korean", "Cantonese"],
|
||||
value="Chinese",
|
||||
label="CustomVoice language",
|
||||
)
|
||||
tts_speaker = gr.Textbox(label="CustomVoice speaker", value="Vivian")
|
||||
|
||||
with gr.Group(visible=True) as base_config_group:
|
||||
tts_ref_audio = gr.Textbox(
|
||||
label="Base reference audio path/URL",
|
||||
value=DEFAULT_TTS_REF_AUDIO,
|
||||
)
|
||||
tts_ref_text = gr.Textbox(
|
||||
label="Base reference transcript",
|
||||
value=DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
lines=3,
|
||||
)
|
||||
tts_x_vector_only_mode = gr.Checkbox(
|
||||
label="Base x-vector only mode (disable ICL)",
|
||||
value=False,
|
||||
)
|
||||
|
||||
tts_task_type.change(
|
||||
_toggle_tts_advanced,
|
||||
inputs=[tts_task_type],
|
||||
outputs=[customvoice_config_group, base_config_group],
|
||||
)
|
||||
button = gr.Button("Generate", variant="primary")
|
||||
with gr.Row():
|
||||
text_output = gr.Textbox(label="Text output", lines=8)
|
||||
audio_output = gr.Audio(label="Audio output", interactive=False)
|
||||
button.click(
|
||||
run,
|
||||
inputs=[
|
||||
audio_input,
|
||||
video_input,
|
||||
prompt,
|
||||
aura_system_prompt,
|
||||
tts_task_type,
|
||||
tts_language,
|
||||
tts_speaker,
|
||||
tts_instruct,
|
||||
tts_ref_audio,
|
||||
tts_ref_text,
|
||||
tts_x_vector_only_mode,
|
||||
tts_pass_token_ids,
|
||||
],
|
||||
outputs=[text_output, audio_output],
|
||||
)
|
||||
demo.queue()
|
||||
return demo
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
client = OpenAI(api_key="EMPTY", base_url=args.api_base)
|
||||
demo = build_interface(client, args.model)
|
||||
demo.launch(server_name=args.ip, server_port=args.port, share=args.share)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,199 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""OpenAI-compatible client for the AURA Omni pipeline."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
|
||||
import soundfile as sf
|
||||
from openai import OpenAI
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
default_qwen3_tts_ref_audio_path,
|
||||
)
|
||||
|
||||
SEED = 42
|
||||
DEFAULT_MODEL = "aurateam/AURA"
|
||||
DEFAULT_VIDEO_URL = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
|
||||
|
||||
def _encode_file(path: str) -> str:
|
||||
with open(path, "rb") as f:
|
||||
return base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
|
||||
def _data_url(path: str, default_mime: str) -> str:
|
||||
suffix = os.path.splitext(path)[1].lower()
|
||||
mime_by_suffix = {
|
||||
".wav": "audio/wav",
|
||||
".mp3": "audio/mpeg",
|
||||
".ogg": "audio/ogg",
|
||||
".flac": "audio/flac",
|
||||
".m4a": "audio/mp4",
|
||||
".mp4": "video/mp4",
|
||||
".webm": "video/webm",
|
||||
".mov": "video/quicktime",
|
||||
".avi": "video/x-msvideo",
|
||||
".mkv": "video/x-matroska",
|
||||
}
|
||||
return f"data:{mime_by_suffix.get(suffix, default_mime)};base64,{_encode_file(path)}"
|
||||
|
||||
|
||||
def media_url(path_or_url: str | None, *, kind: str) -> str:
|
||||
if path_or_url:
|
||||
if path_or_url.startswith(("http://", "https://", "data:")):
|
||||
return path_or_url
|
||||
if not os.path.exists(path_or_url):
|
||||
raise FileNotFoundError(f"{kind} file not found: {path_or_url}")
|
||||
return _data_url(path_or_url, "audio/wav" if kind == "audio" else "video/mp4")
|
||||
if kind == "audio":
|
||||
return AudioAsset("mary_had_lamb").url
|
||||
return DEFAULT_VIDEO_URL
|
||||
|
||||
|
||||
def sampling_params_list() -> list[dict]:
|
||||
return [
|
||||
{"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED},
|
||||
{
|
||||
"temperature": 0.5,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 256,
|
||||
"seed": SEED,
|
||||
"repetition_penalty": 1.0,
|
||||
},
|
||||
{
|
||||
"temperature": 0.9,
|
||||
"top_k": 50,
|
||||
"max_tokens": 4096,
|
||||
"seed": SEED,
|
||||
"detokenize": False,
|
||||
"repetition_penalty": 1.05,
|
||||
"stop_token_ids": [2150],
|
||||
},
|
||||
{
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 65536,
|
||||
"seed": SEED,
|
||||
"repetition_penalty": 1.0,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def parse_modalities(value: str | None) -> list[str] | None:
|
||||
if not value:
|
||||
return None
|
||||
return [item.strip() for item in value.split(",") if item.strip()]
|
||||
|
||||
|
||||
def save_response(response, output_dir: str) -> None:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
for idx, choice in enumerate(response.choices):
|
||||
message = choice.message
|
||||
if message.content:
|
||||
out_txt = os.path.join(output_dir, f"choice_{idx}.txt")
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
f.write(str(message.content).strip() + "\n")
|
||||
print(f"Text saved to {out_txt}")
|
||||
print(message.content)
|
||||
if getattr(message, "audio", None):
|
||||
audio_bytes = base64.b64decode(message.audio.data)
|
||||
audio_np, sample_rate = sf.read(io.BytesIO(audio_bytes))
|
||||
out_wav = os.path.join(output_dir, f"choice_{idx}.wav")
|
||||
sf.write(out_wav, audio_np, int(sample_rate), format="WAV")
|
||||
print(f"Audio saved to {out_wav}")
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
client = OpenAI(base_url=f"http://{args.host}:{args.port}/v1", api_key="EMPTY")
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": media_url(args.audio_path, kind="audio")}},
|
||||
{"type": "video_url", "video_url": {"url": media_url(args.video_path, kind="video")}},
|
||||
{"type": "text", "text": args.prompt},
|
||||
],
|
||||
}
|
||||
]
|
||||
response = client.chat.completions.create(
|
||||
model=args.model,
|
||||
messages=messages,
|
||||
modalities=parse_modalities(args.modalities),
|
||||
extra_body={
|
||||
"sampling_params_list": sampling_params_list(),
|
||||
"additional_information": {
|
||||
"aura_system_prompt": args.aura_system_prompt,
|
||||
"tts_task_type": args.tts_task_type,
|
||||
"tts_language": args.tts_language,
|
||||
"tts_speaker": args.tts_speaker,
|
||||
"tts_instruct": args.tts_instruct,
|
||||
"tts_ref_audio": args.tts_ref_audio,
|
||||
"tts_ref_text": args.tts_ref_text,
|
||||
"tts_x_vector_only_mode": args.tts_x_vector_only_mode,
|
||||
"tts_pass_token_ids": args.tts_pass_token_ids,
|
||||
},
|
||||
},
|
||||
timeout=args.timeout,
|
||||
)
|
||||
save_response(response, args.output_dir)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="AURA Omni online serving client")
|
||||
parser.add_argument("--host", default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8091)
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL)
|
||||
parser.add_argument("--audio-path", default=None, help="Audio file, URL, or data URL.")
|
||||
parser.add_argument("--video-path", default=None, help="Video file, URL, or data URL.")
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.",
|
||||
)
|
||||
parser.add_argument("--modalities", default="text,audio")
|
||||
parser.add_argument("--output-dir", default="output_aura_omni_online")
|
||||
parser.add_argument(
|
||||
"--aura-system-prompt",
|
||||
default=(
|
||||
"You are receiving a live video stream where the final frame is the present moment. "
|
||||
"Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"])
|
||||
parser.add_argument("--tts-language", default="English")
|
||||
parser.add_argument("--tts-speaker", default="Vivian")
|
||||
parser.add_argument("--tts-instruct", default="")
|
||||
parser.add_argument(
|
||||
"--tts-ref-audio",
|
||||
default=default_qwen3_tts_ref_audio_path(),
|
||||
help="Base-mode reference audio path/URL visible to server.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-ref-text",
|
||||
default=DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
help="Base-mode reference audio transcript.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-x-vector-only-mode",
|
||||
action="store_true",
|
||||
help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-pass-token-ids",
|
||||
action="store_true",
|
||||
help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.",
|
||||
)
|
||||
parser.add_argument("--timeout", type=float, default=600.0)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -0,0 +1,60 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
PORT="${PORT:-8091}"
|
||||
MODEL="${MODEL:-aurateam/AURA}"
|
||||
OUTPUT_DIR="${OUTPUT_DIR:-output_aura_omni_online}"
|
||||
TTS_PASS_TOKEN_IDS="${TTS_PASS_TOKEN_IDS:-false}"
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
VLLM_OMNI_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
CLONE_REF_AUDIO="${VLLM_OMNI_ROOT}/tests/assets/qwen3_tts/clone_2.wav"
|
||||
CLONE_REF_TEXT="Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
mkdir -p "$OUTPUT_DIR"
|
||||
|
||||
MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
|
||||
SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
|
||||
request_body=$(cat <<EOF
|
||||
{
|
||||
"model": "$MODEL",
|
||||
"modalities": ["text", "audio"],
|
||||
"sampling_params_list": [
|
||||
{"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42},
|
||||
{"temperature": 0.5, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42, "repetition_penalty": 1.0},
|
||||
{"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "seed": 42, "detokenize": false, "repetition_penalty": 1.05, "stop_token_ids": [2150]},
|
||||
{"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 65536, "seed": 42, "repetition_penalty": 1.0}
|
||||
],
|
||||
"additional_information": {
|
||||
"aura_system_prompt": "You are receiving a live video stream where the final frame is the present moment. Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English.",
|
||||
"tts_task_type": "Base",
|
||||
"tts_ref_audio": "file://${CLONE_REF_AUDIO}",
|
||||
"tts_ref_text": "${CLONE_REF_TEXT}",
|
||||
"tts_language": "English",
|
||||
"tts_speaker": "Vivian",
|
||||
"tts_instruct": "",
|
||||
"tts_pass_token_ids": ${TTS_PASS_TOKEN_IDS}
|
||||
},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": "$MARY_HAD_LAMB_AUDIO_URL"}},
|
||||
{"type": "video_url", "video_url": {"url": "$SAMPLE_VIDEO_URL"}},
|
||||
{"type": "text", "text": "Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English."}
|
||||
]
|
||||
}]
|
||||
}
|
||||
EOF
|
||||
)
|
||||
|
||||
response=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
|
||||
-X POST "http://localhost:${PORT}/v1/chat/completions" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "$request_body")
|
||||
|
||||
echo "$response" | jq '.choices[].message.content'
|
||||
|
||||
audio_b64=$(echo "$response" | jq -r '.choices[]?.message.audio.data // empty' | head -n 1)
|
||||
if [[ -n "$audio_b64" ]]; then
|
||||
echo "$audio_b64" | base64 -d > "${OUTPUT_DIR}/aura_omni_output.wav"
|
||||
echo "Audio saved to ${OUTPUT_DIR}/aura_omni_output.wav"
|
||||
fi
|
||||
+79
@@ -0,0 +1,79 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
MODEL="aurateam/AURA"
|
||||
SERVER_MODEL="aurateam/AURA"
|
||||
DEPLOY_CONFIG="/data/yrr/vllm-omni/vllm_omni/deploy/aura_omni.yaml"
|
||||
SERVER_PORT=8091
|
||||
GRADIO_PORT=7862
|
||||
SERVER_HOST="0.0.0.0"
|
||||
GRADIO_IP="127.0.0.1"
|
||||
GRADIO_SHARE=false
|
||||
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--model) MODEL="$2"; shift 2 ;;
|
||||
--server-model) SERVER_MODEL="$2"; shift 2 ;;
|
||||
--deploy-config) DEPLOY_CONFIG="$2"; shift 2 ;;
|
||||
--server-port) SERVER_PORT="$2"; shift 2 ;;
|
||||
--gradio-port) GRADIO_PORT="$2"; shift 2 ;;
|
||||
--server-host) SERVER_HOST="$2"; shift 2 ;;
|
||||
--gradio-ip) GRADIO_IP="$2"; shift 2 ;;
|
||||
--share) GRADIO_SHARE=true; shift ;;
|
||||
--help)
|
||||
echo "Usage: $0 [--model SERVED_MODEL_NAME] [--server-model MODEL_PATH] [--deploy-config YAML] [--server-port PORT] [--gradio-port PORT] [--share]"
|
||||
exit 0
|
||||
;;
|
||||
*) echo "Unknown option: $1"; exit 1 ;;
|
||||
esac
|
||||
done
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
API_BASE="http://localhost:${SERVER_PORT}/v1"
|
||||
LOG_FILE="/tmp/aura_omni_vllm_${SERVER_PORT}.log"
|
||||
|
||||
cleanup() {
|
||||
echo "Shutting down..."
|
||||
[[ -n "${SERVER_PID:-}" ]] && kill "$SERVER_PID" 2>/dev/null || true
|
||||
[[ -n "${GRADIO_PID:-}" ]] && kill "$GRADIO_PID" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup SIGINT SIGTERM EXIT
|
||||
|
||||
vllm serve "$SERVER_MODEL" \
|
||||
--omni \
|
||||
--host "$SERVER_HOST" \
|
||||
--port "$SERVER_PORT" \
|
||||
--deploy-config "$DEPLOY_CONFIG" \
|
||||
--served-model-name "$MODEL" \
|
||||
--trust-remote-code 2>&1 | tee "$LOG_FILE" &
|
||||
SERVER_PID=$!
|
||||
|
||||
echo "Waiting for server startup..."
|
||||
for _ in $(seq 1 600); do
|
||||
if grep -q "Application startup complete" "$LOG_FILE" 2>/dev/null; then
|
||||
break
|
||||
fi
|
||||
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
|
||||
echo "vLLM server exited before startup completed"
|
||||
wait "$SERVER_PID" || true
|
||||
exit 1
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
# cd "$SCRIPT_DIR"
|
||||
GRADIO_CMD=(python gradio_demo.py --model "$MODEL" --api-base "$API_BASE" --ip "$GRADIO_IP" --port "$GRADIO_PORT")
|
||||
if [[ "$GRADIO_SHARE" == "true" ]]; then
|
||||
GRADIO_CMD+=(--share)
|
||||
fi
|
||||
"${GRADIO_CMD[@]}" &
|
||||
GRADIO_PID=$!
|
||||
|
||||
echo "vLLM server: http://${SERVER_HOST}:${SERVER_PORT}"
|
||||
echo "Gradio demo: http://${GRADIO_IP}:${GRADIO_PORT}"
|
||||
if [[ -n "${SERVER_PID:-}" ]]; then
|
||||
wait "$SERVER_PID" "$GRADIO_PID"
|
||||
else
|
||||
wait "$GRADIO_PID"
|
||||
fi
|
||||
@@ -1,175 +0,0 @@
|
||||
# Helios Video Generation via Online API
|
||||
|
||||
This example demonstrates how to use the `/v1/videos` API with Helios models using the generic `extra_params` field.
|
||||
|
||||
## Overview
|
||||
|
||||
The `/v1/videos` API now supports model-specific parameters through the `extra_params` field, which accepts a JSON object containing any model-specific configuration. This allows supporting new models like Helios without modifying the API.
|
||||
|
||||
## Helios Model Variants
|
||||
|
||||
- **Helios-Base**: Basic T2V/I2V/V2V generation (Stage 1 only)
|
||||
- **Helios-Mid**: Advanced generation with CFG-Zero* (Stage 2)
|
||||
- **Helios-Distilled**: Fast generation with DMD (Stage 2)
|
||||
|
||||
## API Usage Examples
|
||||
|
||||
### 1. T2V (Text-to-Video) - Helios-Base
|
||||
|
||||
Basic text-to-video generation:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A serene lakeside sunrise with mist over the water." \
|
||||
-F "model=BestWishYsh/Helios-Base" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "num_frames=99" \
|
||||
-F "num_inference_steps=50" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42"
|
||||
```
|
||||
|
||||
### 2. I2V (Image-to-Video) - Helios-Base
|
||||
|
||||
Generate video from an input image:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=The lake water gently ripples as morning mist rises." \
|
||||
-F "model=BestWishYsh/Helios-Base" \
|
||||
-F "input_reference=@/path/to/image.jpg" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "num_frames=99" \
|
||||
-F "guidance_scale=5.0"
|
||||
```
|
||||
|
||||
### 3. Helios-Mid with Stage 2 + CFG-Zero*
|
||||
|
||||
Advanced generation with pyramid multi-stage denoising:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A serene lakeside sunrise with mist over the water." \
|
||||
-F "model=BestWishYsh/Helios-Mid" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F 'extra_params={
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"use_cfg_zero_star": true,
|
||||
"use_zero_init": true,
|
||||
"zero_steps": 1
|
||||
}'
|
||||
```
|
||||
|
||||
### 4. Helios-Distilled with DMD
|
||||
|
||||
Fast generation with Distribution Matching Distillation:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A serene lakeside sunrise with mist over the water." \
|
||||
-F "model=BestWishYsh/Helios-Distilled" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "guidance_scale=1.0" \
|
||||
-F 'extra_params={
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [2, 2, 2],
|
||||
"is_amplify_first_chunk": true
|
||||
}'
|
||||
```
|
||||
|
||||
## Model-Specific Parameters
|
||||
|
||||
The `extra_params` field accepts a JSON object with model-specific parameters. For Helios models, supported parameters include:
|
||||
|
||||
### Stage 2 Parameters
|
||||
- `is_enable_stage2` (bool): Enable pyramid multi-stage denoising
|
||||
- `pyramid_num_stages` (int): Number of pyramid stages (default: 3)
|
||||
- `pyramid_num_inference_steps_list` (array): Steps per stage, e.g., `[20, 20, 20]`
|
||||
|
||||
### CFG Zero Star Parameters (Helios-Mid)
|
||||
- `use_cfg_zero_star` (bool): Enable CFG Zero Star guidance
|
||||
- `use_zero_init` (bool): Use zero initialization for first steps
|
||||
- `zero_steps` (int): Number of initial zero prediction steps (default: 1)
|
||||
|
||||
### DMD Parameters (Helios-Distilled)
|
||||
- `is_amplify_first_chunk` (bool): Enable DMD amplification for first chunk
|
||||
|
||||
### Video Input (V2V mode)
|
||||
For video-to-video generation, upload a video file via `input_reference`:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=Transform the video into a watercolor painting style." \
|
||||
-F "model=BestWishYsh/Helios-Base" \
|
||||
-F "input_reference=@/path/to/video.mp4" \
|
||||
-F "guidance_scale=5.0"
|
||||
```
|
||||
|
||||
## Python Example
|
||||
|
||||
```python
|
||||
import requests
|
||||
import json
|
||||
|
||||
url = "http://localhost:8000/v1/videos"
|
||||
|
||||
# Helios-Mid with Stage 2
|
||||
data = {
|
||||
"prompt": "A serene lakeside sunrise with mist over the water.",
|
||||
"model": "BestWishYsh/Helios-Mid",
|
||||
"width": 640,
|
||||
"height": 384,
|
||||
"guidance_scale": 5.0,
|
||||
"extra_params": json.dumps({
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"use_cfg_zero_star": True,
|
||||
"use_zero_init": True,
|
||||
"zero_steps": 1
|
||||
})
|
||||
}
|
||||
|
||||
response = requests.post(url, data=data)
|
||||
video_job = response.json()
|
||||
print(f"Video job created: {video_job['id']}")
|
||||
|
||||
# Poll for completion
|
||||
import time
|
||||
while True:
|
||||
status_response = requests.get(f"{url}/{video_job['id']}")
|
||||
status = status_response.json()
|
||||
|
||||
if status['status'] == 'completed':
|
||||
print(f"Video generated: {status['file_name']}")
|
||||
break
|
||||
elif status['status'] == 'failed':
|
||||
print(f"Generation failed: {status.get('error')}")
|
||||
break
|
||||
|
||||
print(f"Progress: {status['progress']}%")
|
||||
time.sleep(2)
|
||||
```
|
||||
|
||||
## Benefits of This Approach
|
||||
|
||||
1. **No API Changes**: New models can be supported without modifying the API
|
||||
2. **Backward Compatible**: Existing clients continue to work
|
||||
3. **Flexible**: Any model-specific parameter can be passed
|
||||
4. **Type-Safe**: Parameters are validated at the model level
|
||||
5. **Future-Proof**: Supports models that don't exist yet
|
||||
|
||||
## Notes
|
||||
|
||||
- The `extra_params` field must be a valid JSON object
|
||||
- Parameters are passed directly to the model's `extra_args`
|
||||
- Invalid parameters will be caught by the model implementation
|
||||
- Tensor inputs (images/videos) should use `input_reference`, not `extra_params`
|
||||
@@ -1,212 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Helios video generation client example using the /v1/videos API.
|
||||
|
||||
This example demonstrates how to use the generic extra_params field
|
||||
to pass Helios-specific parameters without modifying the API.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def create_video_job(
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
model: str,
|
||||
width: int = 640,
|
||||
height: int = 384,
|
||||
num_frames: int = 99,
|
||||
guidance_scale: float = 5.0,
|
||||
seed: int = 42,
|
||||
extra_params: dict | None = None,
|
||||
input_image: Path | None = None,
|
||||
) -> dict:
|
||||
"""Create a video generation job."""
|
||||
data = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_frames": num_frames,
|
||||
"guidance_scale": guidance_scale,
|
||||
"seed": seed,
|
||||
}
|
||||
|
||||
files = {}
|
||||
if input_image:
|
||||
files["input_reference"] = open(input_image, "rb")
|
||||
|
||||
if extra_params:
|
||||
data["extra_params"] = json.dumps(extra_params)
|
||||
|
||||
response = requests.post(f"{api_url}/v1/videos", data=data, files=files)
|
||||
response.raise_for_status()
|
||||
|
||||
if files:
|
||||
files["input_reference"].close()
|
||||
|
||||
return response.json()
|
||||
|
||||
|
||||
def poll_video_status(api_url: str, video_id: str, poll_interval: int = 2) -> dict:
|
||||
"""Poll video generation status until completion."""
|
||||
print(f"Polling video job: {video_id}")
|
||||
|
||||
while True:
|
||||
response = requests.get(f"{api_url}/v1/videos/{video_id}")
|
||||
response.raise_for_status()
|
||||
status_data = response.json()
|
||||
|
||||
status = status_data["status"]
|
||||
progress = status_data.get("progress", 0)
|
||||
|
||||
print(f"Status: {status}, Progress: {progress}%")
|
||||
|
||||
if status == "completed":
|
||||
print("Video generation completed!")
|
||||
return status_data
|
||||
elif status == "failed":
|
||||
error = status_data.get("error", {})
|
||||
raise RuntimeError(f"Video generation failed: {error}")
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Helios video generation client")
|
||||
parser.add_argument(
|
||||
"--api-url",
|
||||
default="http://localhost:8000",
|
||||
help="API server URL",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="BestWishYsh/Helios-Base",
|
||||
help="Model name (Helios-Base, Helios-Mid, Helios-Distilled)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="A serene lakeside sunrise with mist over the water.",
|
||||
help="Text prompt",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=640,
|
||||
help="Video width",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=384,
|
||||
help="Video height",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=99,
|
||||
help="Number of frames",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--guidance-scale",
|
||||
type=float,
|
||||
default=5.0,
|
||||
help="Guidance scale",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Random seed",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-image",
|
||||
type=Path,
|
||||
help="Input image for I2V mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preset",
|
||||
choices=["base", "mid-stage2", "distilled"],
|
||||
default="base",
|
||||
help="Helios preset configuration",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Define preset configurations
|
||||
presets = {
|
||||
"base": None, # No extra params for base model
|
||||
"mid-stage2": {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"use_cfg_zero_star": True,
|
||||
"use_zero_init": True,
|
||||
"zero_steps": 1,
|
||||
},
|
||||
"distilled": {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [2, 2, 2],
|
||||
"is_amplify_first_chunk": True,
|
||||
},
|
||||
}
|
||||
|
||||
extra_params = presets[args.preset]
|
||||
|
||||
print("=" * 50)
|
||||
print("Helios Video Generation")
|
||||
print("=" * 50)
|
||||
print(f"API URL: {args.api_url}")
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Preset: {args.preset}")
|
||||
print(f"Prompt: {args.prompt}")
|
||||
print(f"Size: {args.width}x{args.height}")
|
||||
print(f"Frames: {args.num_frames}")
|
||||
print(f"Guidance Scale: {args.guidance_scale}")
|
||||
if extra_params:
|
||||
print(f"Extra Params: {json.dumps(extra_params, indent=2)}")
|
||||
print()
|
||||
|
||||
# Create video job
|
||||
print("Creating video generation job...")
|
||||
job = create_video_job(
|
||||
api_url=args.api_url,
|
||||
prompt=args.prompt,
|
||||
model=args.model,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
num_frames=args.num_frames,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
extra_params=extra_params,
|
||||
input_image=args.input_image,
|
||||
)
|
||||
|
||||
video_id = job["id"]
|
||||
print(f"Video job created: {video_id}")
|
||||
print()
|
||||
|
||||
# Poll for completion
|
||||
result = poll_video_status(args.api_url, video_id)
|
||||
|
||||
print()
|
||||
print("=" * 50)
|
||||
print("Generation Complete")
|
||||
print("=" * 50)
|
||||
print(f"Video ID: {result['id']}")
|
||||
print(f"File: {result.get('file_name', 'N/A')}")
|
||||
print(f"Inference Time: {result.get('inference_time_s', 'N/A')}s")
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,79 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Helios-Distilled with DMD Example
|
||||
# Fast generation with Distribution Matching Distillation
|
||||
|
||||
API_URL="${API_URL:-http://localhost:8000}"
|
||||
MODEL="${MODEL:-BestWishYsh/Helios-Distilled}"
|
||||
PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
|
||||
|
||||
echo "==================================="
|
||||
echo "Helios-Distilled with DMD"
|
||||
echo "==================================="
|
||||
echo "API URL: $API_URL"
|
||||
echo "Model: $MODEL"
|
||||
echo "Prompt: $PROMPT"
|
||||
echo ""
|
||||
|
||||
# Model-specific parameters for Helios-Distilled
|
||||
extra_params='{
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [2, 2, 2],
|
||||
"is_amplify_first_chunk": true
|
||||
}'
|
||||
|
||||
echo "Model extra params: $extra_params"
|
||||
echo ""
|
||||
|
||||
# Create video generation job
|
||||
echo "Creating video generation job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
|
||||
-F "prompt=$PROMPT" \
|
||||
-F "model=$MODEL" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "guidance_scale=1.0" \
|
||||
-F "seed=42" \
|
||||
-F "extra_params=$extra_params")
|
||||
|
||||
echo "Response: $RESPONSE"
|
||||
echo ""
|
||||
|
||||
# Extract video ID
|
||||
VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
|
||||
|
||||
if [ -z "$VIDEO_ID" ]; then
|
||||
echo "Error: Failed to create video job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Video job created: $VIDEO_ID"
|
||||
echo ""
|
||||
|
||||
# Poll for completion
|
||||
echo "Polling for completion..."
|
||||
while true; do
|
||||
STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
|
||||
STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
|
||||
PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
|
||||
|
||||
echo "Status: $STATUS, Progress: $PROGRESS%"
|
||||
|
||||
if [ "$STATUS" = "completed" ]; then
|
||||
echo ""
|
||||
echo "Video generation completed!"
|
||||
echo "Full response:"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
break
|
||||
elif [ "$STATUS" = "failed" ]; then
|
||||
echo ""
|
||||
echo "Video generation failed!"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sleep 2
|
||||
done
|
||||
@@ -1,81 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Helios-Mid with Stage 2 + CFG-Zero* Example
|
||||
# This demonstrates advanced generation with pyramid multi-stage denoising
|
||||
|
||||
API_URL="${API_URL:-http://localhost:8000}"
|
||||
MODEL="${MODEL:-BestWishYsh/Helios-Mid}"
|
||||
PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
|
||||
|
||||
echo "==================================="
|
||||
echo "Helios-Mid Stage 2 + CFG-Zero*"
|
||||
echo "==================================="
|
||||
echo "API URL: $API_URL"
|
||||
echo "Model: $MODEL"
|
||||
echo "Prompt: $PROMPT"
|
||||
echo ""
|
||||
|
||||
# Model-specific parameters for Helios-Mid
|
||||
extra_params='{
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"use_cfg_zero_star": true,
|
||||
"use_zero_init": true,
|
||||
"zero_steps": 1
|
||||
}'
|
||||
|
||||
echo "Model extra params: $extra_params"
|
||||
echo ""
|
||||
|
||||
# Create video generation job
|
||||
echo "Creating video generation job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
|
||||
-F "prompt=$PROMPT" \
|
||||
-F "model=$MODEL" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42" \
|
||||
-F "extra_params=$extra_params")
|
||||
|
||||
echo "Response: $RESPONSE"
|
||||
echo ""
|
||||
|
||||
# Extract video ID
|
||||
VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
|
||||
|
||||
if [ -z "$VIDEO_ID" ]; then
|
||||
echo "Error: Failed to create video job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Video job created: $VIDEO_ID"
|
||||
echo ""
|
||||
|
||||
# Poll for completion
|
||||
echo "Polling for completion..."
|
||||
while true; do
|
||||
STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
|
||||
STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
|
||||
PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
|
||||
|
||||
echo "Status: $STATUS, Progress: $PROGRESS%"
|
||||
|
||||
if [ "$STATUS" = "completed" ]; then
|
||||
echo ""
|
||||
echo "Video generation completed!"
|
||||
echo "Full response:"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
break
|
||||
elif [ "$STATUS" = "failed" ]; then
|
||||
echo ""
|
||||
echo "Video generation failed!"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sleep 2
|
||||
done
|
||||
@@ -1,69 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Helios Text-to-Video Generation Example
|
||||
# This script demonstrates how to use the /v1/videos API with Helios models
|
||||
|
||||
API_URL="${API_URL:-http://localhost:8000}"
|
||||
MODEL="${MODEL:-BestWishYsh/Helios-Base}"
|
||||
PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
|
||||
|
||||
echo "==================================="
|
||||
echo "Helios T2V Generation"
|
||||
echo "==================================="
|
||||
echo "API URL: $API_URL"
|
||||
echo "Model: $MODEL"
|
||||
echo "Prompt: $PROMPT"
|
||||
echo ""
|
||||
|
||||
# Create video generation job
|
||||
echo "Creating video generation job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
|
||||
-F "prompt=$PROMPT" \
|
||||
-F "model=$MODEL" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "num_frames=99" \
|
||||
-F "num_inference_steps=50" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42")
|
||||
|
||||
echo "Response: $RESPONSE"
|
||||
echo ""
|
||||
|
||||
# Extract video ID
|
||||
VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
|
||||
|
||||
if [ -z "$VIDEO_ID" ]; then
|
||||
echo "Error: Failed to create video job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Video job created: $VIDEO_ID"
|
||||
echo ""
|
||||
|
||||
# Poll for completion
|
||||
echo "Polling for completion..."
|
||||
while true; do
|
||||
STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
|
||||
STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
|
||||
PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
|
||||
|
||||
echo "Status: $STATUS, Progress: $PROGRESS%"
|
||||
|
||||
if [ "$STATUS" = "completed" ]; then
|
||||
echo ""
|
||||
echo "Video generation completed!"
|
||||
echo "Full response:"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
break
|
||||
elif [ "$STATUS" = "failed" ]; then
|
||||
echo ""
|
||||
echo "Video generation failed!"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sleep 2
|
||||
done
|
||||
@@ -0,0 +1,131 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import struct
|
||||
import wave
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
_HEADER = struct.Struct(">iii")
|
||||
_SAMPLE_RATE = 16000
|
||||
_DEFAULT_PROMPT = "Transcribe the audio."
|
||||
_ASR_TEXT_TAG = "<asr_text>"
|
||||
|
||||
|
||||
def _clean_transcript(raw: str) -> str:
|
||||
"""Qwen3-ASR emits ``language <Lang><asr_text><transcript>``; keep the transcript."""
|
||||
text = raw or ""
|
||||
marker = text.rfind(_ASR_TEXT_TAG)
|
||||
if marker != -1:
|
||||
text = text[marker + len(_ASR_TEXT_TAG) :]
|
||||
return text.replace("</asr_text>", "").strip()
|
||||
|
||||
|
||||
def _pcm_to_wav_b64(pcm: bytes, sample_rate: int) -> str:
|
||||
buf = io.BytesIO()
|
||||
with wave.open(buf, "wb") as wav:
|
||||
wav.setnchannels(1)
|
||||
wav.setsampwidth(2)
|
||||
wav.setframerate(sample_rate)
|
||||
wav.writeframes(pcm)
|
||||
return base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def _result(text: str, *, final: bool) -> dict:
|
||||
return {
|
||||
"asr_response": {
|
||||
"event_type": "IS_FINAL" if final else "IS_PARTIAL",
|
||||
"recognition_result": {"hypothesis": [{"text": text, "confidence": None}]},
|
||||
},
|
||||
"mid": "",
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
}
|
||||
|
||||
|
||||
async def _transcribe(cfg: dict, pcm: bytes) -> str:
|
||||
if not pcm:
|
||||
return ""
|
||||
wav_b64 = _pcm_to_wav_b64(pcm, cfg["sample_rate"])
|
||||
body = {
|
||||
"model": cfg["model"],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_audio", "input_audio": {"data": wav_b64, "format": "wav"}},
|
||||
{"type": "text", "text": cfg["prompt"]},
|
||||
],
|
||||
}
|
||||
],
|
||||
"temperature": 0.0,
|
||||
"max_tokens": cfg["max_tokens"],
|
||||
}
|
||||
url = cfg["backend_url"].rstrip("/") + "/v1/chat/completions"
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, json=body) as resp:
|
||||
payload = await resp.json()
|
||||
choices = payload.get("choices") or []
|
||||
if not choices:
|
||||
return ""
|
||||
return _clean_transcript(choices[0].get("message", {}).get("content") or "")
|
||||
|
||||
|
||||
async def _handle(request: web.Request) -> web.WebSocketResponse:
|
||||
ws = web.WebSocketResponse(max_msg_size=0)
|
||||
await ws.prepare(request)
|
||||
cfg = request.app["cfg"]
|
||||
pcm = bytearray()
|
||||
try:
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.BINARY and len(msg.data) >= _HEADER.size:
|
||||
seqid, _, _ = _HEADER.unpack(msg.data[: _HEADER.size])
|
||||
pcm.extend(msg.data[_HEADER.size :])
|
||||
if seqid < 0:
|
||||
text = await _transcribe(cfg, bytes(pcm))
|
||||
await ws.send_str(json.dumps(_result(text, final=True), ensure_ascii=False))
|
||||
pcm.clear()
|
||||
elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
|
||||
break
|
||||
finally:
|
||||
if not ws.closed:
|
||||
await ws.close()
|
||||
return ws
|
||||
|
||||
|
||||
def create_app(backend_url: str, model: str, sample_rate: int, prompt: str, max_tokens: int) -> web.Application:
|
||||
app = web.Application()
|
||||
app["cfg"] = {
|
||||
"backend_url": backend_url,
|
||||
"model": model,
|
||||
"sample_rate": sample_rate,
|
||||
"prompt": prompt,
|
||||
"max_tokens": max_tokens,
|
||||
}
|
||||
app.router.add_get("/v1/asr", _handle)
|
||||
return app
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="JoyVL webui <-> Qwen3-ASR (chat input_audio) bridge")
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=8093)
|
||||
parser.add_argument("--backend-url", default="http://127.0.0.1:8094")
|
||||
parser.add_argument("--model", default="Qwen/Qwen3-ASR-1.7B")
|
||||
parser.add_argument("--sample-rate", type=int, default=_SAMPLE_RATE)
|
||||
parser.add_argument("--prompt", default=_DEFAULT_PROMPT)
|
||||
parser.add_argument("--max-tokens", type=int, default=256)
|
||||
args = parser.parse_args()
|
||||
app = create_app(args.backend_url, args.model, args.sample_rate, args.prompt, args.max_tokens)
|
||||
web.run_app(app, host=args.host, port=args.port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,117 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
try:
|
||||
import cn2an # optional: Arabic-digit -> Chinese number normalization
|
||||
except ImportError:
|
||||
cn2an = None
|
||||
|
||||
logger = logging.getLogger("tts_bridge")
|
||||
|
||||
|
||||
def _normalize(text: str) -> str:
|
||||
"""Spell out Arabic digits in Chinese so short counting outputs synthesize
|
||||
intelligibly: bare "1次" garbles on Qwen3-TTS (and others), while "一次"
|
||||
is clean. ``cn2an.transform`` rewrites digits in place ("第3次"->"第三次")
|
||||
and is pure-Python (~microseconds). No-op if cn2an is unavailable."""
|
||||
if not cn2an or not text:
|
||||
return text
|
||||
try:
|
||||
return cn2an.transform(text, "an2cn")
|
||||
except Exception as err: # never let TN break synthesis
|
||||
logger.debug("cn2an normalize skipped: %s", err)
|
||||
return text
|
||||
|
||||
|
||||
async def _pump_backend_to_front(back: aiohttp.ClientWebSocketResponse, front: web.WebSocketResponse) -> None:
|
||||
async for msg in back:
|
||||
if msg.type == aiohttp.WSMsgType.BINARY:
|
||||
await front.send_bytes(msg.data)
|
||||
elif msg.type == aiohttp.WSMsgType.TEXT:
|
||||
event = json.loads(msg.data).get("type")
|
||||
if event == "session.done":
|
||||
await front.send_str(json.dumps({"type": "response.done"}))
|
||||
return
|
||||
if event == "error":
|
||||
await front.send_str(json.dumps({"type": "error", "error": "tts backend error"}))
|
||||
return
|
||||
elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
|
||||
return
|
||||
|
||||
|
||||
async def _handle(request: web.Request) -> web.WebSocketResponse:
|
||||
front = web.WebSocketResponse(heartbeat=20, max_msg_size=0)
|
||||
await front.prepare(request)
|
||||
cfg = request.app["cfg"]
|
||||
|
||||
session = aiohttp.ClientSession()
|
||||
back: aiohttp.ClientWebSocketResponse | None = None
|
||||
pump: asyncio.Task | None = None
|
||||
parts: list[str] = []
|
||||
try:
|
||||
async for msg in front:
|
||||
if msg.type != web.WSMsgType.TEXT:
|
||||
continue
|
||||
data = json.loads(msg.data)
|
||||
|
||||
if "config" in data and back is None:
|
||||
voice = (data["config"] or {}).get("voice") or cfg["voice"]
|
||||
back = await session.ws_connect(cfg["backend_url"], max_msg_size=0)
|
||||
await back.send_str(
|
||||
json.dumps({"type": "session.config", "response_format": cfg["response_format"], "voice": voice})
|
||||
)
|
||||
pump = asyncio.create_task(_pump_backend_to_front(back, front))
|
||||
elif data.get("type") == "input_text.append" and back is not None:
|
||||
parts.append(data.get("text", ""))
|
||||
elif data.get("type") == "input_text.commit" and back is not None:
|
||||
# Normalize the full utterance once (digits may span chunks), then send.
|
||||
await back.send_str(json.dumps({"type": "input.text", "text": _normalize("".join(parts))}))
|
||||
parts = []
|
||||
await back.send_str(json.dumps({"type": "input.done"}))
|
||||
if pump is not None:
|
||||
await pump
|
||||
break
|
||||
except Exception as err:
|
||||
logger.warning("tts bridge error: %s", err)
|
||||
finally:
|
||||
if pump is not None and not pump.done():
|
||||
pump.cancel()
|
||||
if back is not None and not back.closed:
|
||||
await back.close()
|
||||
await session.close()
|
||||
if not front.closed:
|
||||
await front.close()
|
||||
return front
|
||||
|
||||
|
||||
def create_app(backend_url: str, voice: str, response_format: str) -> web.Application:
|
||||
app = web.Application()
|
||||
app["cfg"] = {"backend_url": backend_url, "voice": voice, "response_format": response_format}
|
||||
app.router.add_get("/v1/tts", _handle)
|
||||
return app
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="JoyVL webui <-> Qwen3-TTS websocket bridge")
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=8092)
|
||||
parser.add_argument("--backend-url", default="ws://127.0.0.1:8091/v1/audio/speech/stream")
|
||||
parser.add_argument("--voice", default="vivian")
|
||||
parser.add_argument("--response-format", default="pcm")
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
web.run_app(create_app(args.backend_url, args.voice, args.response_format), host=args.host, port=args.port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
from stream_client import stream
|
||||
|
||||
ICON = {"response": "SPEAK ", "silence": "silence ", "delegate": "DELEGATE"}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("video")
|
||||
parser.add_argument("--server", default="http://127.0.0.1:8070")
|
||||
parser.add_argument("--query", default=None)
|
||||
parser.add_argument("--fps", type=float, default=1.0)
|
||||
parser.add_argument("--session-id", default="cli")
|
||||
args = parser.parse_args()
|
||||
|
||||
for tick in stream(args.video, args.server, args.session_id, args.query, args.fps):
|
||||
if tick.action == "silence":
|
||||
continue
|
||||
line = f"[{tick.t:6.1f}s] {ICON.get(tick.action)} {tick.text}"
|
||||
if tick.delegation:
|
||||
line += f" (-> {tick.delegation.get('question', '')})"
|
||||
print(line, flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,79 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
|
||||
import cv2
|
||||
import requests
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tick:
|
||||
index: int
|
||||
t: float
|
||||
action: str
|
||||
text: str
|
||||
skipped: bool
|
||||
latency_ms: float
|
||||
delegation: dict | None
|
||||
|
||||
|
||||
def iter_frames(video_path: str, fps: float) -> Iterator[tuple[float, bytes]]:
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
step = max(1, round(src_fps / fps))
|
||||
idx = 0
|
||||
while True:
|
||||
ok, frame = cap.read()
|
||||
if not ok:
|
||||
break
|
||||
if idx % step == 0:
|
||||
ok, buf = cv2.imencode(".jpg", frame)
|
||||
if ok:
|
||||
yield idx / src_fps, buf.tobytes()
|
||||
idx += 1
|
||||
cap.release()
|
||||
|
||||
|
||||
def stream(
|
||||
video_path: str,
|
||||
server: str,
|
||||
session_id: str,
|
||||
query: str | None = None,
|
||||
fps: float = 1.0,
|
||||
realtime: bool = False,
|
||||
) -> Iterator[Tick]:
|
||||
"""Send each sampled frame as one chat completion; yield the parsed action.
|
||||
|
||||
``query`` is attached to the first tick to arm the session; inject later
|
||||
queries by posting another frame turn with a new text part.
|
||||
"""
|
||||
url = server.rstrip("/") + "/v1/chat/completions"
|
||||
requests.post(server.rstrip("/") + "/reset", json={"session_id": session_id}, timeout=10)
|
||||
for i, (t, jpeg) in enumerate(iter_frames(video_path, fps)):
|
||||
data_url = "data:image/jpeg;base64," + base64.b64encode(jpeg).decode()
|
||||
content = [{"type": "image_url", "image_url": {"url": data_url}}]
|
||||
if i == 0 and query:
|
||||
content.insert(0, {"type": "text", "text": query})
|
||||
resp = requests.post(
|
||||
url,
|
||||
json={"session_id": session_id, "messages": [{"role": "user", "content": content}]},
|
||||
timeout=120,
|
||||
).json()
|
||||
info = resp.get("interaction", {})
|
||||
yield Tick(
|
||||
index=info.get("frame_index", i + 1),
|
||||
t=t,
|
||||
action=info.get("action", "?"),
|
||||
text=resp["choices"][0]["message"]["content"],
|
||||
skipped=info.get("inference_skipped", False),
|
||||
latency_ms=info.get("latency_ms", 0.0),
|
||||
delegation=info.get("delegation"),
|
||||
)
|
||||
if realtime:
|
||||
time.sleep(1.0 / fps)
|
||||
@@ -0,0 +1,46 @@
|
||||
# RTSP local streaming
|
||||
|
||||
Simulate an RTSP camera from a local video file so the WebUI (or any RTSP client) can
|
||||
test the RTSP input path without a physical IP camera. RTSP is a **webui-side input** —
|
||||
the browser pulls the stream and feeds frames to the orchestrator over the normal API;
|
||||
no serving-layer code is involved.
|
||||
|
||||
Needs [`ffmpeg`](https://ffmpeg.org/) and [MediaMTX](https://github.com/bluenviron/mediamtx/releases)
|
||||
(download the binary + `mediamtx.yml` and keep them next to these scripts, or point
|
||||
`MEDIAMTX_BIN` / `MEDIAMTX_CONFIG` at them).
|
||||
|
||||
## 1. Start the RTSP server
|
||||
|
||||
```bash
|
||||
bash ./mediamtx.sh # MediaMTX, listens on :8554
|
||||
```
|
||||
|
||||
## 2. Push a local video as an RTSP stream
|
||||
|
||||
```bash
|
||||
bash ./rtsp.sh ./videos/example.mp4 rtsp://127.0.0.1:8554/fire1
|
||||
```
|
||||
|
||||
`rtsp.sh` takes `[video-path] [rtsp-url]` (or set `VIDEO_PATH` / `RTSP_URL`). To stream a
|
||||
whole directory of clips, one stream per file:
|
||||
|
||||
```bash
|
||||
bash ./rtsp_all.sh ./videos rtsp://127.0.0.1:8554
|
||||
```
|
||||
|
||||
## 3. Use it in the WebUI
|
||||
|
||||
Enter the URL in the WebUI RTSP box:
|
||||
|
||||
```
|
||||
rtsp://127.0.0.1:8554/fire1
|
||||
```
|
||||
|
||||
If the WebUI runs on another machine, replace `127.0.0.1` with the MediaMTX host's IP.
|
||||
|
||||
## Notes
|
||||
|
||||
- `rtsp.sh` transcodes audio to AAC; replace `-c:a aac -b:a 128k -ar 44100` with `-an`
|
||||
if the source clip has no audio track and your `ffmpeg` errors on it.
|
||||
- If a client cannot connect, check that port `8554` is reachable and not blocked by a
|
||||
firewall, container network, or security group.
|
||||
@@ -0,0 +1,11 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Start a local MediaMTX RTSP server (listens on :8554 by default).
|
||||
# Download the binary + mediamtx.yml from https://github.com/bluenviron/mediamtx/releases
|
||||
# and either keep them next to this script or point the vars below at them.
|
||||
|
||||
MEDIAMTX_BIN="${MEDIAMTX_BIN:-./mediamtx}"
|
||||
MEDIAMTX_CONFIG="${MEDIAMTX_CONFIG:-./mediamtx.yml}"
|
||||
|
||||
exec "$MEDIAMTX_BIN" "$MEDIAMTX_CONFIG"
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Loop a local video file and push it to an RTSP URL, simulating an RTSP camera.
|
||||
# bash ./rtsp.sh [video-path] [rtsp-output-url]
|
||||
# Defaults stream ./videos/example.mp4 to rtsp://127.0.0.1:8554/fire1
|
||||
|
||||
VIDEO_PATH="${1:-${VIDEO_PATH:-./videos/example.mp4}}"
|
||||
RTSP_URL="${2:-${RTSP_URL:-rtsp://127.0.0.1:8554/fire1}}"
|
||||
|
||||
if [[ ! -f "$VIDEO_PATH" ]]; then
|
||||
echo "Video file not found: $VIDEO_PATH" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Input video: $VIDEO_PATH"
|
||||
echo "RTSP output: $RTSP_URL"
|
||||
|
||||
# Drop "-c:a aac -b:a 128k -ar 44100" for "-an" if the source has no audio track.
|
||||
exec ffmpeg \
|
||||
-hide_banner \
|
||||
-loglevel info \
|
||||
-re \
|
||||
-stream_loop -1 \
|
||||
-i "$VIDEO_PATH" \
|
||||
-vf "scale='min(1280,iw)':-2" \
|
||||
-c:v libx264 \
|
||||
-preset veryfast \
|
||||
-tune zerolatency \
|
||||
-b:v 2500k \
|
||||
-maxrate 2500k \
|
||||
-bufsize 5000k \
|
||||
-c:a aac \
|
||||
-b:a 128k \
|
||||
-ar 44100 \
|
||||
-f rtsp \
|
||||
-rtsp_transport udp \
|
||||
"$RTSP_URL"
|
||||
@@ -0,0 +1,64 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Push every .mp4 in a directory as its own RTSP stream (one ffmpeg per file),
|
||||
# named rtsp://<base>/<filename-without-extension>.
|
||||
# bash ./rtsp_all.sh [video-dir] [rtsp-base-url]
|
||||
|
||||
VIDEO_DIR="${1:-./videos}"
|
||||
RTSP_BASE="${2:-rtsp://127.0.0.1:8554}"
|
||||
|
||||
if [[ ! -d "$VIDEO_DIR" ]]; then
|
||||
echo "Video directory not found: $VIDEO_DIR" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
shopt -s nullglob
|
||||
VIDEO_FILES=("$VIDEO_DIR"/*.mp4)
|
||||
shopt -u nullglob
|
||||
|
||||
if [[ ${#VIDEO_FILES[@]} -eq 0 ]]; then
|
||||
echo "No video files found in: $VIDEO_DIR" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Found ${#VIDEO_FILES[@]} video(s) in $VIDEO_DIR"
|
||||
|
||||
PIDS=()
|
||||
cleanup() {
|
||||
echo "Stopping all streams..."
|
||||
for pid in "${PIDS[@]}"; do
|
||||
kill "$pid" 2>/dev/null || true
|
||||
done
|
||||
wait 2>/dev/null
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
for video in "${VIDEO_FILES[@]}"; do
|
||||
stream_name="$(basename "${video%.*}")"
|
||||
rtsp_url="${RTSP_BASE}/${stream_name}"
|
||||
echo "Starting stream: $rtsp_url <- $(basename "$video")"
|
||||
ffmpeg \
|
||||
-hide_banner \
|
||||
-loglevel error \
|
||||
-re \
|
||||
-stream_loop -1 \
|
||||
-i "$video" \
|
||||
-vf "format=yuv420p,scale='min(1280,iw)':'min(720,ih)':force_original_aspect_ratio=decrease:force_divisible_by=2" \
|
||||
-c:v libx264 \
|
||||
-preset veryfast \
|
||||
-tune zerolatency \
|
||||
-b:v 4000k \
|
||||
-maxrate 4500k \
|
||||
-bufsize 9000k \
|
||||
-c:a aac \
|
||||
-b:a 128k \
|
||||
-ar 44100 \
|
||||
-f rtsp \
|
||||
-rtsp_transport tcp \
|
||||
"$rtsp_url" &
|
||||
PIDS+=($!)
|
||||
done
|
||||
|
||||
echo "All ${#VIDEO_FILES[@]} streams started. Press Ctrl+C to stop."
|
||||
wait
|
||||
@@ -0,0 +1,108 @@
|
||||
#!/bin/bash
|
||||
# One-shot JoyVL demo: model -> interaction orchestrator -> JD webui.
|
||||
#
|
||||
# The model is the only fixed component; every external module is pluggable and
|
||||
# configured here via env (unset = disabled), matching JoyVL's design:
|
||||
#
|
||||
# TTS_URL voice out (point at our tts_bridge, or any OpenAI-Realtime-style TTS WS)
|
||||
# ASR_URL voice in (point at our asr_bridge, or any compatible ASR WS)
|
||||
# DELEGATION_BACKEND_URL orchestrator background brain for </delegation> (OpenAI-compatible; DELEGATION_MODEL/DELEGATION_API_KEY pick model+key)
|
||||
# DELEGATION_KIND=chat (default) text/VL brain answers; =image a text-to-image model generates a picture
|
||||
# chat+Claude: DELEGATION_BACKEND_URL=https://api.anthropic.com/v1/ DELEGATION_MODEL=claude-... DELEGATION_API_KEY=sk-ant-...
|
||||
# image+Qwen-Image: DELEGATION_KIND=image DELEGATION_BACKEND_URL=http://127.0.0.1:8091/v1
|
||||
# BACKGROUND_CODEX_API_URL webui-side background agent
|
||||
# ENABLE_MEMORY (default 1) 3-tier memory; SUMMARIZER_BACKEND_URL/MODEL pick the summarizer (default: reuse main)
|
||||
#
|
||||
# Usage: bash start_all.sh
|
||||
# TTS_URL=ws://host/v1/tts ASR_URL=ws://host/v1/asr GPU=0 bash start_all.sh
|
||||
set -o pipefail
|
||||
|
||||
CODE="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)"
|
||||
VENV="${VENV:-$CODE/.venv}"
|
||||
WEBUI_DIR="${WEBUI_DIR:-$CODE/joyvl-interaction/joyvl-interaction-webui}"
|
||||
MODEL="${MODEL:-jdopensource/JoyAI-VL-Interaction-Preview}"
|
||||
SERVED_NAME="${SERVED_NAME:-JoyAI-VL-Interaction-Preview}"
|
||||
GPU="${GPU:-0}"
|
||||
MODEL_PORT="${MODEL_PORT:-8061}"
|
||||
ORCH_PORT="${ORCH_PORT:-8070}"
|
||||
WEBUI_PORT="${WEBUI_PORT:-8999}"
|
||||
WEBUI_USERNAME="${WEBUI_USERNAME:-admin}"
|
||||
WEBUI_PASSWORD="${WEBUI_PASSWORD:-changeme}" # override via env for any exposed deployment
|
||||
ENABLE_MEMORY="${ENABLE_MEMORY:-1}"
|
||||
RESPONSE_DEDUP="${RESPONSE_DEDUP:-0.85}" # demo drops near-duplicate narration; set 1.0 for reference (exact-only)
|
||||
TTS_URL="${TTS_URL:-}"
|
||||
ASR_URL="${ASR_URL:-}"
|
||||
BACKGROUND_CODEX_API_URL="${BACKGROUND_CODEX_API_URL:-}"
|
||||
|
||||
LOG=/tmp/joyvl_demo
|
||||
mkdir -p "$LOG"
|
||||
|
||||
[ -d "$WEBUI_DIR" ] || { echo "! WEBUI_DIR not found: $WEBUI_DIR (set WEBUI_DIR=...)"; exit 1; }
|
||||
|
||||
kill_port() {
|
||||
local p
|
||||
p=$(ss -tlnp 2>/dev/null | grep ":$1 " | grep -oE 'pid=[0-9]+' | head -1 | cut -d= -f2)
|
||||
[ -n "${p:-}" ] && kill "$p" 2>/dev/null
|
||||
return 0
|
||||
}
|
||||
wait_http() {
|
||||
for _ in $(seq 1 150); do
|
||||
curl -sk -m2 "$1" -o /dev/null 2>/dev/null && return 0
|
||||
sleep 2
|
||||
done
|
||||
echo " ! timeout waiting for $1"
|
||||
return 1
|
||||
}
|
||||
|
||||
echo "Resetting ports ${MODEL_PORT} ${ORCH_PORT} ${WEBUI_PORT}…"
|
||||
for p in "$MODEL_PORT" "$ORCH_PORT" "$WEBUI_PORT"; do kill_port "$p"; done
|
||||
sleep 3
|
||||
|
||||
[ -f "$VENV/bin/activate" ] && source "$VENV/bin/activate"
|
||||
|
||||
# detach into a new session so servers outlive this launcher
|
||||
spawn() { local log="$1"; shift; setsid "$@" > "$log" 2>&1 < /dev/null & }
|
||||
|
||||
echo "[1/3] model on GPU ${GPU}…"
|
||||
spawn "$LOG/model.log" env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES="$GPU" \
|
||||
vllm serve "$MODEL" --served-model-name "$SERVED_NAME" --port "$MODEL_PORT" \
|
||||
--gpu-memory-utilization 0.85 --max-model-len 131072 --enable-prefix-caching \
|
||||
--limit-mm-per-prompt '{"image":256,"video":1}'
|
||||
wait_http "http://127.0.0.1:${MODEL_PORT}/health" && echo " up :${MODEL_PORT}"
|
||||
|
||||
echo "[2/3] orchestrator (memory=${ENABLE_MEMORY})…"
|
||||
ORCH_ARGS=(--port "$ORCH_PORT" --main-backend-url "http://127.0.0.1:${MODEL_PORT}/v1"
|
||||
--main-model "$SERVED_NAME"
|
||||
--response-dedup-threshold "$RESPONSE_DEDUP")
|
||||
# delegation: point at a stronger background brain to enable it; otherwise keep it off
|
||||
if [ -n "${DELEGATION_BACKEND_URL:-}" ]; then
|
||||
ORCH_ARGS+=(--delegation-backend-url "$DELEGATION_BACKEND_URL")
|
||||
ORCH_ARGS+=(--delegation-kind "${DELEGATION_KIND:-chat}")
|
||||
[ -n "${DELEGATION_MODEL:-}" ] && ORCH_ARGS+=(--delegation-model "$DELEGATION_MODEL")
|
||||
[ -n "${DELEGATION_API_KEY:-}" ] && ORCH_ARGS+=(--delegation-api-key "$DELEGATION_API_KEY")
|
||||
else
|
||||
ORCH_ARGS+=(--no-delegation)
|
||||
fi
|
||||
[ "$ENABLE_MEMORY" = "0" ] && ORCH_ARGS+=(--no-memory)
|
||||
[ -n "${SUMMARIZER_BACKEND_URL:-}" ] && ORCH_ARGS+=(--summarizer-backend-url "$SUMMARIZER_BACKEND_URL")
|
||||
[ -n "${SUMMARIZER_MODEL:-}" ] && ORCH_ARGS+=(--summarizer-model "$SUMMARIZER_MODEL")
|
||||
spawn "$LOG/orch.log" env PYTHONPATH="$CODE" python -m vllm_omni.experimental.fullduplex.joyvl.serving.server "${ORCH_ARGS[@]}"
|
||||
wait_http "http://127.0.0.1:${ORCH_PORT}/health" && echo " up :${ORCH_PORT}"
|
||||
|
||||
echo "[3/3] JD webui…"
|
||||
spawn "$LOG/webui.log" env WEBUI_USERNAME="$WEBUI_USERNAME" WEBUI_PASSWORD="$WEBUI_PASSWORD" \
|
||||
TTS_URL="$TTS_URL" ASR_URL="$ASR_URL" BACKGROUND_CODEX_API_URL="$BACKGROUND_CODEX_API_URL" \
|
||||
bash -c "cd '$WEBUI_DIR' && exec bash scripts/start_server.sh --api-base 'http://127.0.0.1:${ORCH_PORT}/v1'"
|
||||
webui_up=0
|
||||
for _ in $(seq 1 60); do grep -q "Server startup complete" "$LOG/webui.log" 2>/dev/null && { webui_up=1; break; }; sleep 1; done
|
||||
[ "$webui_up" = 1 ] || { echo "! webui did not start — see $LOG/webui.log"; exit 1; }
|
||||
|
||||
IP=$(hostname -I 2>/dev/null | awk '{print $1}')
|
||||
echo
|
||||
echo "=== JoyVL demo ready ==="
|
||||
echo " webui : https://${IP}:${WEBUI_PORT} (user ${WEBUI_USERNAME} / pass ${WEBUI_PASSWORD})"
|
||||
echo " memory : $([ "$ENABLE_MEMORY" = 0 ] && echo off || echo on)"
|
||||
echo " TTS : ${TTS_URL:-<unset>}"
|
||||
echo " ASR : ${ASR_URL:-<unset>}"
|
||||
echo " agent : ${BACKGROUND_CODEX_API_URL:-<unset>}"
|
||||
echo " logs : ${LOG}/{model,orch,webui}.log"
|
||||
@@ -0,0 +1,24 @@
|
||||
#!/bin/bash
|
||||
# Serve Qwen3-ASR (the same audio-LLM AURA uses as its ASR stage) and run the
|
||||
# webui<->ASR bridge. Then start the webui with:
|
||||
# ASR_URL=ws://127.0.0.1:8093/v1/asr bash scripts/start_server.sh
|
||||
set -euo pipefail
|
||||
|
||||
ASR_MODEL="${ASR_MODEL:-Qwen/Qwen3-ASR-1.7B}"
|
||||
ASR_PORT="${ASR_PORT:-8094}"
|
||||
BRIDGE_PORT="${BRIDGE_PORT:-8093}"
|
||||
GPU="${GPU:-2}"
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
# Qwen3-ASR is audio-in / text-out; served as a normal multimodal chat model.
|
||||
CUDA_VISIBLE_DEVICES="${GPU}" vllm serve "${ASR_MODEL}" \
|
||||
--host 0.0.0.0 --port "${ASR_PORT}" \
|
||||
--gpu-memory-utilization 0.3 --trust-remote-code &
|
||||
ASR_PID=$!
|
||||
trap 'kill ${ASR_PID} 2>/dev/null || true' EXIT
|
||||
|
||||
until curl -s "http://127.0.0.1:${ASR_PORT}/health" >/dev/null 2>&1; do sleep 2; done
|
||||
python "${SCRIPT_DIR}/../bridges/asr_bridge.py" \
|
||||
--port "${BRIDGE_PORT}" \
|
||||
--backend-url "http://127.0.0.1:${ASR_PORT}" \
|
||||
--model "${ASR_MODEL}"
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
# Launch the interaction orchestrator in front of a running model server.
|
||||
set -euo pipefail
|
||||
|
||||
PORT="${PORT:-8070}"
|
||||
MAIN_BACKEND_URL="${MAIN_BACKEND_URL:-http://127.0.0.1:8061/v1}"
|
||||
MAIN_MODEL="${MAIN_MODEL:-JoyAI-VL-Interaction-Preview}"
|
||||
PERSONA="${PERSONA:-default}"
|
||||
|
||||
# By default memory reuses the main model as its own summarizer. Set
|
||||
# SUMMARIZER_BACKEND_URL/SUMMARIZER_MODEL to a dedicated Qwen3-VL-4B server,
|
||||
# or pass --no-memory for the lightest setup.
|
||||
EXTRA=()
|
||||
[[ -n "${SUMMARIZER_BACKEND_URL:-}" ]] && EXTRA+=(--summarizer-backend-url "${SUMMARIZER_BACKEND_URL}")
|
||||
[[ -n "${SUMMARIZER_MODEL:-}" ]] && EXTRA+=(--summarizer-model "${SUMMARIZER_MODEL}")
|
||||
[[ "${NO_MEMORY:-0}" == "1" ]] && EXTRA+=(--no-memory)
|
||||
|
||||
python -m vllm_omni.experimental.fullduplex.joyvl.serving.server \
|
||||
--port "${PORT}" \
|
||||
--main-backend-url "${MAIN_BACKEND_URL}" \
|
||||
--main-model "${MAIN_MODEL}" \
|
||||
--persona "${PERSONA}" \
|
||||
"${EXTRA[@]}"
|
||||
@@ -0,0 +1,22 @@
|
||||
#!/bin/bash
|
||||
# Serve the JoyVL interaction model as a plain VLM (NOT --omni: it is a
|
||||
# standard Qwen3-VL autoregressive model, not an omni/diffusion model).
|
||||
set -euo pipefail
|
||||
|
||||
MODEL="${MODEL:-jdopensource/JoyAI-VL-Interaction-Preview}"
|
||||
SERVED_NAME="${SERVED_NAME:-JoyAI-VL-Interaction-Preview}"
|
||||
PORT="${PORT:-8061}"
|
||||
GPU="${GPU:-0}"
|
||||
# Covers the default short-term window (chunk_frames=100) + key_frames=0 (all frames
|
||||
# summarized). Lower MAX_MODEL_LEN / IMAGE_LIMIT (and --chunk-frames) together for ~24GB.
|
||||
MAX_MODEL_LEN="${MAX_MODEL_LEN:-131072}"
|
||||
IMAGE_LIMIT="${IMAGE_LIMIT:-256}"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="${GPU}" vllm serve "${MODEL}" \
|
||||
--served-model-name "${SERVED_NAME}" \
|
||||
--port "${PORT}" \
|
||||
--max-model-len "${MAX_MODEL_LEN}" \
|
||||
--gpu-memory-utilization 0.85 \
|
||||
--enable-prefix-caching \
|
||||
--enable-chunked-prefill \
|
||||
--limit-mm-per-prompt "{\"image\":${IMAGE_LIMIT},\"video\":1}"
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
# Serve Qwen3-TTS (vLLM-Omni) and run the webui<->TTS bridge.
|
||||
# Then start the webui with: TTS_URL=ws://127.0.0.1:8092/v1/tts bash scripts/start_server.sh
|
||||
set -euo pipefail
|
||||
|
||||
TTS_MODEL="${TTS_MODEL:-Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice}"
|
||||
TTS_PORT="${TTS_PORT:-8091}"
|
||||
BRIDGE_PORT="${BRIDGE_PORT:-8092}"
|
||||
GPU="${GPU:-1}"
|
||||
VOICE="${VOICE:-vivian}"
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="${GPU}" vllm-omni serve "${TTS_MODEL}" \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--host 0.0.0.0 --port "${TTS_PORT}" \
|
||||
--gpu-memory-utilization 0.9 --trust-remote-code --omni &
|
||||
TTS_PID=$!
|
||||
trap 'kill ${TTS_PID} 2>/dev/null || true' EXIT
|
||||
|
||||
# wait for the TTS server, then run the bridge in the foreground
|
||||
until curl -s "http://127.0.0.1:${TTS_PORT}/health" >/dev/null 2>&1; do sleep 2; done
|
||||
python "${SCRIPT_DIR}/../bridges/tts_bridge.py" \
|
||||
--port "${BRIDGE_PORT}" \
|
||||
--backend-url "ws://127.0.0.1:${TTS_PORT}/v1/audio/speech/stream" \
|
||||
--voice "${VOICE}"
|
||||
@@ -88,7 +88,7 @@ curl -s http://127.0.0.1:8091/v1/chat/completions \
|
||||
|
||||
### Optional knobs
|
||||
|
||||
Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
|
||||
Pass image-gen overrides as flat keys on the diffusion-stage `sampling_params_list[1].extra_args`:
|
||||
|
||||
| Key | Default | Description |
|
||||
| --- | --- | --- |
|
||||
@@ -97,25 +97,48 @@ Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_
|
||||
| `cfg` | 2.0 | Classifier-free guidance scale. |
|
||||
| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
|
||||
| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
|
||||
| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
|
||||
| `negative_prompt` | empty | Real CFG negative conditioning (set on **stage-0 thinker** `extra_args` so `expand_cfg_prompts` spawns the companion). |
|
||||
|
||||
Example with all the knobs:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Jonathan1909/Ming-flash-omni-2.0",
|
||||
"modalities": ["image"],
|
||||
"sampling_params_list": [
|
||||
{"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
|
||||
"extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
|
||||
{"seed":42,"extra_args":{"image_gen":{
|
||||
"steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
|
||||
"byte5_text":["理解与生成统一"]}}}
|
||||
],
|
||||
"messages": [{"role": "user", "content": "Draw a poster."}]
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Jonathan1909/Ming-flash-omni-2.0",
|
||||
"modalities": ["image"],
|
||||
"sampling_params_list": [
|
||||
{
|
||||
"temperature": 0.4,
|
||||
"top_p": 0.9,
|
||||
"top_k": 1,
|
||||
"max_tokens": 1,
|
||||
"seed": 42,
|
||||
"extra_args": {
|
||||
"negative_prompt": "ugly, blurry, distorted"
|
||||
}
|
||||
},
|
||||
{
|
||||
"seed": 42,
|
||||
"extra_args": {
|
||||
"steps": 6,
|
||||
"cfg": 1.5,
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"seed": 123,
|
||||
"byte5_text": ["理解与生成统一"]
|
||||
}
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Draw a poster."
|
||||
}
|
||||
]
|
||||
}' \
|
||||
| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
|
||||
| base64 -d > ming_imagegen_knobs.png
|
||||
```
|
||||
|
||||
### img2img (reference image)
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
# Streaming Video Generation
|
||||
|
||||
This example uses the custom WebSocket endpoint `WS /v1/realtime/video` to receive a video byte stream as chunks are produced.
|
||||
It covers text-only video generation. Image/reference input is intentionally not included for now.
|
||||
|
||||
## Start The Server
|
||||
|
||||
Start a diffusion video model with streaming output enabled:
|
||||
|
||||
```bash
|
||||
vllm serve BestWishYsh/Helios-Distilled \
|
||||
--omni \
|
||||
--diffusion-streaming-output \
|
||||
--port 8000
|
||||
```
|
||||
|
||||
The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=True` in the default diffusion stage `engine_args`, then loaded by `OmniDiffusionConfig.from_kwargs()`.
|
||||
|
||||
## WebSocket Protocol
|
||||
|
||||
| Direction | Message | Format | Description |
|
||||
| --- | --- | --- | --- |
|
||||
| Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. |
|
||||
| Server to client | `video.start` | JSON text: `{"type":"video.start","request_id":"...","format":"m4s","config":{...}}` | Confirms the session and mirrors the accepted `format`. |
|
||||
| Server to client | Video chunk | Binary WebSocket frame | Fragmented MP4 (`m4s`) video bytes. |
|
||||
| Client to server | `session.stop` | JSON text: `{"type":"session.stop"}` | Requests cancellation of the active session. |
|
||||
| Client to server | `session.ping` | JSON text: `{"type":"session.ping"}` | Optional keepalive; refreshes the server stall clock. |
|
||||
| Server to client | `session.done` | JSON text: `{"type":"session.done","request_id":"...","chunks":3,"stopped":false}` | Ends a completed or stopped session. |
|
||||
| Server to client | `session.pong` | JSON text: `{"type":"session.pong"}` | Reply to `session.ping`. |
|
||||
| Server to client | `error` | JSON text: `{"type":"error","message":"..."}` | Reports invalid input, unsupported formats, generation failures, control-message errors, or stall timeout. |
|
||||
|
||||
During generation the client normally sends only `session.start` and then receives binary chunks; silence on the client socket is expected. The server closes the session with a stall error only when there is no engine progress and no `session.ping` for about 60 seconds.
|
||||
|
||||
## Install Client Dependency
|
||||
|
||||
```bash
|
||||
pip install websockets
|
||||
# For the Gradio demo:
|
||||
pip install 'vllm-omni[demo]' websockets
|
||||
```
|
||||
|
||||
## Run The Client
|
||||
|
||||
```bash
|
||||
python streaming_video_client.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8000 \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--width 640 \
|
||||
--height 384 \
|
||||
--fps 16 \
|
||||
--num-frames 99 \
|
||||
--guidance-scale 1.0 \
|
||||
--seed 42 \
|
||||
--output helios_stream.mp4
|
||||
```
|
||||
|
||||
The client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`.
|
||||
The client remuxes the gathered stream to a regular progressive MP4 file so that local playback knows the video duration.
|
||||
|
||||
## Run The Gradio Demo
|
||||
|
||||
```bash
|
||||
python gradio_demo.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 7860
|
||||
```
|
||||
|
||||
The Gradio demo requests fMP4 (`m4s`) chunks and appends them directly in the browser with a Media Source Extensions player.
|
||||
|
||||
## Model Choice
|
||||
|
||||
### Helios
|
||||
|
||||
The example uses `BestWishYsh/Helios-Distilled` model by default.
|
||||
|
||||
To ensure streaming-level generation speed, `pyramid_num_inference_steps_list` is suggested to be as low as `[1, 1, 1]`. Both example clients uses the following Helios-Distilled preset by default:
|
||||
|
||||
```json
|
||||
{
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"is_amplify_first_chunk": true
|
||||
}
|
||||
```
|
||||
|
||||
Disable it in the CLI example with `--no-helios-distilled-preset`, or override/extend it with `--extra-params`:
|
||||
|
||||
```bash
|
||||
python streaming_video_client.py \
|
||||
--extra-params '{"pyramid_num_inference_steps_list":[2, 2, 2]}'
|
||||
```
|
||||
@@ -0,0 +1,301 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Gradio demo for `/v1/realtime/video`.
|
||||
|
||||
Usage:
|
||||
python gradio_demo.py --host 127.0.0.1 --port 7860
|
||||
|
||||
Start a vLLM-Omni server with streaming output enabled first, for example:
|
||||
vllm serve BestWishYsh/Helios-Distilled --omni --diffusion-streaming-output --port 8000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import gradio as gr
|
||||
except ImportError:
|
||||
raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
|
||||
|
||||
DEFAULT_PROMPT = "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement."
|
||||
VIDEO_STREAM_VIEW_HTML = Path(__file__).with_name("video-stream-view.html")
|
||||
VIDEO_STREAM_VIEW_JS = Path(__file__).with_name("video-stream-view.js")
|
||||
|
||||
|
||||
def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None:
|
||||
if value is None:
|
||||
return
|
||||
if isinstance(value, str) and not value.strip():
|
||||
return
|
||||
payload[key] = value
|
||||
|
||||
|
||||
def _optional_int(value: int | float | None) -> int | None:
|
||||
return None if value is None else int(value)
|
||||
|
||||
|
||||
def _optional_float(value: int | float | None) -> float | None:
|
||||
return None if value is None else float(value)
|
||||
|
||||
|
||||
def _build_session_start(
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
width: int | float | None,
|
||||
height: int | float | None,
|
||||
fps: int | float | None,
|
||||
num_frames: int | float | None,
|
||||
num_inference_steps: int | float | None,
|
||||
guidance_scale: int | float | None,
|
||||
guidance_scale_2: int | float | None,
|
||||
boundary_ratio: int | float | None,
|
||||
flow_shift: int | float | None,
|
||||
true_cfg_scale: int | float | None,
|
||||
seed: int | float | None,
|
||||
) -> dict[str, Any]:
|
||||
payload: dict[str, Any] = {
|
||||
"type": "session.start",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"format": "m4s",
|
||||
}
|
||||
|
||||
_maybe_set(payload, "negative_prompt", negative_prompt)
|
||||
_maybe_set(payload, "width", _optional_int(width))
|
||||
_maybe_set(payload, "height", _optional_int(height))
|
||||
_maybe_set(payload, "fps", _optional_int(fps))
|
||||
_maybe_set(payload, "num_frames", _optional_int(num_frames))
|
||||
_maybe_set(payload, "num_inference_steps", _optional_int(num_inference_steps))
|
||||
_maybe_set(payload, "guidance_scale", _optional_float(guidance_scale))
|
||||
_maybe_set(payload, "guidance_scale_2", _optional_float(guidance_scale_2))
|
||||
_maybe_set(payload, "boundary_ratio", _optional_float(boundary_ratio))
|
||||
_maybe_set(payload, "flow_shift", _optional_float(flow_shift))
|
||||
_maybe_set(payload, "true_cfg_scale", _optional_float(true_cfg_scale))
|
||||
_maybe_set(payload, "seed", _optional_int(seed))
|
||||
|
||||
if "Helios" in model:
|
||||
payload["extra_params"] = {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"is_amplify_first_chunk": True,
|
||||
}
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def _set_optional_field_enabled(enabled: bool) -> Any:
|
||||
return gr.update(interactive=enabled)
|
||||
|
||||
|
||||
def build_browser_stream_config(
|
||||
host: str,
|
||||
port: int | float,
|
||||
model: str,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
width: int | float | None,
|
||||
height: int | float | None,
|
||||
fps: int | float,
|
||||
num_frames: int | float | None,
|
||||
num_inference_steps: int | float | None,
|
||||
include_guidance_scale: bool,
|
||||
guidance_scale: int | float | None,
|
||||
include_guidance_scale_2: bool,
|
||||
guidance_scale_2: int | float | None,
|
||||
include_seed: bool,
|
||||
seed: int | float | None,
|
||||
include_boundary_ratio: bool,
|
||||
boundary_ratio: int | float | None,
|
||||
include_flow_shift: bool,
|
||||
flow_shift: int | float | None,
|
||||
include_true_cfg_scale: bool,
|
||||
true_cfg_scale: int | float | None,
|
||||
) -> tuple[str, Any]:
|
||||
"""Build the browser-side WebSocket request config for the HTML/MSE player."""
|
||||
if not model.strip():
|
||||
raise gr.Error("Model is required.")
|
||||
if not prompt.strip():
|
||||
raise gr.Error("Prompt is required.")
|
||||
|
||||
payload = _build_session_start(
|
||||
model=model.strip(),
|
||||
prompt=prompt.strip(),
|
||||
negative_prompt=negative_prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
fps=fps,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale if include_guidance_scale else None,
|
||||
guidance_scale_2=guidance_scale_2 if include_guidance_scale_2 else None,
|
||||
boundary_ratio=boundary_ratio if include_boundary_ratio else None,
|
||||
flow_shift=flow_shift if include_flow_shift else None,
|
||||
true_cfg_scale=true_cfg_scale if include_true_cfg_scale else None,
|
||||
seed=seed if include_seed else None,
|
||||
)
|
||||
config = {
|
||||
"url": f"ws://{host}:{int(port)}/v1/realtime/video",
|
||||
"payload": payload,
|
||||
}
|
||||
return json.dumps(config, ensure_ascii=False), gr.update(value="Streaming...", interactive=False)
|
||||
|
||||
|
||||
def create_demo() -> gr.Blocks:
|
||||
with gr.Blocks(title="Streaming Video Generation") as demo:
|
||||
gr.Markdown("# Streaming Video Generation")
|
||||
gr.Markdown(
|
||||
"Connects to `WS /v1/realtime/video` in the browser and appends fMP4 chunks "
|
||||
"directly to a Media Source Extensions video player."
|
||||
)
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column(scale=1):
|
||||
host = gr.Textbox(label="Server Host", value="127.0.0.1")
|
||||
port = gr.Number(label="Server Port", value=8000, precision=0)
|
||||
model = gr.Textbox(label="Model", value="BestWishYsh/Helios-Distilled")
|
||||
prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3)
|
||||
negative_prompt = gr.Textbox(label="Negative Prompt", value="", lines=2)
|
||||
|
||||
with gr.Row():
|
||||
width = gr.Number(label="Width", value=640, precision=0)
|
||||
height = gr.Number(label="Height", value=384, precision=0)
|
||||
|
||||
with gr.Row():
|
||||
fps = gr.Number(label="FPS", value=16, precision=0)
|
||||
num_frames = gr.Number(label="Num Frames", value=132, precision=0)
|
||||
num_inference_steps = gr.Number(label="Num Inference Steps", value=50, precision=0)
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
guidance_scale = gr.Number(
|
||||
label="Guidance Scale",
|
||||
value=1.0,
|
||||
interactive=True,
|
||||
scale=4,
|
||||
)
|
||||
include_guidance_scale = gr.Checkbox(label="Send", value=True, scale=1, min_width=70)
|
||||
with gr.Column():
|
||||
guidance_scale_2 = gr.Number(
|
||||
label="Guidance Scale 2",
|
||||
value=None,
|
||||
interactive=False,
|
||||
scale=4,
|
||||
)
|
||||
include_guidance_scale_2 = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
|
||||
with gr.Column():
|
||||
true_cfg_scale = gr.Number(
|
||||
label="True CFG Scale",
|
||||
value=None,
|
||||
interactive=False,
|
||||
scale=4,
|
||||
)
|
||||
include_true_cfg_scale = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
|
||||
# I want it to wrap after 3, but it auto wrap after 2. I'll rather fully delegate to auto wrapping.
|
||||
with gr.Column():
|
||||
boundary_ratio = gr.Number(
|
||||
label="Boundary Ratio",
|
||||
value=None,
|
||||
interactive=False,
|
||||
scale=4,
|
||||
)
|
||||
include_boundary_ratio = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
|
||||
with gr.Column():
|
||||
flow_shift = gr.Number(label="Flow Shift", value=None, interactive=False, scale=4)
|
||||
include_flow_shift = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
|
||||
with gr.Column():
|
||||
seed = gr.Number(label="Seed", value=42, precision=0, interactive=True, scale=4)
|
||||
include_seed = gr.Checkbox(label="Send", value=True, scale=1, min_width=70)
|
||||
|
||||
with gr.Row():
|
||||
start_button = gr.Button("Start", variant="primary", elem_id="streaming-video-start")
|
||||
hidden_stream_config = gr.Textbox(value="", visible=False)
|
||||
|
||||
with gr.Column(scale=1):
|
||||
gr.HTML(
|
||||
VIDEO_STREAM_VIEW_HTML.read_text(encoding="utf-8"),
|
||||
label="Streaming Preview",
|
||||
js_on_load=VIDEO_STREAM_VIEW_JS.read_text(encoding="utf-8"),
|
||||
)
|
||||
|
||||
inputs = [
|
||||
host,
|
||||
port,
|
||||
model,
|
||||
prompt,
|
||||
negative_prompt,
|
||||
width,
|
||||
height,
|
||||
fps,
|
||||
num_frames,
|
||||
num_inference_steps,
|
||||
include_guidance_scale,
|
||||
guidance_scale,
|
||||
include_guidance_scale_2,
|
||||
guidance_scale_2,
|
||||
include_seed,
|
||||
seed,
|
||||
include_boundary_ratio,
|
||||
boundary_ratio,
|
||||
include_flow_shift,
|
||||
flow_shift,
|
||||
include_true_cfg_scale,
|
||||
true_cfg_scale,
|
||||
]
|
||||
for toggle, field in [
|
||||
(include_guidance_scale, guidance_scale),
|
||||
(include_guidance_scale_2, guidance_scale_2),
|
||||
(include_seed, seed),
|
||||
(include_boundary_ratio, boundary_ratio),
|
||||
(include_flow_shift, flow_shift),
|
||||
(include_true_cfg_scale, true_cfg_scale),
|
||||
]:
|
||||
toggle.change(
|
||||
fn=_set_optional_field_enabled,
|
||||
inputs=toggle,
|
||||
outputs=field,
|
||||
)
|
||||
start_button.click(
|
||||
fn=build_browser_stream_config,
|
||||
inputs=inputs,
|
||||
outputs=[hidden_stream_config, start_button],
|
||||
).then(
|
||||
fn=lambda config: config,
|
||||
inputs=hidden_stream_config,
|
||||
outputs=hidden_stream_config,
|
||||
js="""
|
||||
(config) => {
|
||||
if (window.vllmStreamingVideoStart) {
|
||||
window.vllmStreamingVideoStart(config);
|
||||
}
|
||||
return config;
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
return demo
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Streaming video generation Gradio demo")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Host/IP for gradio launch")
|
||||
parser.add_argument("--port", type=int, default=7860, help="Port for gradio launch")
|
||||
parser.add_argument("--share", action="store_true", help="Create a public Gradio link")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
demo = create_demo()
|
||||
demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,279 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Text-only WebSocket client for `/v1/realtime/video`.
|
||||
|
||||
The client prints a line for every binary video chunk it receives and writes
|
||||
the received bytes to disk when the session finishes or is interrupted.
|
||||
Image/reference input is intentionally not supported in this example yet.
|
||||
|
||||
Requirements:
|
||||
pip install websockets
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import av
|
||||
|
||||
try:
|
||||
from websockets.asyncio.client import connect # pyright: ignore[reportMissingImports]
|
||||
except ImportError:
|
||||
try:
|
||||
from websockets import connect # pyright: ignore[reportMissingImports]
|
||||
except ImportError:
|
||||
print("Please install websockets: pip install websockets")
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
HELIOS_DISTILLED_EXTRA_PARAMS = {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [1, 1, 1],
|
||||
"is_amplify_first_chunk": True,
|
||||
}
|
||||
|
||||
|
||||
def _json_object(value: str, *, argument_name: str) -> dict[str, Any]:
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise argparse.ArgumentTypeError(f"{argument_name} must be valid JSON: {exc}") from exc
|
||||
if not isinstance(parsed, dict):
|
||||
raise argparse.ArgumentTypeError(f"{argument_name} must be a JSON object")
|
||||
return parsed
|
||||
|
||||
|
||||
def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None:
|
||||
if value is not None:
|
||||
payload[key] = value
|
||||
|
||||
|
||||
def build_session_start(args: argparse.Namespace) -> dict[str, Any]:
|
||||
payload: dict[str, Any] = {
|
||||
"type": "session.start",
|
||||
"model": args.model,
|
||||
"prompt": args.prompt,
|
||||
}
|
||||
|
||||
_maybe_set(payload, "negative_prompt", args.negative_prompt)
|
||||
_maybe_set(payload, "width", args.width)
|
||||
_maybe_set(payload, "height", args.height)
|
||||
_maybe_set(payload, "size", args.size)
|
||||
_maybe_set(payload, "seconds", args.seconds)
|
||||
_maybe_set(payload, "fps", args.fps)
|
||||
_maybe_set(payload, "num_frames", args.num_frames)
|
||||
_maybe_set(payload, "num_inference_steps", args.num_inference_steps)
|
||||
_maybe_set(payload, "guidance_scale", args.guidance_scale)
|
||||
_maybe_set(payload, "guidance_scale_2", args.guidance_scale_2)
|
||||
_maybe_set(payload, "boundary_ratio", args.boundary_ratio)
|
||||
_maybe_set(payload, "flow_shift", args.flow_shift)
|
||||
_maybe_set(payload, "true_cfg_scale", args.true_cfg_scale)
|
||||
_maybe_set(payload, "seed", args.seed)
|
||||
|
||||
extra_params: dict[str, Any] = {}
|
||||
if "Helios-Distilled" in args.model and args.helios_distilled_preset:
|
||||
extra_params.update(HELIOS_DISTILLED_EXTRA_PARAMS)
|
||||
if args.extra_params:
|
||||
extra_params.update(args.extra_params)
|
||||
if extra_params:
|
||||
payload["extra_params"] = extra_params
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def _count_decoded_frames(video_bytes: bytes, *, stream_format: str) -> int | None:
|
||||
"""Return the number of decodable video frames in a streamed byte sequence, if known."""
|
||||
if not video_bytes:
|
||||
return 0
|
||||
|
||||
av_format = "mp4" if stream_format == "m4s" else None
|
||||
try:
|
||||
with av.open(io.BytesIO(video_bytes), format=av_format, mode="r") as container:
|
||||
stream = container.streams.video[0]
|
||||
return sum(1 for _ in container.decode(stream))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def save_video(
|
||||
chunks: list[bytes],
|
||||
output_path: Path,
|
||||
*,
|
||||
fps: float,
|
||||
stream_format: str,
|
||||
) -> None:
|
||||
if not chunks:
|
||||
print("No video chunks received; nothing to save.")
|
||||
return
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes
|
||||
|
||||
streamed_bytes = b"".join(chunks)
|
||||
playback_bytes = finalize_streaming_video_bytes(
|
||||
streamed_bytes,
|
||||
input_format="m4s",
|
||||
fps=fps,
|
||||
)
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = output_path.with_suffix(output_path.suffix + ".tmp")
|
||||
with tmp_path.open("wb") as f:
|
||||
f.write(playback_bytes)
|
||||
os.replace(tmp_path, output_path)
|
||||
print(
|
||||
f"Saved {len(chunks)} streamed chunk(s), {len(streamed_bytes)} streamed bytes "
|
||||
f"-> {len(playback_bytes)} playback bytes at {output_path}"
|
||||
)
|
||||
|
||||
|
||||
async def stream_video(args: argparse.Namespace) -> None:
|
||||
payload = build_session_start(args)
|
||||
output_path = Path(args.output)
|
||||
url = f"ws://{args.host}:{args.port}/v1/realtime/video"
|
||||
chunks: list[bytes] = []
|
||||
total_bytes = 0
|
||||
total_frames = 0
|
||||
accumulated_stream = bytearray()
|
||||
stream_format = "m4s"
|
||||
done = False
|
||||
ws = None
|
||||
|
||||
try:
|
||||
async with connect(url, max_size=None) as websocket:
|
||||
ws = websocket
|
||||
await websocket.send(json.dumps(payload))
|
||||
print(f"Connected: {url}")
|
||||
print(f"Sent session.start: {json.dumps(payload, ensure_ascii=False)}")
|
||||
started_at = time.perf_counter()
|
||||
last_chunk_at = started_at
|
||||
|
||||
while True:
|
||||
message = await websocket.recv()
|
||||
now = time.perf_counter()
|
||||
|
||||
if isinstance(message, bytes):
|
||||
chunks.append(message)
|
||||
chunk_bytes = len(message)
|
||||
chunk_elapsed = now - last_chunk_at
|
||||
last_chunk_at = now
|
||||
|
||||
total_bytes += chunk_bytes
|
||||
accumulated_stream.extend(message)
|
||||
decoded_total = _count_decoded_frames(bytes(accumulated_stream), stream_format=stream_format)
|
||||
if decoded_total is not None:
|
||||
chunk_frames = max(decoded_total - total_frames, 0)
|
||||
total_frames = decoded_total
|
||||
else:
|
||||
chunk_frames = 0
|
||||
|
||||
total_elapsed = now - started_at
|
||||
print(
|
||||
f"[chunk {len(chunks):03d}] "
|
||||
f"bytes={chunk_bytes} "
|
||||
f"frames={chunk_frames} "
|
||||
f"elapsed={chunk_elapsed:.2f}s "
|
||||
f"total_bytes={total_bytes} "
|
||||
f"total_frames={total_frames} "
|
||||
f"total_elapsed={total_elapsed:.2f}s"
|
||||
)
|
||||
continue
|
||||
|
||||
msg = json.loads(message)
|
||||
msg_type = msg.get("type")
|
||||
if msg_type == "video.start":
|
||||
stream_format = msg.get("format") or stream_format
|
||||
print(f"Video session started: request_id={msg.get('request_id')} format={msg.get('format')}")
|
||||
elif msg_type == "session.done":
|
||||
print(f"Session complete: {json.dumps(msg, ensure_ascii=False)}")
|
||||
done = True
|
||||
break
|
||||
elif msg_type == "error":
|
||||
print(f"ERROR: {msg.get('message', msg)}")
|
||||
else:
|
||||
print(f"Received control message: {msg}")
|
||||
except (KeyboardInterrupt, asyncio.CancelledError):
|
||||
print("\nInterrupted. Asking server to stop the active session...")
|
||||
if ws is not None:
|
||||
try:
|
||||
await ws.send(json.dumps({"type": "session.stop"}))
|
||||
await ws.close(code=1000, reason="client interrupted")
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
finally:
|
||||
print(
|
||||
"Saving video... (May take a while. Remuxing concatenated chunks into one progressive MP4 file with total duration metadata in file header)"
|
||||
)
|
||||
save_video(
|
||||
chunks,
|
||||
output_path,
|
||||
fps=float(payload.get("fps") or args.fps),
|
||||
stream_format=stream_format,
|
||||
)
|
||||
if not done:
|
||||
print("Saved partial output because the stream did not finish normally.")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Streaming video generation WebSocket client")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="API server host")
|
||||
parser.add_argument("--port", type=int, default=8000, help="API server port")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="BestWishYsh/Helios-Distilled",
|
||||
help="Model name served by the API server",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement.",
|
||||
help="Text prompt for video generation",
|
||||
)
|
||||
parser.add_argument("--output", default="streaming_video_output.mp4", help="Path for the received video bytes")
|
||||
parser.add_argument("--width", type=int, default=640, help="Video width")
|
||||
parser.add_argument("--height", type=int, default=384, help="Video height")
|
||||
parser.add_argument("--size", default=None, help="Video size as WIDTHxHEIGHT; alternative to width/height")
|
||||
parser.add_argument("--seconds", default=None, help="Clip duration in seconds")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second")
|
||||
parser.add_argument("--num-frames", type=int, default=99, help="Number of generated frames")
|
||||
|
||||
parser.add_argument("--negative-prompt", default=None, help="Negative prompt")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=None, help="Number of diffusion steps")
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.0, help="Classifier-free guidance scale")
|
||||
parser.add_argument("--guidance-scale-2", type=float, default=None, help="High-noise CFG scale for video models")
|
||||
parser.add_argument("--boundary-ratio", type=float, default=None, help="Boundary split ratio for video models")
|
||||
parser.add_argument("--flow-shift", type=float, default=None, help="Scheduler flow_shift for video models")
|
||||
parser.add_argument("--true-cfg-scale", type=float, default=None, help="True CFG scale for models that support it")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
||||
|
||||
parser.add_argument(
|
||||
"--helios-distilled-preset",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Send the Helios-Distilled extra_params used by the bundled example.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-params",
|
||||
type=lambda value: _json_object(value, argument_name="--extra-params"),
|
||||
default=None,
|
||||
help="JSON object merged into extra_params; overrides preset keys on conflict.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
try:
|
||||
asyncio.run(stream_video(args))
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,11 @@
|
||||
<!-- This file and the corresponding JS file are needed to display a streamable video view in the Gradio demo.
|
||||
Gradio's built-in video player does not support modern M4S streaming format. -->
|
||||
<div id="vllm-streaming-video-view" style="display:flex; flex-direction:column; gap:10px;">
|
||||
<video id="vllm-streaming-video-player" autoplay muted controls playsinline
|
||||
style="width:100%; min-height:320px; background:#111; border-radius:8px;"></video>
|
||||
<div style="display:flex; align-items:center; gap:10px;">
|
||||
<span id="vllm-streaming-video-status" style="font-weight:600;">Ready</span>
|
||||
</div>
|
||||
<pre id="vllm-streaming-video-log" style="height:220px; overflow:auto; margin:0; padding:10px; border-radius:8px;
|
||||
background:#0f172a; color:#dbeafe; font-size:12px; white-space:pre-wrap;">Ready.</pre>
|
||||
</div>
|
||||
@@ -0,0 +1,198 @@
|
||||
/**
|
||||
* This file and the corresponding HTML file are needed to display a streamable video view in the Gradio demo.
|
||||
* Gradio's built-in video player does not support modern M4S streaming format.
|
||||
*/
|
||||
(function () {
|
||||
const playerId = "vllm-streaming-video-player";
|
||||
const statusId = "vllm-streaming-video-status";
|
||||
const logId = "vllm-streaming-video-log";
|
||||
let ws = null;
|
||||
let mediaSource = null;
|
||||
let sourceBuffer = null;
|
||||
let queue = [];
|
||||
let chunkCount = 0;
|
||||
let totalBytes = 0;
|
||||
let done = false;
|
||||
let objectUrl = null;
|
||||
|
||||
function el(id) {
|
||||
return document.getElementById(id);
|
||||
}
|
||||
|
||||
function setStatus(text) {
|
||||
const node = el(statusId);
|
||||
if (node) node.textContent = text;
|
||||
}
|
||||
|
||||
function log(line) {
|
||||
const node = el(logId);
|
||||
if (!node) return;
|
||||
node.textContent += "\n" + line;
|
||||
node.scrollTop = node.scrollHeight;
|
||||
}
|
||||
|
||||
function setStartButton(disabled, label) {
|
||||
const btn = document.querySelector("#streaming-video-start button") || document.getElementById("streaming-video-start");
|
||||
if (!btn) return;
|
||||
btn.disabled = disabled;
|
||||
btn.textContent = label;
|
||||
}
|
||||
|
||||
function chooseMime() {
|
||||
const candidates = [
|
||||
'video/mp4; codecs="avc1.42E01E"',
|
||||
'video/mp4; codecs="avc1.4D401F"',
|
||||
'video/mp4; codecs="avc1.64001F"',
|
||||
"video/mp4"
|
||||
];
|
||||
for (const mime of candidates) {
|
||||
if (window.MediaSource && MediaSource.isTypeSupported(mime)) return mime;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function pump() {
|
||||
if (!sourceBuffer || sourceBuffer.updating || queue.length === 0) return;
|
||||
try {
|
||||
sourceBuffer.appendBuffer(queue.shift());
|
||||
} catch (err) {
|
||||
done = true;
|
||||
setStatus("SourceBuffer error");
|
||||
setStartButton(false, "Restart");
|
||||
log("SourceBuffer append error: " + err.message);
|
||||
finishStream();
|
||||
}
|
||||
}
|
||||
|
||||
function finishStream() {
|
||||
if (!mediaSource || mediaSource.readyState !== "open") return;
|
||||
if (sourceBuffer && sourceBuffer.updating) return;
|
||||
if (queue.length > 0) return;
|
||||
try {
|
||||
mediaSource.endOfStream();
|
||||
} catch (_) { }
|
||||
}
|
||||
|
||||
function stopCurrent() {
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
try { ws.send(JSON.stringify({ type: "session.stop" })); } catch (_) { }
|
||||
try { ws.close(); } catch (_) { }
|
||||
}
|
||||
ws = null;
|
||||
sourceBuffer = null;
|
||||
mediaSource = null;
|
||||
queue = [];
|
||||
if (objectUrl) URL.revokeObjectURL(objectUrl);
|
||||
objectUrl = null;
|
||||
}
|
||||
|
||||
window.vllmStreamingVideoStart = function (configJson) {
|
||||
stopCurrent();
|
||||
chunkCount = 0;
|
||||
totalBytes = 0;
|
||||
done = false;
|
||||
const logNode = el(logId);
|
||||
if (logNode) logNode.textContent = "Starting...";
|
||||
setStatus("Starting...");
|
||||
|
||||
let config;
|
||||
try {
|
||||
config = JSON.parse(configJson);
|
||||
} catch (err) {
|
||||
setStatus("Invalid request");
|
||||
setStartButton(false, "Restart");
|
||||
log("Invalid request config: " + err.message);
|
||||
return configJson;
|
||||
}
|
||||
|
||||
const mime = chooseMime();
|
||||
if (!mime) {
|
||||
setStatus("MSE unsupported");
|
||||
setStartButton(false, "Restart");
|
||||
log("This browser does not support MP4 Media Source Extensions playback.");
|
||||
return configJson;
|
||||
}
|
||||
|
||||
const video = el(playerId);
|
||||
if (!video) return configJson;
|
||||
mediaSource = new MediaSource();
|
||||
objectUrl = URL.createObjectURL(mediaSource);
|
||||
video.src = objectUrl;
|
||||
|
||||
mediaSource.addEventListener("sourceopen", () => {
|
||||
try {
|
||||
sourceBuffer = mediaSource.addSourceBuffer(mime);
|
||||
sourceBuffer.mode = "segments";
|
||||
sourceBuffer.addEventListener("updateend", () => {
|
||||
pump();
|
||||
if (done) finishStream();
|
||||
});
|
||||
} catch (err) {
|
||||
setStatus("SourceBuffer error");
|
||||
setStartButton(false, "Restart");
|
||||
log("SourceBuffer error: " + err.message);
|
||||
return;
|
||||
}
|
||||
|
||||
ws = new WebSocket(config.url);
|
||||
ws.binaryType = "arraybuffer";
|
||||
ws.onopen = () => {
|
||||
setStatus("Streaming...");
|
||||
setStartButton(true, "Streaming...");
|
||||
log("Connected: " + config.url);
|
||||
ws.send(JSON.stringify(config.payload));
|
||||
log("Sent session.start: " + JSON.stringify(config.payload));
|
||||
};
|
||||
ws.onmessage = (event) => {
|
||||
if (typeof event.data === "string") {
|
||||
const msg = JSON.parse(event.data);
|
||||
if (msg.type === "video.start") {
|
||||
log("Video session started: request_id=" + (msg.request_id || "") + " format=" + (msg.format || ""));
|
||||
} else if (msg.type === "session.done") {
|
||||
done = true;
|
||||
setStatus("Done");
|
||||
setStartButton(false, "Restart");
|
||||
log("Session complete: " + JSON.stringify(msg));
|
||||
finishStream();
|
||||
} else if (msg.type === "error") {
|
||||
done = true;
|
||||
setStatus("Error");
|
||||
setStartButton(false, "Restart");
|
||||
log("ERROR: " + (msg.message || JSON.stringify(msg)));
|
||||
finishStream();
|
||||
} else {
|
||||
log("Control message: " + JSON.stringify(msg));
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const data = event.data;
|
||||
queue.push(data);
|
||||
chunkCount += 1;
|
||||
totalBytes += data.byteLength;
|
||||
log(`[chunk ${String(chunkCount).padStart(3, "0")}] bytes=${data.byteLength} total_bytes=${totalBytes}`);
|
||||
pump();
|
||||
};
|
||||
ws.onerror = () => {
|
||||
done = true;
|
||||
setStatus("Error");
|
||||
setStartButton(false, "Restart");
|
||||
log("ERROR: WebSocket error");
|
||||
finishStream();
|
||||
};
|
||||
ws.onclose = () => {
|
||||
if (!done) {
|
||||
done = true;
|
||||
setStatus("Closed");
|
||||
setStartButton(false, "Restart");
|
||||
log("WebSocket closed before session.done");
|
||||
finishStream();
|
||||
}
|
||||
};
|
||||
}, { once: true });
|
||||
|
||||
return configJson;
|
||||
};
|
||||
|
||||
console.log("video-stream-view.html loaded");
|
||||
})();
|
||||
@@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | compat only, non-chunk | uploaded audio voice only; no presets | — |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
|
||||
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
|
||||
@@ -23,6 +24,7 @@ For the full list of supported architectures across all modalities, see
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
|
||||
| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) |
|
||||
|
||||
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead.
|
||||
|
||||
@@ -34,7 +36,7 @@ Launch the server (defaults shown — adjust `--port`, `--gpu-memory-utilization
|
||||
vllm serve <hf-repo-or-local-path> --omni --port 8091
|
||||
```
|
||||
|
||||
Send a TTS request via curl:
|
||||
Send a TTS request via curl. These generic snippets assume a model with a preset/default voice; voice-cloning-only models such as IndexTTS-2 require `ref_audio` or an uploaded audio `voice` (see model-specific sections below).
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
@@ -90,7 +92,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others).
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for IndexTTS-2, and 24 kHz for many others).
|
||||
|
||||
For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](../../../docs/serving/speech_api.md).
|
||||
|
||||
@@ -135,6 +137,40 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2
|
||||
|
||||
2-stage TTS (GPT AR + S2Mel CFM DiT + BigVGAN) at 22.05 kHz. Requests use `ref_audio` for voice cloning, or an uploaded audio `voice` from `/v1/audio/voices`. Supports emotion conditioning via `emo_audio`, `emo_text`, or `emo_vector` passed in `extra_params`.
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
|
||||
# or, to pass the bundled deploy config explicitly:
|
||||
bash examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# Voice cloning (ref_audio required)
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--text "你好,世界!" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
|
||||
# With emotion audio
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
--text "今天心情很好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 22.05 kHz mono WAV.
|
||||
- Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; IndexTTS-2 does not provide a built-in text-only preset voice.
|
||||
- Emotion params (`emo_audio`, `emo_text`, `emo_vector`, `emo_alpha`, `use_emo_text`, `use_random`) are passed via the `extra_params` field. Official precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as the speaker reference.
|
||||
- `stream=true` is accepted as an OpenAI-compatible response path, but IndexTTS-2 is not async-chunk streaming; audio is produced after S2Mel receives the full mel-code sequence.
|
||||
- Deploy config: `vllm_omni/deploy/indextts2.yaml` (auto-loaded).
|
||||
|
||||
---
|
||||
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
|
||||
@@ -504,6 +540,35 @@ python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
To receive word-level timestamps, launch the server with a forced aligner:
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--trust-remote-code \
|
||||
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
|
||||
```
|
||||
Then request PCM JSON sidecar chunks:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "Hello world. How are you?" \
|
||||
--stream-audio \
|
||||
--response-format pcm \
|
||||
--word-timestamps
|
||||
```
|
||||
The client writes one PCM file per sentence and a matching
|
||||
`sentence_XXX_timestamps.json` sidecar.
|
||||
|
||||
To *see* the alignment instead of reading a JSON sidecar, run the
|
||||
word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
|
||||
```bash
|
||||
python qwen3_tts/word_timestamps_demo.py --api-base http://localhost:8091
|
||||
```
|
||||
Each sentence's audio plays in an `<audio>` element while its text is rendered
|
||||
as inline word spans; the current word highlights as `audio.currentTime`
|
||||
crosses each `start_ms`. The **Stop (barge-in)** button cuts playback and
|
||||
reports the last-spoken word, useful for the voice-agent barge-in case.
|
||||
|
||||
### Gradio demos
|
||||
```bash
|
||||
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
|
||||
@@ -597,3 +662,65 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
|
||||
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
|
||||
- Voice cloning is gated upstream and may require a recent `mistral_common`.
|
||||
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
|
||||
|
||||
---
|
||||
|
||||
## SoulX-Singer
|
||||
|
||||
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
# SVS (default)
|
||||
export MODEL=/path/to/SoulX-Singer
|
||||
export PREPROCESS=/path/to/SoulX-Singer-Preprocess
|
||||
bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
|
||||
|
||||
# SVC
|
||||
export MODE=svc
|
||||
export MODEL=/path/to/SoulX-Singer-svc
|
||||
bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
|
||||
```
|
||||
|
||||
Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly:
|
||||
```bash
|
||||
export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS
|
||||
vllm serve $MODEL --omni \
|
||||
--deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \
|
||||
--port 8192 --trust-remote-code --enforce-eager
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`.
|
||||
|
||||
```bash
|
||||
# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
|
||||
python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
|
||||
--prompt-audio /path/on/server/zh_prompt.mp3 \
|
||||
--target-audio /path/on/server/music.mp3 \
|
||||
--preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \
|
||||
-o output.wav
|
||||
```
|
||||
|
||||
Use precomputed metadata to skip online preprocess with following command:
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
|
||||
--prompt-metadata-path /path/on/server/zh_prompt.json \
|
||||
--target-metadata-path /path/on/server/music.json \
|
||||
--audio-path /path/on/server/zh_prompt.mp3 \
|
||||
-o output.wav
|
||||
```
|
||||
|
||||
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 24 kHz mono WAV; batch only.
|
||||
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
|
||||
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for IndexTTS2
|
||||
#
|
||||
# Usage from repository root:
|
||||
# examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
# CUDA_VISIBLE_DEVICES=0 PORT=8092 MODEL=/path/to/IndexTeam/IndexTTS-2 examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
ROOT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)"
|
||||
|
||||
MODEL="${MODEL:-IndexTeam/IndexTTS-2}"
|
||||
PORT="${PORT:-8092}"
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT_DIR/vllm_omni/deploy/indextts2.yaml}"
|
||||
|
||||
echo "Starting IndexTTS2 server with model: $MODEL"
|
||||
|
||||
FLASHINFER_DISABLE_VERSION_CHECK=1 \
|
||||
vllm serve "$MODEL" \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--omni \
|
||||
--trust-remote-code \
|
||||
--deploy-config "$DEPLOY_CONFIG"
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""OpenAI-compatible client for IndexTTS2 TTS via /v1/audio/speech endpoint.
|
||||
|
||||
Examples:
|
||||
# With reference audio for voice cloning
|
||||
python speech_client.py --text "你好,世界!" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
|
||||
# With emotion audio
|
||||
python speech_client.py --text "今天心情很好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
|
||||
Server setup:
|
||||
vllm serve IndexTeam/IndexTTS-2 --omni --host 0.0.0.0 --port 8092
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
DEFAULT_API_BASE = "http://localhost:8092"
|
||||
DEFAULT_API_KEY = "sk-empty"
|
||||
|
||||
|
||||
def encode_audio_to_base64(audio_path: str) -> str:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
ext = audio_path.lower().rsplit(".", 1)[-1]
|
||||
mime = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac"}.get(ext, "audio/wav")
|
||||
with open(audio_path, "rb") as f:
|
||||
b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
return f"data:{mime};base64,{b64}"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="IndexTTS2 OpenAI speech client")
|
||||
parser.add_argument("--text", type=str, required=True)
|
||||
parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio for voice cloning")
|
||||
parser.add_argument("--emo-audio", type=str, default=None, help="Emotion reference audio")
|
||||
parser.add_argument("--emo-text", type=str, default=None, help="Emotion description text")
|
||||
parser.add_argument(
|
||||
"--emo-vector",
|
||||
type=float,
|
||||
nargs=8,
|
||||
default=None,
|
||||
help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm",
|
||||
)
|
||||
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1]")
|
||||
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text")
|
||||
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes")
|
||||
parser.add_argument("--model", type=str, default="IndexTeam/IndexTTS-2")
|
||||
parser.add_argument("--voice", type=str, default=None, help="Uploaded voice name to use instead of --ref-audio")
|
||||
parser.add_argument("--output", type=str, default="output.wav")
|
||||
parser.add_argument("--api-base", type=str, default=DEFAULT_API_BASE)
|
||||
parser.add_argument("--api-key", type=str, default=DEFAULT_API_KEY)
|
||||
parser.add_argument("--response-format", type=str, default="wav")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.ref_audio and not args.voice:
|
||||
parser.error("IndexTTS2 requires --ref-audio or --voice for voice cloning")
|
||||
|
||||
payload: dict = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
if args.voice:
|
||||
payload["voice"] = args.voice
|
||||
|
||||
if args.ref_audio:
|
||||
ref = args.ref_audio
|
||||
if ref.startswith(("http://", "https://", "data:")):
|
||||
payload["ref_audio"] = ref
|
||||
else:
|
||||
payload["ref_audio"] = encode_audio_to_base64(ref)
|
||||
|
||||
extra_params = {}
|
||||
if args.emo_audio:
|
||||
emo = args.emo_audio
|
||||
if emo.startswith(("http://", "https://", "data:")):
|
||||
extra_params["emo_audio"] = emo
|
||||
else:
|
||||
extra_params["emo_audio"] = encode_audio_to_base64(emo)
|
||||
if args.emo_text:
|
||||
extra_params["emo_text"] = args.emo_text
|
||||
if args.emo_vector is not None:
|
||||
extra_params["emo_vector"] = args.emo_vector
|
||||
if args.emo_alpha is not None:
|
||||
extra_params["emo_alpha"] = args.emo_alpha
|
||||
if args.use_emo_text:
|
||||
extra_params["use_emo_text"] = True
|
||||
if args.use_random:
|
||||
extra_params["use_random"] = True
|
||||
if extra_params:
|
||||
payload["extra_params"] = extra_params
|
||||
|
||||
url = f"{args.api_base}/v1/audio/speech"
|
||||
print(f"POST {url}")
|
||||
print(f" text: {args.text}")
|
||||
if args.ref_audio:
|
||||
print(f" ref_audio: {args.ref_audio[:80]}...")
|
||||
|
||||
with httpx.Client(timeout=300) as client:
|
||||
resp = client.post(
|
||||
url,
|
||||
json=payload,
|
||||
headers={"Authorization": f"Bearer {args.api_key}"},
|
||||
)
|
||||
|
||||
if resp.status_code != 200:
|
||||
print(f"Error {resp.status_code}: {resp.text[:500]}")
|
||||
return
|
||||
|
||||
with open(args.output, "wb") as f:
|
||||
f.write(resp.content)
|
||||
print(f"Saved: {args.output} ({len(resp.content):,} bytes)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -18,6 +18,21 @@ Usage:
|
||||
--text "Hello world. How are you? I am fine." \
|
||||
--simulate-stt --stt-delay 0.1
|
||||
|
||||
# Receive JSON sidecar chunks with word-level timestamps
|
||||
python streaming_speech_client.py \
|
||||
--text "Hello world. How are you?" \
|
||||
--stream-audio \
|
||||
--response-format pcm \
|
||||
--word-timestamps
|
||||
|
||||
# Also save every received audio.chunk frame as individual PCM/WAV files
|
||||
python streaming_speech_client.py \
|
||||
--text "Hello world. How are you?" \
|
||||
--stream-audio \
|
||||
--response-format pcm \
|
||||
--word-timestamps \
|
||||
--save-chunks
|
||||
|
||||
# VoiceDesign task
|
||||
python streaming_speech_client.py \
|
||||
--text "Today is a great day. The weather is nice." \
|
||||
@@ -37,8 +52,10 @@ Requirements:
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import wave
|
||||
|
||||
try:
|
||||
import websockets
|
||||
@@ -47,6 +64,37 @@ except ImportError:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str:
|
||||
audio_bytes = b"".join(chunks)
|
||||
if response_format == "pcm":
|
||||
raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm")
|
||||
with open(raw_path, "wb") as f:
|
||||
f.write(audio_bytes)
|
||||
|
||||
wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav")
|
||||
save_pcm_wav(wav_path, audio_bytes)
|
||||
return wav_path
|
||||
|
||||
filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}")
|
||||
with open(filename, "wb") as f:
|
||||
f.write(audio_bytes)
|
||||
return filename
|
||||
|
||||
|
||||
def save_pcm_wav(path: str, audio_bytes: bytes, sample_rate: int = 24000) -> None:
|
||||
with wave.open(path, "wb") as wf:
|
||||
wf.setnchannels(1)
|
||||
wf.setsampwidth(2)
|
||||
wf.setframerate(sample_rate)
|
||||
wf.writeframes(audio_bytes)
|
||||
|
||||
|
||||
def timestamp_words_text(timestamps: list[dict] | None) -> str:
|
||||
if not timestamps:
|
||||
return ""
|
||||
return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word"))
|
||||
|
||||
|
||||
async def stream_tts(
|
||||
url: str,
|
||||
text: str,
|
||||
@@ -54,6 +102,7 @@ async def stream_tts(
|
||||
output_dir: str,
|
||||
simulate_stt: bool = False,
|
||||
stt_delay: float = 0.1,
|
||||
save_chunks: bool = False,
|
||||
) -> None:
|
||||
"""Connect to the streaming TTS endpoint and process audio responses."""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
@@ -99,8 +148,14 @@ async def stream_tts(
|
||||
sender_task = asyncio.create_task(send_text())
|
||||
|
||||
response_format = config.get("response_format", "wav")
|
||||
word_timestamps = bool(config.get("word_timestamps", False))
|
||||
current_sentence_index = 0
|
||||
current_sentence_text = ""
|
||||
current_chunks: list[bytes] = []
|
||||
current_timestamps: list[dict] = []
|
||||
# Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
|
||||
# the sentence is a failure only if no frame ever carried timestamps.
|
||||
alignment_frame_seen = False
|
||||
|
||||
try:
|
||||
while True:
|
||||
@@ -116,22 +171,94 @@ async def stream_tts(
|
||||
|
||||
if msg_type == "audio.start":
|
||||
current_sentence_index = msg["sentence_index"]
|
||||
current_sentence_text = msg.get("sentence_text", "")
|
||||
current_chunks = []
|
||||
print(f" [sentence {msg['sentence_index']}] Generating: {msg['sentence_text']!r}")
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
print(f" [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}")
|
||||
elif msg_type == "audio.chunk":
|
||||
audio = base64.b64decode(msg["audio_b64"])
|
||||
current_chunks.append(audio)
|
||||
chunk_timestamps = msg.get("timestamps")
|
||||
# Sentence-level alignment sends a trailing timestamp-only
|
||||
# frame with empty audio; don't write an empty chunk file.
|
||||
if save_chunks and audio:
|
||||
chunk_base = os.path.join(
|
||||
output_dir,
|
||||
f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}",
|
||||
)
|
||||
chunk_pcm = f"{chunk_base}.pcm"
|
||||
chunk_wav = f"{chunk_base}.wav"
|
||||
chunk_json = f"{chunk_base}_timestamps.json"
|
||||
with open(chunk_pcm, "wb") as f:
|
||||
f.write(audio)
|
||||
save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
|
||||
with open(chunk_json, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"chunk_id": msg["chunk_id"],
|
||||
"chunk_start_ms": msg.get("chunk_start_ms"),
|
||||
"chunk_end_ms": msg.get("chunk_end_ms"),
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(chunk_timestamps),
|
||||
"timestamps": chunk_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
if chunk_timestamps is None:
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
|
||||
f"{len(audio)} bytes, timestamps unavailable"
|
||||
)
|
||||
else:
|
||||
alignment_frame_seen = True
|
||||
current_timestamps.extend(chunk_timestamps)
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
|
||||
f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
|
||||
)
|
||||
if save_chunks and audio:
|
||||
print(f" saved chunk -> {chunk_wav}")
|
||||
elif msg_type == "audio.done":
|
||||
filename = os.path.join(
|
||||
filename = save_audio_file(
|
||||
output_dir,
|
||||
f"sentence_{msg['sentence_index']:03d}.{response_format}",
|
||||
msg["sentence_index"],
|
||||
response_format,
|
||||
current_chunks,
|
||||
)
|
||||
with open(filename, "wb") as f:
|
||||
f.write(b"".join(current_chunks))
|
||||
if word_timestamps:
|
||||
# `null` on failure, else the accumulated list (`[]` for silence).
|
||||
sentence_timestamps = current_timestamps if alignment_frame_seen else None
|
||||
ts_filename = os.path.join(
|
||||
output_dir,
|
||||
f"sentence_{msg['sentence_index']:03d}_timestamps.json",
|
||||
)
|
||||
with open(ts_filename, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(sentence_timestamps),
|
||||
"timestamps": sentence_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] Done"
|
||||
f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
|
||||
f" error={msg.get('error', False)}"
|
||||
f" -> {filename}"
|
||||
)
|
||||
if word_timestamps:
|
||||
print(f" [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
elif msg_type == "session.done":
|
||||
print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
|
||||
break
|
||||
@@ -189,6 +316,16 @@ def main():
|
||||
action="store_true",
|
||||
help="Receive one or more PCM chunks per sentence (requires --response-format pcm)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--word-timestamps",
|
||||
action="store_true",
|
||||
help="Receive JSON sidecar chunks with word-level timestamps (requires --stream-audio --response-format pcm)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-chunks",
|
||||
action="store_true",
|
||||
help="Save each audio.chunk frame as sentence_XXX_chunk_YYY.{pcm,wav,json}",
|
||||
)
|
||||
parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens")
|
||||
|
||||
@@ -238,6 +375,10 @@ def main():
|
||||
config[key] = val
|
||||
if args.stream_audio:
|
||||
config["stream_audio"] = True
|
||||
if args.word_timestamps:
|
||||
config["word_timestamps"] = True
|
||||
config["stream_audio"] = True
|
||||
config["response_format"] = "pcm"
|
||||
if args.x_vector_only_mode:
|
||||
config["x_vector_only_mode"] = True
|
||||
|
||||
@@ -249,6 +390,7 @@ def main():
|
||||
output_dir=args.output_dir,
|
||||
simulate_stt=args.simulate_stt,
|
||||
stt_delay=args.stt_delay,
|
||||
save_chunks=args.save_chunks,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,457 @@
|
||||
"""Gradio demo for streaming TTS word-level timestamps.
|
||||
|
||||
Connects to the WebSocket endpoint ``/v1/audio/speech/stream`` with
|
||||
``word_timestamps: true`` and visualizes the alignment: each sentence's
|
||||
audio plays in a native ``<audio>`` element while its text is rendered as
|
||||
inline word spans, the current word highlighting as ``audio.currentTime``
|
||||
crosses each ``start_ms``.
|
||||
|
||||
A "Stop (barge-in)" button cuts playback and reports the last-spoken word,
|
||||
demonstrating the voice-agent barge-in use case.
|
||||
|
||||
Timestamps are sentence-level: audio streams in real time, then the forced
|
||||
aligner runs once per sentence and the complete word timestamps arrive in a
|
||||
trailing frame. The demo therefore plays one sentence block at a time.
|
||||
|
||||
Usage:
|
||||
# Launch the server with a forced aligner first, e.g.:
|
||||
# vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
|
||||
# --deploy-config vllm_omni/deploy/qwen3_tts.yaml --trust-remote-code \
|
||||
# --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
|
||||
# Then:
|
||||
python word_timestamps_demo.py --api-base http://localhost:8000
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
try:
|
||||
import gradio as gr
|
||||
except ImportError:
|
||||
raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
|
||||
try:
|
||||
import websockets
|
||||
except ImportError:
|
||||
raise ImportError("websockets is required to run this demo. Install it with: pip install websockets") from None
|
||||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
||||
from tts_common import (
|
||||
PCM_SAMPLE_RATE,
|
||||
SUPPORTED_LANGUAGES,
|
||||
TASK_TYPES,
|
||||
add_common_args,
|
||||
fetch_voices,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _upstream_ws_url(api_base: str) -> str:
|
||||
"""Map an http(s) API base to its WebSocket streaming endpoint."""
|
||||
base = api_base.rstrip("/")
|
||||
if base.startswith("https://"):
|
||||
base = "wss://" + base[len("https://") :]
|
||||
elif base.startswith("http://"):
|
||||
base = "ws://" + base[len("http://") :]
|
||||
else:
|
||||
base = "ws://" + base
|
||||
return base + "/v1/audio/speech/stream"
|
||||
|
||||
|
||||
# ── Front-end ────────────────────────────────────────────────────────
|
||||
PLAYER_HTML = """
|
||||
<div id="wt-root">
|
||||
<div style="display:flex; align-items:center; gap:10px;">
|
||||
<div id="wt-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div>
|
||||
<span id="wt-status" style="font-weight:600;font-size:1.05em;">Ready</span>
|
||||
<button id="wt-stop" onclick="window.wtStop()"
|
||||
style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552;
|
||||
background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em;">Stop (barge-in)</button>
|
||||
</div>
|
||||
<div id="wt-barge" style="display:none; margin-top:10px; padding:8px 12px; border-radius:6px;
|
||||
background:#fff5f5; border:1px solid #f3c2c2; color:#b02a2a; font-size:0.9em;"></div>
|
||||
<div id="wt-sentences" style="margin-top:12px; display:flex; flex-direction:column; gap:14px;"></div>
|
||||
</div>
|
||||
"""
|
||||
|
||||
|
||||
def _build_head_js(sample_rate: int) -> str:
|
||||
"""Browser-side player: WebSocket -> per-sentence <audio> + word highlight."""
|
||||
return f"""
|
||||
<script>
|
||||
const WT_SR = {sample_rate};
|
||||
let wtWS = null, wtQueue = [], wtPlaying = false, wtCur = null, wtLastWord = "";
|
||||
|
||||
function wtSet(text, color) {{
|
||||
const s = document.getElementById('wt-status');
|
||||
const d = document.getElementById('wt-dot');
|
||||
if (s) s.textContent = text;
|
||||
if (d) d.style.background = color || '#ccc';
|
||||
}}
|
||||
|
||||
function wtStopBtn(show) {{
|
||||
const b = document.getElementById('wt-stop');
|
||||
if (b) b.style.display = show ? 'inline-block' : 'none';
|
||||
}}
|
||||
|
||||
// Build a WAV Blob (16-bit mono PCM) from concatenated Int16 samples.
|
||||
function wtPcmToWav(int16, sr) {{
|
||||
const n = int16.length;
|
||||
const buf = new ArrayBuffer(44 + n * 2);
|
||||
const v = new DataView(buf);
|
||||
const w = (off, s) => {{ for (let i = 0; i < s.length; i++) v.setUint8(off + i, s.charCodeAt(i)); }};
|
||||
w(0, 'RIFF'); v.setUint32(4, 36 + n * 2, true); w(8, 'WAVE');
|
||||
w(12, 'fmt '); v.setUint32(16, 16, true); v.setUint16(20, 1, true); v.setUint16(22, 1, true);
|
||||
v.setUint32(24, sr, true); v.setUint32(28, sr * 2, true); v.setUint16(32, 2, true); v.setUint16(34, 16, true);
|
||||
w(36, 'data'); v.setUint32(40, n * 2, true);
|
||||
let off = 44;
|
||||
for (let i = 0; i < n; i++) {{ v.setInt16(off, int16[i], true); off += 2; }}
|
||||
return new Blob([buf], {{ type: 'audio/wav' }});
|
||||
}}
|
||||
|
||||
function wtConcat(chunks) {{
|
||||
let total = 0;
|
||||
for (const c of chunks) total += c.length;
|
||||
const out = new Int16Array(total);
|
||||
let off = 0;
|
||||
for (const c of chunks) {{ out.set(c, off); off += c.length; }}
|
||||
return out;
|
||||
}}
|
||||
|
||||
// Render one finished sentence: <audio> + inline word spans.
|
||||
function wtRenderSentence(sentence) {{
|
||||
const wrap = document.getElementById('wt-sentences');
|
||||
if (!wrap) return;
|
||||
const block = document.createElement('div');
|
||||
block.style.cssText = 'border:1px solid #e3e8ef;border-radius:8px;padding:10px 12px;';
|
||||
|
||||
const head = document.createElement('div');
|
||||
head.style.cssText = 'font-size:0.75em;color:#888;margin-bottom:6px;';
|
||||
head.textContent = 'Sentence ' + sentence.index +
|
||||
(sentence.timestamps ? ' (' + sentence.timestamps.length + ' words)' : ' (no timestamps)');
|
||||
block.appendChild(head);
|
||||
|
||||
const textEl = document.createElement('div');
|
||||
textEl.style.cssText = 'font-size:1.15em;line-height:1.8;margin-bottom:8px;';
|
||||
if (sentence.timestamps && sentence.timestamps.length) {{
|
||||
sentence.spans = [];
|
||||
for (const ts of sentence.timestamps) {{
|
||||
const span = document.createElement('span');
|
||||
span.textContent = ts.word;
|
||||
span.dataset.start = ts.start_ms;
|
||||
span.dataset.end = ts.end_ms;
|
||||
span.style.cssText = 'padding:1px 2px;border-radius:3px;transition:background 0.1s;';
|
||||
textEl.appendChild(span);
|
||||
// Keep spacing for spaced languages; CJK chars sit flush.
|
||||
textEl.appendChild(document.createTextNode(' '));
|
||||
sentence.spans.push(span);
|
||||
}}
|
||||
}} else {{
|
||||
textEl.textContent = sentence.text || '(timestamps unavailable)';
|
||||
}}
|
||||
block.appendChild(textEl);
|
||||
|
||||
const audio = document.createElement('audio');
|
||||
audio.controls = true;
|
||||
audio.style.cssText = 'width:100%;';
|
||||
audio.src = URL.createObjectURL(wtPcmToWav(wtConcat(sentence.chunks), WT_SR));
|
||||
block.appendChild(audio);
|
||||
sentence.audio = audio;
|
||||
|
||||
// Highlight the current word. Driven by requestAnimationFrame (~60fps)
|
||||
// rather than the <audio> 'timeupdate' event, which only fires ~4x/sec
|
||||
// and is too coarse for short words. We light the latest word whose
|
||||
// start_ms has passed, so there is always exactly one active word (no
|
||||
// inter-word flicker) and wtLastWord stays well-defined for barge-in.
|
||||
let raf = null;
|
||||
function highlight() {{
|
||||
if (!sentence.spans) return;
|
||||
const t = audio.currentTime * 1000;
|
||||
let idx = -1;
|
||||
for (let i = 0; i < sentence.spans.length; i++) {{
|
||||
if (t >= +sentence.spans[i].dataset.start) idx = i;
|
||||
else break;
|
||||
}}
|
||||
sentence.spans.forEach((sp, i) => {{
|
||||
const on = i === idx;
|
||||
sp.style.background = on ? '#ffe680' : 'transparent';
|
||||
sp.style.fontWeight = on ? '700' : '400';
|
||||
}});
|
||||
if (idx >= 0) wtLastWord = sentence.spans[idx].textContent;
|
||||
}}
|
||||
function loop() {{
|
||||
highlight();
|
||||
if (!audio.paused && !audio.ended) raf = requestAnimationFrame(loop);
|
||||
}}
|
||||
audio.addEventListener('play', () => {{
|
||||
wtCur = audio;
|
||||
wtSet('Playing sentence ' + sentence.index + '...', '#64dd17');
|
||||
cancelAnimationFrame(raf);
|
||||
raf = requestAnimationFrame(loop);
|
||||
}});
|
||||
audio.addEventListener('pause', () => cancelAnimationFrame(raf));
|
||||
audio.addEventListener('seeked', highlight);
|
||||
audio.addEventListener('ended', () => {{
|
||||
cancelAnimationFrame(raf);
|
||||
highlight();
|
||||
wtPlaying = false;
|
||||
wtPlayNext();
|
||||
}});
|
||||
|
||||
wrap.appendChild(block);
|
||||
wtQueue.push(audio);
|
||||
wtPlayNext();
|
||||
}}
|
||||
|
||||
function wtPlayNext() {{
|
||||
if (wtPlaying) return;
|
||||
const next = wtQueue.shift();
|
||||
if (!next) {{
|
||||
if (wtWS === null) {{ wtSet('Done', '#64dd17'); wtStopBtn(false); }}
|
||||
return;
|
||||
}}
|
||||
wtPlaying = true;
|
||||
next.play().catch(() => {{ wtPlaying = false; }});
|
||||
}}
|
||||
|
||||
window.wtStop = function() {{
|
||||
// Barge-in: cut playback, stop generation, report the last-spoken word.
|
||||
if (wtCur) {{ try {{ wtCur.pause(); }} catch (e) {{}} }}
|
||||
wtQueue = [];
|
||||
wtPlaying = false;
|
||||
if (wtWS) {{ try {{ wtWS.close(); }} catch (e) {{}} wtWS = null; }}
|
||||
wtSet('Stopped (barge-in)', '#999');
|
||||
wtStopBtn(false);
|
||||
const b = document.getElementById('wt-barge');
|
||||
if (b) {{
|
||||
b.style.display = 'block';
|
||||
b.textContent = wtLastWord
|
||||
? 'Barge-in at word: "' + wtLastWord + '" (a voice agent would resume the conversation from here)'
|
||||
: 'Barge-in before any word was spoken.';
|
||||
}}
|
||||
}};
|
||||
|
||||
window.wtGenerate = function(payload) {{
|
||||
if (wtWS) {{ try {{ wtWS.close(); }} catch (e) {{}} }}
|
||||
wtQueue = []; wtPlaying = false; wtCur = null; wtLastWord = "";
|
||||
document.getElementById('wt-sentences').innerHTML = '';
|
||||
const barge = document.getElementById('wt-barge');
|
||||
if (barge) {{ barge.style.display = 'none'; barge.textContent = ''; }}
|
||||
wtStopBtn(true);
|
||||
wtSet('Connecting...', '#4A90D9');
|
||||
|
||||
const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
const ws = new WebSocket(proto + '//' + location.host + '/proxy/stream');
|
||||
wtWS = ws;
|
||||
const sentences = {{}};
|
||||
|
||||
ws.onopen = () => {{
|
||||
wtSet('Streaming...', '#4A90D9');
|
||||
ws.send(JSON.stringify(Object.assign({{ type: 'session.config' }}, payload.config)));
|
||||
ws.send(JSON.stringify({{ type: 'input.text', text: payload.text }}));
|
||||
ws.send(JSON.stringify({{ type: 'input.done' }}));
|
||||
}};
|
||||
|
||||
ws.onmessage = (ev) => {{
|
||||
let msg;
|
||||
try {{ msg = JSON.parse(ev.data); }} catch (e) {{ return; }}
|
||||
if (msg.type === 'audio.start') {{
|
||||
sentences[msg.sentence_index] = {{
|
||||
index: msg.sentence_index, text: msg.sentence_text || '',
|
||||
chunks: [], timestamps: null
|
||||
}};
|
||||
}} else if (msg.type === 'audio.chunk') {{
|
||||
const s = sentences[msg.sentence_index];
|
||||
if (!s) return;
|
||||
if (msg.audio_b64) {{
|
||||
const bin = atob(msg.audio_b64);
|
||||
const bytes = new Uint8Array(bin.length);
|
||||
for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i);
|
||||
const usable = bytes.length - (bytes.length % 2);
|
||||
if (usable > 0) s.chunks.push(new Int16Array(bytes.buffer, 0, usable / 2));
|
||||
}}
|
||||
if (msg.timestamps != null) s.timestamps = msg.timestamps;
|
||||
}} else if (msg.type === 'audio.done') {{
|
||||
const s = sentences[msg.sentence_index];
|
||||
if (s) wtRenderSentence(s);
|
||||
}} else if (msg.type === 'session.done') {{
|
||||
wtWS = null;
|
||||
try {{ ws.close(); }} catch (e) {{}}
|
||||
wtPlayNext();
|
||||
}} else if (msg.type === 'error') {{
|
||||
wtSet('Error: ' + (msg.message || 'unknown'), '#EF5552');
|
||||
}}
|
||||
}};
|
||||
|
||||
ws.onerror = () => wtSet('WebSocket error', '#EF5552');
|
||||
ws.onclose = () => {{ if (wtWS === ws) wtWS = null; }};
|
||||
}};
|
||||
</script>
|
||||
"""
|
||||
|
||||
|
||||
def create_app(api_base: str) -> FastAPI:
|
||||
"""FastAPI app with a same-origin WebSocket proxy + Gradio UI."""
|
||||
fastapi_app = FastAPI()
|
||||
upstream_url = _upstream_ws_url(api_base)
|
||||
|
||||
@fastapi_app.websocket("/proxy/stream")
|
||||
async def proxy_stream(client_ws: WebSocket):
|
||||
await client_ws.accept()
|
||||
try:
|
||||
async with websockets.connect(upstream_url, max_size=None) as upstream:
|
||||
|
||||
async def client_to_upstream():
|
||||
try:
|
||||
while True:
|
||||
await upstream.send(await client_ws.receive_text())
|
||||
except (WebSocketDisconnect, Exception):
|
||||
await upstream.close()
|
||||
|
||||
async def upstream_to_client():
|
||||
try:
|
||||
async for msg in upstream:
|
||||
if isinstance(msg, bytes):
|
||||
await client_ws.send_bytes(msg)
|
||||
else:
|
||||
await client_ws.send_text(msg)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
await client_ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
await asyncio.gather(client_to_upstream(), upstream_to_client())
|
||||
except Exception as exc:
|
||||
logger.exception("WS proxy error")
|
||||
try:
|
||||
await client_ws.send_text(json.dumps({"type": "error", "message": f"proxy: {exc}"}))
|
||||
await client_ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
voices = fetch_voices(api_base)
|
||||
|
||||
with gr.Blocks(title="Qwen3-TTS Word Timestamps") as demo:
|
||||
gr.HTML(f"""
|
||||
<div style="margin-bottom:8px;">
|
||||
<h1 style="margin:0; font-size:1.5em;">Qwen3-TTS Streaming Word Timestamps</h1>
|
||||
<span style="font-size:0.85em; color:#666;">
|
||||
Served by <a href="https://github.com/vllm-project/vllm-omni" target="_blank">vLLM-Omni</a>
|
||||
· <code>{api_base}</code> ·
|
||||
launch the server with <code>--forced-aligner</code>.
|
||||
</span>
|
||||
</div>
|
||||
""")
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column(scale=3):
|
||||
text_input = gr.Textbox(
|
||||
label="Text to Synthesize",
|
||||
value="Hello world. How are you today? I am doing great.",
|
||||
lines=3,
|
||||
)
|
||||
with gr.Row():
|
||||
task_type = gr.Radio(choices=TASK_TYPES, value="CustomVoice", label="Task Type", scale=2)
|
||||
language = gr.Dropdown(choices=SUPPORTED_LANGUAGES, value="Auto", label="Language", scale=1)
|
||||
voice = gr.Dropdown(
|
||||
choices=voices,
|
||||
value=voices[0] if voices else None,
|
||||
label="Speaker",
|
||||
allow_custom_value=True,
|
||||
)
|
||||
instructions = gr.Textbox(
|
||||
label="Instructions",
|
||||
placeholder="Optional style/emotion (required for VoiceDesign)",
|
||||
lines=1,
|
||||
)
|
||||
with gr.Column(visible=False) as ref_group:
|
||||
ref_audio_url = gr.Textbox(
|
||||
label="Reference Audio URL (Base task)",
|
||||
placeholder="https://example.com/reference.wav",
|
||||
lines=1,
|
||||
)
|
||||
ref_text = gr.Textbox(label="Reference Transcript", lines=1)
|
||||
generate_btn = gr.Button("Generate + Align", variant="primary", size="lg")
|
||||
|
||||
with gr.Column(scale=2):
|
||||
gr.HTML(value=PLAYER_HTML)
|
||||
|
||||
hidden_payload = gr.Textbox(visible=False, elem_id="wt-payload")
|
||||
|
||||
def on_task_change(tt):
|
||||
return gr.update(visible=(tt == "Base"))
|
||||
|
||||
task_type.change(fn=on_task_change, inputs=[task_type], outputs=[ref_group])
|
||||
|
||||
def on_generate(text, tt, voice_v, lang_v, instr_v, ref_url, ref_t):
|
||||
if not text or not text.strip():
|
||||
raise gr.Error("Please enter text to synthesize.")
|
||||
config: dict = {
|
||||
"task_type": tt,
|
||||
"word_timestamps": True,
|
||||
"stream_audio": True,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
if lang_v:
|
||||
config["language"] = lang_v
|
||||
if tt == "CustomVoice":
|
||||
if voice_v:
|
||||
config["voice"] = voice_v
|
||||
if instr_v and instr_v.strip():
|
||||
config["instructions"] = instr_v.strip()
|
||||
elif tt == "VoiceDesign":
|
||||
if not instr_v or not instr_v.strip():
|
||||
raise gr.Error("VoiceDesign requires voice style instructions.")
|
||||
config["instructions"] = instr_v.strip()
|
||||
elif tt == "Base":
|
||||
if not ref_url or not ref_url.strip():
|
||||
raise gr.Error("Base (voice clone) needs a reference audio URL.")
|
||||
config["ref_audio"] = ref_url.strip()
|
||||
if ref_t and ref_t.strip():
|
||||
config["ref_text"] = ref_t.strip()
|
||||
return json.dumps({"config": config, "text": text.strip()})
|
||||
|
||||
generate_btn.click(
|
||||
fn=on_generate,
|
||||
inputs=[text_input, task_type, voice, language, instructions, ref_audio_url, ref_text],
|
||||
outputs=[hidden_payload],
|
||||
).then(
|
||||
fn=lambda p: p,
|
||||
inputs=[hidden_payload],
|
||||
outputs=[hidden_payload],
|
||||
js="(p) => { if (p && p.trim()) window.wtGenerate(JSON.parse(p)); return p; }",
|
||||
)
|
||||
|
||||
demo.queue()
|
||||
|
||||
return gr.mount_gradio_app(
|
||||
fastapi_app,
|
||||
demo,
|
||||
path="/",
|
||||
head=_build_head_js(PCM_SAMPLE_RATE),
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Gradio demo for streaming TTS word-level timestamps.")
|
||||
add_common_args(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
print(f"Connecting to vLLM server at: {args.api_base}")
|
||||
app = create_app(args.api_base)
|
||||
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host=args.host, port=args.port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,186 @@
|
||||
#!/usr/bin/env python3
|
||||
"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC).
|
||||
|
||||
Sends prompt audio via ``input_audio`` and target accompaniment via
|
||||
``extra_args['target_audio']`` (server-local path). For integrated preprocess,
|
||||
also pass ``preprocess_weights_dir`` in ``extra_args``.
|
||||
|
||||
Usage:
|
||||
python openai_chat_client.py \\
|
||||
--prompt-audio /path/on/server/zh_prompt.mp3 \\
|
||||
--target-audio /path/on/server/music.mp3 \\
|
||||
--preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\
|
||||
-o output.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
import soundfile
|
||||
import torch
|
||||
|
||||
|
||||
def _audio_to_data_url(path: Path) -> str:
|
||||
with path.open("rb") as handle:
|
||||
data = base64.b64encode(handle.read()).decode("ascii")
|
||||
return f"data:audio/mpeg;base64,{data}"
|
||||
|
||||
|
||||
def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
|
||||
audio = audio.to(torch.float32)
|
||||
peak = audio.abs().max().clamp(min=1e-8)
|
||||
audio = audio / peak
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
|
||||
|
||||
|
||||
def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
|
||||
for choice in body.get("choices", []):
|
||||
audio_obj = choice.get("message", {}).get("audio")
|
||||
if isinstance(audio_obj, dict) and audio_obj.get("data"):
|
||||
data, sr = soundfile.read(
|
||||
io.BytesIO(base64.b64decode(audio_obj["data"])),
|
||||
dtype="float32",
|
||||
always_2d=True,
|
||||
)
|
||||
return torch.from_numpy(data).transpose(0, 1), sr
|
||||
brief = {k: v for k, v in body.items() if k != "choices"}
|
||||
raise RuntimeError(f"no audio in response message.audio: {brief}")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
default_assets = repo_root / "tests" / "assets" / "soulxsinger"
|
||||
|
||||
parser = argparse.ArgumentParser(description="SoulX-Singer online chat client")
|
||||
parser.add_argument("--port", type=int, default=8192)
|
||||
parser.add_argument("--model", default="Soul-AILab/SoulX-Singer")
|
||||
parser.add_argument(
|
||||
"--prompt-audio",
|
||||
default=str(default_assets / "zh_prompt.mp3"),
|
||||
help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-audio",
|
||||
default=str(default_assets / "music.mp3"),
|
||||
help="Target accompaniment path on the server (extra_args['target_audio'])",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-metadata-path",
|
||||
default=None,
|
||||
help="SVS precomputed prompt metadata.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-metadata-path",
|
||||
default=None,
|
||||
help="SVS precomputed target metadata.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
default=None,
|
||||
help="SVS prompt vocal wav for precomputed metadata",
|
||||
)
|
||||
parser.add_argument("--preprocess-weights-dir", default=None)
|
||||
parser.add_argument("--output", "-o", default="soulxsinger_out.wav")
|
||||
parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs")
|
||||
parser.add_argument("--language", default="Mandarin")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=32)
|
||||
parser.add_argument("--guidance-scale", type=float, default=3.0)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Optional CFM seed. Omit for non-deterministic sampling.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auto-shift",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Auto pitch shift (default: on, original upstream infer.sh)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--control",
|
||||
default="melody",
|
||||
choices=["melody", "score"],
|
||||
help="SVS control mode",
|
||||
)
|
||||
parser.add_argument("--vocal-sep", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path)
|
||||
if any(meta_paths) and not all(meta_paths):
|
||||
print(
|
||||
"ERROR: precomputed metadata requires --prompt-metadata-path, "
|
||||
"--target-metadata-path, and --audio-path together.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
extra_args: dict = {
|
||||
"vocal_sep": args.vocal_sep,
|
||||
"auto_shift": args.auto_shift,
|
||||
"pitch_shift": 0,
|
||||
}
|
||||
if all(meta_paths):
|
||||
extra_args.update(
|
||||
{
|
||||
"prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()),
|
||||
"target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()),
|
||||
"audio_path": str(Path(args.audio_path).expanduser().resolve()),
|
||||
}
|
||||
)
|
||||
content = [{"type": "text", "text": "soulx-singer"}]
|
||||
else:
|
||||
prompt_path = Path(args.prompt_audio).expanduser().resolve()
|
||||
if not prompt_path.is_file():
|
||||
print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr)
|
||||
return 2
|
||||
extra_args["prompt_audio"] = str(prompt_path)
|
||||
extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve())
|
||||
if args.preprocess_weights_dir:
|
||||
extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve())
|
||||
content = [
|
||||
{"type": "text", "text": "soulx-singer"},
|
||||
{
|
||||
"type": "input_audio",
|
||||
"input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"},
|
||||
},
|
||||
]
|
||||
if not args.svc:
|
||||
extra_args["language"] = args.language
|
||||
extra_args["control"] = args.control
|
||||
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"modalities": ["audio"],
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
"extra_args": extra_args,
|
||||
}
|
||||
if args.seed is not None:
|
||||
payload["seed"] = args.seed
|
||||
|
||||
print(f"POST http://localhost:{args.port}/v1/chat/completions")
|
||||
response = requests.post(
|
||||
f"http://localhost:{args.port}/v1/chat/completions",
|
||||
headers={"Content-Type": "application/json"},
|
||||
json=payload,
|
||||
timeout=1800,
|
||||
)
|
||||
response.raise_for_status()
|
||||
audio, sample_rate = _decode_audio_from_response(response.json())
|
||||
_save_wav(audio, Path(args.output), sample_rate)
|
||||
duration = audio.shape[-1] / sample_rate
|
||||
print(f"saved {args.output} sr={sample_rate}Hz duration={duration:.2f}s")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,39 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline).
|
||||
#
|
||||
# Usage:
|
||||
# MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \
|
||||
#
|
||||
# Audio paths in client extra_args must be readable on the server host.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
|
||||
|
||||
MODEL="${MODEL:-Soul-AILab/SoulX-Singer}"
|
||||
MODE="${MODE:-svs}"
|
||||
PORT="${PORT:-8192}"
|
||||
GPUS="${GPUS:-0}"
|
||||
|
||||
if [[ "$MODE" == "svc" ]]; then
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}"
|
||||
else
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}"
|
||||
fi
|
||||
|
||||
echo "Starting SoulX-Singer server"
|
||||
echo " MODEL=$MODEL"
|
||||
echo " MODE=$MODE"
|
||||
echo " PORT=$PORT"
|
||||
echo " DEPLOY_CONFIG=$DEPLOY_CONFIG"
|
||||
echo " CUDA_VISIBLE_DEVICES=$GPUS"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="$GPUS" \
|
||||
vllm serve "$MODEL" \
|
||||
--omni \
|
||||
--deploy-config "$DEPLOY_CONFIG" \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
@@ -10,6 +10,7 @@ This example demonstrates how to deploy text-to-video models for online video ge
|
||||
| Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` |
|
||||
| Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
|
||||
| LTX-2 | `Lightricks/LTX-2` |
|
||||
| Helios (Base / Mid / Distilled) | `BestWishYsh/Helios-Base`, `Helios-Mid`, `Helios-Distilled` |
|
||||
|
||||
## Wan2.2 T2V
|
||||
|
||||
@@ -323,3 +324,45 @@ curl -sS -X POST http://localhost:8098/v1/videos \
|
||||
-F "guidance_scale=3.0" \
|
||||
-F "seed=42"
|
||||
```
|
||||
|
||||
## Helios
|
||||
|
||||
Helios ships three variants (`Helios-Base`, `Helios-Mid`, `Helios-Distilled`) that
|
||||
share the same server launch. Variant-specific knobs (declared in
|
||||
`vllm_omni/model_extras/helios.py`) are sent per request through the generic
|
||||
`extra_params` JSON form field — no per-model server flags required.
|
||||
|
||||
### Start Server
|
||||
|
||||
```bash
|
||||
vllm serve BestWishYsh/Helios-Base --omni --port 8098
|
||||
# or: MODEL=BestWishYsh/Helios-Mid bash run_server_helios.sh
|
||||
```
|
||||
|
||||
### Send Requests (curl)
|
||||
|
||||
```bash
|
||||
# Helios-Base (Stage 1 only)
|
||||
bash run_curl_helios.sh
|
||||
|
||||
# Helios-Mid (Stage 2 pyramid + CFG-Zero*)
|
||||
PRESET=mid-stage2 MODEL=BestWishYsh/Helios-Mid bash run_curl_helios.sh
|
||||
|
||||
# Helios-Distilled (Stage 2 pyramid + DMD, few-step)
|
||||
PRESET=distilled MODEL=BestWishYsh/Helios-Distilled bash run_curl_helios.sh
|
||||
```
|
||||
|
||||
The `mid-stage2` and `distilled` presets attach an `extra_params` field, e.g. for Helios-Distilled:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:8098/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A dynamic time-lapse of scenery rushing past the window of a speeding train." \
|
||||
-F "model=BestWishYsh/Helios-Distilled" \
|
||||
-F "size=640x384" \
|
||||
-F "num_frames=99" \
|
||||
-F "fps=16" \
|
||||
-F "guidance_scale=1.0" \
|
||||
-F "seed=42" \
|
||||
-F 'extra_params={"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}'
|
||||
```
|
||||
Loaded 100 of 478 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user