From 1e2a6446f79da72b66e15f684d0256244c725916 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:22:05 +0800 Subject: [PATCH] [CI] Remove AudioX and MammothModa2 support (#6353) Signed-off-by: wangyu <410167048@qq.com> --- docs/api/README.md | 4 - docs/models/supported_models.md | 3 - docs/user_guide/diffusion_features.md | 1 - examples/offline_inference/audex/README.md | 2 +- .../offline_inference/audex/text_to_audio.py | 2 +- .../run_mammothmoda2_image_summarize.py | 117 -- .../offline_inference/text_to_audio/README.md | 39 +- .../text_to_audio/text_to_audio.py | 110 +- .../text_to_image/text_to_image.py | 11 - .../offline_inference/x_to_text/README.md | 34 +- .../offline_inference/x_to_text/x_to_text.py | 5 +- examples/online_serving/audiox/README.md | 64 - .../audiox/openai_chat_client.py | 119 -- examples/online_serving/audiox/run_server.sh | 16 - recipes/MammothModa2/MammothModa2.md | 212 ---- recipes/README.md | 2 - recipes/audiox/AudioX.md | 130 --- tests/config/test_config_factory.py | 2 - .../test_diffusion_output_metadata.py | 2 - .../test_audiox_model_expansion.py | 71 -- .../test_mammoth_moda2_expansion.py | 204 ---- .../test_stable_audio_expansion.py | 3 +- .../online_serving/test_audiox_expansion.py | 63 - tests/helpers/runtime.py | 2 +- .../test_mammoth_moda2_config.py | 119 -- .../models/test_kv_cache_scale_mapper.py | 5 - tests/model_extras/test_model_extras.py | 119 -- tests/model_tests/diffusion/test_alignment.py | 1 - tests/worker/test_omni_connector_mixin.py | 1 - vllm_omni/config/pipeline_registry.py | 6 - vllm_omni/deploy/mammoth_moda2.yaml | 27 - vllm_omni/deploy/mammoth_moda2_ar.yaml | 14 - vllm_omni/diffusion/models/audiox/__init__.py | 11 - .../models/audiox/audiox_transformer.py | 487 -------- .../models/audiox/pipeline_audiox.py | 933 --------------- .../mammoth_moda2/mammothmoda2_dit_model.py | 839 -------------- .../pipeline_mammothmoda2_dit.py | 405 ------- .../models/mammoth_moda2/rope_real.py | 252 ---- .../models/mammoth_moda2/schedulers.py | 137 --- .../stable_audio/pipeline_stable_audio.py | 3 +- vllm_omni/diffusion/registry.py | 7 - vllm_omni/entrypoints/openai/serving_chat.py | 2 +- .../models/mammoth_moda2/__init__.py | 8 - .../models/mammoth_moda2/mammoth_moda2.py | 1032 ----------------- .../models/mammoth_moda2/pipeline.py | 68 -- .../pipeline_mammothmoda2_dit.py | 15 - vllm_omni/model_executor/models/registry.py | 35 - .../stage_input_processors/mammoth_moda2.py | 115 -- vllm_omni/model_extras/audiox.py | 27 - .../model_extras/mammothmodal2_preview.py | 101 -- vllm_omni/model_extras/registry.py | 34 - vllm_omni/outputs/output_metadata.py | 1 - .../npu/worker/npu_generation_model_runner.py | 6 +- .../tokenizers/mammoth_moda2_tokenizer.py | 401 ------- .../transformers_utils/configs/__init__.py | 9 - .../configs/mammoth_moda2.py | 377 ------ .../transformers_utils/processors/audiox.py | 240 ---- .../worker/gpu_generation_model_runner.py | 6 +- vllm_omni/worker/gpu_model_runner.py | 10 +- 59 files changed, 42 insertions(+), 7029 deletions(-) delete mode 100644 examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py delete mode 100644 examples/online_serving/audiox/README.md delete mode 100755 examples/online_serving/audiox/openai_chat_client.py delete mode 100755 examples/online_serving/audiox/run_server.sh delete mode 100644 recipes/MammothModa2/MammothModa2.md delete mode 100644 recipes/audiox/AudioX.md delete mode 100644 tests/e2e/offline_inference/test_audiox_model_expansion.py delete mode 100644 tests/e2e/offline_inference/test_mammoth_moda2_expansion.py delete mode 100644 tests/e2e/online_serving/test_audiox_expansion.py delete mode 100644 tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py delete mode 100644 vllm_omni/deploy/mammoth_moda2.yaml delete mode 100644 vllm_omni/deploy/mammoth_moda2_ar.yaml delete mode 100644 vllm_omni/diffusion/models/audiox/__init__.py delete mode 100644 vllm_omni/diffusion/models/audiox/audiox_transformer.py delete mode 100644 vllm_omni/diffusion/models/audiox/pipeline_audiox.py delete mode 100644 vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py delete mode 100644 vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py delete mode 100644 vllm_omni/diffusion/models/mammoth_moda2/rope_real.py delete mode 100644 vllm_omni/diffusion/models/mammoth_moda2/schedulers.py delete mode 100644 vllm_omni/model_executor/models/mammoth_moda2/__init__.py delete mode 100644 vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py delete mode 100644 vllm_omni/model_executor/models/mammoth_moda2/pipeline.py delete mode 100644 vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py delete mode 100644 vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py delete mode 100644 vllm_omni/model_extras/audiox.py delete mode 100644 vllm_omni/model_extras/mammothmodal2_preview.py delete mode 100644 vllm_omni/tokenizers/mammoth_moda2_tokenizer.py delete mode 100644 vllm_omni/transformers_utils/configs/mammoth_moda2.py delete mode 100644 vllm_omni/transformers_utils/processors/audiox.py diff --git a/docs/api/README.md b/docs/api/README.md index 4a34e3f31..62eecb3a4 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -114,10 +114,6 @@ Configuration classes. - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][] - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][] - [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][] -- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][] -- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][] -- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][] -- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][] - [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][] - [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][] diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 672a10174..f1d18ea48 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -68,15 +68,12 @@ th { | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — | | `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — | -| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — | -| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — | | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — | | `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — | | `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — | | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — | | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — | -| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — | | `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 4a46c6b25..538c3ee71 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -130,7 +130,6 @@ The following tables show which models support each feature: | **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | diff --git a/examples/offline_inference/audex/README.md b/examples/offline_inference/audex/README.md index 793b69eb9..4b570b06a 100644 --- a/examples/offline_inference/audex/README.md +++ b/examples/offline_inference/audex/README.md @@ -17,7 +17,7 @@ Shared plumbing (corpus/WAV IO, tokenizer loading, the CFG cond/uncond pair contract) lives in `common.py`; each task script keeps only its own prompt recipe and flow. Note `text_to_audio.py` here drives the Audex two-stage AR pipeline — for diffusion text-to-audio models (Stable Audio -Open, AudioX) use the generic +Open) use the generic `examples/offline_inference/text_to_audio/text_to_audio.py` instead. ## text_to_speech.py — text → speech diff --git a/examples/offline_inference/audex/text_to_audio.py b/examples/offline_inference/audex/text_to_audio.py index c357819c3..e81a53bb3 100644 --- a/examples/offline_inference/audex/text_to_audio.py +++ b/examples/offline_inference/audex/text_to_audio.py @@ -7,7 +7,7 @@ thinker generates interleaved 4-codebook RVQ tokens under an RVQ phase mask, and the external XCodec1 checkpoint decodes them to a 16 kHz mono WAV. -(For DIFFUSION text-to-audio models — Stable Audio Open, AudioX — use +(For DIFFUSION text-to-audio models — Stable Audio Open — use ``examples/offline_inference/text_to_audio/text_to_audio.py`` instead; Audex TTA is a two-stage autoregressive pipeline with a different interface.) diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py deleted file mode 100644 index 5363948dd..000000000 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py +++ /dev/null @@ -1,117 +0,0 @@ -""" -Offline inference example: MammothModa2 image summarization (single AR stage). - -Example: - uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \ - --model path/to/MammothModa2-Preview \ - --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \ - --image /path/to/input.jpg \ - --question "Please summarize the content of this image." -""" - -from __future__ import annotations - -import argparse -import os - -from PIL import Image -from vllm import SamplingParams -from vllm.multimodal.image import convert_image_mode - -from vllm_omni import Omni - -DEFAULT_SYSTEM = "You are a helpful assistant." -DEFAULT_QUESTION = "Please summarize the content of this image." - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).") - parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.") - parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.") - parser.add_argument("--image", type=str, required=True, help="Path to input image.") - parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.") - parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.") - parser.add_argument( - "--max-tokens", - type=int, - default=512, - help="Max new tokens to generate.", - ) - parser.add_argument("--temperature", type=float, default=0.2) - parser.add_argument("--top-p", type=float, default=0.9) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--trust-remote-code", action="store_true") - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable diffusion pipeline profiler to display stage durations.", - ) - return parser.parse_args() - - -def build_prompt(system: str, question: str) -> str: - return ( - f"<|im_start|>system\n{system}<|im_end|>\n" - "<|im_start|>user\n" - "<|vision_start|><|image_pad|><|vision_end|>" - f"{question}<|im_end|>\n" - "<|im_start|>assistant\n" - ) - - -def main() -> None: - args = parse_args() - - if not os.path.exists(args.image): - raise FileNotFoundError(f"Image file not found: {args.image}") - - pil_image = Image.open(args.image) - image_data = convert_image_mode(pil_image, "RGB") - prompt = build_prompt(args.system, args.question) - - omni = Omni( - model=args.model, - deploy_config=args.deploy_config, - trust_remote_code=args.trust_remote_code, - enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, - ) - try: - sp = SamplingParams( - temperature=float(args.temperature), - top_p=float(args.top_p), - top_k=-1, - max_tokens=int(args.max_tokens), - seed=int(args.seed), - detokenize=True, - ) - # NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None]. - # Consume it inside the try block so the worker isn't closed early. - outputs = list( - omni.generate( - [ - { - "prompt": prompt, - "multi_modal_data": {"image": image_data}, - "additional_information": {"omni_task": ["chat"]}, - } - ], - [sp], - ) - ) - finally: - omni.close() - - lines: list[str] = [] - for stage_outputs in outputs: - ro = stage_outputs - text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro) - lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n") - lines.append("answer:\n") - lines.append(text.strip() + "\n") - lines.append("\n") - - print("\n".join(lines)) - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/text_to_audio/README.md b/examples/offline_inference/text_to_audio/README.md index 89ffc8bd8..f096d96c8 100644 --- a/examples/offline_inference/text_to_audio/README.md +++ b/examples/offline_inference/text_to_audio/README.md @@ -1,11 +1,10 @@ # Text-To-Audio -A unified script for text/video-to-audio generation. Supported models: +A unified script for text-to-audio generation. Supported models: | Model | Tasks | Notes | |-------|-------|-------| | `stabilityai/stable-audio-open-1.0` | text-to-audio | gated; uses `--audio-length` | -| `zhangj1an/AudioX` | `t2a` / `t2m` / `v2a` / `v2m` / `tv2a` / `tv2m` | pass `--task`; video tasks need `--video` | The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts. @@ -50,42 +49,12 @@ python text_to_audio.py \ --output stable_audio_output.wav ``` -### AudioX - -AudioX supports six tasks. Sampler and reference knobs (declared in -`vllm_omni/model_extras/audiox.py`) are passed via the generic `--extra-body` -JSON flag, routed into sampling `extra_args`. - -Text tasks (`t2a` / `t2m`): - -```bash -python text_to_audio.py \ - --model zhangj1an/AudioX --task t2a \ - --prompt "Fireworks burst twice, followed by a clock ticking." \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 --seed 42 \ - --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \ - --output t2a.wav -``` - -Video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) require `--video`: - -```bash -python text_to_audio.py \ - --model zhangj1an/AudioX --task tv2a \ - --prompt "drum beating sound and human talking" \ - --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ - --output tv2a.wav -``` - Key arguments: - `--prompt`: text description (string). -- `--task`: [AudioX] one of `t2a`/`t2m`/`v2a`/`v2m`/`tv2a`/`tv2m`. -- `--video`: [AudioX `v2*`/`tv2*`] video file/URL for conditioning (→ `video_path`). -- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s` for Stable Audio, `seconds_start` for AudioX). -- `--audio-length`: audio duration in seconds (audio length for Stable Audio, `seconds_total` for AudioX). -- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. For AudioX, sampler/reference knobs go here, e.g. `'{"sigma_min": 0.03, "sigma_max": 1000.0, "cfg_rescale": 0.0, "audio_path": "ref.wav"}'`. +- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s`). +- `--audio-length`: audio duration in seconds (audio length for Stable Audio). +- `--extra-body`: JSON dict of model-specific knobs, merged into sampling `extra_args`. - `--negative-prompt`: negative prompt for classifier-free guidance. - `--seed`: integer seed for deterministic generation. - `--guidance-scale`: classifier-free guidance scale. diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 3a67c662a..094a1c66e 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -2,15 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Example script for text/video-to-audio generation with diffusion audio models. +Example script for text-to-audio generation with diffusion audio models. This script supports: * Stable Audio Open (`stabilityai/stable-audio-open-1.0`) — text-to-audio. - * AudioX (`zhangj1an/AudioX`) — six tasks: t2a / t2m / v2a / v2m / tv2a / tv2m. -Model-specific generation knobs are declared in `vllm_omni/model_extras` and -routed into `sampling_params.extra_args` via `apply_declared_extra_args`, so -the same script serves both models without per-model branching in the engine. +Model-specific generation knobs can be passed via `--extra-body` and are +merged into `sampling_params.extra_args`. Usage: # Stable Audio Open @@ -18,15 +16,6 @@ Usage: python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0 python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality" python text_to_audio.py --prompt "A soft synth pad" --cache-backend tea_cache - - # AudioX (text-to-audio / text-to-music) - python text_to_audio.py --model zhangj1an/AudioX --task t2a \ - --prompt "Fireworks burst twice, then a clock ticking." \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 - # AudioX video-conditioned (v2*/tv2*) — pass a video file/URL - python text_to_audio.py --model zhangj1an/AudioX --task tv2a \ - --prompt "drum beating sound and human talking" --video clip.mp4 \ - --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' """ import argparse @@ -38,16 +27,11 @@ from typing import Any import numpy as np import torch -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 -AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m") -AUDIOX_VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"}) - def parse_extra_body(value: str) -> dict[str, Any]: """Parse a JSON object string for --extra-body.""" @@ -65,7 +49,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--model", default="stabilityai/stable-audio-open-1.0", - help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0, zhangj1an/AudioX.", + help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0.", ) parser.add_argument( "--prompt", @@ -96,19 +80,13 @@ def parse_args() -> argparse.Namespace: "--audio-start", type=float, default=0.0, - help=( - "Audio start offset in seconds. Applies to both models " - "(audio_start_in_s for Stable Audio, seconds_start for AudioX)." - ), + help="Audio start offset in seconds (maps to audio_start_in_s).", ) parser.add_argument( "--audio-length", type=float, default=10.0, - help=( - "Audio duration in seconds. Maps to audio length for Stable Audio " - "(max ~47s for stable-audio-open-1.0) and seconds_total for AudioX." - ), + help=("Audio duration in seconds. Maps to audio length for Stable Audio (max ~47s for stable-audio-open-1.0)."), ) parser.add_argument( "--num-inference-steps", @@ -122,41 +100,20 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of audio waveforms to generate for the given prompt.", ) - parser.add_argument( - "--task", - type=str, - default=None, - choices=list(AUDIOX_TASKS), - help=( - "[AudioX only] Generation task: t2a/t2m/v2a/v2m/tv2a/tv2m. " - "v2*/tv2* require --video. Ignored for Stable Audio Open." - ), - ) - parser.add_argument( - "--video", - type=str, - default=None, - help="[AudioX v2*/tv2*] Video file path or URL used for conditioning.", - ) parser.add_argument( "--extra-body", type=parse_extra_body, default=None, help=( - "Model-specific generation params as a JSON object, e.g. " - '\'{"sigma_min": 0.03, "sigma_max": 1000.0}\'. Each key is filtered ' - "against the model's declared extra_body_params (see vllm_omni/model_extras), " - "so unknown keys are silently dropped. Values here override the equivalent flags." + "Model-specific generation params as a JSON object, merged into " + "sampling extra_args. Values here override the equivalent flags." ), ) parser.add_argument( "--output", type=str, default=None, - help=( - "Path to save the generated audio (WAV format). Defaults to " - "audiox_output.wav for AudioX and stable_audio_output.wav otherwise." - ), + help=("Path to save the generated audio (WAV format). Defaults to stable_audio_output.wav."), ) parser.add_argument( "--sample-rate", @@ -275,8 +232,6 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410 def main(): args = parse_args() - model_name = str(args.model).lower() if args.model else "" - is_audiox = "audiox" in model_name generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) cache_config = None if args.cache_backend == "tea_cache": @@ -285,11 +240,9 @@ def main(): } print(f"\n{'=' * 60}") - print("Text/Video-to-Audio Generation") + print("Text-to-Audio Generation") print(f"{'=' * 60}") print(f" Model: {args.model}") - if args.task is not None: - print(f" Task: {args.task}") print(f" Prompt: {args.prompt}") print(f" Negative prompt: {args.negative_prompt}") print(f" Duration: {args.audio_length}s") @@ -305,8 +258,6 @@ def main(): print(f" Seed: {args.seed}") print(f"{'=' * 60}\n") - # Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio - # is resolved from its config. omni_kwargs: dict[str, Any] = dict( model=args.model, use_hsdp=args.use_hsdp, @@ -318,11 +269,7 @@ def main(): enable_cpu_offload=args.enable_cpu_offload, enable_layerwise_offload=args.enable_layerwise_offload, ) - if is_audiox: - omni_kwargs["model_class_name"] = "AudioXPipeline" omni = Omni(**omni_kwargs) - model_class_name = get_model_class_name(omni) - declared_extra_body_params = get_extra_body_params(model_class_name) diffusion_params = OmniDiffusionSamplingParams( generator=generator, @@ -332,38 +279,19 @@ def main(): seed=args.seed, ) - # Negative prompt is opt-in for both models (off by default). Including a + # Negative prompt is opt-in (off by default). Including a # non-empty negative prompt enables the separate CFG conditioning branch. prompt: dict[str, Any] = {"prompt": args.prompt} if args.negative_prompt: prompt["negative_prompt"] = args.negative_prompt - if model_class_name == "AudioXPipeline": - if args.task is None: - raise ValueError("AudioX requires --task (one of t2a/t2m/v2a/v2m/tv2a/tv2m).") - if args.task in AUDIOX_VIDEO_TASKS and not args.video: - raise ValueError(f"AudioX task {args.task!r} requires --video.") - # AudioX core knobs are surfaced as shared flags; sampler/reference knobs - # (sigma_min/sigma_max/cfg_rescale/audio_path) arrive via --extra-body and - # are filtered against the declared extra_body contract. - user_extra: dict[str, Any] = { - "audiox_task": args.task, - "seconds_start": args.audio_start, - "seconds_total": args.audio_length, - "video_path": args.video, - } - if args.extra_body: - user_extra.update(args.extra_body) - apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra) - else: - # Stable Audio Open path (backward-compatible). - audio_end_in_s = args.audio_start + args.audio_length - diffusion_params.extra_args = { - "audio_start_in_s": args.audio_start, - "audio_end_in_s": audio_end_in_s, - } - if args.extra_body: - diffusion_params.extra_args.update(args.extra_body) + audio_end_in_s = args.audio_start + args.audio_length + diffusion_params.extra_args = { + "audio_start_in_s": args.audio_start, + "audio_end_in_s": audio_end_in_s, + } + if args.extra_body: + diffusion_params.extra_args.update(args.extra_body) generation_start = time.perf_counter() @@ -374,7 +302,7 @@ def main(): print(f"Total generation time: {generation_time:.2f} seconds") - output = args.output or ("audiox_output.wav" if model_class_name == "AudioXPipeline" else "stable_audio_output.wav") + output = args.output or "stable_audio_output.wav" output_path = Path(output) output_path.parent.mkdir(parents=True, exist_ok=True) suffix = output_path.suffix or ".wav" diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 2ee09e178..a08f97cfd 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -599,17 +599,6 @@ def main(): if args.seed is not None and hasattr(params, "seed"): params.seed = args.seed - # MammothModa2's AR stage emits one visual token per grid cell, - # one EOL token per row, and one final look-ahead token whose hidden - # state is unavailable. Size the first stage from the prompt metadata - # instead of SamplingParams' default of 16 tokens. - prompt_info = prompt_dict.get("additional_information", {}) - if idx == 0 and prompt_info.get("omni_task") == ["t2i"]: - ar_width = int(prompt_info.get("ar_width", [0])[0]) - ar_height = int(prompt_info.get("ar_height", [0])[0]) - if ar_width > 0 and ar_height > 0: - params.max_tokens = ar_height * (ar_width + 1) + 1 - if not diffusion_replaced and len(sampling_params_list) == 1: sampling_params_list = [diffusion_params] diff --git a/examples/offline_inference/x_to_text/README.md b/examples/offline_inference/x_to_text/README.md index de668b6ee..0cbf9731d 100644 --- a/examples/offline_inference/x_to_text/README.md +++ b/examples/offline_inference/x_to_text/README.md @@ -10,12 +10,10 @@ Generate text from text or image inputs with vLLM-Omni's shared offline entrypoi | --- | --- | --- | --- | | `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` | | `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` | -| `bytedance-research/MammothModa2-Preview` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` | -| `bytedance-research/MammothModa2-Dev` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` | The script recognizes these model families from `config.json`, applies the -model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3 -and MammothModa2. BAGEL uses its registered default deploy. +model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3. +BAGEL uses its registered default deploy. ## Text-To-Text @@ -35,22 +33,6 @@ python examples/offline_inference/x_to_text/x_to_text.py \ --prompt "Explain multimodal inference in three concise sentences." ``` -### MammothModa2 - -```bash -python examples/offline_inference/x_to_text/x_to_text.py \ - --model bytedance-research/MammothModa2-Preview \ - --prompt "Explain multimodal inference in three concise sentences." -``` - -The same command supports the Dev checkpoint: - -```bash -python examples/offline_inference/x_to_text/x_to_text.py \ - --model bytedance-research/MammothModa2-Dev \ - --prompt "Explain multimodal inference in three concise sentences." -``` - ## Image-To-Text Pass one image with `--image`. The default question can be replaced with any @@ -72,16 +54,6 @@ python examples/offline_inference/x_to_text/x_to_text.py \ --prompt "Describe this image in detail." ``` -```bash -python examples/offline_inference/x_to_text/x_to_text.py \ - --model bytedance-research/MammothModa2-Preview \ - --image image.png \ - --prompt "Describe this image in detail." -``` - -Use `bytedance-research/MammothModa2-Dev` in the same command to run I2T with -the Dev checkpoint. - ## Key Arguments | Argument | Default | Description | @@ -102,4 +74,4 @@ the Dev checkpoint. The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a different topology, pass a custom AR-only YAML through `--deploy-config`. -BAGEL and MammothModa2 use their repository deploy defaults unless overridden. +BAGEL uses its repository deploy default unless overridden. diff --git a/examples/offline_inference/x_to_text/x_to_text.py b/examples/offline_inference/x_to_text/x_to_text.py index 9937a57ff..258bc79df 100644 --- a/examples/offline_inference/x_to_text/x_to_text.py +++ b/examples/offline_inference/x_to_text/x_to_text.py @@ -15,7 +15,6 @@ from vllm_omni.model_extras import build_x_to_text_prompt, get_x_to_text_model_f _REPO_ROOT = Path(__file__).resolve().parents[3] _HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml" -_MAMMOTH_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "mammoth_moda2_ar.yaml" def parse_args() -> argparse.Namespace: @@ -62,7 +61,7 @@ def main() -> None: omni_kwargs: dict[str, Any] = { "model": args.model, "mode": "image-to-text" if image is not None else "text-to-text", - "trust_remote_code": args.trust_remote_code or family in {"hunyuan_image3", "mammoth_moda2"}, + "trust_remote_code": args.trust_remote_code or family == "hunyuan_image3", "log_stats": args.log_stats, } if args.enforce_eager is not None: @@ -71,8 +70,6 @@ def main() -> None: omni_kwargs["deploy_config"] = args.deploy_config elif family == "hunyuan_image3": omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG) - elif family == "mammoth_moda2": - omni_kwargs["deploy_config"] = str(_MAMMOTH_AR_DEPLOY_CONFIG) omni = Omni(**omni_kwargs) try: diff --git a/examples/online_serving/audiox/README.md b/examples/online_serving/audiox/README.md deleted file mode 100644 index bf068b3eb..000000000 --- a/examples/online_serving/audiox/README.md +++ /dev/null @@ -1,64 +0,0 @@ -# AudioX online serving - -Launches the `AudioXPipeline` behind vLLM-Omni's OpenAI-compatible chat endpoint and provides a -minimal Python client that covers all six tasks (`t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`). - -## Start the server - -```bash -cd examples/online_serving/audiox -bash run_server.sh # defaults: MODEL=zhangj1an/AudioX, PORT=8099 -``` - -Environment overrides: `MODEL`, `PORT`, `DIFFUSION_ATTENTION_BACKEND`. - -## Call from Python - -```bash -# text-to-audio -python openai_chat_client.py --task t2a \ - --prompt "Fireworks burst twice, followed by a period of silence before a clock begins ticking." \ - --output t2a.wav - -# text-to-music -python openai_chat_client.py --task t2m \ - --prompt "Uplifting ukulele tune for a travel vlog" \ - --output t2m.wav - -# video-to-audio (no text) -python openai_chat_client.py --task v2a --video path/to/clip.mp4 --output v2a.wav - -# text+video-to-audio -python openai_chat_client.py --task tv2a \ - --prompt "drum beating sound and human talking" \ - --video path/to/clip.mp4 \ - --output tv2a.wav -``` - -The client sends: - -- `num_inference_steps`, `guidance_scale`, `seed`, `audiox_task`, `seconds_start`, - `seconds_total`, `sigma_min`, `sigma_max` as flattened keys under `extra_body`. The server - routes the declared knobs (see `vllm_omni/model_extras/audiox.py`) into the pipeline's - `sampling_params.extra_args`. -- For `v2*` / `tv2*` tasks, the video as a `video_url` content item (data URI for local files) - -## curl - -```bash -curl -sS -X POST http://localhost:8099/v1/chat/completions \ - -H 'Content-Type: application/json' \ - -d '{ - "model": "zhangj1an/AudioX", - "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}], - "extra_body": { - "num_inference_steps": 250, - "guidance_scale": 7.0, - "seed": 42, - "audiox_task": "t2m", - "seconds_total": 10.0, - "sigma_min": 0.3, - "sigma_max": 500.0 - } - }' > t2m.json -``` diff --git a/examples/online_serving/audiox/openai_chat_client.py b/examples/online_serving/audiox/openai_chat_client.py deleted file mode 100755 index fd9aa8cc5..000000000 --- a/examples/online_serving/audiox/openai_chat_client.py +++ /dev/null @@ -1,119 +0,0 @@ -#!/usr/bin/env python3 -"""AudioX OpenAI-compatible chat client. - -AudioX supports 6 tasks (t2a, t2m, v2a, v2m, tv2a, tv2m). Text-only tasks send the prompt as the -chat message; video-conditioned tasks additionally attach the video via a ``video_url`` content -item (data URI for local files). Task + generation knobs (steps, cfg, sigma range, seconds, seed) -are sent as flattened keys under ``extra_body``; the server routes the declared knobs into -``sampling_params.extra_args`` (see ``vllm_omni/model_extras/audiox.py``). - -Usage: - python openai_chat_client.py --task t2a --prompt "Fireworks burst twice..." --output t2a.wav - python openai_chat_client.py --task tv2a --prompt "drum beating" --video clip.mp4 -o tv2a.wav -""" - -from __future__ import annotations - -import argparse -import base64 -import io -import mimetypes -import sys -from pathlib import Path - -import requests -import soundfile -import torch - -VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"}) -TEXT_TASKS = frozenset({"t2a", "t2m", "tv2a", "tv2m"}) - - -def _to_data_url(path: str) -> str: - mime, _ = mimetypes.guess_type(path) - mime = mime or "video/mp4" - with open(path, "rb") as f: - data = base64.b64encode(f.read()).decode("ascii") - return f"data:{mime};base64,{data}" - - -def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None: - audio = audio.to(torch.float32) - audio = audio / audio.abs().max().clamp(min=1e-8) - path.parent.mkdir(parents=True, exist_ok=True) - # soundfile expects channels-last (T, C); project convention is (C, T). - 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 not (isinstance(audio_obj, dict) and audio_obj.get("data")): - continue - 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: - p = argparse.ArgumentParser(description="AudioX OpenAI chat client") - p.add_argument("--task", required=True, choices=["t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m"]) - p.add_argument("--prompt", "-p", default="", help="Text prompt (required for t2*/tv2*).") - p.add_argument("--video", help="Video path or URL (required for v2*/tv2*).") - p.add_argument("--output", "-o", default="audiox_out.wav") - p.add_argument("--server", "-s", default="http://localhost:8099") - p.add_argument("--model", default="zhangj1an/AudioX") - p.add_argument("--steps", type=int, default=250) - p.add_argument("--guidance-scale", type=float, default=7.0) - p.add_argument("--seed", type=int, default=42) - p.add_argument("--seconds-total", type=float, default=10.0) - p.add_argument("--seconds-start", type=float, default=0.0) - p.add_argument("--sigma-min", type=float, default=0.03) - p.add_argument("--sigma-max", type=float, default=1000.0) - args = p.parse_args() - - if args.task in VIDEO_TASKS and not args.video: - print(f"ERROR: task {args.task!r} requires --video", file=sys.stderr) - return 2 - if args.task in TEXT_TASKS and not args.prompt.strip() and args.task not in {"v2a", "v2m"}: - print(f"ERROR: task {args.task!r} requires --prompt", file=sys.stderr) - return 2 - - content: list[dict] = [{"type": "text", "text": args.prompt}] - if args.task in VIDEO_TASKS: - vurl = args.video if args.video.startswith(("http://", "https://")) else _to_data_url(args.video) - content.append({"type": "video_url", "video_url": {"url": vurl}}) - - payload = { - "model": args.model, - "messages": [{"role": "user", "content": content}], - "extra_body": { - "num_inference_steps": args.steps, - "guidance_scale": args.guidance_scale, - "seed": args.seed, - "audiox_task": args.task, - "seconds_start": args.seconds_start, - "seconds_total": args.seconds_total, - "sigma_min": args.sigma_min, - "sigma_max": args.sigma_max, - }, - } - - print(f"POST {args.server}/v1/chat/completions task={args.task} steps={args.steps}") - r = requests.post( - f"{args.server}/v1/chat/completions", - headers={"Content-Type": "application/json"}, - json=payload, - timeout=600, - ) - r.raise_for_status() - audio, sr = _decode_audio_from_response(r.json()) - _save_wav(audio, Path(args.output), sr) - dur = audio.shape[-1] / sr - print(f"saved {args.output} sr={sr}Hz duration={dur:.2f}s channels={audio.shape[0]}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/examples/online_serving/audiox/run_server.sh b/examples/online_serving/audiox/run_server.sh deleted file mode 100755 index d1a259bfd..000000000 --- a/examples/online_serving/audiox/run_server.sh +++ /dev/null @@ -1,16 +0,0 @@ -#!/bin/bash -# AudioX online serving startup script. - -MODEL="${MODEL:-zhangj1an/AudioX}" -PORT="${PORT:-8099}" -DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-FLASH_ATTN}" - -echo "Starting AudioX server..." -echo "Model: $MODEL" -echo "Port: $PORT" -echo "Diffusion attention backend: $DIFFUSION_ATTENTION_BACKEND" - -DIFFUSION_ATTENTION_BACKEND="$DIFFUSION_ATTENTION_BACKEND" \ - vllm serve "$MODEL" --omni \ - --model-class-name AudioXPipeline \ - --port "$PORT" diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md deleted file mode 100644 index 6dc5f12e9..000000000 --- a/recipes/MammothModa2/MammothModa2.md +++ /dev/null @@ -1,212 +0,0 @@ -# MammothModa2 - -> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation - -## Summary - -- Vendor: ByteDance Research -- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev` -- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding -- Mode: Offline inference -- Maintainer: Community - -## When to use this recipe - -Use this recipe to run MammothModa2-Preview through shared task-oriented -offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script. -The generic example formats the AR prompt, drives the AR → DiT stage pipeline, -and forwards MammothModa2-specific generation parameters through the -pipeline-declared `extra_body` contract. - -MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs -interface (not `OmniDiffusionRequest`), so its generation knobs -(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via -`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale` -flags. Image size uses the standard `--height` / `--width` flags. - -## References - -- Upstream model: - [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview) -- Dev model: - [`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev) -- Related offline example: - [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) -- Related T2T/I2T example: - [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py) -- Declared parameters: - [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py) -- Deploy config: - [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml) - -## Hardware Support - -The default deploy config runs both the AR and DiT stages on a single GPU -(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`, -stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU -after rebalancing the split so the AR weights (~23 GB) leave room for the KV -cache — see the note under *1x L40S 48GB*. - -## GPU - -### 1x L40S 48GB - -> **48 GB config adjustment:** the committed -> `vllm_omni/deploy/mammoth_moda2.yaml` uses -> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S, -> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16` -> before running. (On an ~80 GB GPU, leave the defaults unchanged.) - -#### Environment - -- OS: Linux -- Python: Match the repository requirements for your checkout -- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an - ~80 GB GPU for the default config -- vLLM version: Match the repository requirements for your checkout -- vLLM-Omni version or commit: Use the commit you are deploying from - -#### Offline Commands - -Download the model: - -```bash -hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview -``` - -Run text-to-image with the shared offline example from the repository root. The -deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the -MammothModa2 generation parameters as a JSON object through `--extra-body`: - -```bash -python examples/offline_inference/text_to_image/text_to_image.py \ - --model ./MammothModa2-Preview \ - --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \ - --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ - --height 1024 \ - --width 1024 \ - --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \ - --output mammoth_t2i.png -``` - -The `--extra-body` JSON forwards MammothModa2-specific parameters into -`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's -declared `extra_body_params` (see -[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)), -so unknown keys for MammothModa2 are silently dropped: - -- `text_guidance_scale` — classifier-free guidance scale for the DiT stage - (default `9.0`; CFG is active only when `> 1.0`). -- `cfg_range` — relative step range `[start, end]` over which CFG is applied - (default `[0.0, 1.0]`). -- `num_inference_steps` — number of DiT denoising steps (default `50`). - -`--height` and `--width` must be multiples of 16. - -Run text-to-text through the shared understanding example. It recognizes the -MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`: - -```bash -python examples/offline_inference/x_to_text/x_to_text.py \ - --model ./MammothModa2-Preview \ - --prompt "Explain multimodal generation in three sentences." -``` - -Add an image for image-to-text or image summarization. The shared example -uses MammothModa2's chat and vision-token template: - -```bash -python examples/offline_inference/x_to_text/x_to_text.py \ - --model ./MammothModa2-Preview \ - --image /path/to/input.jpg \ - --prompt "Please summarize the content of this image." -``` - -#### Verification - -The example writes the generated image to the `--output` path. Confirm the file -exists and is a valid image: - -```bash -ls -lh mammoth_t2i.png -python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)" -``` - -### 1x AMD MI300X, MammothModa2 Preview - -#### Environment - -- OS: Linux 6.8.0-134-generic, x86_64 -- Container: official ROCm image built from `docker/Dockerfile.rocm` -- Python: 3.12.13 -- PyTorch: 2.11.0+gitd0c8b1f -- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211 -- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM -- vLLM version: 0.27.0+rocm723 -- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0` -- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm` - -#### Offline Commands - -The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT: - -```bash -python3 examples/offline_inference/text_to_image/text_to_image.py \ - --model bytedance-research/MammothModa2-Preview \ - --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \ - --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ - --height 1024 \ - --width 1024 \ - --seed 42 \ - --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \ - --enable-diffusion-pipeline-profiler \ - --log-stats \ - --output mammoth_t2i.png -``` - -#### Verification - -The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting. - -The output was a valid 1024 by 1024 RGB PNG. - -## MammothModa2-Dev unified inference - -MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses -Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested -`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are -required. - -Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the -Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head, -then sends the generated visual tokens and hidden states to the DiT stage. - -Download the checkpoint: - -```bash -hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev -``` - -Run text-to-text through the shared understanding example. It recognizes the -Dev checkpoint as MammothModa2 and automatically selects -`mammoth_moda2_ar.yaml`: - -```bash -python examples/offline_inference/x_to_text/x_to_text.py \ - --model ./MammothModa2-Dev \ - --prompt "Explain multimodal generation in three sentences." -``` - -Add an image for image-to-text or image summarization: - -```bash -python examples/offline_inference/x_to_text/x_to_text.py \ - --model ./MammothModa2-Dev \ - --image ./image.png \ - --prompt "Please summarize the content of this image." -``` - -The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only -run, loaded model weights used approximately 16.97 GiB of GPU memory before KV -and encoder caches. Allow additional GPU memory for those caches and the input -image. diff --git a/recipes/README.md b/recipes/README.md index 1a384d08e..1bf35745a 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -39,7 +39,6 @@ recipes/ | Recipe | Task | Hardware | |--------|------|----------| -| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB | | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB | | [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs | | [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ | @@ -58,7 +57,6 @@ recipes/ | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU | | [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified | -| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom | | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) | diff --git a/recipes/audiox/AudioX.md b/recipes/audiox/AudioX.md deleted file mode 100644 index a2e260736..000000000 --- a/recipes/audiox/AudioX.md +++ /dev/null @@ -1,130 +0,0 @@ -# AudioX - -> AudioX MMDiT for unified audio + music generation: t2a / t2m / v2a / v2m / tv2a / tv2m - -## Summary - -- Vendor: HKUSTAudio (project), `zhangj1an/AudioX` weight bundle -- Model: `zhangj1an/AudioX` -- Task: Text/video → audio or music. Six tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. -- Mode: Offline inference + online serving (pure diffusion) -- Maintainer: Community - -## When to use this recipe - -Use this recipe to run AudioX for sound-effect (`*2a`) or music (`*2m`) generation -from a text prompt and/or video clip. AudioX is a unified diffusion transformer -that produces stereo 44.1 kHz audio up to ~10 s per call. - -## References - -- Project page: -- vLLM-Omni weight bundle: -- Pipeline: `vllm_omni.diffusion.models.audiox.pipeline_audiox.AudioXPipeline` -- Input transforms: `vllm_omni.transformers_utils.processors.audiox` -- Param contract: `vllm_omni/model_extras/audiox.py` (declared `extra_body` knobs) -- Offline example: [`examples/offline_inference/text_to_audio/text_to_audio.py`](../../examples/offline_inference/text_to_audio/text_to_audio.py) -- Online: standard OpenAI chat-completions endpoint (see commands below) - -## Hardware Support - -## GPU - -### 1x L4 24GB - -#### Environment - -- OS: Ubuntu 22.04 -- Python: 3.10+ -- Driver / runtime: CUDA 12.4 -- vLLM version: 0.20.0 -- vLLM-Omni version: 0.1.x - -#### Command - -AudioX uses the **standard** `text_to_audio` example offline and the standard -OpenAI chat-completions endpoint online. All model-specific knobs -(`audiox_task`, `seconds_start`, `seconds_total`, `sigma_min`, `sigma_max`, -`cfg_rescale`, `video_path`, `audio_path`) are declared in -`vllm_omni/model_extras/audiox.py` and routed through `extra_body` (online) / -`--extra-body` (offline). The offline example also exposes `--task`, `--video`, -`--audio-start`, and `--audio-length` shortcuts for the most common knobs. - -Offline — text tasks (`t2a` / `t2m`): - -```bash -huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights -python examples/offline_inference/text_to_audio/text_to_audio.py \ - --model ./audiox_weights --task t2a \ - --prompt "Fireworks burst twice, followed by a clock ticking." \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ - --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \ - --output t2a.wav -``` - -Offline — video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) need `--video`: - -```bash -python examples/offline_inference/text_to_audio/text_to_audio.py \ - --model ./audiox_weights --task tv2a \ - --prompt "drum beating sound and human talking" \ - --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ - --output tv2a.wav -``` - -Online: - -```bash -DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN \ - vllm serve zhangj1an/AudioX --omni --model-class-name AudioXPipeline --port 8099 - -curl -sS -X POST http://localhost:8099/v1/chat/completions \ - -H 'Content-Type: application/json' \ - -d '{ - "model": "zhangj1an/AudioX", - "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}], - "extra_body": { - "num_inference_steps": 250, - "guidance_scale": 7.0, - "seed": 42, - "audiox_task": "t2m", - "seconds_total": 10.0, - "sigma_min": 0.3, - "sigma_max": 500.0 - } - }' > t2m.json -``` - -For `v2*` / `tv2*` online, attach the video as a `video_url` content item -(data URI for local files) in the message `content`. - -#### Verification - -```bash -# Health check -curl http://localhost:8099/health - -# Listen to the saved file (stereo, 44.1 kHz, sigma_min=0.03, sigma_max=1000 — upstream defaults) -ffprobe t2a.wav -``` - -#### Notes - -- Memory usage: ~10 GB peak with `num_inference_steps=250`, 10 s of audio. -- Output rate: 44.1 kHz stereo regardless of requested rate. -- Supported tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. Pass via the - `audiox_task` knob (in `extra_body` online, or `--task` offline). -- Video conditioning: `v2*` and `tv2*` require a video; offline via `--video` - (→ `video_path` knob), online via a `video_url` content item. -- Cache acceleration is **not** supported (AudioXPipeline is in `_NO_CACHE_ACCELERATION`). -- Tensor parallelism is supported via `--tensor-parallel-size` (DiT QKV is sharded with - `QKVParallelLinear`); cross-attention K/V is also TP-sharded. - -### Known limitations - -- Inference uses an inlined DPM-Solver++(3M) SDE sampler (k-diffusion port). Replacing it with - diffusers' `EDMDPMSolverMultistepScheduler` introduces a fixed ~861 Hz resonance and is not - recommended. -- Generation is fixed at 10 s (configured by the bundle's `sample_size`); longer outputs require - a different bundle. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 708018dd0..ac5e5aea9 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1267,8 +1267,6 @@ stages: @pytest.mark.parametrize( ("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"), [ - ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"), - ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"), ("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"), ("mimo_audio.yaml", "mimo_audio", 2, "audio", None), ("step_audio_2.yaml", "step_audio_2", 2, "audio", None), diff --git a/tests/diffusion/test_diffusion_output_metadata.py b/tests/diffusion/test_diffusion_output_metadata.py index ee3559f47..b1228cebd 100644 --- a/tests/diffusion/test_diffusion_output_metadata.py +++ b/tests/diffusion/test_diffusion_output_metadata.py @@ -25,7 +25,6 @@ def test_validate_accepts_canonical_metadata_fields() -> None: }, "audio": { "sample_rate": 48000, - "audiox_task": "tts", }, "actions": { "raw_action_dim": 7, @@ -55,7 +54,6 @@ def test_validate_accepts_canonical_metadata_fields() -> None: ({"video": {"video_fps_multiplier": 0}}, ValueError), ({"video": {"shape": [1, 0, 64]}}, ValueError), ({"audio": {"sample_rate": 0}}, ValueError), - ({"audio": {"audiox_task": ""}}, ValueError), ({"actions": {"raw_action_dim": 0}}, ValueError), ({"actions": {"action_mode": ""}}, ValueError), ({"actions": {"domain_id": "7"}}, TypeError), diff --git a/tests/e2e/offline_inference/test_audiox_model_expansion.py b/tests/e2e/offline_inference/test_audiox_model_expansion.py deleted file mode 100644 index 5739cc463..000000000 --- a/tests/e2e/offline_inference/test_audiox_model_expansion.py +++ /dev/null @@ -1,71 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import os - -import numpy as np -import pytest -import torch - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner -from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform - -# Prefer a tiny/random checkpoint for CI. -# Override in CI if needed: AUDIOX_TEST_MODEL= -models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")] - -# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize -_OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models] - -pytestmark = [ - pytest.mark.slow, - pytest.mark.diffusion, - pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True, ids=models), -] - - -@hardware_test(res={"cuda": "L4", "xpu": "B60"}) -def test_audiox_model(omni_runner: OmniRunner) -> None: - # Keep runtime short for CI. - seconds_total = 2.0 - # AudioXPipeline always emits 44.1 kHz stereo (advertised via class-level - # ``audio_sample_rate``); the trimmed output should match this rate. - sample_rate = 44100 - - outputs = omni_runner.omni.generate( - prompts={"prompt": "A dog barking in a quiet park."}, - sampling_params_list=OmniDiffusionSamplingParams( - num_inference_steps=4, - guidance_scale=6.0, - generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), - num_outputs_per_prompt=1, - extra_args={ - "audiox_task": "t2a", - "seconds_start": 0.0, - "seconds_total": seconds_total, - }, - ), - ) - - assert outputs is not None - first_output = outputs[0] - assert first_output.final_output_type == "audio" - assert isinstance(first_output, OmniRequestOutput) and first_output - - req_out = first_output - assert isinstance(req_out, OmniRequestOutput) - assert req_out.final_output_type == "audio" - assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output - - audio = req_out.multimodal_output.get("audio") - assert isinstance(audio, np.ndarray) - # audio shape: (batch, channels, samples) - assert audio.ndim == 3 - assert audio.shape[0] == 1 - assert audio.shape[1] == 2 - assert audio.shape[2] > 0 - expected_samples = int(seconds_total * sample_rate) - assert abs(audio.shape[2] - expected_samples) <= 2 * 1024 diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py deleted file mode 100644 index f570a6ea7..000000000 --- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py +++ /dev/null @@ -1,204 +0,0 @@ -""" -End-to-end test for MammothModa2 text-to-image generation. - -Verifies that the AR->DiT pipeline produces an image tensor whose pixel values -match a golden reference. - -Model Hub repo id: ``bytedance-research/MammothModa2-Preview``. -Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml`` - -Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json`` - Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py`` -""" - -from __future__ import annotations - -import json -import os -from pathlib import Path - -import pytest -import torch -from huggingface_hub import snapshot_download -from vllm.sampling_params import SamplingParams - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import get_deploy_config_path - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- -_IMAGE_TOKEN_ID = 151655 # "<|image_pad|>" -_VIDEO_TOKEN_ID = 151656 # "<|video_pad|>" -_VISION_START_TOKEN_ID = 151652 # "<|vision_start|>" -_VISION_END_TOKEN_ID = 151653 # "<|vision_end|>" -_AR_PATCH_SIZE = 16 - -MODEL_PATH = "bytedance-research/MammothModa2-Preview" -T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml") - -_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG) - -# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate. -_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json" -# Fixed sampling coordinates: (channel, row_fraction, col_fraction) -# Covers corners, centre, and mid-edges across all 3 channels. -_PIXEL_SAMPLE_COORDS = [ - (0, 0.0, 0.0), - (0, 0.5, 0.5), - (0, 1.0, 1.0), - (0, 0.25, 0.75), - (1, 0.0, 1.0), - (1, 0.5, 0.0), - (1, 0.75, 0.25), - (1, 1.0, 0.5), - (2, 0.0, 0.5), - (2, 0.5, 1.0), - (2, 0.75, 0.75), - (2, 1.0, 0.0), -] - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- -def _load_t2i_gen_config(repo_id: str) -> dict: - weights_dir = Path(snapshot_download(repo_id)) - cfg_path = weights_dir / "t2i_generation_config.json" - if not cfg_path.exists(): - pytest.skip(f"t2i_generation_config.json not found at {cfg_path}") - with cfg_path.open() as f: - return json.load(f) - - -def _format_t2i_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str: - return ( - "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n" - f"<|im_start|>user\n{user_prompt}<|im_end|>\n" - "<|im_start|>assistant\n" - f"<|image start|>{ar_width}*{ar_height}<|image token|>" - ) - - -def _sample_pixels(img_tensor: torch.Tensor) -> list[float]: - """Sample pixel values at fixed fractional coordinates from a (C, H, W) tensor.""" - t = img_tensor.float().clamp(0.0, 1.0) - if t.ndim == 4: - t = t[0] # unbatch - C, H, W = t.shape - values = [] - for c, rh, rw in _PIXEL_SAMPLE_COORDS: - ri = min(int(rh * (H - 1)), H - 1) - ci = min(int(rw * (W - 1)), W - 1) - values.append(round(float(t[c, ri, ci]), 6)) - return values - - -# --------------------------------------------------------------------------- -# End-to-end test -# --------------------------------------------------------------------------- -pytestmark = [ - pytest.mark.slow, - pytest.mark.diffusion, - pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), -] - - -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3201") -@hardware_test(res={"cuda": "H100"}) -def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner): - """ - End-to-end text-to-image generation with MammothModa2 (AR -> DiT). - - Verifies: - - Omni pipeline initialises with the two-stage YAML config. - - DiT stage outputs an image tensor with the correct shape. - - A fixed set of pixel values matches a golden reference - (regenerate with ``UPDATE_GOLDEN=1``). - """ - gen_cfg = _load_t2i_gen_config(MODEL_PATH) - eol_token_id = int(gen_cfg["eol_token_id"]) - visual_start = int(gen_cfg["visual_token_start_id"]) - visual_end = int(gen_cfg["visual_token_end_id"]) - - height, width = 256, 256 # small for CI speed - ar_height, ar_width = height // _AR_PATCH_SIZE, width // _AR_PATCH_SIZE - expected_grid_tokens = ar_height * (ar_width + 1) - - prompt_text = "A cat sitting on a laptop keyboard" - formatted_prompt = _format_t2i_prompt(prompt_text, ar_width, ar_height) - - omni = omni_runner.omni - ar_sampling = SamplingParams( - temperature=0.0, - top_k=1, - max_tokens=max(1, expected_grid_tokens + 1), - detokenize=False, - ) - dit_sampling = SamplingParams(temperature=0.0, max_tokens=1, detokenize=False) - - outputs = list( - omni.generate( - [ - { - "prompt": formatted_prompt, - "additional_information": { - "omni_task": ["t2i"], - "ar_width": [ar_width], - "ar_height": [ar_height], - "eol_token_id": [eol_token_id], - "visual_token_start_id": [visual_start], - "visual_token_end_id": [visual_end], - "image_height": [height], - "image_width": [width], - "num_inference_steps": [2], - "text_guidance_scale": [1.0], - "cfg_range": [0.0, 1.0], - "visual_ids": [ - _IMAGE_TOKEN_ID, - _VIDEO_TOKEN_ID, - _VISION_START_TOKEN_ID, - _VISION_END_TOKEN_ID, - ], - }, - } - ], - [ar_sampling, dit_sampling], - ) - ) - - assert len(outputs) > 0, "Pipeline produced no outputs" - - found_image = False - for out in outputs: - ro_list = out - if not isinstance(ro_list, list): - ro_list = [ro_list] - for ro in ro_list: - completion_outputs = getattr(ro, "outputs", None) - if not isinstance(completion_outputs, list): - continue - for completion in completion_outputs: - mm = getattr(completion, "multimodal_output", None) - if not (isinstance(mm, dict) and "image" in mm): - continue - img_list = mm["image"] if isinstance(mm["image"], list) else [mm["image"]] - for img_tensor in img_list: - assert isinstance(img_tensor, torch.Tensor), f"Expected image tensor, got {type(img_tensor)}" - assert img_tensor.ndim in (3, 4), f"Expected 3D or 4D image tensor, got {img_tensor.ndim}D" - - sampled = _sample_pixels(img_tensor) - - if os.environ.get("UPDATE_GOLDEN"): - _GOLDEN_T2I_PATH.parent.mkdir(parents=True, exist_ok=True) - _GOLDEN_T2I_PATH.write_text(json.dumps({"pixels": sampled}, indent=2)) - print(f"\nGolden file written to {_GOLDEN_T2I_PATH}") - elif _GOLDEN_T2I_PATH.exists(): - golden = json.loads(_GOLDEN_T2I_PATH.read_text())["pixels"] - for i, (got, exp) in enumerate(zip(sampled, golden)): - assert abs(got - exp) < 1e-4, f"Pixel {i} mismatch: got {got}, expected {exp}" - - found_image = True - - assert found_image, "No image tensor found in pipeline output" diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py index 00ecd09d9..92f5bbfb3 100644 --- a/tests/e2e/offline_inference/test_stable_audio_expansion.py +++ b/tests/e2e/offline_inference/test_stable_audio_expansion.py @@ -84,8 +84,7 @@ def test_stable_audio_quantization_and_teacache() -> None: # ``model_class_name`` must be passed explicitly: the default-stage-cfg # factory in ``async_omni_engine.py`` reads it out of ``kwargs`` when # deciding ``final_output_type`` (#2077), and at construction time the - # auto-resolution from ``model_index.json`` has not run yet. AudioX's - # offline test follows the same pattern. + # auto-resolution from ``model_index.json`` has not run yet. m = Omni( model="stabilityai/stable-audio-open-1.0", model_class_name="StableAudioPipeline", diff --git a/tests/e2e/online_serving/test_audiox_expansion.py b/tests/e2e/online_serving/test_audiox_expansion.py deleted file mode 100644 index 612f397e0..000000000 --- a/tests/e2e/online_serving/test_audiox_expansion.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -E2E online serving test for AudioX text-to-audio diffusion. - -Mirrors `tests/e2e/online_serving/test_sd3_expansion.py` for image diffusion: spin up -`vllm-omni` with `--model-class-name AudioXPipeline` and validate via the standard -`send_diffusion_request` helper (now audio-aware). -""" - -import os - -import pytest - -from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data - -# Tiny / random checkpoint usable in CI; override to a real bundle locally. -AUDIOX_TEST_MODEL = os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random") -T2A_PROMPT = "A quiet living room with soft fabric rustle and gentle cat breathing." - -SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}) - - -def _audiox_server_cases(model: str): - return [ - pytest.param( - OmniServerParams( - model=model, - server_args=["--model-class-name", "AudioXPipeline"], - ), - id="t2a", - marks=SINGLE_CARD_FEATURE_MARKS, - ), - ] - - -@pytest.mark.slow -@pytest.mark.diffusion -@pytest.mark.parametrize("omni_server", _audiox_server_cases(AUDIOX_TEST_MODEL), indirect=True) -def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - """AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`. - - AudioX is registered in ``vllm_omni/model_extras`` (``AudioXPipeline``), so the - ``audiox_task`` / ``seconds_*`` / ``sigma_*`` keys below flow through the standard - ``apply_declared_extra_args`` server path into ``sampling_params.extra_args`` rather - than the old top-level ``extra_args`` escape hatch. - """ - request_config = { - "model": omni_server.model, - "messages": dummy_messages_from_mix_data(content_text=T2A_PROMPT), - "extra_body": { - "num_inference_steps": 4, - "guidance_scale": 6.0, - "seed": 42, - "audiox_task": "t2a", - "seconds_start": 0.0, - "seconds_total": 2.0, - "sigma_min": 0.03, - "sigma_max": 1000.0, - }, - } - openai_client.send_diffusion_request(request_config) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index d61dcc25f..0554a8fe7 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -1170,7 +1170,7 @@ class OpenAIClientHandler: b64_data = image_url.split(",", 1)[1] images.append(decode_b64_image(b64_data)) - # OpenAI audio responses (e.g. AudioX text-to-audio) populate `message.audio`. + # OpenAI audio responses populate `message.audio`. audio_obj = getattr(choice.message, "audio", None) audio_b64 = getattr(audio_obj, "data", None) if audio_obj is not None else None if audio_b64: diff --git a/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py b/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py deleted file mode 100644 index 05d5dbdde..000000000 --- a/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py +++ /dev/null @@ -1,119 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for MammothModa2 config classes. - -These tests verify the correct behavior of the Mammothmoda2Config class, -particularly the get_text_config() method which is called during parent -class initialization (PretrainedConfig.__init__ -> validation -> get_text_config). - -Regression test for: AttributeError: 'Mammothmoda2Config' object has no attribute 'llm_config' -""" - -import pytest - -from vllm_omni.transformers_utils.configs.mammoth_moda2 import ( - Mammothmoda2Config, - Mammothmoda2Qwen2_5_VLConfig, - Mammothmoda2Qwen2_5_VLTextConfig, - Mammothmoda2Qwen3VLConfig, - Mammothmoda2Qwen3VLTextConfig, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class TestMammothmoda2Config: - """Tests for Mammothmoda2Config class.""" - - def test_get_text_config_returns_nested_text_config(self): - """ - Verify that get_text_config() returns the nested text_config from llm_config. - - This is a regression test for the AttributeError that occurred when - get_text_config() was called during parent class initialization before - llm_config was set. - """ - config = Mammothmoda2Config( - llm_config={ - "model_type": "mammothmoda2_qwen2_5_vl", - } - ) - - assert config.llm_config is not None - assert isinstance(config.llm_config, Mammothmoda2Qwen2_5_VLConfig) - - text_config = config.get_text_config() - assert text_config is not None - assert isinstance(text_config, Mammothmoda2Qwen2_5_VLTextConfig) - - # The text_config should be the same object as llm_config.text_config - assert text_config is config.llm_config.text_config - - def test_get_text_config_returns_none_when_llm_config_is_none(self): - """Verify that get_text_config() returns None when llm_config is None.""" - config = Mammothmoda2Config(llm_config=None) - - assert config.llm_config is None - assert config.get_text_config() is None - - def test_init_sets_llm_config_before_parent_init(self): - """ - Verify that llm_config is set before parent __init__ is called. - - This is critical because the parent PretrainedConfig.__init__ calls - validation which in turn calls get_text_config(). If llm_config is not - set before super().__init__(), an AttributeError will be raised. - """ - # This should not raise an AttributeError - config = Mammothmoda2Config( - llm_config={ - "model_type": "mammothmoda2_qwen2_5_vl", - } - ) - assert config.llm_config is not None - - -class TestMammothmoda2Qwen2_5_VLConfig: - """Tests for Mammothmoda2Qwen2_5_VLConfig class.""" - - def test_has_text_config_attribute(self): - """Verify that Mammothmoda2Qwen2_5_VLConfig has text_config attribute.""" - config = Mammothmoda2Qwen2_5_VLConfig() - assert hasattr(config, "text_config") - assert isinstance(config.text_config, Mammothmoda2Qwen2_5_VLTextConfig) - - def test_text_config_has_expected_attributes(self): - """Verify that the nested text_config has expected attributes for token validation.""" - config = Mammothmoda2Qwen2_5_VLConfig() - text_config = config.text_config - - assert hasattr(text_config, "vocab_size") - assert hasattr(text_config, "image_token_id") - assert hasattr(text_config, "video_token_id") - - -class TestMammothmoda2Qwen3VLConfig: - def test_dev_config_uses_qwen3_vl_subconfigs(self): - config = Mammothmoda2Config( - llm_config={ - "model_type": "mammothmoda2_qwen3_vl", - "text_config": { - "model_type": "mammothmoda2_qwen3_vl_text", - "vocab_size": 151936, - "gen_vocab_size": 32800, - "gen_vocab_start_index": 152064, - }, - "vision_config": { - "model_type": "mammothmoda2_qwen3_vl_vision", - "deepstack_visual_indexes": [8, 16, 24], - }, - } - ) - - assert isinstance(config.llm_config, Mammothmoda2Qwen3VLConfig) - assert isinstance(config.get_text_config(), Mammothmoda2Qwen3VLTextConfig) - # vLLM validates sampling against the complete base + visual vocabulary. - assert config.get_text_config().vocab_size == 152064 + 32800 - assert config.get_text_config().gen_vocab_start_index == 152064 - assert config.llm_config.vision_config.deepstack_visual_indexes == [8, 16, 24] diff --git a/tests/model_executor/models/test_kv_cache_scale_mapper.py b/tests/model_executor/models/test_kv_cache_scale_mapper.py index 98f016bb4..218b4bfdf 100644 --- a/tests/model_executor/models/test_kv_cache_scale_mapper.py +++ b/tests/model_executor/models/test_kv_cache_scale_mapper.py @@ -14,7 +14,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] _REPO_ROOT = Path(__file__).resolve().parents[3] _STALE_API_FILES = [ "vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py", - "vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py", "vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py", "vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py", ] @@ -118,10 +117,6 @@ def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale( @pytest.mark.parametrize( ("module_name", "class_name"), [ - ( - "vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2", - "MammothModa2Qwen2ForCausalLM", - ), ( "vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm", "MiMoAudioLLMForConditionalGeneration", diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index b7f2decf9..5f48b6eb4 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -31,7 +31,6 @@ from vllm_omni.model_extras import ( [ ({"model_type": "bagel"}, "bagel"), ({"architectures": ["HunyuanImage3ForConditionalGeneration"]}, "hunyuan_image3"), - ({"model_type": "mammothmoda2"}, "mammoth_moda2"), ({"model_type": "qwen2_vl"}, "generic"), ], ) @@ -51,22 +50,6 @@ def test_bagel_x_to_text_prompt_builder() -> None: assert stop_token_ids is None -@pytest.mark.core_model -@pytest.mark.cpu -def test_mammoth_x_to_text_prompt_builder() -> None: - prompt, stop_token_ids = build_x_to_text_prompt("mammoth_moda2", "unused", "Hello.", has_image=False) - assert prompt == { - "prompt": ( - "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" - "<|im_start|>user\nHello.<|im_end|>\n" - "<|im_start|>assistant\n" - ), - "modalities": ["text"], - "additional_information": {"omni_task": ["chat"]}, - } - assert stop_token_ids is None - - @pytest.mark.core_model @pytest.mark.cpu def test_generic_x_to_text_prompt_builder() -> None: @@ -386,47 +369,6 @@ def test_ming_flash_omni_image_to_image_prompt_builder() -> None: assert "negative_prompt" not in result -@pytest.mark.core_model -@pytest.mark.cpu -def test_audiox_extra_registry_declares_request_and_response_params() -> None: - assert get_extra_body_params("AudioXPipeline") == frozenset( - { - "audiox_task", - "seconds_start", - "seconds_total", - "sigma_min", - "sigma_max", - "cfg_rescale", - "video_path", - "audio_path", - } - ) - assert get_extra_output_params("AudioXPipeline") == frozenset({"audiox_task"}) - assert should_init_extra_args_for_non_diffusion_stages("AudioXPipeline") is False - - -@pytest.mark.core_model -@pytest.mark.cpu -def test_audiox_declared_extra_args_route_into_sampling_params() -> None: - params = OmniDiffusionSamplingParams() - declared = get_extra_body_params("AudioXPipeline") - apply_declared_extra_args( - params, - declared, - { - "audiox_task": "t2a", - "seconds_total": 10.0, - "sigma_min": 0.03, - "unknown": "ignored", - }, - ) - assert params.extra_args == { - "audiox_task": "t2a", - "seconds_total": 10.0, - "sigma_min": 0.03, - } - - @pytest.mark.core_model @pytest.mark.cpu def test_helios_extra_registry_declares_request_and_response_params() -> None: @@ -630,64 +572,3 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None: "cfg_text_scale": 4.0, "think": False, } - - -@pytest.mark.core_model -@pytest.mark.cpu -def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None: - for model_class_name in ( - "MammothModa2DiTPipeline", - "MammothModa2ForConditionalGeneration", - "Mammothmoda2Model", - ): - assert get_extra_body_params(model_class_name) == frozenset( - { - "text_guidance_scale", - "cfg_range", - "num_inference_steps", - } - ) - assert get_extra_output_params(model_class_name) == frozenset() - assert should_init_extra_args_for_non_diffusion_stages(model_class_name) is True - - wrapper_prompt = build_text_to_image_prompt( - "MammothModa2ForConditionalGeneration", - {"prompt": "a cat", "modalities": ["image"]}, - height=256, - width=256, - ) - assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"] - - -@pytest.mark.core_model -@pytest.mark.cpu -def test_mammothmoda2_text_to_image_prompt_builder() -> None: - # Image dims are converted to the AR grid (width/16 x height/16); the negative - # prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path). - assert build_text_to_image_prompt( - "MammothModa2DiTPipeline", - { - "prompt": "a cat", - "modalities": ["image"], - "negative_prompt": "blurry", - }, - height=512, - width=768, - ) == { - "prompt": ( - "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n" - "<|im_start|>user\na cat<|im_end|>\n" - "<|im_start|>assistant\n" - "<|image start|>48*32<|image token|>" - ), - "additional_information": { - "omni_task": ["t2i"], - "ar_width": [48], - "ar_height": [32], - "image_height": [512], - "image_width": [768], - "eol_token_id": [152064], - "visual_token_start_id": [152072], - "visual_token_end_id": [168456], - }, - } diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 650429bf8..63e989668 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -59,7 +59,6 @@ EXCLUDED_MODELS = [ "HeliosPyramidPipeline", "DreamIDOmniPipeline", "SenseNovaU1Pipeline", - "AudioXPipeline", "HunyuanVideo15Pipeline", "HunyuanVideo15ImageToVideoPipeline", "LingBotVideoPipeline", diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index af472fd95..956228db5 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -379,7 +379,6 @@ class TestLoadCustomFuncSelection(unittest.TestCase): def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self): incompatible_paths = [ "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav", - "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit", "vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow", "vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion", ] diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index cecb952eb..47b649e85 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -69,10 +69,6 @@ from vllm_omni.model_executor.models.indextts2.pipeline import ( INDEXTTS25_PIPELINE, ) from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE -from vllm_omni.model_executor.models.mammoth_moda2.pipeline import ( - MAMMOTH_MODA2_AR_PIPELINE, - MAMMOTH_MODA2_PIPELINE, -) from vllm_omni.model_executor.models.mimo_audio.pipeline import MIMO_AUDIO_PIPELINE from vllm_omni.model_executor.models.ming_flash_omni.pipeline import ( MING_FLASH_OMNI_IMAGE_PIPELINE, @@ -168,8 +164,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE, "moss_tts_nano": MOSS_TTS_NANO_PIPELINE, "omnivoice": OMNIVOICE_PIPELINE, - "mammoth_moda2": MAMMOTH_MODA2_PIPELINE, - "mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE, "moss_tts_delay": MOSS_TTS_PIPELINE, "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE, "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE, diff --git a/vllm_omni/deploy/mammoth_moda2.yaml b/vllm_omni/deploy/mammoth_moda2.yaml deleted file mode 100644 index afae38ee0..000000000 --- a/vllm_omni/deploy/mammoth_moda2.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# MammothModa2 AR→DiT (text-to-image). - -async_chunk: false -pipeline: mammoth_moda2 - -stages: - - stage_id: 0 - devices: "0" - max_num_seqs: 100 - max_model_len: 8192 - gpu_memory_utilization: 0.5 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - - - stage_id: 1 - devices: "0" - max_num_seqs: 1 - gpu_memory_utilization: 0.3 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - default_sampling_params: - extra_args: - text_guidance_scale: 9.0 - cfg_range: [0.0, 1.0] - num_inference_steps: 50 diff --git a/vllm_omni/deploy/mammoth_moda2_ar.yaml b/vllm_omni/deploy/mammoth_moda2_ar.yaml deleted file mode 100644 index 45e2580d9..000000000 --- a/vllm_omni/deploy/mammoth_moda2_ar.yaml +++ /dev/null @@ -1,14 +0,0 @@ -# MammothModa2 AR-only (image/video understanding and summarization). - -async_chunk: false -pipeline: mammoth_moda2_ar - -stages: - - stage_id: 0 - devices: "0" - max_num_seqs: 16 - max_model_len: 8192 - gpu_memory_utilization: 0.5 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false diff --git a/vllm_omni/diffusion/models/audiox/__init__.py b/vllm_omni/diffusion/models/audiox/__init__.py deleted file mode 100644 index cec1470f2..000000000 --- a/vllm_omni/diffusion/models/audiox/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from .audiox_transformer import MMDiffusionTransformer -from .pipeline_audiox import AudioXPipeline, get_audiox_post_process_func - -__all__ = [ - "AudioXPipeline", - "MMDiffusionTransformer", - "get_audiox_post_process_func", -] diff --git a/vllm_omni/diffusion/models/audiox/audiox_transformer.py b/vllm_omni/diffusion/models/audiox/audiox_transformer.py deleted file mode 100644 index 9377f2b03..000000000 --- a/vllm_omni/diffusion/models/audiox/audiox_transformer.py +++ /dev/null @@ -1,487 +0,0 @@ -from __future__ import annotations - -import logging -from collections.abc import Iterable -from typing import Any - -import torch -import torch.nn as nn -from einops import rearrange -from torch.nn import functional as F -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_gather, - tensor_model_parallel_all_reduce, -) -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - MergedColumnParallelLinear, - QKVParallelLinear, -) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader, sharded_weight_loader - -from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.layers.fourier import GaussianFourierProjection -from vllm_omni.diffusion.layers.rope import RotaryEmbedding - - -class AudioXCrossAttention(nn.Module): - def __init__(self, dim: int, nheads: int, prefix: str = ""): - super().__init__() - self.dim = dim - self.nheads = nheads - head_dim = dim // nheads - self.head_dim = head_dim - - # to_kv bundle weights arrive in (head, dim, VK-index) interleaved order; pipeline's - # load_weights restacks them to [V|K] before MergedColumnParallelLinear consumes them. - self.to_q = ColumnParallelLinear(dim, dim, bias=False, gather_output=False, prefix=f"{prefix}.to_q") - self.to_kv = MergedColumnParallelLinear( - input_size=dim, output_sizes=[dim, dim], bias=False, gather_output=False, prefix=f"{prefix}.to_kv" - ) - self.q_norm = AudioXRMSNorm(head_dim) - self.k_norm = AudioXRMSNorm(head_dim) - local_nheads = nheads // get_tensor_model_parallel_world_size() - self.attn = Attention( - num_heads=local_nheads, - head_size=head_dim, - softmax_scale=head_dim**-0.5, - causal=False, - ) - - def forward(self, x: torch.Tensor, context: torch.Tensor | None = None) -> torch.Tensor: - tp = get_tensor_model_parallel_world_size() - local_h = self.nheads // tp - d = self.head_dim - - # ROCm's ``wvSplitK`` GEMM bypasses autocast and rejects mismatched dtypes, - # so pre-match the input to the linear's weight dtype. - weight_dtype = self.to_q.weight.dtype - q_flat, _ = self.to_q(x.to(weight_dtype)) - q = rearrange(q_flat, "b n (h d) -> b n h d", h=local_h, d=d) - kv_flat, _ = self.to_kv(context.to(weight_dtype)) - v_flat, k_flat = kv_flat.chunk(2, dim=-1) - # Upstream `CrossAttention.forward` unpacks `pre_attention()` as `q, v, k = ...` (K/V - # swapped), AND only normalizes the first chunk of to_kv via k_norm. Trained weights - # depend on this quirk: first chunk normalized -> V; second chunk unnormalized -> K. - v = rearrange(v_flat, "b n (h d) -> b n h d", h=local_h, d=d) - k = rearrange(k_flat, "b n (h d) -> b n h d", h=local_h, d=d) - q = self.q_norm(q) - v = self.k_norm(v) - - out = self.attn(q.contiguous(), k.contiguous(), v.contiguous(), attn_metadata=None) - out = rearrange(out, "b n h d -> b n (h d)").contiguous() - if tp > 1: - out = tensor_model_parallel_all_gather(out, dim=-1) - return out - - -logger = logging.getLogger(__name__) - -__all__ = [ - "AudioXMMChannelLastConv1d", - "AudioXMMConvFeedForward", - "AudioXMMDiTSelfAttention", - "AudioXMMDiTBlock", - "ContinuousMMDiTTransformer", - "MMDiffusionTransformer", -] - - -class AudioXMMChannelLastConv1d(nn.Conv1d): - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = rearrange(x, "b n c -> b c n") - x = super().forward(x) - x = rearrange(x, "b c n -> b n c") - return x - - -class _ColumnParallelChannelLastConv1d(AudioXMMChannelLastConv1d): - def __init__(self, in_channels: int, out_channels_total: int, **kwargs: Any): - tp_size = get_tensor_model_parallel_world_size() - assert out_channels_total % tp_size == 0, (out_channels_total, tp_size) - super().__init__(in_channels, out_channels_total // tp_size, **kwargs) - self.weight.weight_loader = sharded_weight_loader(0) - - -class _RowParallelChannelLastConv1d(AudioXMMChannelLastConv1d): - def __init__(self, in_channels_total: int, out_channels: int, **kwargs: Any): - tp_size = get_tensor_model_parallel_world_size() - assert in_channels_total % tp_size == 0, (in_channels_total, tp_size) - super().__init__(in_channels_total // tp_size, out_channels, **kwargs) - self._tp_size = tp_size - self.weight.weight_loader = sharded_weight_loader(1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - y = super().forward(x) - if self._tp_size > 1: - y = tensor_model_parallel_all_reduce(y) - return y - - -class AudioXMMConvFeedForward(nn.Module): - def __init__( - self, - dim: int, - hidden_dim: int, - multiple_of: int = 256, - kernel_size: int = 3, - padding: int = 1, - ): - super().__init__() - hidden_dim = int(2 * hidden_dim / 3) - hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) - - self.w1 = _ColumnParallelChannelLastConv1d( - dim, hidden_dim, bias=False, kernel_size=kernel_size, padding=padding - ) - self.w2 = _RowParallelChannelLastConv1d(hidden_dim, dim, bias=False, kernel_size=kernel_size, padding=padding) - self.w3 = _ColumnParallelChannelLastConv1d( - dim, hidden_dim, bias=False, kernel_size=kernel_size, padding=padding - ) - - def forward(self, x): - return self.w2(F.silu(self.w1(x)) * self.w3(x)) - - -class AudioXRMSNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: - super().__init__() - self.eps = eps - - def forward(self, x: torch.Tensor) -> torch.Tensor: - mean_sq = torch.mean(x**2, dim=-1, keepdim=True) - scale = torch.rsqrt(mean_sq + self.eps) - return x * scale - - -class AudioXMMDiTSelfAttention(nn.Module): - def __init__(self, dim: int, nheads: int, prefix: str = ""): - super().__init__() - self.dim = dim - self.nheads = nheads - - head_dim = dim // nheads - self.head_dim = head_dim - # Bundle weights arrive in interleaved (h, d, qkv) layout; pipeline's load_weights - # restacks to [Q|K|V] before QKVParallelLinear's weight_loader consumes them. - self.qkv = QKVParallelLinear( - hidden_size=dim, - head_size=head_dim, - total_num_heads=nheads, - bias=True, - prefix=f"{prefix}.qkv", - ) - self.q_norm = AudioXRMSNorm(head_dim) - self.k_norm = AudioXRMSNorm(head_dim) - - self.rope = RotaryEmbedding(is_neox_style=False) - self.attn = Attention( - num_heads=self.qkv.num_heads, - head_size=head_dim, - softmax_scale=head_dim**-0.5, - causal=False, - ) - - def apply_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - out = self.attn(q.contiguous(), k.contiguous(), v.contiguous(), attn_metadata=None) - out = rearrange(out, "b n h d -> b n (h d)").contiguous() - # Downstream linear1/ffn are replicated and expect full hidden dim. - if get_tensor_model_parallel_world_size() > 1: - out = tensor_model_parallel_all_gather(out, dim=-1) - return out - - def pre_attention(self, x: torch.Tensor, rot: tuple[torch.Tensor, torch.Tensor] | None = None): - # ROCm GEMM dtype cast — see AudioXCrossAttention.forward. - qkv, _ = self.qkv(x.to(self.qkv.weight.dtype)) - local_h = self.qkv.num_heads - d = self.head_dim - q_size = local_h * d - q, k, v = qkv.split([q_size, q_size, q_size], dim=-1) - q = rearrange(q, "b n (h d) -> b n h d", h=local_h, d=d) - k = rearrange(k, "b n (h d) -> b n h d", h=local_h, d=d) - v = rearrange(v, "b n (h d) -> b n h d", h=local_h, d=d) - q = self.q_norm(q) - k = self.k_norm(k) - - if rot is not None: - cos, sin = rot - cos = cos.to(dtype=q.dtype) - sin = sin.to(dtype=q.dtype) - q = self.rope(q, cos, sin) - k = self.rope(k, cos, sin) - - return q, k, v - - def forward( - self, - x: torch.Tensor, - rot: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> torch.Tensor: - q, k, v = self.pre_attention(x, rot=rot) - return self.apply_attention(q, k, v) - - -class AudioXMMDiTBlock(nn.Module): - def __init__( - self, - dim: int, - nhead: int, - mlp_ratio: float = 4.0, - prefix: str = "", - ): - super().__init__() - self.norm1 = nn.LayerNorm(dim, elementwise_affine=False) - self.attn = AudioXMMDiTSelfAttention(dim, nhead, prefix=f"{prefix}.attn") - self.cross_attn = AudioXCrossAttention(dim, nhead, prefix=f"{prefix}.cross_attn") - self.linear1 = AudioXMMChannelLastConv1d(dim, dim, kernel_size=3, padding=1) - self.norm2 = nn.LayerNorm(dim, elementwise_affine=False) - self.ffn = AudioXMMConvFeedForward(dim, int(dim * mlp_ratio), kernel_size=3, padding=1) - self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True)) - - def pre_attention(self, x: torch.Tensor, c: torch.Tensor, rot: tuple[torch.Tensor, torch.Tensor] | None): - modulation = self.adaLN_modulation(c) - (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = modulation.chunk(6, dim=-1) - x = self.norm1(x) * (1 + scale_msa) + shift_msa - q, k, v = self.attn.pre_attention(x, rot) - return (q, k, v), (gate_msa, shift_mlp, scale_mlp, gate_mlp) - - def post_attention(self, x: torch.Tensor, attn_out: torch.Tensor, c: tuple[torch.Tensor, ...], context=None): - (gate_msa, shift_mlp, scale_mlp, gate_mlp) = c - x = x + self.linear1(attn_out) * gate_msa - - x = x + self.cross_attn(x, context=context) - - r = self.norm2(x) * (1 + scale_mlp) + shift_mlp - x = x + self.ffn(r) * gate_mlp - return x - - def forward( - self, - x: torch.Tensor, - cond: torch.Tensor, - rot: tuple[torch.Tensor, torch.Tensor] | None, - context: torch.Tensor = None, - ) -> torch.Tensor: - x_qkv, x_conditions = self.pre_attention(x, cond, rot) - attn_out = self.attn.apply_attention(*x_qkv) - x = self.post_attention(x, attn_out, x_conditions, context=context) - return x - - -class MMDiffusionTransformer(nn.Module): - """AudioX MMDiT, specialized for the published bundle (`zhangj1an/AudioX`). - - The bundle fixes patch_size=1, transformer_type="continuous_transformer", - cond_token_dim=768 (>0, project_cond_tokens=False), and never sets - prepend_cond_dim or input_concat_dim, so those code paths are removed. - """ - - def __init__( - self, - io_channels: int, - embed_dim: int, - cond_token_dim: int, - global_cond_dim: int, - depth: int, - num_heads: int, - project_cond_tokens: bool = False, - project_global_cond: bool = True, - **kwargs, - ): - super().__init__() - if kwargs: - logger.debug("MMDiffusionTransformer ignoring unused config keys: %s", sorted(kwargs.keys())) - if project_cond_tokens: - raise ValueError("AudioX bundle requires project_cond_tokens=False to match official checkpoints.") - - self.cond_token_dim = cond_token_dim - - timestep_features_dim = 256 - self.timestep_features = GaussianFourierProjection( - in_features=1, - embedding_size=timestep_features_dim // 2, - scale=1.0, - trainable=False, - ) - self.to_timestep_embed = nn.Sequential( - nn.Linear(timestep_features_dim, embed_dim, bias=True), - nn.SiLU(), - nn.Linear(embed_dim, embed_dim, bias=True), - ) - - cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim - self.to_cond_embed = nn.Sequential( - nn.Linear(cond_token_dim, cond_embed_dim, bias=False), - nn.SiLU(), - nn.Linear(cond_embed_dim, cond_embed_dim, bias=False), - ) - - # ``to_global_embed`` weights live in the bundle but global conditioning is always None - # at inference; kept so AutoWeightsLoader has a slot to load them into. - global_embed_dim = global_cond_dim if not project_global_cond else embed_dim - self.to_global_embed = nn.Sequential( - nn.Linear(global_cond_dim, global_embed_dim, bias=False), - nn.SiLU(), - nn.Linear(global_embed_dim, global_embed_dim, bias=False), - ) - - self.transformer = ContinuousMMDiTTransformer( - dim=embed_dim, - depth=depth, - dim_heads=embed_dim // num_heads, - dim_in=io_channels, - dim_out=io_channels, - ) - - self.preprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) - nn.init.zeros_(self.preprocess_conv.weight) - self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) - nn.init.zeros_(self.postprocess_conv.weight) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - def _forward(self, x, t, cross_attn_cond): - timestep_embed = self.to_timestep_embed(self.timestep_features(t[:, None])) - prepend_inputs = timestep_embed.unsqueeze(1) - prepend_length = prepend_inputs.shape[1] - - x = self.preprocess_conv(x) + x - x = rearrange(x, "b c n -> b n c") - output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond) - output = rearrange(output, "b n c -> b c n")[:, :, prepend_length:] - return self.postprocess_conv(output) + output - - def forward( - self, - x, - t, - cross_attn_cond, - negative_cross_attn_cond=None, - negative_cross_attn_mask=None, - cfg_scale: float = 1.0, - scale_phi: float = 0.0, - **kwargs, - ): - if cfg_scale == 1.0: - return self._forward(x, t, cross_attn_cond) - - # Classifier-free guidance: batch the conditional + unconditional pass. - null_embed = torch.zeros_like(cross_attn_cond) - if negative_cross_attn_cond is not None and negative_cross_attn_mask is not None: - mask = negative_cross_attn_mask.to(torch.bool).unsqueeze(2) - negative_cross_attn_cond = torch.where(mask, negative_cross_attn_cond, null_embed) - uncond = negative_cross_attn_cond if negative_cross_attn_cond is not None else null_embed - - batch_output = self._forward( - torch.cat([x, x], dim=0), - torch.cat([t, t], dim=0), - torch.cat([cross_attn_cond, uncond], dim=0), - ) - cond_output, uncond_output = torch.chunk(batch_output, 2, dim=0) - cfg_output = uncond_output + (cond_output - uncond_output) * cfg_scale - - if scale_phi == 0.0: - return cfg_output - cond_std = cond_output.std(dim=1, keepdim=True) - cfg_std = cfg_output.std(dim=1, keepdim=True) - return scale_phi * (cfg_output * (cond_std / cfg_std)) + (1 - scale_phi) * cfg_output - - -class ContinuousMMDiTTransformer(nn.Module): - def __init__( - self, - dim, - depth, - *, - dim_in=None, - dim_out=None, - dim_heads=64, - _latent_seq_len=237, - ): - super().__init__() - - self.dim = dim - self.depth = depth - - self.project_in = nn.Linear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity() - self.project_out = nn.Linear(dim, dim_out, bias=False) if dim_out is not None else nn.Identity() - - hidden_dim = dim - num_heads = dim_heads - mlp_ratio = 4.0 - self._latent_seq_len = _latent_seq_len - - self.layers = nn.ModuleList( - [ - AudioXMMDiTBlock( - hidden_dim, - num_heads, - mlp_ratio=mlp_ratio, - prefix=f"layers.{i}", - ) - for i in range(depth) - ] - ) - self.proj_mm_tokens = nn.Linear(768, hidden_dim) if dim != 768 else nn.Identity() - self.proj_mm_seq_len = nn.Linear(384, self._latent_seq_len) if self._latent_seq_len != 384 else nn.Identity() - - # AudioX RoPE: interleaved (GPT-J) pair layout, theta=10000. - head_dim = hidden_dim // num_heads - pos = torch.arange(self._latent_seq_len, dtype=torch.float32, device=self.device) - inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=self.device) / head_dim)) - ang = torch.outer(pos, inv_freq) - self.register_buffer("latent_rope_cos", torch.cos(ang), persistent=False) - self.register_buffer("latent_rope_sin", torch.sin(ang), persistent=False) - - @property - def device(self): - return next(self.parameters()).device - - def forward( - self, - x, - prepend_embeds=None, - context=None, - ): - x = self.project_in(x) - - if prepend_embeds is not None: - prepend_dim = prepend_embeds.shape[-1] - assert prepend_dim == x.shape[-1], "prepend dimension must match sequence dimension" - x = torch.cat((prepend_embeds, x), dim=-2) - - time_cond = prepend_embeds.squeeze(1) - mm_tokens = context - - mm_tokens = self.proj_mm_tokens(mm_tokens) - mm_tokens = rearrange(mm_tokens, "b s d -> b d s") - mm_tokens = self.proj_mm_seq_len(mm_tokens) - mm_tokens = rearrange(mm_tokens, "b d s -> b s d") - - time_cond = time_cond.unsqueeze(1) - rot = ( - self.latent_rope_cos.to(device=x.device, dtype=x.dtype), - self.latent_rope_sin.to(device=x.device, dtype=x.dtype), - ) - for block in self.layers: - x = block(x, mm_tokens, rot, context=time_cond) - - x = self.project_out(x) - return x - - -def __getattr__(name: str): - if name in __all__: - return globals()[name] - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py b/vllm_omni/diffusion/models/audiox/pipeline_audiox.py deleted file mode 100644 index 395e06dc8..000000000 --- a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py +++ /dev/null @@ -1,933 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import json -import math -import os -from collections.abc import Iterable -from typing import Any, ClassVar - -import torch -import torch.nn.functional as F -import torchsde -from diffusers import AutoencoderOobleck -from einops import rearrange -from torch import einsum, nn -from torchvision import transforms -from transformers import AutoConfig, CLIPVisionModelWithProjection, T5EncoderModel, T5TokenizerFast -from vllm.logger import init_logger -from vllm.model_executor.models.utils import AutoWeightsLoader - -from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.models.audiox.audiox_transformer import MMDiffusionTransformer -from vllm_omni.diffusion.models.interface import SupportAudioOutput -from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from vllm_omni.transformers_utils.processors import audiox as _audiox_transforms - -_VIDEO_ONLY_TASKS = _audiox_transforms.VIDEO_ONLY_TASKS -_TEXT_VIDEO_TASKS = _audiox_transforms.TEXT_VIDEO_TASKS -_VIDEO_CONDITIONED_TASKS = _audiox_transforms.VIDEO_CONDITIONED_TASKS -_normalize_prompts = _audiox_transforms.normalize_prompts -prepare_audio_reference = _audiox_transforms.prepare_audio_reference -prepare_video_reference = _audiox_transforms.prepare_video_reference - -# Polyexponential sigma schedule defaults; mirror upstream AudioX gradio interface -# (``audiox/interface/gradio.py`` ``generate_cond`` defaults: ``sigma_min=0.03``, ``sigma_max=1000``). -_DEFAULT_UPSTREAM_SIGMA_MIN = 0.03 -_DEFAULT_UPSTREAM_SIGMA_MAX = 1000.0 - -logger = init_logger(__name__) - - -def _load_audiox_bundle_config(model_root: str) -> dict[str, Any]: - with open(os.path.join(os.path.abspath(model_root), "config.json"), encoding="utf-8") as f: - return json.load(f) - - -def _audio_conditioning_input_samples(model_config: dict[str, Any]) -> int: - configs = model_config["model"]["conditioning"]["configs"] - cfg = next(c["config"] for c in configs if c["id"] == "audio_prompt") - return int(cfg["latent_seq_len"]) * int(cfg["pretransform_config"]["config"]["downsampling_ratio"]) - - -def get_audiox_post_process_func(od_config: OmniDiffusionConfig): - """Convert the pipeline's float audio tensor to a CPU numpy array for serving.""" - - def post_process_func(audio: torch.Tensor) -> Any: - if isinstance(audio, dict) and isinstance(audio.get("payload"), dict): - payload = dict(audio["payload"]) - audio_payload = payload.get("audio") - if isinstance(audio_payload, torch.Tensor): - payload["audio"] = audio_payload.detach().cpu().float().numpy() - return { - "payload": payload, - "metadata": audio.get("metadata") or {}, - } - if isinstance(audio, torch.Tensor): - return audio.detach().cpu().float().numpy() - return audio - - return post_process_func - - -class SA_PreNorm(nn.Module): - def __init__(self, dim, fn): - super().__init__() - self.norm = nn.LayerNorm(dim) - self.fn = fn - - def forward(self, x, **kwargs): - return self.fn(self.norm(x), **kwargs) - - -class SA_FeedForward(nn.Module): - def __init__(self, dim, hidden_dim): - super().__init__() - # Dropout p=0 preserves upstream ``net.{2,4}`` state-dict keys so the upstream weights - # load into the right slots; inference is identical to no dropout. - self.net = nn.Sequential( - nn.Linear(dim, hidden_dim), - nn.GELU(), - nn.Dropout(0.0), - nn.Linear(hidden_dim, dim), - nn.Dropout(0.0), - ) - - def forward(self, x): - return self.net(x) - - -# Manual einsum+softmax only. SDPA/diffusion Attention here degrades conditioning vs upstream. -class SA_Attention(nn.Module): - def __init__(self, dim, heads=8, dim_head=64): - super().__init__() - inner_dim = dim_head * heads - project_out = not (heads == 1 and dim_head == dim) - - self.heads = heads - self.scale = dim_head**-0.5 - - self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False) - - self.to_out = ( - nn.Sequential( - nn.Linear(inner_dim, dim), - nn.Dropout(0.0), - ) - if project_out - else nn.Identity() - ) - - def forward(self, x): - h = self.heads - qkv = self.to_qkv(x).chunk(3, dim=-1) - q, k, v = (rearrange(t, "b n (h d) -> b h n d", h=h) for t in qkv) - - dots = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale - - attn = dots.softmax(dim=-1) - - out = einsum("b h i j, b h j d -> b h i d", attn, v) - out = rearrange(out, "b h n d -> b n (h d)") - out = self.to_out(out) - return out - - -class SA_Transformer(nn.Module): - def __init__(self, dim, depth, heads, dim_head, mlp_dim): - super().__init__() - self.layers = nn.ModuleList([]) - self.norm = nn.LayerNorm(dim) - for _ in range(depth): - self.layers.append( - nn.ModuleList( - [ - SA_PreNorm(dim, SA_Attention(dim, heads=heads, dim_head=dim_head)), - SA_PreNorm(dim, SA_FeedForward(dim, mlp_dim)), - ] - ) - ) - - def forward(self, x): - for attn, ff in self.layers: - x = attn(x) + x - x = ff(x) + x - return self.norm(x) - - -_AUDIOX_OOBLECK_CONFIG = { - "audio_channels": 2, - "channel_multiples": [1, 2, 4, 8, 16], - "decoder_channels": 128, - "decoder_input_channels": 64, - "downsampling_ratios": [2, 4, 4, 8, 8], - "encoder_hidden_size": 128, - "sampling_rate": 44100, -} - - -def _build_audiox_oobleck(scaling_factor: float = 1.0) -> AutoencoderOobleck: - vae = AutoencoderOobleck(**_AUDIOX_OOBLECK_CONFIG) - vae.audiox_scaling_factor = float(scaling_factor) # type: ignore[attr-defined] - return vae.eval().requires_grad_(False) - - -class AudioVaePromptAdapter(nn.Module): - def __init__(self, *, cond_dim: int, latent_seq_len: int = 215): - super().__init__() - self.pretransform = _build_audiox_oobleck() - in_ch = int(self.pretransform.config.decoder_input_channels) - self.proj_features_128 = nn.Linear(latent_seq_len, 128) - self.proj_out = nn.Linear(in_ch, cond_dim) if in_ch != cond_dim else nn.Identity() - - def forward(self, audio: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - z = self.pretransform.encode(audio, return_dict=True).latent_dist.sample() - latents = z / float(self.pretransform.audiox_scaling_factor) - latents = rearrange(self.proj_features_128(latents), "b c s -> b s c") - latents = self.proj_out(latents) - ones = torch.ones(latents.shape[0], latents.shape[2], device=latents.device) - return latents, ones - - -class _MAFCrossAttentionBlock(nn.Module): - def __init__(self, dim: int, num_heads: int): - super().__init__() - assert dim % num_heads == 0, "dim must be divisible by num_heads" - self._num_heads = num_heads - self._head_dim = dim // num_heads - self._scale = self._head_dim**-0.5 - self.to_q = nn.Linear(dim, dim, bias=True) - self.to_kv = nn.Linear(dim, dim * 2, bias=True) - self.to_out = nn.Linear(dim, dim, bias=True) - nn.init.zeros_(self.to_out.weight) - if self.to_out.bias is not None: - nn.init.zeros_(self.to_out.bias) - - def forward(self, experts: torch.Tensor, full_context: torch.Tensor) -> torch.Tensor: - nh, hd = self._num_heads, self._head_dim - q = rearrange(self.to_q(experts), "b n (h d) -> b h n d", h=nh, d=hd) - k, v = (rearrange(t, "b n (h d) -> b h n d", h=nh, d=hd) for t in self.to_kv(full_context).chunk(2, dim=-1)) - out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False, scale=self._scale) - return self.to_out(rearrange(out, "b h n d -> b n (h d)")) - - -class _MAFFusionBlock(nn.Module): - def __init__(self, dim: int, num_heads: int, mlp_ratio: float): - super().__init__() - assert dim % num_heads == 0, "dim must be divisible by num_heads" - head_dim = dim // num_heads - self._num_heads = num_heads - hidden = int(dim * mlp_ratio) - self.norm1 = nn.LayerNorm(dim) - self.to_qkv = nn.Linear(dim, dim * 3, bias=True) - self.self_attn = Attention( - num_heads=num_heads, - head_size=head_dim, - softmax_scale=head_dim**-0.5, - causal=False, - ) - self.out_proj = nn.Linear(dim, dim, bias=True) - self.norm2 = nn.LayerNorm(dim) - self.ff = nn.Sequential( - nn.Linear(dim, hidden), - nn.GELU(), - nn.Linear(hidden, dim), - ) - nn.init.zeros_(self.out_proj.weight) - if self.out_proj.bias is not None: - nn.init.zeros_(self.out_proj.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - h = self.norm1(x) - q, k, v = self.to_qkv(h).chunk(3, dim=-1) - nh = self._num_heads - q_bsn = rearrange(q, "b n (h d) -> b n h d", h=nh).contiguous() - k_bsn = rearrange(k, "b n (h d) -> b n h d", h=nh).contiguous() - v_bsn = rearrange(v, "b n (h d) -> b n h d", h=nh).contiguous() - out = self.self_attn(q_bsn, k_bsn, v_bsn, attn_metadata=None) - out = rearrange(out, "b n h d -> b n (h d)") - out = self.out_proj(out) - x = x + out - x = x + self.ff(self.norm2(x)) - return x - - -class MAF_Block(nn.Module): - DIM = 768 - MLP_RATIO = 4.0 - - def __init__( - self, - *, - dim: int = 768, - num_experts_per_modality: int = 64, - num_heads: int = 24, - num_fusion_layers: int = 8, - mlp_ratio: float = 4.0, - ): - super().__init__() - total_experts = num_experts_per_modality * 3 - - self.gating_network = nn.Sequential( - nn.Linear(dim * 3, dim), - nn.GELU(), - nn.Linear(dim, 3), - nn.Sigmoid(), - ) - - self.unified_experts = nn.Parameter(torch.randn(total_experts, dim)) - - self.cross_block = _MAFCrossAttentionBlock(dim, num_heads) - self.norm1 = nn.LayerNorm(dim) - self.fusion_blocks = nn.ModuleList( - [_MAFFusionBlock(dim, num_heads, mlp_ratio) for _ in range(num_fusion_layers)] - ) - - self.norm_v2 = nn.LayerNorm(dim) - self.norm_t2 = nn.LayerNorm(dim) - self.norm_a2 = nn.LayerNorm(dim) - self.bypass_gate_v = nn.Parameter(torch.tensor(-10.0)) - self.bypass_gate_t = nn.Parameter(torch.tensor(-10.0)) - self.bypass_gate_a = nn.Parameter(torch.tensor(-10.0)) - - def forward( - self, - video_tokens: torch.Tensor, - text_tokens: torch.Tensor, - audio_tokens: torch.Tensor, - ) -> dict[str, torch.Tensor]: - batch_size = video_tokens.shape[0] - - v_global = video_tokens.mean(dim=1) - t_global = text_tokens.mean(dim=1) - a_global = audio_tokens.mean(dim=1) - - all_global = torch.cat([v_global, t_global, a_global], dim=1) - gates = self.gating_network(all_global) - w_v, w_t, w_a = gates.chunk(3, dim=-1) - - gated_v = video_tokens * w_v.unsqueeze(-1) - gated_t = text_tokens * w_t.unsqueeze(-1) - gated_a = audio_tokens * w_a.unsqueeze(-1) - - full_context = torch.cat([gated_v, gated_t, gated_a], dim=1) - - experts = self.unified_experts.unsqueeze(0).expand(batch_size, -1, -1) - cross_out = self.cross_block(experts, full_context) - updated_experts = self.norm1(experts + cross_out) - - for blk in self.fusion_blocks: - updated_experts = blk(updated_experts) - - fused_v_experts, fused_t_experts, fused_a_experts = updated_experts.chunk(3, dim=1) - - refinement_v = fused_v_experts.mean(dim=1) - refinement_t = fused_t_experts.mean(dim=1) - refinement_a = fused_a_experts.mean(dim=1) - - alpha_v = torch.sigmoid(self.bypass_gate_v) - alpha_t = torch.sigmoid(self.bypass_gate_t) - alpha_a = torch.sigmoid(self.bypass_gate_a) - - final_v = video_tokens + alpha_v * self.norm_v2(refinement_v).unsqueeze(1) - final_t = text_tokens + alpha_t * self.norm_t2(refinement_t).unsqueeze(1) - final_a = audio_tokens + alpha_a * self.norm_a2(refinement_a).unsqueeze(1) - - return { - "video": final_v, - "text": final_t, - "audio": final_a, - } - - -class _BrownianTreeNoiseSampler: - """Brownian-tree noise sampler for DPM-Solver++ SDE, ported from k-diffusion. - - Returns a scaled Brownian increment between two sigma levels; the tree is indexed by - transformed ``sigma`` (here linear, same as k-diffusion's default), so re-querying - identical (sigma, sigma_next) pairs returns the same noise. Entropy must come from - the user's seed — otherwise the unseeded global RNG would make sampling non-deterministic. - """ - - def __init__(self, x: torch.Tensor, sigma_min: torch.Tensor, sigma_max: torch.Tensor, entropy: int): - t0 = torch.as_tensor(sigma_min) - t1 = torch.as_tensor(sigma_max) - self._tree = torchsde.BrownianTree(t0, torch.zeros_like(x), t1, entropy=entropy) - - def __call__(self, sigma: torch.Tensor, sigma_next: torch.Tensor) -> torch.Tensor: - # AudioX denoises from sigma_max → 0, so sigma > sigma_next throughout the loop. - t0 = torch.as_tensor(sigma_next) - t1 = torch.as_tensor(sigma) - return -self._tree(t0, t1) / (t1 - t0).sqrt() - - -class AudioXPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMixin): - supports_request_batch = False - support_audio_output: ClassVar[bool] = True - audio_sample_rate: ClassVar[int] = 44100 - audio_channels: ClassVar[int] = 2 - _PROFILER_TARGETS: ClassVar[list[str]] = ["diffuse"] - _CLIP_SYNC_DURATION_SEC: ClassVar[float] = 10.0 - _VIDEO_SYNC_FRAME_COUNT: ClassVar[int] = 240 - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - self.device = get_local_device() - if od_config.model is None: - raise ValueError( - "AudioXPipeline requires od_config.model (directory with unified safetensors; " - "see https://huggingface.co/zhangj1an/AudioX)." - ) - - if os.path.exists(od_config.model): - self._model_root = os.path.abspath(od_config.model) - else: - from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific - - self._model_root = download_weights_from_hf_specific(od_config.model, None, ["*"]) - self._model_config = _load_audiox_bundle_config(self._model_root) - - model_config = self._model_config["model"] - diffusion_config = model_config["diffusion"] - - self.model = MMDiffusionTransformer(**dict(diffusion_config["config"])) - - cond_configs = {c["id"]: c["config"] for c in model_config["conditioning"]["configs"]} - self.audio_vae_adapter = AudioVaePromptAdapter( - cond_dim=int(model_config["conditioning"]["cond_dim"]), - latent_seq_len=int(cond_configs["audio_prompt"]["latent_seq_len"]), - ) - - t5_name = cond_configs["text_prompt"]["t5_model_name"] - self._t5_max_length = int(cond_configs["text_prompt"]["max_length"]) - self.tokenizer = T5TokenizerFast.from_pretrained(t5_name) - t5_config = AutoConfig.from_pretrained(t5_name) - self.text_encoder = T5EncoderModel(t5_config).train(False).requires_grad_(False).to(torch.float16) - - clip_name = cond_configs["video_prompt"]["clip_model_name"] - clip_config = AutoConfig.from_pretrained(clip_name) - self.clip_encoder = CLIPVisionModelWithProjection(clip_config.vision_config) - _CLIP_PATCH_TOKENS, _VIDEO_FPS, _DURATION_SEC, _DIM = 50, 5, 10, 768 - _in_features = _CLIP_PATCH_TOKENS * _VIDEO_FPS * _DURATION_SEC - self._clip_in_features = _in_features - self._clip_out_features = 128 - self.clip_proj = nn.Linear(_in_features, self._clip_out_features) - self.clip_proj_sync = nn.Linear(240, self._clip_out_features) - self.clip_sync_weight = nn.Parameter(torch.tensor(0.0)) - self.clip_temp_transformer = SA_Transformer(_DIM, depth=4, heads=16, dim_head=64, mlp_dim=_DIM * 4) - self.clip_temp_pos_embedding = nn.Parameter(torch.randn(1, _VIDEO_FPS * _DURATION_SEC, _DIM)) - self.clip_empty_visual_feat = nn.Parameter(torch.zeros(1, self._clip_out_features, _DIM), requires_grad=False) - _CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) - _CLIP_STD = (0.26862954, 0.26130258, 0.27577711) - self._clip_normalize = transforms.Compose([transforms.Normalize(mean=list(_CLIP_MEAN), std=list(_CLIP_STD))]) - - self.pretransform = _build_audiox_oobleck( - scaling_factor=float(model_config["pretransform"].get("scale", 1.0)), - ) - - self.io_channels = model_config["io_channels"] - self.diffusion_objective = "v" - - gate_type_config = diffusion_config["gate_type_config"] - self.maf_block = MAF_Block( - dim=768, - num_experts_per_modality=int(gate_type_config["num_experts_per_modality"]), - num_heads=int(gate_type_config["num_heads"]), - num_fusion_layers=int(gate_type_config["num_fusion_layers"]), - ) - - logger.debug("AudioX model built from %s", self._model_root) - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=self._model_root, - subfolder="transformer", - revision=getattr(od_config, "revision", None), - prefix="", - ), - ] - sample_rate = int(self._model_config.get("sample_rate", 48000)) - self._sample_rate = sample_rate - self._sample_size = int(self._model_config.get("sample_size", sample_rate * 10)) - self._target_fps = int(self._model_config.get("video_fps", 5)) - self._audio_conditioning_samples = _audio_conditioning_input_samples(self._model_config) - - self.setup_diffusion_pipeline_profiler( - profiler_targets=list(self._PROFILER_TARGETS), - enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - _legacy_prefix = "conditioner.conditioners.audio_prompt." - - # DiT self-attn QKV: bundle stores weights with last dim (h, d, qkv) interleaved - # (i.e. Q/K/V rows for the same head sit next to each other), whereas - # QKVParallelLinear expects stacked [Q|K|V] blocks along the output axis. - # Reshape once at load time so the parallel loader can consume them. - nheads = int(self._model_config["model"]["diffusion"]["config"]["num_heads"]) - embed_dim = int(self._model_config["model"]["diffusion"]["config"]["embed_dim"]) - total_heads = embed_dim // nheads - head_dim = embed_dim // total_heads - qkv_mid = ".attn.qkv." - to_kv_mid = ".cross_attn.to_kv." - - def _restack_interleaved(tensor: torch.Tensor, n_slots: int) -> torch.Tensor: - """Turn bundle's (h, d, slot) interleaved-last-dim layout into stacked [slot|slot|...].""" - if tensor.dim() == 2: # weight: (out, in) - out, inp = tensor.shape - assert out == n_slots * total_heads * head_dim, (out, n_slots, total_heads, head_dim) - return ( - tensor.view(total_heads, head_dim, n_slots, inp).permute(2, 0, 1, 3).reshape(out, inp).contiguous() - ) - out = tensor.shape[0] - assert out == n_slots * total_heads * head_dim - return tensor.view(total_heads, head_dim, n_slots).permute(2, 0, 1).reshape(out).contiguous() - - def _remap(items): - for name, tensor in items: - if name.startswith(_legacy_prefix): - name = "audio_vae_adapter." + name[len(_legacy_prefix) :] - if qkv_mid in name and (name.endswith(".weight") or name.endswith(".bias")): - tensor = _restack_interleaved(tensor, 3) - elif to_kv_mid in name and (name.endswith(".weight") or name.endswith(".bias")): - tensor = _restack_interleaved(tensor, 2) - yield name, tensor - - loaded = AutoWeightsLoader(self).load_weights(_remap(weights)) - - self.to(torch.float32) - self.eval().requires_grad_(False) - - return loaded - - def _conditioning_dtype(self) -> torch.dtype: - p = next(self.model.parameters()) - return p.dtype if p.dtype.is_floating_point else torch.float32 - - @staticmethod - def _normalize_task(task: str | None) -> str | None: - if task is None: - return None - t = str(task).strip().lower() - return t or None - - @staticmethod - def _text_for_task(task_norm: str | None, prompt: str) -> str: - if task_norm in _VIDEO_ONLY_TASKS: - return "" - return prompt - - @staticmethod - def _ensure_text_video_prompts(task_norm: str | None, prompts: list[str]) -> None: - if task_norm not in _TEXT_VIDEO_TASKS: - return - for i, p in enumerate(prompts): - if not str(p).strip(): - raise ValueError( - f"audiox_task={task_norm!r} requires a non-empty text prompt for item {i}; " - "use v2a/v2m for video-only generation." - ) - - def _audio_prompt_tensors( - self, - *, - raw_prompts: list[Any], - extra: dict[str, Any], - device: torch.device, - cond_dtype: torch.dtype = torch.float32, - ) -> list[torch.Tensor]: - target_len = self._audio_conditioning_samples - sample_rate = self._sample_rate - seconds_start = float(extra.get("seconds_start", 0.0)) - seconds_total = float(target_len) / float(sample_rate) - out: list[torch.Tensor] = [] - for raw in raw_prompts: - src = extra.get("audio_path") - if src is None: - out.append(torch.zeros(2, target_len, device=device, dtype=cond_dtype)) - continue - wav = prepare_audio_reference( - src, - model_sample_rate=sample_rate, - seconds_start=seconds_start, - seconds_total=seconds_total, - device=device, - ) - out.append(wav.to(dtype=cond_dtype)) - return out - - def _video_feature_tensors( - self, - *, - task_norm: str | None, - raw_prompts: list[Any], - extra: dict[str, Any], - seconds_start: float, - target_fps: int, - device: torch.device, - cond_dtype: torch.dtype = torch.float32, - ) -> list[torch.Tensor]: - clip_frames = int(round(self._CLIP_SYNC_DURATION_SEC * target_fps)) - if task_norm not in _VIDEO_CONDITIONED_TASKS: - empty = torch.zeros(clip_frames, 3, 224, 224, device=device, dtype=cond_dtype) - return [empty for _ in raw_prompts] - - tensors: list[torch.Tensor] = [] - for _ in raw_prompts: - src = extra.get("video_path") - if src is None: - raise ValueError(f"audiox_task={task_norm!r} requires video input: set extra_args['video_path'].") - vt = prepare_video_reference( - src, - duration=float(self._CLIP_SYNC_DURATION_SEC), - target_fps=target_fps, - seek_time=seconds_start, - ) - tensors.append(vt.to(device=device, dtype=cond_dtype)) - return tensors - - def get_conditioning_inputs(self, conditioning_tensors: dict[str, Any], negative: bool = False) -> dict[str, Any]: - video_feature, video_mask = conditioning_tensors["video_prompt"] - text_feature, text_mask = conditioning_tensors["text_prompt"] - audio_feature, audio_mask = conditioning_tensors["audio_prompt"] - - refined = self.maf_block(text_feature, video_feature, audio_feature) - fused = torch.cat(list(refined.values()), dim=1) - masks = torch.cat([video_mask, text_mask, audio_mask], dim=1) - - if negative: - return {"negative_cross_attn_cond": fused, "negative_cross_attn_mask": masks} - return {"cross_attn_cond": fused} - - def diffuse( - self, - *, - steps: int, - guidance_scale: float, - conditioning_tensors: dict[str, Any], - negative_conditioning_tensors: dict[str, Any] | None, - batch_size: int, - sigma_min: float, - sigma_max: float, - generator: torch.Generator, - cfg_rescale: float, - ) -> torch.Tensor: - device = self.device - model_dtype = next(self.model.parameters()).dtype - - # Match upstream AudioX: disable TF32 matmul + fp16 reduced precision + cudnn benchmark - # for numerical parity with audiox/inference/generation.py:152-156. - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False - torch.backends.cudnn.benchmark = False - - latent_len = self._sample_size // int(self.pretransform.hop_length) - noise = torch.randn( - [batch_size, self.io_channels, latent_len], device=device, generator=generator, dtype=model_dtype - ) - - def _cast(d: dict[str, Any]) -> dict[str, Any]: - return {k: (v.type(model_dtype) if isinstance(v, torch.Tensor) else v) for k, v in d.items()} - - cond = _cast(self.get_conditioning_inputs(conditioning_tensors)) - neg = ( - _cast(self.get_conditioning_inputs(negative_conditioning_tensors, negative=True)) - if negative_conditioning_tensors is not None - else {} - ) - - # Inlined k-diffusion VDenoiser + sample_dpmpp_3m_sde, matching upstream AudioX exactly. - # diffusers' EDMDPMSolverMultistepScheduler uses different v-prediction preconditioning - # and a different stochastic update rule, which here produces a fixed ~861 Hz resonance - # in the decoded audio regardless of conditioning. - def denoise(x_in: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: - s2 = sigma * sigma - c_skip = 1.0 / (s2 + 1.0) - c_out = -sigma / (s2 + 1.0).sqrt() - c_in = 1.0 / (s2 + 1.0).sqrt() - t_cond = sigma.atan() * (2.0 / math.pi) - v = self.model( - x_in * c_in, - t_cond, - cross_attn_cond=cond["cross_attn_cond"], - negative_cross_attn_cond=neg.get("negative_cross_attn_cond"), - negative_cross_attn_mask=neg.get("negative_cross_attn_mask"), - cfg_scale=guidance_scale, - scale_phi=cfg_rescale, - ) - return v * c_out + x_in * c_skip - - ramp = torch.linspace(1.0, 0.0, steps, device=device) - sigmas = torch.cat( - [ - torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min)), - torch.zeros(1, device=device), - ] - ) - x = noise * sigmas[0] - - # Match upstream AudioX: sampler runs under fp16 autocast (see audiox/inference/sampling.py:184). - with torch.cuda.amp.autocast(): - if steps <= 1: - s_in = x.new_ones([x.shape[0]]) - sampled = denoise(x, sigmas[0] * s_in) - else: - # DPM-Solver++(3M) SDE loop (k-diffusion sample_dpmpp_3m_sde), eta=1.0, s_noise=1.0. - noise_sampler = _BrownianTreeNoiseSampler( - x, - sigmas[sigmas > 0].min(), - sigmas.max(), - entropy=generator.initial_seed(), - ) - s_in = x.new_ones([x.shape[0]]) - denoised_1 = denoised_2 = None - h_1 = h_2 = None - eta = 1.0 - for i in range(len(sigmas) - 1): - denoised = denoise(x, sigmas[i] * s_in) - if sigmas[i + 1] == 0: - x = denoised - else: - t_, s_ = -sigmas[i].log(), -sigmas[i + 1].log() - h = s_ - t_ - h_eta = h * (eta + 1) - x = torch.exp(-h_eta) * x + (-h_eta).expm1().neg() * denoised - if h_2 is not None: - r0 = h_1 / h - r1 = h_2 / h - d1_0 = (denoised - denoised_1) / r0 - d1_1 = (denoised_1 - denoised_2) / r1 - d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1) - d2 = (d1_0 - d1_1) / (r0 + r1) - phi_2 = h_eta.neg().expm1() / h_eta + 1 - phi_3 = phi_2 / h_eta - 0.5 - x = x + phi_2 * d1 - phi_3 * d2 - elif h_1 is not None: - r = h_1 / h - d = (denoised - denoised_1) / r - phi_2 = h_eta.neg().expm1() / h_eta + 1 - x = x + phi_2 * d - x = ( - x - + noise_sampler(sigmas[i], sigmas[i + 1]) - * sigmas[i + 1] - * (-2 * h * eta).expm1().neg().sqrt() - ) - denoised_1, denoised_2 = denoised, denoised_1 - h_1, h_2 = h, h_1 - sampled = x - - vae = self.pretransform.to(device=sampled.device, dtype=torch.float32).eval() - return vae.decode(sampled.to(torch.float32) * float(vae.audiox_scaling_factor), return_dict=True).sample - - def _encode_text(self, texts: list[str], device: torch.device) -> list[torch.Tensor]: - self.text_encoder.to(device) - encoded = self.tokenizer( - texts, - truncation=True, - max_length=self._t5_max_length, - padding="max_length", - return_tensors="pt", - ) - input_ids = encoded["input_ids"].to(device) - attention_mask = encoded["attention_mask"].to(device).to(torch.bool) - - self.text_encoder.eval() - with torch.no_grad(): - embeddings = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask)["last_hidden_state"] - - embeddings = embeddings.float() * attention_mask.unsqueeze(-1).float() - return [embeddings, attention_mask] - - def _encode_video(self, video_list: list[dict], device: torch.device) -> list[torch.Tensor]: - self.clip_encoder.to(device).eval() - - video_tensors = [item["video_tensors"] for item in video_list] - video_sync_frames = torch.cat([item["video_sync_frames"] for item in video_list], dim=0).to(device) - - original_videos = torch.cat(video_tensors, dim=0).to(device) - batch_size, time_length, _, _, _ = original_videos.size() - is_zero = torch.all(original_videos == 0, dim=(1, 2, 3, 4)) - - frames = original_videos.flatten(0, 1) - pixel_values = self._clip_normalize(frames).to(device) - - with torch.no_grad(): - outputs = self.clip_encoder(pixel_values=pixel_values) - hidden = outputs.last_hidden_state - hidden = rearrange(hidden, "(b t) p d -> (b p) t d", b=batch_size, t=time_length) - hidden = hidden + self.clip_temp_pos_embedding - hidden = self.clip_temp_transformer(hidden) - hidden = rearrange(hidden, "(b p) t d -> b (t p) d", b=batch_size) - hidden = self.clip_proj(hidden.view(-1, self._clip_in_features)) - hidden = hidden.view(batch_size, self._clip_out_features, -1) - - sync = self.clip_proj_sync(video_sync_frames.view(-1, 240)) - sync = sync.view(batch_size, self._clip_out_features, -1) - hidden = hidden + self.clip_sync_weight * sync - - empty = self.clip_empty_visual_feat.expand(batch_size, -1, -1) - hidden = torch.where(is_zero.view(batch_size, 1, 1), empty, hidden) - return [hidden, torch.ones(batch_size, 1, device=device)] - - def _encode_conditioning_tensors(self, batch_metadata: list[dict[str, Any]]) -> dict[str, Any]: - device = self.device - audio = torch.cat([item["audio_prompt"] for item in batch_metadata], dim=0).to(device) - return { - "audio_prompt": list(self.audio_vae_adapter(audio)), - "text_prompt": self._encode_text([item["text_prompt"] for item in batch_metadata], device), - "video_prompt": self._encode_video([item["video_prompt"] for item in batch_metadata], device), - } - - def _build_conditioning_batch( - self, - *, - texts: list[str], - video_tensors_list: list[torch.Tensor], - audio_prompt_list: list[torch.Tensor], - sync_features: torch.Tensor, - seconds_start: float, - seconds_model: float, - num_outputs_per_prompt: int, - task_norm: str | None, - ) -> list[dict[str, Any]]: - batch: list[dict[str, Any]] = [] - for i, text in enumerate(texts): - for _ in range(num_outputs_per_prompt): - batch.append( - { - "video_prompt": { - "video_tensors": video_tensors_list[i].unsqueeze(0), - "video_sync_frames": sync_features, - }, - "text_prompt": self._text_for_task(task_norm, text), - "audio_prompt": audio_prompt_list[i].unsqueeze(0), - "seconds_start": seconds_start, - "seconds_total": seconds_model, - } - ) - return batch - - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - if req.prompts is None or len(req.prompts) == 0: - raise ValueError("AudioXPipeline requires at least one prompt.") - normalized_prompts = _normalize_prompts(list(req.prompts)) - prompts = [p["prompt"] for p in normalized_prompts] - - sampling_params = req.sampling_params - if sampling_params.num_inference_steps is None: - raise ValueError("AudioXPipeline requires sampling_params.num_inference_steps.") - num_inference_steps = int(sampling_params.num_inference_steps) - extra_args = sampling_params.extra_args or {} - task_norm = self._normalize_task(extra_args.get("audiox_task")) - self._ensure_text_video_prompts(task_norm, prompts) - - neg: list[str] | None = None - if not all(p.get("negative_prompt") is None for p in normalized_prompts): - neg = [str(p.get("negative_prompt") or "") for p in normalized_prompts] - - guidance_scale = float(sampling_params.guidance_scale) - if sampling_params.num_outputs_per_prompt <= 0: - raise ValueError("AudioXPipeline requires sampling_params.num_outputs_per_prompt > 0.") - num_outputs_per_prompt = int(sampling_params.num_outputs_per_prompt) - batch_size = len(prompts) * num_outputs_per_prompt - - seconds_start = float(extra_args.get("seconds_start", 0.0)) - seconds_model = self._sample_size / self._sample_rate - seconds_total = float(extra_args.get("seconds_total", seconds_model)) - sigma_min = float(extra_args.get("sigma_min", _DEFAULT_UPSTREAM_SIGMA_MIN)) - sigma_max = float(extra_args.get("sigma_max", _DEFAULT_UPSTREAM_SIGMA_MAX)) - cfg_rescale = float(extra_args.get("cfg_rescale", 0.0)) - device = self.device - generator = sampling_params.generator - if generator is None: - raise ValueError("AudioXPipeline requires sampling_params.generator.") - target_fps = self._target_fps - cond_dtype = self._conditioning_dtype() - - sync_features = torch.zeros(1, self._VIDEO_SYNC_FRAME_COUNT, 768, device=device, dtype=cond_dtype) - - audio_prompt_list = self._audio_prompt_tensors( - raw_prompts=normalized_prompts, - extra=extra_args, - device=device, - cond_dtype=cond_dtype, - ) - - video_tensors_list = self._video_feature_tensors( - task_norm=task_norm, - raw_prompts=normalized_prompts, - extra=extra_args, - seconds_start=seconds_start, - target_fps=target_fps, - device=device, - cond_dtype=cond_dtype, - ) - - conditioning_batch = self._build_conditioning_batch( - texts=prompts, - video_tensors_list=video_tensors_list, - audio_prompt_list=audio_prompt_list, - sync_features=sync_features, - seconds_start=seconds_start, - seconds_model=seconds_model, - num_outputs_per_prompt=num_outputs_per_prompt, - task_norm=task_norm, - ) - - negative_conditioning_batch: list[dict[str, Any]] | None = None - if neg is not None and guidance_scale > 1.0: - negative_conditioning_batch = self._build_conditioning_batch( - texts=neg, - video_tensors_list=video_tensors_list, - audio_prompt_list=audio_prompt_list, - sync_features=sync_features, - seconds_start=seconds_start, - seconds_model=seconds_model, - num_outputs_per_prompt=num_outputs_per_prompt, - task_norm=task_norm, - ) - - conditioning_tensors = self._encode_conditioning_tensors(conditioning_batch) - negative_conditioning_tensors: dict[str, Any] | None = None - if negative_conditioning_batch is not None: - negative_conditioning_tensors = self._encode_conditioning_tensors(negative_conditioning_batch) - - audio = self.diffuse( - steps=num_inference_steps, - guidance_scale=guidance_scale, - conditioning_tensors=conditioning_tensors, - negative_conditioning_tensors=negative_conditioning_tensors, - batch_size=batch_size, - sigma_min=sigma_min, - sigma_max=sigma_max, - generator=generator, - cfg_rescale=cfg_rescale, - ) - - # Trim decoded audio to the requested duration (matches upstream AudioX sample script). - if 0.0 < seconds_total < seconds_model: - target_samples = int(seconds_total * self._sample_rate) - audio = audio[..., :target_samples] - - stage_durations = self.stage_durations if getattr(self, "enable_diffusion_pipeline_profiler", False) else None - return DiffusionOutput( - output={ - "payload": {"audio": audio}, - "metadata": {"audio": {"audiox_task": task_norm}}, - }, - stage_durations=stage_durations, - ) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py deleted file mode 100644 index 000abe051..000000000 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ /dev/null @@ -1,839 +0,0 @@ -import math - -import torch -import torch.nn.functional as F -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.models.attention_processor import Attention -from diffusers.models.embeddings import TimestepEmbedding, Timesteps -from diffusers.models.modeling_utils import ModelMixin -from einops import rearrange, repeat -from torch import nn -from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm - -from .rope_real import RotaryPosEmbedReal - -try: - from transformers.modeling_flash_attention_utils import ( # type: ignore - flash_attn_varlen_func, # pyright: ignore[reportAttributeAccessIssue] - is_flash_attn_available, - ) -except Exception: # pragma: no cover - best-effort compatibility - flash_attn_varlen_func = None # type: ignore[assignment] - - def is_flash_attn_available() -> bool: # type: ignore[override] - return False - - -from .rope_real import apply_real_rotary_emb - -_HAS_FLASH_ATTN_VARLEN = bool(is_flash_attn_available()) and flash_attn_varlen_func is not None - - -class LuminaRMSNormZero(nn.Module): - """ - Norm layer adaptive RMS normalization zero. - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - """ - - def __init__( - self, - embedding_dim: int, - norm_eps: float, - norm_elementwise_affine: bool, - ): - super().__init__() - self.silu = nn.SiLU() - self.linear = nn.Linear( - min(embedding_dim, 1024), - 4 * embedding_dim, - bias=True, - ) - - self.norm = Qwen2RMSNorm(embedding_dim, eps=norm_eps) - - def forward( - self, - x: torch.Tensor, - emb: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - emb = self.linear(self.silu(emb)) - scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1) - x = self.norm(x) * (1 + scale_msa[:, None]) - return x, gate_msa, scale_mlp, gate_mlp - - -class LuminaFeedForward(nn.Module): - def __init__( - self, - dim: int, - inner_dim: int, - multiple_of: int | None = 256, - ffn_dim_multiplier: float | None = None, - ): - super().__init__() - - # custom hidden_size factor multiplier - if ffn_dim_multiplier is not None: - inner_dim = int(ffn_dim_multiplier * inner_dim) - inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of) - - self.linear_1 = nn.Linear( - dim, - inner_dim, - bias=False, - ) - self.linear_2 = nn.Linear( - inner_dim, - dim, - bias=False, - ) - self.linear_3 = nn.Linear( - dim, - inner_dim, - bias=False, - ) - - def swiglu(self, x, y): - return F.silu(x.float(), inplace=False).to(x.dtype) * y - - def forward(self, x): - h1, h2 = self.linear_1(x), self.linear_3(x) - return self.linear_2(self.swiglu(h1, h2)) - - -class LuminaLayerNormContinuous(nn.Module): - def __init__( - self, - embedding_dim: int, - conditioning_embedding_dim: int, - # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters - # because the output is immediately scaled and shifted by the projected conditioning embeddings. - # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. - # However, this is how it was implemented in the original code, and it's rather likely you should - # set `elementwise_affine` to False. - elementwise_affine=True, - eps=1e-5, - bias=True, - norm_type="layer_norm", - out_dim: int | None = None, - ): - super().__init__() - - self.silu = nn.SiLU() - self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias) - - if norm_type == "layer_norm": - self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias) - elif norm_type == "rms_norm": - self.norm = Qwen2RMSNorm(embedding_dim, eps=eps) - else: - raise ValueError(f"unknown norm_type {norm_type}") - - self.linear_2 = None - if out_dim is not None: - self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias) - - def forward( - self, - x: torch.Tensor, - conditioning_embedding: torch.Tensor, - ) -> torch.Tensor: - scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype)) - x = self.norm(x) * (1 + scale)[:, None, :] - - if self.linear_2 is not None: - x = self.linear_2(x) - - return x - - -class Lumina2CombinedTimestepCaptionEmbedding(nn.Module): - def __init__( - self, - hidden_size: int = 2520, - text_feat_dim: int = 3584, - frequency_embedding_size: int = 256, - norm_eps: float = 1e-5, - timestep_scale: float = 1.0, - ) -> None: - super().__init__() - - self.time_proj = Timesteps( - num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=timestep_scale - ) - - self.timestep_embedder = TimestepEmbedding( - in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024) - ) - - self.caption_embedder = nn.Sequential( - Qwen2RMSNorm(text_feat_dim, eps=norm_eps), - nn.Linear(text_feat_dim, hidden_size, bias=True), - ) - - def forward( - self, - timestep: torch.Tensor, - text_hidden_states: torch.Tensor, - dtype: torch.dtype, - ) -> tuple[torch.Tensor, torch.Tensor]: - timestep_proj = self.time_proj(timestep).to(dtype=dtype) - time_embed = self.timestep_embedder(timestep_proj) - caption_embed = self.caption_embedder(text_hidden_states) - return time_embed, caption_embed - - -class SimpleQFormerImageRefiner(nn.Module): - def __init__( - self, - hidden_size: int, - output_hidden_size: int | None = None, - num_queries: int = 128, - num_layers: int = 2, - num_heads: int | None = None, - dropout: float = 0.0, - norm_eps: float = 1e-5, - ) -> None: - super().__init__() - input_hidden_size = hidden_size - hidden_size = output_hidden_size or hidden_size - self.hidden_size = hidden_size - self.num_queries = num_queries - # ensure num_heads divides hidden_size - if num_heads is None: - num_heads = max(1, hidden_size // 128) - self.num_heads = self._choose_valid_num_heads(hidden_size, num_heads) - self.input_proj = nn.Sequential( - Qwen2RMSNorm(input_hidden_size, eps=norm_eps), - nn.Linear(input_hidden_size, hidden_size, bias=True), - ) - - # Learnable query embeddings - scale = hidden_size**-0.5 - self.query = nn.Parameter(scale * torch.randn(1, num_queries, hidden_size)) - - # Decoder layers - self.layers = nn.ModuleList() - for _ in range(num_layers): - self.layers.append( - nn.ModuleDict( - dict( - ln_q1=Qwen2RMSNorm(hidden_size, eps=norm_eps), - self_attn=nn.MultiheadAttention( - embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True - ), - ln_q2=Qwen2RMSNorm(hidden_size, eps=norm_eps), - cross_attn=nn.MultiheadAttention( - embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True - ), - ln_ffn=Qwen2RMSNorm(hidden_size, eps=norm_eps), - ffn=LuminaFeedForward(dim=hidden_size, inner_dim=4 * hidden_size), - ) - ) - ) - - @staticmethod - def _choose_valid_num_heads(hidden_size: int, proposed_heads: int, preferred_head_dim: int = 128) -> int: - """Pick a number of heads that divides hidden_size, close to proposed or preferred.""" - # If proposed is valid, use it - if proposed_heads > 0 and hidden_size % proposed_heads == 0: - return proposed_heads - # target based on preferred head dim - target = max(1, round(hidden_size / preferred_head_dim)) - # collect divisors up to 128 heads (more than enough) - max_heads_cap = min(128, hidden_size) - divisors = [d for d in range(1, max_heads_cap + 1) if hidden_size % d == 0] - # choose closest to target - best = min(divisors, key=lambda d: (abs(d - target), -d)) - return best - - def forward(self, x: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - """ - Args: - x: Tensor of shape (batch, seq_len, input_dim) - Returns: - Tensor of shape (batch, num_queries, hidden_size) - """ - batch, _, _ = x.shape - kv = self.input_proj(x) - q = self.query.repeat(batch, 1, 1).to(kv.dtype) - - for layer in self.layers: - # Self-attention on queries - q_norm = layer["ln_q1"](q) - attn_out, _ = layer["self_attn"](q_norm, q_norm, q_norm, need_weights=False) - q = q + attn_out - - # Cross-attention: queries attend to inputs - q_norm = layer["ln_q2"](q) - cross_out, _ = layer["cross_attn"](q_norm, kv, kv, need_weights=False, key_padding_mask=attention_mask) - q = q + cross_out - - # Feed-forward - q = q + layer["ffn"](layer["ln_ffn"](q)) - - return q - - -class AttnProcessor: - def __init__(self) -> None: - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor requires PyTorch 2.0+ (F.scaled_dot_product_attention).") - - def __call__( - self, - attn: Attention, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - image_rotary_emb: torch.Tensor | None = None, - base_sequence_length: int | None = None, - ) -> torch.Tensor: - batch_size, sequence_length, _ = hidden_states.shape - - # Get Query-Key-Value Pair - query = attn.to_q(hidden_states) - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query_dim = query.shape[-1] - inner_dim = key.shape[-1] - head_dim = query_dim // attn.heads - dtype = query.dtype - - # Get key-value heads - kv_heads = inner_dim // head_dim - - # Reshape tensors for attention computation - query = query.view(batch_size, -1, attn.heads, head_dim) - key = key.view(batch_size, -1, kv_heads, head_dim) - value = value.view(batch_size, -1, kv_heads, head_dim) - - # Apply Query-Key normalization - if attn.norm_q is not None: - query = attn.norm_q(query) - if attn.norm_k is not None: - key = attn.norm_k(key) - - # Apply Rotary Position Embeddings - if image_rotary_emb is not None: - query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1]) - key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1]) - - query, key = query.to(dtype), key.to(dtype) - - # Calculate attention scale - if base_sequence_length is not None: - softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale - else: - softmax_scale = attn.scale - - if _HAS_FLASH_ATTN_VARLEN and attention_mask is not None and hidden_states.is_cuda: - # Flash-Attn varlen expects packed tokens + cu_seqlens. Here we only need - # the self-attention case (q/k/v share the same padding mask). - attention_mask = attention_mask.to(torch.bool) - seqlens = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() - max_seqlen = int(seqlens.max().item()) - cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) - - query_states = query.reshape(batch_size * sequence_length, attn.heads, head_dim)[indices] - key_states = key.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices] - value_states = value.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices] - - if kv_heads < attn.heads: - key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) - value_states = repeat(value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) - - attn_output_unpad = flash_attn_varlen_func( - query_states, - key_states, - value_states, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - dropout_p=0.0, - causal=False, - softmax_scale=softmax_scale, - ) - - out = torch.zeros( - (batch_size * sequence_length, attn.heads, head_dim), - device=attn_output_unpad.device, - dtype=attn_output_unpad.dtype, - ) - out[indices] = attn_output_unpad - hidden_states = out.view(batch_size, sequence_length, attn.heads, head_dim).flatten(-2) - hidden_states = hidden_states.type_as(query) - else: - # PyTorch SDPA path. - attn_mask = None - if attention_mask is not None: - attention_mask = attention_mask.to(torch.bool) - attn_mask = attention_mask.view(batch_size, 1, 1, -1) - - query = query.transpose(1, 2) # [B, H, S, D] - key = key.transpose(1, 2) # [B, H_kv, S, D] - value = value.transpose(1, 2) - - if kv_heads < attn.heads: - key = key.repeat_interleave(attn.heads // kv_heads, dim=1) - value = value.repeat_interleave(attn.heads // kv_heads, dim=1) - - hidden_states = F.scaled_dot_product_attention( - query, - key, - value, - attn_mask=attn_mask, - dropout_p=0.0, - is_causal=False, - scale=softmax_scale, - ) - - if attention_mask is not None: - # Keep padding tokens consistent with the flash-varlen path (zero output). - hidden_states = hidden_states * attention_mask[:, None, :, None] - - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.type_as(query) - - # Apply output projection - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - return hidden_states - - -class TransformerBlock(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - num_kv_heads: int, - multiple_of: int, - ffn_dim_multiplier: float, - norm_eps: float, - modulation: bool = True, - ) -> None: - """Initialize the transformer block.""" - super().__init__() - self.head_dim = dim // num_attention_heads - self.modulation = modulation - - processor = AttnProcessor() - - # Initialize attention layer - self.attn = Attention( - query_dim=dim, - cross_attention_dim=None, - dim_head=dim // num_attention_heads, - qk_norm=None, - heads=num_attention_heads, - kv_heads=num_kv_heads, - eps=1e-5, - bias=False, - out_bias=False, - processor=processor, - ) - # 显式使用 transformers 的 Qwen2RMSNorm，避免依赖 diffusers 内部创建的 `RMSNorm` 再做递归替换。 - self.attn.norm_q = Qwen2RMSNorm(self.head_dim, eps=1e-5) - self.attn.norm_k = Qwen2RMSNorm(self.head_dim, eps=1e-5) - - # Initialize feed-forward network - self.feed_forward = LuminaFeedForward( - dim=dim, inner_dim=4 * dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier - ) - - # Initialize normalization layers - if modulation: - self.norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True) - else: - self.norm1 = Qwen2RMSNorm(dim, eps=norm_eps) - - self.ffn_norm1 = Qwen2RMSNorm(dim, eps=norm_eps) - self.norm2 = Qwen2RMSNorm(dim, eps=norm_eps) - self.ffn_norm2 = Qwen2RMSNorm(dim, eps=norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - image_rotary_emb: torch.Tensor, - temb: torch.Tensor | None = None, - ) -> torch.Tensor: - if self.modulation: - if temb is None: - raise ValueError("temb must be provided when modulation is enabled") - - norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) - attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output) - mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1))) - hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output) - else: - norm_hidden_states = self.norm1(hidden_states) - attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, - ) - hidden_states = hidden_states + self.norm2(attn_output) - mlp_output = self.feed_forward(self.ffn_norm1(hidden_states)) - hidden_states = hidden_states + self.ffn_norm2(mlp_output) - - return hidden_states - - -class Transformer2DModel(ModelMixin, ConfigMixin): - """MammothModa2 DiT transformer""" - - @register_to_config - def __init__( - self, - patch_size: int = 2, - in_channels: int = 16, - out_channels: int | None = None, - hidden_size: int = 2304, - num_layers: int = 26, - num_refiner_layers: int = 2, - num_attention_heads: int = 24, - num_kv_heads: int = 8, - multiple_of: int = 256, - ffn_dim_multiplier: float | None = None, - norm_eps: float = 1e-5, - axes_dim_rope: tuple[int, int, int] = (32, 32, 32), - axes_lens: tuple[int, int, int] = (300, 512, 512), - text_feat_dim: int = 1024, - timestep_scale: float = 1.0, - ) -> None: - """Initialize the transformer model.""" - super().__init__() - self.hidden_size = hidden_size - - # Validate configuration - if (hidden_size // num_attention_heads) != sum(axes_dim_rope): - raise ValueError( - f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) " - f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})" - ) - - self.out_channels = out_channels or in_channels - # Initialize embeddings - self.rope_embedder = RotaryPosEmbedReal( - theta=10000, - axes_dim=axes_dim_rope, - axes_lens=axes_lens, - patch_size=patch_size, - ) - - self.x_embedder = nn.Linear( - in_features=patch_size * patch_size * in_channels, - out_features=hidden_size, - ) - - self.ref_image_patch_embedder = nn.Linear( - in_features=patch_size * patch_size * in_channels, - out_features=hidden_size, - ) - - self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding( - hidden_size=hidden_size, - text_feat_dim=text_feat_dim, - norm_eps=norm_eps, - timestep_scale=timestep_scale, - ) - - # Initialize transformer blocks - self.noise_refiner = nn.ModuleList( - [ - TransformerBlock( - hidden_size, - num_attention_heads, - num_kv_heads, - multiple_of, - ffn_dim_multiplier, - norm_eps, - modulation=True, - ) - for _ in range(num_refiner_layers) - ] - ) - - self.ref_image_refiner = nn.ModuleList( - [ - TransformerBlock( - hidden_size, - num_attention_heads, - num_kv_heads, - multiple_of, - ffn_dim_multiplier, - norm_eps, - modulation=True, - ) - for _ in range(num_refiner_layers) - ] - ) - - self.context_refiner = nn.ModuleList( - [ - TransformerBlock( - hidden_size, - num_attention_heads, - num_kv_heads, - multiple_of, - ffn_dim_multiplier, - norm_eps, - modulation=False, - ) - for _ in range(num_refiner_layers) - ] - ) - - # 3. Transformer blocks - self.layers = nn.ModuleList( - [ - TransformerBlock( - hidden_size, - num_attention_heads, - num_kv_heads, - multiple_of, - ffn_dim_multiplier, - norm_eps, - modulation=True, - ) - for _ in range(num_layers) - ] - ) - - # 4. Output norm & projection - self.norm_out = LuminaLayerNormContinuous( - embedding_dim=hidden_size, - conditioning_embedding_dim=min(hidden_size, 1024), - elementwise_affine=False, - eps=1e-6, - bias=True, - out_dim=patch_size * patch_size * self.out_channels, - ) - - # Add learnable embeddings to distinguish different images - self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size)) # support max 5 ref images - - def _validate_inputs( - self, - hidden_states: torch.Tensor, - text_hidden_states: torch.Tensor, - text_attention_mask: torch.Tensor, - ref_image_hidden_states: list[list[torch.Tensor]] | None, - return_dict: bool, - ) -> tuple[int, int, int]: - if return_dict: - raise ValueError("return_dict=True is not supported in vLLM inference.") - if ref_image_hidden_states is not None: - raise ValueError("ref_image_hidden_states is not supported in vLLM inference.") - if hidden_states.ndim != 4: - raise ValueError(f"Expected hidden_states to be 4D [B,C,H,W], got shape={tuple(hidden_states.shape)}") - - batch_size, _channels, height, width = hidden_states.shape - if batch_size != text_hidden_states.shape[0] or batch_size != text_attention_mask.shape[0]: - raise ValueError( - "Batch size mismatch: " - f"hidden_states={batch_size}, text_hidden_states={text_hidden_states.shape[0]}, " - f"text_attention_mask={text_attention_mask.shape[0]}" - ) - - p = self.config.patch_size - if height % p != 0 or width % p != 0: - raise ValueError(f"Input latent H/W must be divisible by patch_size={p}, got {height}x{width}") - return batch_size, height, width - - def _prepare_embeddings( - self, - hidden_states: torch.Tensor, - timestep: torch.Tensor, - text_hidden_states: torch.Tensor, - text_attention_mask: torch.Tensor, - freqs_cis: torch.Tensor, - batch_size: int, - height: int, - width: int, - ar_image_hidden_states: torch.Tensor | None = None, - ar_image_attention_mask: torch.Tensor | None = None, - ): - device = hidden_states.device - p = self.config.patch_size - - temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states.dtype) - image_embedder = getattr(self.time_caption_embed, "image_embedder", None) - if image_embedder is not None and ar_image_hidden_states is not None: - if ar_image_attention_mask is None: - ar_image_attention_mask = torch.ones( - ar_image_hidden_states.shape[:2], - dtype=torch.bool, - device=ar_image_hidden_states.device, - ) - image_hidden_states = image_embedder( - ar_image_hidden_states, - ~ar_image_attention_mask.bool(), - ) - image_attention_mask = torch.ones( - image_hidden_states.shape[:2], - dtype=torch.bool, - device=image_hidden_states.device, - ) - text_hidden_states = torch.cat([text_hidden_states, image_hidden_states], dim=1) - text_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1) - - img_tokens = rearrange(hidden_states, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p) - img_tokens = self.x_embedder(img_tokens) - - img_len = (height // p) * (width // p) - img_mask = torch.ones((batch_size, img_len), dtype=torch.bool, device=device) - l_effective_img_len = [img_len for _ in range(batch_size)] - img_sizes = [(height, width) for _ in range(batch_size)] - - l_effective_ref_img_len = [[] for _ in range(batch_size)] - ref_img_sizes = [None for _ in range(batch_size)] - - ( - context_rotary_emb, - _ref_img_rotary_emb, - noise_rotary_emb, - rotary_emb, - encoder_seq_lengths, - seq_lengths, - ) = self.rope_embedder( - freqs_cis, - text_attention_mask, - l_effective_ref_img_len, - l_effective_img_len, - ref_img_sizes, - img_sizes, - device, - ) - - return ( - temb, - text_hidden_states, - text_attention_mask, - img_tokens, - img_mask, - img_len, - context_rotary_emb, - noise_rotary_emb, - rotary_emb, - encoder_seq_lengths, - seq_lengths, - ) - - def _apply_refiners( - self, - text_hidden_states: torch.Tensor, - text_attention_mask: torch.Tensor, - context_rotary_emb: torch.Tensor, - img_tokens: torch.Tensor, - img_mask: torch.Tensor, - noise_rotary_emb: torch.Tensor, - temb: torch.Tensor, - ): - for layer in self.context_refiner: - text_hidden_states = layer(text_hidden_states, text_attention_mask, context_rotary_emb) - - for layer in self.noise_refiner: - img_tokens = layer(img_tokens, img_mask, noise_rotary_emb, temb) - - return text_hidden_states, img_tokens - - def _apply_transformer_layers(self, hidden_states, attention_mask, rotary_emb, temb): - for layer in self.layers: - hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb) - return hidden_states - - def forward( - self, - hidden_states: torch.Tensor, - timestep: torch.Tensor, - text_hidden_states: torch.Tensor, - freqs_cis: torch.Tensor, - text_attention_mask: torch.Tensor, - ref_image_hidden_states: list[list[torch.Tensor]] | None = None, - ar_image_hidden_states: torch.Tensor | None = None, - ar_image_attention_mask: torch.Tensor | None = None, - return_dict: bool = False, - ) -> torch.Tensor: - batch_size, height, width = self._validate_inputs( - hidden_states, text_hidden_states, text_attention_mask, ref_image_hidden_states, return_dict - ) - - ( - temb, - text_hidden_states, - text_attention_mask, - img_tokens, - img_mask, - img_len, - context_rotary_emb, - noise_rotary_emb, - rotary_emb, - encoder_seq_lengths, - seq_lengths, - ) = self._prepare_embeddings( - hidden_states, - timestep, - text_hidden_states, - text_attention_mask, - freqs_cis, - batch_size, - height, - width, - ar_image_hidden_states, - ar_image_attention_mask, - ) - - text_hidden_states, img_tokens = self._apply_refiners( - text_hidden_states, - text_attention_mask, - context_rotary_emb, - img_tokens, - img_mask, - noise_rotary_emb, - temb, - ) - - max_seq_len = max(seq_lengths) - attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool) - joint_hidden_states = hidden_states.new_zeros(batch_size, max_seq_len, self.config.hidden_size) - for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)): - attention_mask[i, :seq_len] = True - joint_hidden_states[i, :encoder_seq_len] = text_hidden_states[i, :encoder_seq_len] - joint_hidden_states[i, encoder_seq_len : encoder_seq_len + img_len] = img_tokens[i, :img_len] - - hidden_states = self._apply_transformer_layers(joint_hidden_states, attention_mask, rotary_emb, temb) - - hidden_states = self.norm_out(hidden_states, temb) - - p = self.config.patch_size - img_hidden_states = torch.stack( - [ - hidden_states[i, encoder_seq_len : encoder_seq_len + img_len] - for i, encoder_seq_len in enumerate(encoder_seq_lengths) - ], - dim=0, - ) - output = rearrange( - img_hidden_states, - "b (h w) (p1 p2 c) -> b c (h p1) (w p2)", - h=height // p, - w=width // p, - p1=p, - p2=p, - ) - return output diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py deleted file mode 100644 index 5298d5514..000000000 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ /dev/null @@ -1,405 +0,0 @@ -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any, ClassVar - -import torch -from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL -from diffusers.utils.torch_utils import randn_tensor -from torch import nn -from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -from vllm.config import VllmConfig -from vllm.logger import init_logger -from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper - -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config - -from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel -from .rope_real import RotaryPosEmbedReal -from .schedulers import FlowMatchEulerDiscreteScheduler - -logger = init_logger(__name__) - - -class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): - """ - MammothModa2 DiT + VAE generation stage (non-autoregressive). - - This stage expects "image condition token hidden states" from the upstream AR stage, - and outputs image tensors via diffusion transformer + VAE decode. - - """ - - _dit_modules: ClassVar[list[str]] = ["gen_transformer"] - _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"] - _vae_modules: ClassVar[list[str]] = ["gen_vae"] - - have_multimodal_outputs = True - - # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage. - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={ - "llm_model.": None, - "gen_tokenizer.": None, - } - ) - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - del prefix - - hf_config = vllm_config.model_config.hf_config - if not isinstance(hf_config, Mammothmoda2Config): - raise TypeError(f"Expected Mammothmoda2Config, got {type(hf_config)}") - - self.config = hf_config - - # --- Build DiT / VAE modules (names must match checkpoint keys) --- - if self.config.gen_vae_config is None or self.config.gen_dit_config is None: - raise ValueError("Mammothmoda2Config.gen_vae_config / gen_dit_config must not be None") - - self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config) - self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config) - - # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config - llm_hidden_size = 0 - text_config = self.config.get_text_config() - if text_config is None: - logger.warning("No text config; failed to infer llm_hidden_size.") - elif not hasattr(text_config, "hidden_size"): - logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.") - else: - llm_hidden_size = int(text_config.hidden_size or 0) - if llm_hidden_size <= 0: - raise ValueError( - "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size" - ) - self._reinit_caption_embedder(llm_hidden_size) - - # Optional image condition Q-Former. Preview stores it as a standalone - # module; Dev stores it under the DiT timestep/caption embedder. - llm_model_type = getattr(self.config.llm_config, "model_type", "") - refiner_config = self.config.gen_image_condition_refiner_config - if refiner_config is not None and llm_model_type == "mammothmoda2_qwen3_vl": - dit_hidden_size = int(self.gen_transformer.hidden_size) - self.gen_transformer.time_caption_embed.image_embedder = SimpleQFormerImageRefiner( - hidden_size=llm_hidden_size, - output_hidden_size=dit_hidden_size, - num_heads=max(1, dit_hidden_size // 128), - **refiner_config, - ) - self.gen_image_condition_refiner = None - elif refiner_config is not None: - self.gen_image_condition_refiner = SimpleQFormerImageRefiner( - hidden_size=llm_hidden_size, - **refiner_config, - ) - else: - self.gen_image_condition_refiner = None - - # Precompute rotary freqs for diffusion transformer - # IMPORTANT: follow upstream mammothmoda: use top-level `config.gen_axes_*` - # (the checkpoint's `gen_dit_config.axes_lens` can be as small as 1024, - # which is insufficient for vLLM dummy-run/cudagraph warmup). - self.gen_freqs_cis = RotaryPosEmbedReal.get_freqs_real( - tuple(self.config.gen_axes_dim_rope), - tuple(self.config.gen_axes_lens), - theta=10000, - ) - - # vLLM PP interface compatibility - self.make_empty_intermediate_tensors = lambda: None - - self._llm_hidden_size = llm_hidden_size - - def _reinit_caption_embedder(self, in_features: int) -> None: - # Align with upstream Mammothmoda2Model's `reinit_caption_embedder`: - # Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features). - out_features = int(getattr(self.gen_transformer, "hidden_size", 0) or self.gen_transformer.config.hidden_size) - self.gen_transformer.time_caption_embed.caption_embedder = nn.Sequential( - Qwen2RMSNorm(in_features, eps=1e-5), - nn.Linear(in_features, out_features, bias=True), - ) - - def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]: - if num_reqs <= 0: - raise ValueError(f"num_reqs must be positive, got {num_reqs}") - if num_reqs > 1: - raise NotImplementedError( - f"get_dummy_runtime_additional_information does not support num_reqs > 1, got {num_reqs}" - ) - text_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32) - image_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32) - negative_prompt_embeds = torch.zeros((0, self._llm_hidden_size), dtype=torch.float32) - info = { - "text_prompt_embeds": text_prompt_embeds, - "image_prompt_embeds": image_prompt_embeds, - "negative_prompt_embeds": negative_prompt_embeds, - "negative_prompt_attention_mask": [], - "image_height": [512], - "image_width": [512], - "text_guidance_scale": [1.0], - "cfg_range": [0.0, 1.0], - "num_inference_steps": [1], - } - return [info for _ in range(num_reqs)] - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - # DiT stage does not consume token embeddings; return a dummy tensor. - try: - dtype = next(self.parameters()).dtype - except StopIteration: - dtype = torch.float32 - return torch.zeros( - (input_ids.numel(), self._llm_hidden_size), - device=input_ids.device, - dtype=dtype, - ) - - def _split_ar_conditions( - self, - *, - full_hidden_states: torch.Tensor, - full_token_ids: list[int], - answer_start_index: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Split AR-stage hidden states into text / image condition embeds. - - The token ids that distinguish question (text) tokens, generated visual - tokens, and multi-modal placeholder tokens are read from the model config - (``gen_vocab_start_index`` and the vision placeholder token ids), so the - caller no longer needs to pass them. Mirrors the masking the bespoke - MammothModa2 example performed via ar2dit. - """ - gen_vocab_start_index = int(self.config.llm_config.gen_vocab_start_index) - visual_ids = [ - int(self.config.image_token_id), - int(self.config.video_token_id), - int(self.config.vision_start_token_id), - int(self.config.vision_end_token_id), - ] - - device = full_hidden_states.device - token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device) - positions = torch.arange(token_ids.shape[0], device=device) - questions_mask = positions < answer_start_index - answers_mask = ~questions_mask - gen_token_mask = token_ids >= gen_vocab_start_index - visual_token_mask = torch.isin(token_ids, torch.tensor(visual_ids, dtype=torch.long, device=device)) - text_mask = questions_mask & ~(visual_token_mask | gen_token_mask) - image_mask = answers_mask & gen_token_mask - - text_cond = full_hidden_states[text_mask].to(dtype=torch.float32).contiguous() - image_cond = full_hidden_states[image_mask].to(dtype=torch.float32).contiguous() - return text_cond, image_cond - - @torch.inference_mode() - def forward( - self, - *, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, # noqa: ARG002 - ) -> OmniOutput: - runtime_addi = kwargs.get("runtime_additional_information", None) - info = runtime_addi[0] - - # Sampling knobs are declared in vllm_omni/model_extras/mammothmodal2_preview.py - # and routed via extra_body -> sampling_params.extra_args (surfaced here as - # ``sampling_extra_args``). Fall back to runtime_additional_information for the - # legacy bespoke-example path during the transition. - extra_args_list = kwargs.get("sampling_extra_args") or [] - extra_args = extra_args_list[0] if extra_args_list else {} - text_guidance_scale = float(extra_args.get("text_guidance_scale", info["text_guidance_scale"][0])) - cfg_range_val = extra_args.get("cfg_range", info["cfg_range"]) - cfg_range = float(cfg_range_val[0]), float(cfg_range_val[1]) - num_inference_steps = int(extra_args.get("num_inference_steps", info["num_inference_steps"][0])) - - negative_cond = info.get("negative_prompt_embeds") - negative_attention_mask = info.get("negative_prompt_attention_mask") - image_hw = info["image_height"][0], info["image_width"][0] - - # Split the AR hidden states into text / image conditions. The token ids that - # drive the split are sourced from the model config (see _split_ar_conditions), - # formerly supplied by the bespoke example via additional_information. Legacy - # fallback: ar2dit may have already produced the split conditions. - if "text_prompt_embeds" in info: - text_cond = info["text_prompt_embeds"] - image_cond = info["image_prompt_embeds"] - else: - text_cond, image_cond = self._split_ar_conditions( - full_hidden_states=info["full_hidden_states"], - full_token_ids=info["full_token_ids"], - answer_start_index=int(info["answer_start_index"][0]), - ) - - # Move to model device/dtype. - model_device = next(self.parameters()).device - if self.gen_image_condition_refiner is not None: - target_dtype = next(self.gen_image_condition_refiner.parameters()).dtype - else: - target_dtype = next(self.gen_transformer.parameters()).dtype - - def _ensure_2d(x: torch.Tensor, name: str) -> torch.Tensor: - if x.ndim == 3 and x.shape[0] == 1: - x = x[0] - if x.ndim != 2: - raise ValueError(f"Expected {name} to be 2D [T,H], got shape={tuple(x.shape)}") - return x - - text_cond = _ensure_2d(text_cond, "text_prompt_embeds") - image_cond = _ensure_2d(image_cond, "image_prompt_embeds") - if image_cond.shape[0] == 0: - answer_token_ids = info.get("full_token_ids", [])[int(info.get("answer_start_index", [0])[0]) :] - raise ValueError( - "MammothModa2 AR stage produced no visual-token hidden states; " - "the DiT stage requires at least one generated visual token. " - f"Generated token ids: {answer_token_ids[:32]}" - ) - text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous() - image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous() - - text_embeds = text_cond.unsqueeze(0) # [1, T_text, H] - text_attention_mask = torch.ones( - (1, text_embeds.shape[1]), - dtype=torch.bool, - device=text_embeds.device, - ) - - image_embeds = image_cond.unsqueeze(0) # [1, T_img, H] - image_attention_mask = torch.ones( - (1, image_embeds.shape[1]), - dtype=torch.bool, - device=image_embeds.device, - ) - - # Apply optional refiner ONLY on image condition tokens. - if self.gen_image_condition_refiner is not None and image_embeds.shape[1] > 0: - image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_attention_mask.bool()) - image_attention_mask = torch.ones( - image_embeds.shape[:2], - dtype=torch.bool, - device=image_embeds.device, - ) - - nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None) - if nested_image_embedder is None: - prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1) - prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1) - ar_image_embeds = None - ar_image_attention_mask = None - else: - prompt_embeds = text_embeds - prompt_attention_mask = text_attention_mask - ar_image_embeds = image_embeds - ar_image_attention_mask = image_attention_mask - - # Prepare negative prompt (for CFG). If none provided, fall back to unconditional. - negative_prompt_embeds = None - negative_prompt_attention_mask = None - if text_guidance_scale > 1.0: - if negative_cond is not None: - negative_cond = _ensure_2d(negative_cond, "negative_prompt_embeds") - negative_prompt_embeds = ( - negative_cond.to(device=model_device, dtype=target_dtype, non_blocking=True) - .contiguous() - .unsqueeze(0) - ) - if isinstance(negative_attention_mask, torch.Tensor): - neg_mask = negative_attention_mask - elif isinstance(negative_attention_mask, list): - neg_mask = torch.tensor(negative_attention_mask, dtype=torch.bool) - else: - neg_mask = None - if neg_mask is None: - negative_prompt_attention_mask = torch.ones( - (1, negative_prompt_embeds.shape[1]), - dtype=torch.bool, - device=negative_prompt_embeds.device, - ) - else: - neg_mask = neg_mask.to(device=negative_prompt_embeds.device, dtype=torch.bool) - if neg_mask.ndim == 1: - neg_mask = neg_mask.unsqueeze(0) - negative_prompt_attention_mask = neg_mask - else: - hidden_size = int(prompt_embeds.shape[-1]) - negative_prompt_embeds = torch.zeros( - (1, 0, hidden_size), - dtype=target_dtype, - device=prompt_embeds.device, - ) - negative_prompt_attention_mask = torch.zeros( - (1, 0), - dtype=torch.bool, - device=prompt_embeds.device, - ) - - # Output image size (px), passed from stage input processor. - height, width = image_hw - if height <= 0 or width <= 0: - raise ValueError(f"Invalid image size: {height}x{width}") - if height % 16 != 0 or width % 16 != 0: - raise ValueError(f"Image size must be multiples of 16, got {height}x{width}") - vae_scale_factor = 16 - - latent_channels = int(self.gen_transformer.config.in_channels) - shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor) - latents = randn_tensor(shape, device=prompt_embeds.device, dtype=prompt_embeds.dtype) - - scheduler = FlowMatchEulerDiscreteScheduler() - - scheduler.set_timesteps( - num_inference_steps=num_inference_steps, - device=prompt_embeds.device, - num_tokens=latents.shape[-2] * latents.shape[-1], - ) - - # Run diffusion loop (CFG supported when text_guidance_scale > 1.0) - total_steps = max(1, len(scheduler.timesteps)) - for i, t in enumerate(scheduler.timesteps): - timestep = t.expand(latents.shape[0]).to(latents.dtype) - model_pred = self.gen_transformer( - hidden_states=latents, - timestep=timestep, - text_hidden_states=prompt_embeds, - text_attention_mask=prompt_attention_mask, - ref_image_hidden_states=None, - ar_image_hidden_states=ar_image_embeds, - ar_image_attention_mask=ar_image_attention_mask, - freqs_cis=self.gen_freqs_cis, - ) - guidance_scale = text_guidance_scale if cfg_range[0] <= i / total_steps <= cfg_range[1] else 1.0 - if guidance_scale > 1.0 and negative_prompt_embeds is not None: - model_pred_uncond = self.gen_transformer( - hidden_states=latents, - timestep=timestep, - text_hidden_states=negative_prompt_embeds, - text_attention_mask=negative_prompt_attention_mask, - ref_image_hidden_states=None, - freqs_cis=self.gen_freqs_cis, - ) - model_pred = model_pred_uncond + guidance_scale * (model_pred - model_pred_uncond) - latents = scheduler.step(model_pred, t, latents, return_dict=False)[0] - latents = latents.to(dtype=prompt_embeds.dtype) - - # VAE decode - if self.gen_vae.config.scaling_factor is not None: - latents = latents / self.gen_vae.config.scaling_factor - if self.gen_vae.config.shift_factor is not None: - latents = latents + self.gen_vae.config.shift_factor - image = self.gen_vae.decode(latents, return_dict=False)[0] - - return OmniOutput( - text_hidden_states=inputs_embeds, # placeholder, not used by runner - multimodal_outputs=image, - intermediate_tensors=None, - ) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: # noqa: ARG002 - return None - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/rope_real.py b/vllm_omni/diffusion/models/mammoth_moda2/rope_real.py deleted file mode 100644 index 64cc43248..000000000 --- a/vllm_omni/diffusion/models/mammoth_moda2/rope_real.py +++ /dev/null @@ -1,252 +0,0 @@ -# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import numpy as np -import torch -from diffusers.models.embeddings import get_1d_rotary_pos_embed as _get_1d_rotary_pos_embed -from einops import repeat -from torch import nn - -from vllm_omni.platforms import current_omni_platform - - -def apply_real_rotary_emb(x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> torch.Tensor: - """ - Apply real-valued rotary embeddings to input tensor. - - Args: - x: Input tensor of shape [..., seq_len, num_heads, dim] or [..., seq_len, dim] - freqs_cos: Cosine frequencies of shape [batch, seq_len, dim] or [seq_len, dim] - freqs_sin: Sine frequencies of shape [batch, seq_len, dim] or [seq_len, dim] - - Returns: - Tensor with rotary embeddings applied - """ - # x: [batch, seq_len, num_heads, dim] or [batch, seq_len, dim] - # freqs_cos: [batch, seq_len, dim] or [seq_len, dim] - # freqs_sin: [batch, seq_len, dim] or [seq_len, dim] - - x_shape = x.shape - if len(x_shape) == 4: - batch, seq_len, num_heads, dim = x_shape - x_reshaped = x.view(batch, seq_len, num_heads, dim // 2, 2) - elif len(x_shape) == 3: - batch, seq_len, dim = x_shape - num_heads = None - x_reshaped = x.view(batch, seq_len, dim // 2, 2) - else: - raise ValueError(f"Unsupported x shape: {x.shape}") - - # freqs_cos/sin: [batch, seq_len, dim] or [seq_len, dim] - # Expand freqs_cos/sin to [batch, seq_len, dim] if needed - if freqs_cos.dim() == 2: - # [seq_len, dim] -> [1, seq_len, dim] - freqs_cos = freqs_cos.unsqueeze(0) - freqs_sin = freqs_sin.unsqueeze(0) - if freqs_cos.shape[0] == 1 and batch > 1: - freqs_cos = freqs_cos.expand(batch, -1, -1) - freqs_sin = freqs_sin.expand(batch, -1, -1) - - # Reshape freqs to [batch, seq_len, dim//2, 2] - freqs_cos_reshaped = freqs_cos.view(batch, seq_len, dim // 2, 2) - freqs_sin_reshaped = freqs_sin.view(batch, seq_len, dim // 2, 2) - - cos_1 = freqs_cos_reshaped[..., 0] # [batch, seq_len, dim//2] - cos_2 = freqs_cos_reshaped[..., 1] # [batch, seq_len, dim//2] - sin_1 = freqs_sin_reshaped[..., 0] # [batch, seq_len, dim//2] - sin_2 = freqs_sin_reshaped[..., 1] # [batch, seq_len, dim//2] - - # Broadcast cos/sin to match x_reshaped - if len(x_shape) == 4: - # [batch, seq_len, 1, dim//2] - cos_1 = cos_1.unsqueeze(2) - cos_2 = cos_2.unsqueeze(2) - sin_1 = sin_1.unsqueeze(2) - sin_2 = sin_2.unsqueeze(2) - - x1 = x_reshaped[..., 0] # [..., seq_len, num_heads, dim//2] or [..., seq_len, dim//2] - x2 = x_reshaped[..., 1] # same - - out1 = x1 * cos_1 - x2 * sin_1 - out2 = x1 * sin_2 + x2 * cos_2 - - out = torch.stack([out1, out2], dim=-1) - return out.view(*x_shape) - - -def get_1d_rotary_pos_embed_real( - dim: int, - pos: np.ndarray | int, - theta: float = 10000.0, - linear_factor: float = 1.0, - ntk_factor: float = 1.0, - freqs_dtype: torch.dtype = torch.float32, -): - freqs_cos, freqs_sin = _get_1d_rotary_pos_embed( - dim, - pos, - theta=theta, - use_real=True, - linear_factor=linear_factor, - ntk_factor=ntk_factor, - repeat_interleave_real=True, - freqs_dtype=freqs_dtype, - ) - return freqs_cos, freqs_sin - - -class RotaryPosEmbedReal(nn.Module): - def __init__( - self, theta: int, axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], patch_size: int = 2 - ): - super().__init__() - self.theta = theta - self.axes_dim = axes_dim - self.axes_lens = axes_lens - self.patch_size = patch_size - - @staticmethod - def get_freqs_real( - axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int - ) -> list[tuple[torch.Tensor, torch.Tensor]]: - freqs_real = [] - freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32 - for i, (d, e) in enumerate(zip(axes_dim, axes_lens)): - cos_emb, sin_emb = get_1d_rotary_pos_embed_real(d, e, theta=theta, freqs_dtype=freqs_dtype) - freqs_real.append((cos_emb, sin_emb)) - return freqs_real - - def _get_freqs_real( - self, freqs_real: list[tuple[torch.Tensor, torch.Tensor]], ids: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - device = ids.device - if ids.device.type == "mps": - ids = ids.to("cpu") - - cos_result = [] - sin_result = [] - for i in range(len(self.axes_dim)): - freqs_cos, freqs_sin = freqs_real[i] - freqs_cos = freqs_cos.to(ids.device) - freqs_sin = freqs_sin.to(ids.device) - index = ids[:, :, i : i + 1].repeat(1, 1, freqs_cos.shape[-1]).to(torch.int64) - cos_result.append(torch.gather(freqs_cos.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index)) - sin_result.append(torch.gather(freqs_sin.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index)) - - combined_cos = torch.cat(cos_result, dim=-1).to(device) - combined_sin = torch.cat(sin_result, dim=-1).to(device) - return combined_cos, combined_sin - - def forward( - self, freqs_real, attention_mask, l_effective_ref_img_len, l_effective_img_len, ref_img_sizes, img_sizes, device - ): - batch_size = len(attention_mask) - p = self.patch_size - - encoder_seq_len = attention_mask.shape[1] - l_effective_cap_len = attention_mask.sum(dim=1).tolist() - l_effective_cap_len = [int(len) for len in l_effective_cap_len] - seq_lengths = [ - int(cap_len + sum(ref_img_len) + img_len) - for cap_len, ref_img_len, img_len in zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len) - ] - - max_seq_len = max(seq_lengths) - max_ref_img_len = max([int(sum(ref_img_len)) for ref_img_len in l_effective_ref_img_len]) - max_img_len = int(max(l_effective_img_len)) - - # Create position IDs - position_ids = torch.zeros(batch_size, int(max_seq_len), 3, dtype=torch.int32, device=device) - - for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)): - # add text position ids - position_ids[i, :cap_seq_len] = repeat( - torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3" - ) - pe_shift = cap_seq_len - pe_shift_len = cap_seq_len - - if ref_img_sizes[i] is not None: - for ref_img_size, ref_img_len in zip(ref_img_sizes[i], l_effective_ref_img_len[i]): - H, W = ref_img_size - ref_H_tokens, ref_W_tokens = H // p, W // p - assert ref_H_tokens * ref_W_tokens == ref_img_len - # add image position ids - - row_ids = repeat( - torch.arange(ref_H_tokens, dtype=torch.int32, device=device), "h -> h w", w=ref_W_tokens - ).flatten() - col_ids = repeat( - torch.arange(ref_W_tokens, dtype=torch.int32, device=device), "w -> h w", h=ref_H_tokens - ).flatten() - position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = pe_shift - position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = row_ids - position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = col_ids - - pe_shift += max(ref_H_tokens, ref_W_tokens) - pe_shift_len += ref_img_len - - H, W = img_sizes[i] - H_tokens, W_tokens = H // p, W // p - assert H_tokens * W_tokens == l_effective_img_len[i] - - row_ids = repeat(torch.arange(H_tokens, dtype=torch.int32, device=device), "h -> h w", w=W_tokens).flatten() - col_ids = repeat(torch.arange(W_tokens, dtype=torch.int32, device=device), "w -> h w", h=H_tokens).flatten() - - assert pe_shift_len + l_effective_img_len[i] == seq_len - position_ids[i, pe_shift_len:seq_len, 0] = pe_shift - position_ids[i, pe_shift_len:seq_len, 1] = row_ids - position_ids[i, pe_shift_len:seq_len, 2] = col_ids - - # Get combined rotary embeddings (real version) - freqs_cos, freqs_sin = self._get_freqs_real(freqs_real, position_ids) - - # create separate rotary embeddings for captions and images - cap_freqs_cos = torch.zeros( - batch_size, encoder_seq_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype - ) - cap_freqs_sin = torch.zeros( - batch_size, encoder_seq_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype - ) - ref_img_freqs_cos = torch.zeros( - batch_size, max_ref_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype - ) - ref_img_freqs_sin = torch.zeros( - batch_size, max_ref_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype - ) - img_freqs_cos = torch.zeros(batch_size, max_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype) - img_freqs_sin = torch.zeros(batch_size, max_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype) - - for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate( - zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths) - ): - cap_freqs_cos[i, :cap_seq_len] = freqs_cos[i, :cap_seq_len] - cap_freqs_sin[i, :cap_seq_len] = freqs_sin[i, :cap_seq_len] - ref_img_freqs_cos[i, : sum(ref_img_len)] = freqs_cos[i, cap_seq_len : cap_seq_len + sum(ref_img_len)] - ref_img_freqs_sin[i, : sum(ref_img_len)] = freqs_sin[i, cap_seq_len : cap_seq_len + sum(ref_img_len)] - img_freqs_cos[i, :img_len] = freqs_cos[ - i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len - ] - img_freqs_sin[i, :img_len] = freqs_sin[ - i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len - ] - - return ( - (cap_freqs_cos, cap_freqs_sin), - (ref_img_freqs_cos, ref_img_freqs_sin), - (img_freqs_cos, img_freqs_sin), - (freqs_cos, freqs_sin), - l_effective_cap_len, - seq_lengths, - ) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/schedulers.py b/vllm_omni/diffusion/models/mammoth_moda2/schedulers.py deleted file mode 100644 index ac73301ce..000000000 --- a/vllm_omni/diffusion/models/mammoth_moda2/schedulers.py +++ /dev/null @@ -1,137 +0,0 @@ -# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates -# SPDX-License-Identifier: Apache-2.0 -# -# This file has been modified by ByteDance Ltd. on 2025-09-30. -# -# Original file was released under Apache-2.0, with the full license text -# available at https://www.apache.org/licenses/LICENSE-2.0 -# -# This modified file is released under the same license. -# -# --- Upstream header preserved below --- -# -# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np -import torch - - -@dataclass -class FlowMatchEulerDiscreteSchedulerOutput: - prev_sample: torch.FloatTensor - - -class FlowMatchEulerDiscreteScheduler: - order = 1 - - def __init__(self, num_train_timesteps: int = 1000, dynamic_time_shift: bool = True): - self.num_train_timesteps = int(num_train_timesteps) - self.dynamic_time_shift = bool(dynamic_time_shift) - - timesteps = torch.linspace(0, 1, self.num_train_timesteps + 1, dtype=torch.float32)[:-1] - self.timesteps = timesteps - self._timesteps = torch.cat([timesteps, torch.ones(1, dtype=timesteps.dtype)]) - - self._step_index: int | None = None - self._begin_index: int | None = None - - @property - def step_index(self) -> int | None: - return self._step_index - - @property - def begin_index(self) -> int | None: - return self._begin_index - - def set_begin_index(self, begin_index: int = 0) -> None: - self._begin_index = int(begin_index) - - def index_for_timestep(self, timestep: torch.Tensor, schedule_timesteps: torch.Tensor | None = None) -> int: - schedule_timesteps = self._timesteps if schedule_timesteps is None else schedule_timesteps - indices = (schedule_timesteps == timestep).nonzero() - pos = 1 if len(indices) > 1 else 0 - return int(indices[pos].item()) - - def set_timesteps( - self, - num_inference_steps: int | None = None, - device: str | torch.device | None = None, - timesteps: list[float] | None = None, - num_tokens: int | None = None, - ) -> None: - if timesteps is None: - if num_inference_steps is None: - raise ValueError("`num_inference_steps` must be provided when `timesteps` is None.") - timesteps_np = np.linspace(0, 1, num_inference_steps + 1, dtype=np.float32)[:-1] - if self.dynamic_time_shift and num_tokens is not None: - m = np.sqrt(float(num_tokens)) / 40.0 - timesteps_np = timesteps_np / (m - m * timesteps_np + timesteps_np) - else: - timesteps_np = np.asarray(timesteps, dtype=np.float32) - - timesteps_t = torch.from_numpy(timesteps_np).to(dtype=torch.float32, device=device) - self.timesteps = timesteps_t - self._timesteps = torch.cat([timesteps_t, torch.ones(1, device=timesteps_t.device, dtype=timesteps_t.dtype)]) - - self._step_index = None - self._begin_index = None - - def _init_step_index(self, timestep: float | torch.Tensor) -> None: - if self.begin_index is None: - if isinstance(timestep, torch.Tensor): - timestep = timestep.to(self.timesteps.device) - else: - timestep = torch.tensor(timestep, device=self.timesteps.device, dtype=self.timesteps.dtype) - self._step_index = self.index_for_timestep(timestep) - else: - self._step_index = self._begin_index - - def step( - self, - model_output: torch.FloatTensor, - timestep: float | torch.FloatTensor, - sample: torch.FloatTensor, - generator: torch.Generator | None = None, # noqa: ARG002 - kept for API compatibility - return_dict: bool = True, - ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor]: - if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): - raise ValueError( - "Integer timesteps (e.g. a loop index from enumerate) are not supported; " - "pass a value from `scheduler.timesteps` instead.", - ) - - if self.step_index is None: - self._init_step_index(timestep) - assert self._step_index is not None - - sample_fp32 = sample.to(torch.float32) - t = self._timesteps[self._step_index] - t_next = self._timesteps[self._step_index + 1] - - prev_sample = sample_fp32 + (t_next - t) * model_output - prev_sample = prev_sample.to(model_output.dtype) - - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) - - def __len__(self) -> int: - return self.num_train_timesteps diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py index 41f12f654..f8681408a 100644 --- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py +++ b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py @@ -81,8 +81,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov # Picked up by ``supports_audio_output`` in the diffusion engine so the # default stage metadata reports ``final_output_type="audio"`` and the - # ``multimodal_output`` payload includes the sample rate (mirrors the - # contract introduced for AudioX in #2077). + # ``multimodal_output`` payload includes the sample rate. support_audio_output: ClassVar[bool] = True audio_sample_rate: ClassVar[int] = 44100 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 63f9eb366..5178371fc 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -271,11 +271,6 @@ _DIFFUSION_MODELS = { "pipeline_sensenova_u1", "SenseNovaU1Pipeline", ), - "AudioXPipeline": ( - "audiox", - "pipeline_audiox", - "AudioXPipeline", - ), "HunyuanVideo15Pipeline": ( "hunyuan_video", "pipeline_hunyuan_video_1_5", @@ -362,7 +357,6 @@ DiffusionModelRegistry = _ModelRegistry( _NO_CACHE_ACCELERATION = { # Pipelines that do not support cache acceleration (cache_dit / tea_cache). "NextStep11Pipeline", - "AudioXPipeline", # π0 is a flow-matching VLA with a self-contained sample_actions loop and no # DiT-style ``.transformer`` block list, so cache_dit / tea_cache cannot apply # to it; list it here so a stray cache_backend override disables gracefully @@ -556,7 +550,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "MiniMaxH3Pipeline": "get_minimax_h3_post_process_func", "MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", - "AudioXPipeline": "get_audiox_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", "WanS2VPipeline": "get_wan22_s2v_post_process_func", "WanT2VDMD2Pipeline": "get_wan22_post_process_func", diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 2bcbecc33..48812324f 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -3649,7 +3649,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if reference_audios: gen_params.extra_args["audio_path"] = reference_audios[0] - # Generate image or audio (e.g. AudioX) via AsyncOmni + # Generate image or audio via AsyncOmni diffusion_engine = cast(AsyncOmni, self._diffusion_engine) stage_configs = list(getattr(diffusion_engine, "stage_configs", []) or []) default_sampling_params_list = get_default_sampling_params_list(diffusion_engine) diff --git a/vllm_omni/model_executor/models/mammoth_moda2/__init__.py b/vllm_omni/model_executor/models/mammoth_moda2/__init__.py deleted file mode 100644 index fb6738a29..000000000 --- a/vllm_omni/model_executor/models/mammoth_moda2/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# NOTE: Do not import model classes in this file. Importing any -# submodule in this package triggers __init__.py execution, and -# both the model registry and pipeline registry import submodules -# directly — heavy imports here would be loaded as a side effect -# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py deleted file mode 100644 index 68c4df022..000000000 --- a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py +++ /dev/null @@ -1,1032 +0,0 @@ -from __future__ import annotations - -from collections.abc import Callable, Iterable, Mapping -from itertools import islice -from typing import Any - -import torch -from torch import nn -from transformers import Qwen2Config, Qwen3VLProcessor -from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_kv_scale_name, -) -from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP -from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP -from vllm.model_executor.models.qwen2_5_vl import ( - Qwen2_5_VLDummyInputsBuilder, - Qwen2_5_VLForConditionalGeneration, - Qwen2_5_VLMultiModalProcessor, - Qwen2_5_VLProcessingInfo, -) -from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser -from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer, Qwen3MLP -from vllm.model_executor.models.qwen3_vl import ( - Qwen3VLDummyInputsBuilder, - Qwen3VLForConditionalGeneration, - Qwen3VLMultiModalProcessor, - Qwen3VLProcessingInfo, -) -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - PPMissingLayer, - WeightsMapper, - init_vllm_registered_model, - is_pp_missing_parameter, - make_layers, - maybe_prefix, -) -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.config import ( - patch_rope_parameters, - set_default_rope_theta, -) - -from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, is_interleaved -from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config - - -def _runtime_meta(runtime_info: dict[str, Any]) -> dict[str, Any]: - """Accept both legacy flat metadata and the canonical nested payload.""" - meta = runtime_info.get("meta") - return meta if isinstance(meta, dict) else runtime_info - - -def _first_int(value: Any) -> int: - if isinstance(value, (list, tuple)): - value = value[0] - return int(value) - - -def moe_enable(moe_type, layer_type, layer_idx) -> bool: - """Determine if MoE should be enabled for a specific layer type and index. - - Args: - moe_type (str): The MoE configuration string (e.g., "ffn", "ffn_attention-14:28"). - layer_type (str): The type of layer being checked (e.g., "ffn", "attention"). - layer_idx (int): The index of the current layer. - - Returns: - bool: True if MoE should be enabled for this layer, False otherwise. - """ - if ":" in moe_type: - # moe_type like ffn_attention-14:28 - moe_type, layers = moe_type.split("-") - start, end = [int(n) for n in layers.split(":")] - else: - start, end = 0, float("inf") - assert moe_type in ["none", "attention", "ffn", "ffn_attention"] - return layer_type in moe_type and start <= layer_idx < end - - -def moe_forward( - hidden_states: torch.Tensor, - und_expert: Callable[[torch.Tensor], torch.Tensor], - gen_expert: Callable[[torch.Tensor], torch.Tensor] | None, - gen_token_mask: torch.Tensor | None = None, -) -> torch.Tensor: - """Perform Mixture-of-Experts (MoE) routing and forward pass. - - This function routes tokens to either the understanding expert (`und_expert`) or the - generation expert (`gen_expert`) based on the `gen_token_mask`. - - Routing Logic: - - If `gen_expert` is None: All tokens go to `und_expert`. - - If `gen_token_mask` is None or all False: All tokens go to `und_expert`. - - If `gen_token_mask` is all True: All tokens go to `gen_expert`. - - Otherwise (mixed batch): - - Tokens where `gen_token_mask` is True go to `gen_expert`. - - Tokens where `gen_token_mask` is False go to `und_expert`. - - Results are concatenated and reordered to match the original input order. - - Args: - hidden_states (torch.Tensor): Input hidden states. Shape: `(batch_size, seq_len, hidden_size)` - or `(num_tokens, hidden_size)`. - und_expert (Callable): The expert module for understanding/text tokens. - Takes `(N, D)` tensor, returns `(N, D_out)`. - gen_expert (Callable | None): The expert module for generation/image tokens, or None. - If provided, takes `(N, D)` tensor, returns `(N, D_out)`. - gen_token_mask (torch.Tensor | None): Boolean mask indicating generation tokens. - Shape matches `hidden_states` (excluding feature dim). True for generation tokens. - - Returns: - torch.Tensor: The processed hidden states with the same shape as input (except potentially - different feature dimension if `D_out != D`). - """ - if gen_expert is None: - return und_expert(hidden_states) - - if gen_token_mask is None or not gen_token_mask.any(): - return und_expert(hidden_states) - if gen_token_mask.all(): - return gen_expert(hidden_states) - - if hidden_states.ndim == 2: - flat_hid = hidden_states - d_model = hidden_states.shape[-1] - total_tokens = hidden_states.shape[0] - elif hidden_states.ndim == 3: - d_model = hidden_states.shape[-1] - flat_hid = hidden_states.reshape(-1, d_model) # (B*L, D) - total_tokens = flat_hid.shape[0] - else: - raise ValueError(f"Unexpected hidden_states shape: {tuple(hidden_states.shape)}") - - # Validate before reshape to catch shape mismatches where numel() would - # coincidentally match after flattening dimensions of different sizes. - if gen_token_mask.numel() != total_tokens: # type: ignore[union-attr] - raise ValueError( - "gen_token_mask shape mismatch: " - f"mask={tuple(gen_token_mask.shape)}, hidden_states={tuple(hidden_states.shape)}" - ) - # mask: [num_tokens] or [B, L] -> flatten to [total_tokens] - flat_mask = gen_token_mask.reshape(-1) # type: ignore[union-attr] - gen_pos = torch.where(flat_mask)[0] - und_pos = torch.where(~flat_mask)[0] - permute_order = torch.cat([gen_pos, und_pos], dim=0) - inverse_order = torch.argsort(permute_order) - gen_token_num = int(flat_mask.sum().item()) - gen_hid, und_hid = flat_hid[permute_order].split([gen_token_num, total_tokens - gen_token_num], dim=0) - - # 1.1 Generation tokens (True) - gen_out = gen_expert(gen_hid) # (N_gen, D) - - # 1.2 Understanding tokens (False) - und_out = und_expert(und_hid) # (N_und, D) - out_dim = und_out.shape[-1] - - merged = torch.cat([gen_out, und_out], dim=0) - merged = merged[inverse_order] - - if hidden_states.ndim == 2: - return merged.view(total_tokens, out_dim).contiguous() - return merged.view(*hidden_states.shape[:-1], out_dim).contiguous() - - -class Mammothmoda2Processor(Qwen2_5_VLProcessor): - """Qwen2.5-VL Processor with MammothU tokenizer.""" - - tokenizer_class = ("MammothUTokenizer", None) - - -class MammothModa2ARProcessingInfo(Qwen2_5_VLProcessingInfo): - """Processes multi-modal information for MammothModa2 AR, returning the VL sub-configuration.""" - - def get_hf_config(self): - mammoth_cfg: Mammothmoda2Config = self.ctx.get_hf_config(Mammothmoda2Config) - llm_cfg = getattr(mammoth_cfg, "llm_config", None) - return llm_cfg - - def get_hf_processor(self, **kwargs: object) -> Mammothmoda2Processor: - return self.ctx.get_hf_processor( - Mammothmoda2Processor, - use_fast=kwargs.pop("use_fast", True), - **kwargs, - ) - - def get_supported_mm_limits(self) -> Mapping[str, int | None]: - # MammothModa2 currently supports only image input, not video. - return {"image": None} - - def get_data_parser(self) -> Qwen2VLMultiModalDataParser: - # vLLM >=0.16 expects the parser to be provided by ProcessingInfo, - # not by BaseMultiModalProcessor._get_data_parser. - return Qwen2VLMultiModalDataParser( - spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size, - expected_hidden_size=self._get_expected_hidden_size(), - ) - - -class MammothModa2ARDummyInputsBuilder(Qwen2_5_VLDummyInputsBuilder): - """Reuse Qwen2.5-VL's dummy input generation logic.""" - - -class MammothModa2ARMultiModalProcessor(Qwen2_5_VLMultiModalProcessor): - """Reuse Qwen2.5-VL's multi-modal processing,""" - - -class Mammoth2DecoderLayer(Qwen2DecoderLayer): - def __init__( - self, - config: Qwen2Config, - layer_idx: int, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - # Must patch rope parameters on config BEFORE calling super().__init__, - # because Qwen2DecoderLayer.__init__ creates Qwen2Attention using the - # config's rope settings. Patching afterwards would leave the attention - # module initialised with wrong parameters, causing shape mismatches in - # rotary_embedding at runtime. - patch_rope_parameters(config) - set_default_rope_theta(config, default_theta=1000000) - super().__init__(config, cache_config, quant_config, prefix) - - self.moe_enable = moe_enable(config.moe_type, "ffn", layer_idx) - if self.moe_enable: - self.gen_mlp = Qwen2MLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.gen_mlp", - ) - else: - self.gen_mlp = None - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - gen_token_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Self Attention - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - hidden_states = self.self_attn( - positions=positions, - hidden_states=hidden_states, - ) - - # Fully Connected - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = moe_forward(hidden_states, self.mlp, self.gen_mlp, gen_token_mask) - return hidden_states, residual - - -class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): - def __init__( - self, *, vllm_config: VllmConfig, prefix: str = "", decoder_layer_type: type[nn.Module] = Mammoth2DecoderLayer - ): - super().__init__() - - hf_config = vllm_config.model_config.hf_config - if hasattr(hf_config, "get_text_config"): - config = hf_config.get_text_config() - elif hasattr(hf_config, "text_config"): - config = hf_config.text_config - else: - config = hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - self.prefix = prefix - - if is_interleaved(vllm_config.model_config.hf_text_config): - assert config.max_window_layers == config.num_hidden_layers, ( - "Sliding window for some but all layers is not supported. " - f"This model uses sliding window but `max_window_layers` = {config.max_window_layers} " - f"is less than `num_hidden_layers` = {config.num_hidden_layers}. Please open an issue " - "to discuss this feature." - ) - - self.config = config - self.quant_config = quant_config - # NOTE: MammothModa2 supports extra generation vocabulary (for image tokens). - # Token ID range: [gen_vocab_start_index, gen_vocab_start_index + gen_vocab_size). - # vLLM sampler/processor expects "last dimension of logits == model_config.get_vocab_size()", - # so we output base+gen logits in compute_logits, and embeddings must accept these IDs. - self.extra_gen_vocab = bool(getattr(config, "extra_gen_vocab", False)) - # Starting index for generation tokens (used for gen_token_mask). - self.gen_vocab_start_index = getattr(hf_config, "gen_vocab_start_index", None) or getattr( - config, "gen_vocab_start_index", None - ) - self.gen_vocab_size = int(getattr(config, "gen_vocab_size", 0) or 0) - - self.base_vocab_size = int(self.gen_vocab_start_index) if self.extra_gen_vocab else int(config.vocab_size) - self.org_base_vocab_size = int(getattr(config, "base_vocab_size", self.base_vocab_size)) - # The configuration level (hf_text_config.vocab_size) has been extended to base+gen - # by the upstream config class. Use config.vocab_size as the total vocab size. - self.total_vocab_size = int(getattr(config, "vocab_size", self.base_vocab_size)) - - if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): - self.embed_tokens = VocabParallelEmbedding( - self.base_vocab_size, - config.hidden_size, - org_num_embeddings=self.org_base_vocab_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) - else: - self.embed_tokens = PPMissingLayer() - - if self.extra_gen_vocab: - if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): - self.gen_embed_tokens = VocabParallelEmbedding( - self.gen_vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.gen_embed_tokens", - ) - else: - self.gen_embed_tokens = PPMissingLayer() - else: - self.gen_embed_tokens = None - - if get_pp_group().is_last_rank: - self.lm_head = ParallelLMHead( - self.base_vocab_size, - config.hidden_size, - org_num_embeddings=self.org_base_vocab_size, - quant_config=quant_config, - prefix=f"{prefix}.lm_head", - ) - else: - self.lm_head = PPMissingLayer() - - if self.extra_gen_vocab: - if get_pp_group().is_last_rank: - self.gen_head = ParallelLMHead( - self.gen_vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.gen_head", - ) - else: - self.gen_head = PPMissingLayer() - else: - self.gen_head = None - - self.logits_processor = LogitsProcessor(self.base_vocab_size) - self.gen_logits_processor = LogitsProcessor(self.gen_vocab_size) if self.extra_gen_vocab else None - - decoder_layer_type = decoder_layer_type or Mammoth2DecoderLayer - - def _make_decoder_layer(*, prefix: str) -> nn.Module: - try: - layer_idx = int(prefix.rsplit(".", 1)[-1]) - except Exception: - layer_idx = 0 - return decoder_layer_type( - config=config, - layer_idx=layer_idx, - cache_config=cache_config, - quant_config=quant_config, - prefix=prefix, - ) - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - _make_decoder_layer, - prefix=f"{prefix}.layers", - ) - - def _make_empty_intermediate_tensors( - batch_size: int, - dtype: torch.dtype, - device: torch.device, - ) -> IntermediateTensors: - return IntermediateTensors( - { - "hidden_states": torch.zeros((batch_size, config.hidden_size), dtype=dtype, device=device), - "residual": torch.zeros((batch_size, config.hidden_size), dtype=dtype, device=device), - "gen_token_mask": torch.zeros((batch_size,), dtype=torch.bool, device=device), - } - ) - - self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors - if get_pp_group().is_last_rank: - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - else: - self.norm = PPMissingLayer() - - @property - def model(self) -> MammothModa2Qwen2ForCausalLM: - return self - - def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: - if not self.extra_gen_vocab or self.gen_embed_tokens is None: - return self.embed_tokens(input_ids) - - gen_mask = input_ids >= int(self.gen_vocab_start_index) - if not gen_mask.any(): - return self.embed_tokens(input_ids) - if gen_mask.all(): - gen_ids = input_ids - int(self.gen_vocab_start_index) - return self.gen_embed_tokens(gen_ids) - - flat_ids = input_ids.reshape(-1) - flat_mask = gen_mask.reshape(-1) - out = torch.empty( - (flat_ids.shape[0], self.config.hidden_size), - dtype=self.embed_tokens.weight.dtype, # type: ignore[attr-defined] - device=flat_ids.device, - ) - - base_pos = torch.where(~flat_mask)[0] - gen_pos = torch.where(flat_mask)[0] - if base_pos.numel() > 0: - out[base_pos] = self.embed_tokens(flat_ids[base_pos]) - if gen_pos.numel() > 0: - gen_ids = flat_ids[gen_pos] - int(self.gen_vocab_start_index) - out[gen_pos] = self.gen_embed_tokens(gen_ids) - return out.view(*input_ids.shape, -1).contiguous() - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.get_input_embeddings(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - deepstack_input_embeds: IntermediateTensors | None = None, - ) -> torch.Tensor | IntermediateTensors: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - assert input_ids is not None - hidden_states = self.get_input_embeddings(input_ids) - # gen_token_mask: True indicates image generation tokens, which use gen_mlp. - # In vLLM v1 path, only inputs_embeds might be provided, with input_ids set to None. - # In this case, gen tokens cannot be distinguished by ID, falling back to und_expert. - if self.gen_vocab_start_index is None or input_ids is None: - gen_token_mask = None - else: - gen_token_mask = input_ids >= self.gen_vocab_start_index - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - gen_token_mask = intermediate_tensors.tensors.get("gen_token_mask") - - for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer), - start=self.start_layer, - ): - hidden_states, residual = layer(positions, hidden_states, residual, gen_token_mask) - # The vision encoder's deepstack_visual_indexes select which vision - # layers produce these features. The resulting tensors are keyed - # 0..N-1 and, as in Qwen3-Omni, are injected into the first N - # language-model layers. - if deepstack_input_embeds is not None and idx < len(deepstack_input_embeds): - hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{idx}"] - - if not get_pp_group().is_last_rank: - tensors = {"hidden_states": hidden_states, "residual": residual} - if gen_token_mask is not None: - tensors["gen_token_mask"] = gen_token_mask - return IntermediateTensors(tensors) - - hidden_states, _ = self.norm(hidden_states, residual) - - return hidden_states - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: - if isinstance(self.lm_head, PPMissingLayer): - return None - base_logits = self.logits_processor(self.lm_head, hidden_states) - if not self.extra_gen_vocab: - return base_logits - if self.gen_head is None or isinstance(self.gen_head, PPMissingLayer): - return base_logits - assert self.gen_logits_processor is not None - gen_logits = self.gen_logits_processor(self.gen_head, hidden_states) - if base_logits is None or gen_logits is None: - return None - return torch.cat([base_logits, gen_logits], dim=-1) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - if self.quant_config is not None: - cache_scale_mapper = self.quant_config.get_cache_scale_mapper() - if cache_scale_mapper is not None: - weights = cache_scale_mapper.apply(weights) - - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - - name = name.replace(weight_name, param_name) - - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - if name.endswith("scale"): - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - if weight_loader == default_weight_loader: - weight_loader(param, loaded_weight) - else: - weight_loader(param, loaded_weight, shard_id) - break - else: - if name.endswith(".bias") and name not in params_dict: - continue - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - - loaded_params.add(name) - - return loaded_params - - -class Mammoth3DecoderLayer(Qwen3DecoderLayer): - """Qwen3 decoder layer with MammothModa2's generation MLP expert.""" - - def __init__( - self, - config, - layer_idx: int, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__(config, cache_config, quant_config, prefix) - self.moe_enable = moe_enable(config.moe_type, "ffn", layer_idx) - self.gen_mlp = ( - Qwen3MLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.gen_mlp", - ) - if self.moe_enable - else None - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - gen_token_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = moe_forward( - hidden_states, - self.mlp, - self.gen_mlp, - gen_token_mask if self.moe_enable else None, - ) - return hidden_states, residual - - -class MammothModa2Qwen3ForCausalLM(MammothModa2Qwen2ForCausalLM): - """MammothModa2 generation-aware language model using Qwen3 attention.""" - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__( - vllm_config=vllm_config, - prefix=prefix, - decoder_layer_type=Mammoth3DecoderLayer, - ) - - -@MULTIMODAL_REGISTRY.register_processor( - MammothModa2ARMultiModalProcessor, - info=MammothModa2ARProcessingInfo, - dummy_inputs=MammothModa2ARDummyInputsBuilder, -) -class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration): - """Replaces the language backbone with MoE within the Qwen2_5_VLForConditionalGeneration multi-modal framework.""" - - have_multimodal_outputs = True - - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={ - # Skip generation-side (DiT/VAE) weights as they do not belong to the AR stage. - "gen_image_condition_refiner.": None, - "gen_transformer.": None, - "gen_vae.": None, - # LLM backbone: checkpoint uses the llm_model.* prefix. - # Extra generation vocab (image tokens) weights: mapped separately to the vLLM language_model submodule. - "llm_model.model.language_model.gen_embed_tokens.": "language_model.gen_embed_tokens.", - "llm_model.gen_head.": "language_model.gen_head.", - "llm_model.model.language_model.": "language_model.", - "llm_model.model.visual.": "visual.", - "llm_model.lm_head.": "language_model.lm_head.", - } - ) - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - # Switch hf_config to the AR sub-config to ensure the Qwen2.5-VL path receives the correct type. - mammoth_cfg = vllm_config.model_config.hf_config - ar_hf_config = getattr(mammoth_cfg, "llm_config", mammoth_cfg) - ar_vllm_config = vllm_config.with_hf_config(ar_hf_config, architectures=vllm_config.model_config.architectures) - # Initialize multi-modal components like the vision tower first. - super().__init__(vllm_config=ar_vllm_config, prefix=prefix) - # Replace with the custom MoE language model. - lm_hf_config = getattr( - ar_vllm_config.model_config.hf_config, "text_config", ar_vllm_config.model_config.hf_config - ) - self.language_model = init_vllm_registered_model( - vllm_config=ar_vllm_config, - prefix=maybe_prefix(prefix, "language_model"), - hf_config=lm_hf_config, - architectures=["MammothModa2Qwen2ForCausalLM"], - ) - self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors - - # -------- t2i (AR grid) token constraints -------- - # Constraint logic depends on per-step sampling_metadata + runtime_additional_information. - # These are passed by the vllm-omni runner via kwargs, so caching them in the model is sufficient. - self._last_runtime_additional_information: list[dict[str, Any]] | None = None - - def _apply_t2i_token_constraints(self, logits: torch.Tensor) -> torch.Tensor: - """Applies per-request token constraints. - - - For T2I requests: constrain AR grid tokens (force EOL at row end and - restrict intra-row sampling to visual token range). - - For non-T2I (text/understanding/chat) requests: disallow sampling - from the extra generation vocabulary (image tokens) to prevent - accidentally emitting visual-token sequences. - """ - if logits is None or not isinstance(logits, torch.Tensor): - return logits - - runtime_infos = self._last_runtime_additional_information - - if runtime_infos is None: - # There is no runtime info in dummy/profile run - return logits - - neg_inf = -float("inf") - num_reqs = int(logits.shape[0]) - for i in range(num_reqs): - runtime_info = runtime_infos[i] if isinstance(runtime_infos[i], dict) else {} - meta = _runtime_meta(runtime_info) - omni_task = meta.get("omni_task") - if not isinstance(omni_task, list) or not omni_task or omni_task[0] != "t2i": - # Text/understanding/chat: forbid sampling from the extra gen vocab. - logits[i, self.language_model.base_vocab_size :] = neg_inf - continue - - ar_width = _first_int(meta["ar_width"]) - eol_token_id = _first_int(meta["eol_token_id"]) - visual_start = _first_int(meta["visual_token_start_id"]) - visual_end = _first_int(meta["visual_token_end_id"]) - generated_len = runtime_info["generated_len"] - - row = logits[i] - column_id = generated_len % (ar_width + 1) - if column_id == ar_width: - # End-of-row token: only allow eol. - eol_logit = row[eol_token_id].clone() - row.fill_(neg_inf) - row[eol_token_id] = eol_logit - else: - # Intra-row tokens: only allow visual tokens (explicitly forbid eol). - row[:visual_start] = neg_inf - row[visual_end + 1 :] = neg_inf - row[eol_token_id] = neg_inf - - return logits - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, - ): - # vllm-omni runner passes sampling_metadata and runtime_additional_information - # in each forward step. compute_logits is called immediately after - # forward, so caching here enables step-by-step dynamic token constraints. - runtime_infos = kwargs.get("runtime_additional_information") - self._last_runtime_additional_information = runtime_infos if isinstance(runtime_infos, list) else None - hidden_states = super().forward( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **kwargs, - ) - if isinstance(hidden_states, IntermediateTensors) and not get_pp_group().is_last_rank: - return hidden_states - # NOTE: gpu_model_runner._dummy_run performs hidden_states[logit_indices] after forward. - # We must ensure text_hidden_states is a torch.Tensor to avoid errors when - # indexing (which happens if it's a list/tuple). - if isinstance(hidden_states, IntermediateTensors): - text_hidden_states = hidden_states["hidden_states"] - out_intermediate_tensors = hidden_states - elif isinstance(hidden_states, list): - text_hidden_states = hidden_states[0] - out_intermediate_tensors = None - else: - text_hidden_states = hidden_states - out_intermediate_tensors = None - - return OmniOutput( - text_hidden_states=text_hidden_states, - multimodal_outputs={}, - intermediate_tensors=out_intermediate_tensors, - ) - - def compute_logits(self, hidden_states: torch.Tensor | OmniOutput): - if isinstance(hidden_states, OmniOutput): - hidden_states = hidden_states.text_hidden_states - logits = super().compute_logits(hidden_states) - if isinstance(logits, torch.Tensor): - logits = self._apply_t2i_token_constraints(logits) - return logits - - -class MammothModa2Qwen3ProcessingInfo(Qwen3VLProcessingInfo): - """Qwen3-VL processing using MammothModa2's nested AR configuration.""" - - def get_hf_config(self): - mammoth_cfg: Mammothmoda2Config = self.ctx.get_hf_config(Mammothmoda2Config) - return mammoth_cfg.llm_config - - def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor: - # The MammothModa2 repository does not ship tokenizer_config.json, so - # AutoProcessor otherwise resolves only the image processor and leaves - # Qwen3-VL's image placeholder undefined. - return self.ctx.get_hf_processor( - Qwen3VLProcessor, - use_fast=kwargs.pop("use_fast", True), - **kwargs, - ) - - def get_supported_mm_limits(self) -> Mapping[str, int | None]: - return {"image": None} - - -@MULTIMODAL_REGISTRY.register_processor( - Qwen3VLMultiModalProcessor, - info=MammothModa2Qwen3ProcessingInfo, - dummy_inputs=Qwen3VLDummyInputsBuilder, -) -class MammothModa2Qwen3ARForConditionalGeneration(Qwen3VLForConditionalGeneration): - """MammothModa2-Dev Qwen3-VL stage for understanding and generation.""" - - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={ - "llm_model.model.visual.": "visual.", - "llm_model.lm_head.": "language_model.lm_head.", - "llm_model.model.language_model.gen_embed_tokens.": "language_model.gen_embed_tokens.", - "llm_model.model.language_model.": "language_model.", - "llm_model.gen_head.": "language_model.gen_head.", - "gen_image_condition_refiner.": None, - "gen_transformer.": None, - "gen_vae.": None, - "gen_tokenizer.": None, - } - ) - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - mammoth_cfg = vllm_config.model_config.hf_config - ar_hf_config = getattr(mammoth_cfg, "llm_config", mammoth_cfg) - ar_vllm_config = vllm_config.with_hf_config( - ar_hf_config, - architectures=vllm_config.model_config.architectures, - ) - super().__init__(vllm_config=ar_vllm_config, prefix=prefix) - text_config = ar_hf_config.text_config - self.language_model = MammothModa2Qwen3ForCausalLM( - vllm_config=ar_vllm_config.with_hf_config(text_config), - prefix=maybe_prefix(prefix, "language_model"), - ) - self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors - self._last_runtime_additional_information: list[dict[str, Any]] | None = None - - def forward(self, *args, **kwargs): - runtime_infos = kwargs.get("runtime_additional_information") - self._last_runtime_additional_information = runtime_infos if isinstance(runtime_infos, list) else None - return super().forward(*args, **kwargs) - - def compute_logits(self, hidden_states: torch.Tensor): - logits = super().compute_logits(hidden_states) - if not isinstance(logits, torch.Tensor) or self._last_runtime_additional_information is None: - return logits - neg_inf = -float("inf") - for i, runtime_info in enumerate(self._last_runtime_additional_information): - meta = _runtime_meta(runtime_info) if isinstance(runtime_info, dict) else {} - if (meta.get("omni_task") or [None])[0] != "t2i": - logits[i, self.language_model.base_vocab_size :] = neg_inf - continue - ar_width = _first_int(meta["ar_width"]) - eol_token_id = _first_int(meta["eol_token_id"]) - visual_start = _first_int(meta["visual_token_start_id"]) - visual_end = _first_int(meta["visual_token_end_id"]) - row = logits[i] - if int(runtime_info["generated_len"]) % (ar_width + 1) == ar_width: - eol_logit = row[eol_token_id].clone() - row.fill_(neg_inf) - row[eol_token_id] = eol_logit - else: - row[:visual_start] = neg_inf - row[visual_end + 1 :] = neg_inf - row[eol_token_id] = neg_inf - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - - -@MULTIMODAL_REGISTRY.register_processor( - MammothModa2ARMultiModalProcessor, - info=MammothModa2ARProcessingInfo, - dummy_inputs=MammothModa2ARDummyInputsBuilder, -) -class MammothModa2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): - # Ensure vllm_omni/worker/gpu_model_runner.py's `extract_multimodal_outputs` follows - # the OmniOutput branch to retrieve text_hidden_states as a pure torch.Tensor, - # preventing errors in `hidden_states[logit_indices]` due to type mismatch (list/tuple). - have_multimodal_outputs = True - - multimodal_cpu_fields = {"image_grid_thw", "video_grid_thw"} - merge_by_field_config = True - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - # Consistent with Qwen2_5OmniForConditionalGeneration: instance-level flag. - self.have_multimodal_outputs = True - self.vllm_config = vllm_config - cfg = vllm_config.model_config.hf_config - self.model_stage = vllm_config.model_config.model_stage - self.multimodal_config = vllm_config.model_config.multimodal_config - - # For debugging/alignment with qwen2.5-omni: explicitly nullify unused stages. - self.ar = None - self.dit = None - self.vae = None - - if self.model_stage == "ar": - # AR stage: multi-modal + MoE text. - self.ar = init_vllm_registered_model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "ar"), - hf_config=cfg.llm_config if hasattr(cfg, "llm_config") else cfg.text_config, - architectures=[ - "MammothModa2Qwen3ARForConditionalGeneration" - if getattr(cfg.llm_config, "model_type", "") == "mammothmoda2_qwen3_vl" - else "MammothModa2ARForConditionalGeneration" - ], - ) - self.model = self.ar - elif self.model_stage == "dit": - self.dit = init_vllm_registered_model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "dit"), - # NOTE: init_vllm_registered_model -> VllmConfig.with_hf_config requires a - # transformers.PretrainedConfig; however, Mammothmoda2Config.gen_dit_config - # is a dict (diffusers config). The DiT stage hf_config still uses the - # top-level Mammothmoda2Config, and the DiT module reads its own - # gen_dit_config / gen_vae_config dicts. - hf_config=cfg, - architectures=["MammothModa2DiTPipeline"], - ) - self.model = self.dit - elif self.model_stage == "vae": - # Reserved: VAEs not implemented yet; raise explicit error. - raise NotImplementedError("MammothModa2 VAE stage not implemented yet.") - else: - raise ValueError(f"Unsupported model_stage: {self.model_stage}") - - # Expose intermediate tensor factory for PP if provided by the submodule. - self.make_empty_intermediate_tensors = getattr(self.model, "make_empty_intermediate_tensors", lambda: None) - - @classmethod - def get_placeholder_str(cls, modality: str, i: int): # noqa: ARG003 - return Qwen2_5_VLForConditionalGeneration.get_placeholder_str(modality, i) - - def get_language_model(self) -> nn.Module: - if hasattr(self.model, "get_language_model"): - return self.model.get_language_model() - return self.model - - def embed_input_ids( - self, - input_ids: torch.Tensor, - multimodal_embeddings=None, - *, - is_multimodal: torch.Tensor | None = None, - ) -> torch.Tensor: - if isinstance(self.model, MammothModa2Qwen3ARForConditionalGeneration): - return self.model.embed_input_ids( - input_ids, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - ) - return SupportsMultiModal.embed_input_ids( - self, - input_ids, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - ) - - def get_multimodal_embeddings(self, **kwargs: object): - # Backward compatibility: route through embed_multimodal. - return self.embed_multimodal(**kwargs) - - def embed_multimodal(self, **kwargs: object): - if hasattr(self.model, "embed_multimodal"): - return self.model.embed_multimodal(**kwargs) - if hasattr(self.model, "get_multimodal_embeddings"): - return self.model.get_multimodal_embeddings(**kwargs) - return [] - - def get_input_embeddings(self, input_ids: torch.Tensor, multimodal_embeddings=None) -> torch.Tensor: - if hasattr(self.model, "get_input_embeddings"): - return self.model.get_input_embeddings(input_ids, multimodal_embeddings=multimodal_embeddings) - # DiT stage does not consume token embeddings from `input_ids`; it uses - # condition embeddings passed via additional_information. - # However, vLLM's generation runner may still request token embeddings - # to populate `inputs_embeds` buffers, so we provide a dummy tensor. - if self.model_stage == "dit": - hidden_size = int(self.vllm_config.model_config.get_hidden_size()) - try: - target_dtype = next(self.model.parameters()).dtype - except StopIteration: - target_dtype = self.vllm_config.model_config.dtype - return torch.zeros( - (input_ids.numel(), hidden_size), - device=input_ids.device, - dtype=target_dtype, - ) - raise NotImplementedError("Underlying model does not implement get_input_embeddings") - - def forward(self, *args, **kwargs) -> OmniOutput | torch.Tensor: - out = self.model(*args, **kwargs) - if isinstance(out, OmniOutput): - return out - if isinstance(out, list): - out = out[0] - return OmniOutput(text_hidden_states=out, multimodal_outputs={}, intermediate_tensors=None) - - def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, *args, **kwargs): - if isinstance(hidden_states, OmniOutput): - hidden_states = hidden_states.text_hidden_states - if hasattr(self.model, "compute_logits"): - return self.model.compute_logits(hidden_states) - return None - - def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]: - if self.model_stage != "dit": - raise RuntimeError( - f"get_dummy_runtime_additional_information only valid for dit stage, got {self.model_stage}" - ) - if self.dit is None: - raise RuntimeError("dit stage model is not initialized") - if not hasattr(self.dit, "get_dummy_runtime_additional_information"): - raise AttributeError("dit model missing get_dummy_runtime_additional_information") - return self.dit.get_dummy_runtime_additional_information(num_reqs) - - def load_weights(self, weights): - if self.model_stage == "ar": - if self.ar is None or not hasattr(self.ar, "load_weights"): - return set() - loaded = self.ar.load_weights(weights) - return add_prefix_to_loaded_weights(loaded, "ar") - if self.model_stage == "dit": - if self.dit is None or not hasattr(self.dit, "load_weights"): - return set() - loaded = self.dit.load_weights(weights) - return add_prefix_to_loaded_weights(loaded, "dit") - if self.model_stage == "vae": - return set() - raise ValueError(f"Unsupported model_stage: {self.model_stage}") diff --git a/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py b/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py deleted file mode 100644 index 830c63aff..000000000 --- a/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py +++ /dev/null @@ -1,68 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MammothModa2 pipeline topology (frozen). - -Stage 0: AR — multimodal understanding + latent generation -Stage 1: DiT — latent → image - -For text/image understanding tasks (text output only), use MAMMOTH_MODA2_AR_PIPELINE. -""" - -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) - -_PROC = "vllm_omni.model_executor.stage_input_processors.mammoth_moda2" - -MAMMOTH_MODA2_PIPELINE = PipelineConfig( - model_type="mammoth_moda2", - default_deploy_config_name="mammoth_moda2.yaml", - model_arch="MammothModa2ForConditionalGeneration", - hf_architectures=("Mammothmoda2Model", "MammothModa2ForConditionalGeneration"), - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="ar", - execution_type=StageExecutionType.LLM_AR, - input_sources=(), - final_output=False, - owns_tokenizer=True, - requires_multimodal_data=True, - engine_output_type="latent", - ), - StagePipelineConfig( - stage_id=1, - model_stage="dit", - execution_type=StageExecutionType.LLM_GENERATION, - input_sources=(0,), - final_output=True, - final_output_type="image", - owns_tokenizer=False, - requires_multimodal_data=False, - engine_output_type="image", - custom_process_input_func=f"{_PROC}.ar2dit", - ), - ), -) - -# Single-stage AR variant for understanding / summarization tasks -MAMMOTH_MODA2_AR_PIPELINE = PipelineConfig( - model_type="mammoth_moda2_ar", - default_deploy_config_name="mammoth_moda2_ar.yaml", - model_arch="MammothModa2ForConditionalGeneration", - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="ar", - execution_type=StageExecutionType.LLM_AR, - input_sources=(), - final_output=True, - final_output_type="text", - owns_tokenizer=True, - requires_multimodal_data=True, - engine_output_type="text", - ), - ), -) diff --git a/vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py deleted file mode 100644 index afe38f015..000000000 --- a/vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ /dev/null @@ -1,15 +0,0 @@ -""" -Compatibility shim. - -The MammothModa2 DiT implementation lives under `vllm_omni.diffusion` to align -with other ARDiT structured models. We keep this module path so existing -OmniModelRegistry entries (and downstream code) keep working. -""" - -from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import ( # noqa: F401 - MammothModa2DiTPipeline, -) - -__all__ = [ - "MammothModa2DiTPipeline", -] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 0064f1c4a..3d2777ada 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -114,41 +114,6 @@ _OMNI_MODELS = { "omnivoice", "OmniVoiceModel", ), - "MammothModa2Qwen2ForCausalLM": ( - "mammoth_moda2", - "mammoth_moda2", - "MammothModa2Qwen2ForCausalLM", - ), - "MammothModa2ARForConditionalGeneration": ( - "mammoth_moda2", - "mammoth_moda2", - "MammothModa2ARForConditionalGeneration", - ), - "MammothModa2Qwen3ARForConditionalGeneration": ( - "mammoth_moda2", - "mammoth_moda2", - "MammothModa2Qwen3ARForConditionalGeneration", - ), - "MammothModa2Qwen3ForCausalLM": ( - "mammoth_moda2", - "mammoth_moda2", - "MammothModa2Qwen3ForCausalLM", - ), - "MammothModa2DiTPipeline": ( - "mammoth_moda2", - "pipeline_mammothmoda2_dit", - "MammothModa2DiTPipeline", - ), - "MammothModa2ForConditionalGeneration": ( - "mammoth_moda2", - "mammoth_moda2", - "MammothModa2ForConditionalGeneration", - ), - "Mammothmoda2Model": ( - "mammoth_moda2", - "mammoth_moda2", - "MammothModa2ForConditionalGeneration", - ), "Qwen3TTSForConditionalGeneration": ( "qwen3_tts", "qwen3_tts_talker", diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py deleted file mode 100644 index 7b23f8e49..000000000 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Stage input processor for MammothModa2 (AR -> DiT).""" - -from collections.abc import Mapping -from typing import Any - -from vllm.inputs import TextPrompt - -from vllm_omni.inputs.data import OmniTokensPrompt - - -def _as_dict(prompt: Any) -> dict[str, Any]: - """Coerce an original-stage prompt to a dict. - - It may arrive as a dict, a NamedTuple/object, or a bare string depending on - the calling flow (the shared text_to_image example vs the bespoke script). - """ - if isinstance(prompt, dict): - return prompt - if hasattr(prompt, "_asdict"): - return prompt._asdict() - if hasattr(prompt, "__dict__"): - return vars(prompt) - return {} - - -def _coerce_dim(value: Any, default: int) -> int: - try: - iv = int(value) - except (TypeError, ValueError): - return default - return iv if iv > 0 else default - - -def ar2dit( - source_outputs: list[Any], - prompts: OmniTokensPrompt | TextPrompt | list[OmniTokensPrompt | TextPrompt] | None = None, - _requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - """Convert AR stage outputs to DiT stage inputs.""" - ar_outputs = source_outputs - - # The shared text_to_image example forwards a single prompt (not a list); normalize - # so a lone dict isn't iterated as its keys. Mirrors glm_image.ar2diffusion. - if not isinstance(prompts, list): - prompts = [prompts] if prompts is not None else [{}] - - dit_inputs: list[OmniTokensPrompt] = [] - for i, ar_output in enumerate(ar_outputs): - prompt_dict = _as_dict(prompts[i] if i < len(prompts) else {}) - addi_info = prompt_dict.get("additional_information") or {} - mm_kwargs = prompt_dict.get("mm_processor_kwargs") or {} - - # Image size: prefer mm_processor_kwargs target_h/target_w (set by the serving - # layer), fall back to additional_information, then a 1024 default. - image_height = _coerce_dim( - mm_kwargs.get("target_h"), - _coerce_dim((addi_info.get("image_height") or [None])[0], 1024), - ) - image_width = _coerce_dim( - mm_kwargs.get("target_w"), - _coerce_dim((addi_info.get("image_width") or [None])[0], 1024), - ) - - # Sampling knobs arrive on the DiT stage via extra_body -> extra_args; these are - # defensive fallbacks (defaults mirror the former bespoke script's argparse). - text_guidance_scale = (addi_info.get("text_guidance_scale") or [9.0])[0] - cfg_range = addi_info.get("cfg_range") or [0.0, 1.0] - num_inference_steps = (addi_info.get("num_inference_steps") or [50])[0] - - prompt_token_ids = ar_output.prompt_token_ids - # exclude the last token because it has no corresponding hidden state - completion_output = ar_output.outputs[0] - gen_token_ids = completion_output.cumulative_token_ids[:-1] - full_token_ids = prompt_token_ids + gen_token_ids - - mm_output = getattr(completion_output, "multimodal_output", None) - if not isinstance(mm_output, Mapping) or "latent" not in mm_output: - raise ValueError( - "AR stage output missing latent multimodal output. " - f"request_id={getattr(ar_output, 'request_id', None)}, " - f"completion_has_mm={hasattr(completion_output, 'multimodal_output')}" - ) - full_hidden_states = mm_output["latent"] - hidden_total = int(full_hidden_states.shape[0]) - assert hidden_total == len(prompt_token_ids) + len(gen_token_ids), ( - f"Hidden states length mismatch: expected {len(prompt_token_ids) + len(gen_token_ids)}, got {hidden_total}" - ) - - # The text/image condition split is performed in the DiT pipeline, which sources - # the distinguishing token ids (gen_vocab_start_index, vision placeholder ids) - # from the model config. Pass through the raw AR hidden states + token ids and - # the question/answer boundary so the pipeline can reconstruct the masks. - additional_information = { - # float32 so the tensor crosses the stage boundary (the serializer uses - # numpy, which has no bf16); the DiT re-casts to the model dtype. - "full_hidden_states": full_hidden_states.float().contiguous(), - "full_token_ids": full_token_ids, - "answer_start_index": [len(prompt_token_ids)], - "image_height": [int(image_height)], - "image_width": [int(image_width)], - "text_guidance_scale": [float(text_guidance_scale)], - "cfg_range": [float(cfg_range[0]), float(cfg_range[1])], - "num_inference_steps": [int(num_inference_steps)], - } - - dit_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[0], - additional_information=additional_information, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - - return dit_inputs diff --git a/vllm_omni/model_extras/audiox.py b/vllm_omni/model_extras/audiox.py deleted file mode 100644 index 220972148..000000000 --- a/vllm_omni/model_extras/audiox.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -# Request params the AudioX pipeline reads from ``sampling_params.extra_args`` -# (see ``pipeline_audiox.py`` ``forward``). Declaring them here routes -# ``extra_body`` (online) / ``--extra-body`` (offline) keys into ``extra_args`` -# instead of silently dropping them. -AUDIOX_EXTRA_BODY_PARAMS = frozenset( - { - "audiox_task", # t2a/t2m/v2a/v2m/tv2a/tv2m - "seconds_start", - "seconds_total", - "sigma_min", - "sigma_max", - "cfg_rescale", - "video_path", # v2*/tv2* conditioning - "audio_path", # optional reference-audio conditioning - } -) -# Echoed back in diffusion output metadata (pipeline_audiox.py). -AUDIOX_EXTRA_OUTPUT_PARAMS = frozenset( - { - "audiox_task", - } -) diff --git a/vllm_omni/model_extras/mammothmodal2_preview.py b/vllm_omni/model_extras/mammothmodal2_preview.py deleted file mode 100644 index 96f3fb939..000000000 --- a/vllm_omni/model_extras/mammothmodal2_preview.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -from typing import Any - -# AR image-grid patch size (pixels per visual token). Mirrors the constant the -# bespoke MammothModa2 example used to convert image dimensions -> AR grid dims. -_PATCH_SIZE = 16 - -# Shared by Preview and Dev (see ``t2i_generation_config.json``). -_EOL_TOKEN_ID = 152064 -_VISUAL_TOKEN_START_ID = 152072 -_VISUAL_TOKEN_END_ID = 168456 - -_AR_SYSTEM_PROMPT = "You are a helpful image generator." - - -def build_text_to_image_prompt( - prompt: str, - negative_prompt: str | None, - height: int | None = None, - width: int | None = None, -) -> dict[str, Any]: - """Build the MammothModa2 AR-stage prompt for text-to-image generation. - - Reproduces the prompt string and the structural ``additional_information`` - that the former bespoke example (``run_mammothmoda2_t2i.py``) constructed, - deriving the AR image grid from ``height`` / ``width``. - - Model-specific sampling knobs (``text_guidance_scale``, ``cfg_range``, - ``num_inference_steps``) flow separately via ``extra_body`` -> ``extra_args``. - The structural token ids are included in ``additional_information`` so the - AR sampler can apply the same constraints for both Preview and Dev. - - MammothModa2 t2i uses classifier-free guidance via ``text_guidance_scale`` - and has no explicit negative-prompt path, so ``negative_prompt`` is accepted - for signature compatibility but not injected. - """ - h = height or 1024 - w = width or 1024 - ar_height = h // _PATCH_SIZE - ar_width = w // _PATCH_SIZE - - return { - "prompt": ( - f"<|im_start|>system\n{_AR_SYSTEM_PROMPT}<|im_end|>\n" - f"<|im_start|>user\n{prompt}<|im_end|>\n" - "<|im_start|>assistant\n" - f"<|image start|>{ar_width}*{ar_height}<|image token|>" - ), - "additional_information": { - "omni_task": ["t2i"], - "ar_width": [ar_width], - "ar_height": [ar_height], - "image_height": [h], - "image_width": [w], - "eol_token_id": [_EOL_TOKEN_ID], - "visual_token_start_id": [_VISUAL_TOKEN_START_ID], - "visual_token_end_id": [_VISUAL_TOKEN_END_ID], - }, - } - - -MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS = frozenset( - { - "text_guidance_scale", - "cfg_range", - # MammothModa2's DiT stage consumes inputs via the kwargs interface rather - # than OmniDiffusionRequest, so the standard --num-inference-steps flag does - # not reach it; it is routed through extra_body like the CFG knobs. - "num_inference_steps", - } -) - - -def build_x_to_text_prompt( - model: str, - prompt: str, - has_image: bool, -) -> tuple[dict[str, Any], list[int] | None]: - """Build MammothModa2's chat prompt for text-output tasks.""" - del model - vision = "<|vision_start|><|image_pad|><|vision_end|>" if has_image else "" - return ( - { - "prompt": ( - "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" - f"<|im_start|>user\n{vision}{prompt}<|im_end|>\n" - "<|im_start|>assistant\n" - ), - "modalities": ["text"], - "additional_information": {"omni_task": ["chat"]}, - }, - None, - ) - - -MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS = frozenset() -MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index c394b8999..0bd57cb02 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -8,10 +8,6 @@ from typing import Any, Literal, Protocol from PIL import Image -from vllm_omni.model_extras.audiox import ( - AUDIOX_EXTRA_BODY_PARAMS, - AUDIOX_EXTRA_OUTPUT_PARAMS, -) from vllm_omni.model_extras.bagel import ( BAGEL_EXTRA_BODY_PARAMS, BAGEL_EXTRA_OUTPUT_PARAMS, @@ -44,17 +40,6 @@ from vllm_omni.model_extras.magi_human import ( MAGI_HUMAN_EXTRA_BODY_PARAMS, MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, ) -from vllm_omni.model_extras.mammothmodal2_preview import ( - MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS, - MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS, - MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, -) -from vllm_omni.model_extras.mammothmodal2_preview import ( - build_text_to_image_prompt as build_mammothmoda2_text_to_image_prompt, -) -from vllm_omni.model_extras.mammothmodal2_preview import ( - build_x_to_text_prompt as build_mammothmoda2_x_to_text_prompt, -) from vllm_omni.model_extras.ming_flash_omni import ( MING_FLASH_OMNI_EXTRA_BODY_PARAMS, MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS, @@ -131,7 +116,6 @@ def default_x_to_text_prompt( _X_TO_TEXT_SPECS: dict[str, XToTextPromptBuilder] = { "bagel": build_bagel_x_to_text_prompt, "hunyuan_image3": build_hunyuan_x_to_text_prompt, - "mammoth_moda2": build_mammothmoda2_x_to_text_prompt, } @@ -146,8 +130,6 @@ def get_x_to_text_model_family(model: str) -> str: return "bagel" if model_type == "hunyuan_image_3_moe" or any("hunyuanimage3" in value for value in architectures): return "hunyuan_image3" - if "mammoth" in model_type or any("mammothmoda2" in value for value in architectures): - return "mammoth_moda2" return "generic" @@ -179,10 +161,6 @@ def default_image_to_image_prompt( _EXTRA_SPECS: dict[str, dict[str, Any]] = { - "AudioXPipeline": { - "extra_body_params": AUDIOX_EXTRA_BODY_PARAMS, - "extra_output_params": AUDIOX_EXTRA_OUTPUT_PARAMS, - }, "BagelPipeline": { "extra_body_params": BAGEL_EXTRA_BODY_PARAMS, "extra_output_params": BAGEL_EXTRA_OUTPUT_PARAMS, @@ -240,12 +218,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS, "image_to_video_prompt_builder": build_vace_image_to_video_prompt, }, - "MammothModa2DiTPipeline": { - "extra_body_params": MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS, - "extra_output_params": MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS, - "init_extra_args_for_non_diffusion_stages": MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, - "text_to_image_prompt_builder": build_mammothmoda2_text_to_image_prompt, - }, "MingImagePipeline": { "extra_body_params": MING_FLASH_OMNI_EXTRA_BODY_PARAMS, "extra_output_params": MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS, @@ -259,12 +231,6 @@ for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"): _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder -# Multi-stage discovery reports the top-level wrapper rather than its DiT -# submodule, so both names must resolve to the same request builders. -_EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] -_EXTRA_SPECS["Mammothmoda2Model"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] - - def _get_spec(model_class_name: str | None) -> dict[str, Any] | None: if not model_class_name: return None diff --git a/vllm_omni/outputs/output_metadata.py b/vllm_omni/outputs/output_metadata.py index f3e364a34..82faed499 100644 --- a/vllm_omni/outputs/output_metadata.py +++ b/vllm_omni/outputs/output_metadata.py @@ -130,7 +130,6 @@ def _validate_audio_metadata(metadata: DiffusionMetadataMapping | None) -> None: if metadata is None: return _check_optional_int(metadata, "sample_rate", group="audio", min_value=1) - _check_optional_non_empty_str(metadata, "audiox_task", group="audio") def _validate_actions_metadata(metadata: DiffusionMetadataMapping | None) -> None: diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 94c6f8916..0f9c155ee 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -850,9 +850,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin inputs_embeds = None # -------------------------------------- Omni-new ------------------------------------------------- - # Some generation-stage models (e.g. MammothModa2DiTPipeline) require - # model-specific runtime information (such as image size and conditioning - # embeddings) even during the dummy profiling run that vLLM uses to + # Some generation-stage models require model-specific runtime + # information (such as image size and conditioning embeddings) + # even during the dummy profiling run that vLLM uses to # estimate KV-cache capacity. get_dummy_runtime_additional_information # provides placeholder values of the correct shape so that the profiling # run does not raise an error due to missing inputs. diff --git a/vllm_omni/tokenizers/mammoth_moda2_tokenizer.py b/vllm_omni/tokenizers/mammoth_moda2_tokenizer.py deleted file mode 100644 index a419b2a92..000000000 --- a/vllm_omni/tokenizers/mammoth_moda2_tokenizer.py +++ /dev/null @@ -1,401 +0,0 @@ -# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates -# SPDX-License-Identifier: Apache-2.0 -# -# This file has been modified by ByteDance Ltd. on 2025-09-30. -# -# Original file was released under Apache-2.0, with the full license text -# available at https://www.apache.org/licenses/LICENSE-2.0 -# -# This modified file is released under the same license. -# -# --- Upstream header preserved below --- -# -# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -import base64 -import os -import unicodedata -from collections.abc import Collection - -import tiktoken -from loguru import logger -from transformers import AddedToken, PreTrainedTokenizer - -VOCAB_FILES_NAMES = { - "vocab_file": "mammothu.tiktoken", - "special_tokens_file": "mammothu_vision_tokens.txt", -} - -PAT_STR = ( - r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?""" - r"""[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""" -) -ENDOFTEXT = "<|endoftext|>" -IMSTART = "<|im_start|>" -IMEND = "<|im_end|>" -# as the default behavior is changed to allow special tokens in -# regular texts, the surface forms of special tokens need to be -# as different as possible to minimize the impact -QWEN_SPECIAL_TOKENS = ( - "<|object_ref_start|>", - "<|object_ref_end|>", - "<|box_start|>", - "<|box_end|>", - "<|quad_start|>", - "<|quad_end|>", - "<|vision_start|>", - "<|vision_end|>", - "<|vision_pad|>", - "<|image_pad|>", - "<|video_pad|>", - "", - "", - "<|fim_prefix|>", - "<|fim_middle|>", - "<|fim_suffix|>", - "<|fim_pad|>", - "<|repo_name|>", - "<|file_sep|>", -) - -# align to qwen2.5 tokenizer length (151846) -EXTRAS = [f"<|extra_{i}|>" for i in range(181)] # 205 - 19[len(QWEN_SPECIAL_TOKENS)] - 5 -# align to qwen2.5 embedding size (152064) -EXTRAS += [f"<|extra_margin_{i}|>" for i in range(152064 - 151846)] -# append new token in gen embedding range -EXTRAS += ["<|endofline|>", "<|endoffile|>", "<|gen_placeholder|>", "<|useless token|>", "<|beginoftext|>"] -EXTRAS = tuple(EXTRAS) -# changed to use actual index to avoid misconfiguration with vocabulary expansion -SPECIAL_START_ID = 151643 - - -def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> dict[bytes, int]: - with open(tiktoken_bpe_file, "rb") as f: - contents = f.read() - return { - base64.b64decode(token): int(rank) for token, rank in (line.split() for line in contents.splitlines() if line) - } - - -class MammothUTokenizer(PreTrainedTokenizer): - """MammothU tokenizer.""" - - vocab_files_names = VOCAB_FILES_NAMES - - def __init__( - self, - vocab_file: str, - special_tokens_file: str, - errors: str = "replace", - bos_token: str = "<|beginoftext|>", - eos_token: str = "<|endoftext|>", - pad_token: str = "<|endoftext|>", - img_token: str = "<|image token|>", - boi_token: str = "<|image start|>", - eoi_token: str = "<|image end|>", - eol_token: str = "<|endofline|>", - eof_token: str = "<|endoffile|>", - **kwargs, - ) -> None: - super().__init__(**kwargs) - - # how to handle errors in decoding UTF-8 byte sequences - # use ignore if you are in streaming inference - self.errors = errors - self.mergeable_ranks = _load_tiktoken_bpe(vocab_file) - with open(special_tokens_file) as f: - vision_tokens = [t.strip() for t in f.readlines() if len(t.strip()) > 0] - SPECIAL_TOKENS = tuple( - enumerate( - ( - ( - ENDOFTEXT, - IMSTART, - IMEND, - ) - + QWEN_SPECIAL_TOKENS - + EXTRAS - + tuple(vision_tokens) - ), - start=SPECIAL_START_ID, - ) - ) - - self.special_tokens = {token: index for index, token in SPECIAL_TOKENS} - self.special_tokens_set = set(t for _, t in SPECIAL_TOKENS) - - enc = tiktoken.Encoding( - "mammothu", - pat_str=PAT_STR, - mergeable_ranks=self.mergeable_ranks, - special_tokens=self.special_tokens, - ) - - assert len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab, ( - f"{len(self.mergeable_ranks) + len(self.special_tokens)} != {enc.n_vocab} in encoding" - ) - - self.decoder = {v: k for k, v in self.mergeable_ranks.items()} - self.decoder.update({v: k for k, v in self.special_tokens.items()}) - self.tokenizer = enc - self.eod_id = self.tokenizer.eot_token - self.bos_token = bos_token - self.eos_token = eos_token - self.pad_token = pad_token - self.img_token = img_token - self.boi_token = boi_token - self.eoi_token = eoi_token - self.eol_token = eol_token - self.eof_token = eof_token - self.image_content_token = "<|image_pad|>" # come from Qwen2.5-VL - self.gen_image_token = "<|gen_image_pad|>" - self.gen_image_placeholder_token = "<|gen_placeholder|>" - self.visual_tokens = ["<|image_pad|>", "<|video_pad|>", "<|vision_start|>", "<|vision_end|>"] - self.visual_tokens_ids = [self.get_vocab()[token] for token in self.visual_tokens] - - self.vision_range = (self.get_vocab()[self.boi_token], self.tokenizer.n_vocab - 1) - logger.info(f"MammothUTokeniser Vision range: {self.vision_range}") - - def __getstate__(self): - # for pickle lovers - state = self.__dict__.copy() - del state["tokenizer"] - return state - - def __setstate__(self, state): - # tokenizer is not python native; don't pass it; rebuild it - self.__dict__.update(state) - enc = tiktoken.Encoding( - "mammothu", - pat_str=PAT_STR, - mergeable_ranks=self.mergeable_ranks, - special_tokens=self.special_tokens, - ) - self.tokenizer = enc - - def __len__(self) -> int: - return self.tokenizer.n_vocab - - def get_vocab(self) -> dict[bytes | str, int]: - vocab = self.mergeable_ranks.copy() - vocab.update(self.special_tokens) - return vocab - - @property - def gen_placeholder_id(self): - return self.get_vocab()[self.gen_image_placeholder_token] - - def convert_tokens_to_ids(self, tokens: bytes | str | list[bytes | str]) -> list[int]: - if isinstance(tokens, (str, bytes)): - if tokens in self.special_tokens: - return self.special_tokens[tokens] - else: - return self.mergeable_ranks.get(tokens) - - ids = [] - for token in tokens: - if token in self.special_tokens: - ids.append(self.special_tokens[token]) - else: - ids.append(self.mergeable_ranks.get(token)) - return ids - - def _add_tokens( - self, - new_tokens: list[str] | list[AddedToken], - special_tokens: bool = False, - ) -> int: - if not special_tokens and new_tokens: - raise ValueError("Adding regular tokens is not supported") - - added_tokens = 0 - for token in new_tokens: - surface_form = token.content if isinstance(token, AddedToken) else token - if surface_form in self.special_tokens_set: - # Token already exists in our special tokens set - added_tokens += 1 - else: - logger.warning(f"Token {surface_form} is not in the predefined special tokens set and cannot be added") - - return added_tokens - - def add_special_tokens(self, special_tokens_dict: dict[str, str | AddedToken]) -> int: - """ - Add special tokens to the tokenizer and update the special tokens mapping. - Only adds tokens that are already in the special_tokens_set. - - Args: - special_tokens_dict: dictionary of special tokens to add. - The key is the token type and the value is the token to add. - - Returns: - Number of tokens added to the vocabulary. - """ - added_tokens = 0 - - for token_type, token in special_tokens_dict.items(): - if token_type == "additional_special_tokens" and isinstance(token, list): - added_tokens += self._add_tokens(token, special_tokens=True) - else: - token_value = token.content if isinstance(token, AddedToken) else token - if token_value in self.special_tokens_set: - setattr(self, token_type, token_value) - added_tokens += 1 - else: - logger.warning( - f"Token {token_value} is not in the predefined special tokens set and cannot be added" - ) - - return added_tokens - - def save_vocabulary(self, save_directory: str, **kwargs) -> tuple[str]: - """ - Save only the vocabulary of the tokenizer (vocabulary). - - Returns: - `tuple(str)`: Paths to the files saved. - """ - regular_file_path = os.path.join(save_directory, self.vocab_files_names["vocab_file"]) - with open(regular_file_path, "w", encoding="utf8") as w: - for k, v in self.mergeable_ranks.items(): - line = base64.b64encode(k).decode("utf8") + " " + str(v) + "\n" - w.write(line) - - excluded_special_tokens = set( - ( - ENDOFTEXT, - IMSTART, - IMEND, - ) - + EXTRAS - ) - special_file_path = os.path.join(save_directory, self.vocab_files_names["special_tokens_file"]) - with open(special_file_path, "w", encoding="utf8") as w: - for k in self.special_tokens: - if k not in excluded_special_tokens: - print(k, file=w) - - return (regular_file_path, special_file_path) - - def tokenize( - self, - text: str, - allowed_special: set | str = "all", - disallowed_special: Collection | str = (), - **kwargs, - ) -> list[bytes | str]: - """ - Converts a string in a sequence of tokens. - - Args: - text (`str`): - The sequence to be encoded. - allowed_special (`Literal["all"]` or `set`): - The surface forms of the tokens to be encoded as special tokens in regular texts. - Default to "all". - disallowed_special (`Literal["all"]` or `Collection`): - The surface forms of the tokens that should not be in regular texts and trigger errors. - Default to an empty tuple. - - kwargs (additional keyword arguments, *optional*): - Will be passed to the underlying model specific encode method. - - Returns: - `list[bytes|str]`: The list of tokens. - """ - tokens = [] - text = unicodedata.normalize("NFC", text) - - # this implementation takes a detour: text -> token id -> token surface forms - for t in self.tokenizer.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special): - tokens.append(self.decoder[t]) - - return tokens - - def convert_tokens_to_string(self, tokens: list[bytes | str]) -> str: - """ - Converts a sequence of tokens in a single string. - """ - text = "" - temp = b"" - for t in tokens: - if isinstance(t, str): - if temp: - text += temp.decode("utf-8", errors=self.errors) - temp = b"" - text += t - elif isinstance(t, bytes): - temp += t - else: - raise TypeError("token should only be of type types or str") - if temp: - text += temp.decode("utf-8", errors=self.errors) - return text - - @property - def vocab_size(self) -> int: - return self.tokenizer.n_vocab - - def _convert_id_to_token(self, index: int) -> bytes | str: - """Converts an id to a token, special tokens included""" - if index in self.decoder: - return self.decoder[index] - raise ValueError("unknown ids") - - def _convert_token_to_id(self, token: bytes | str) -> int: - """Converts a token to an id using the vocab, special tokens included""" - if token in self.special_tokens: - return self.special_tokens[token] - if token in self.mergeable_ranks: - return self.mergeable_ranks[token] - raise ValueError("unknown token") - - def _decode( - self, - token_ids: int | list[int], - skip_special_tokens: bool = False, - errors: str | None = None, - **kwargs, - ) -> str: - if isinstance(token_ids, int): - token_ids = [token_ids] - - if skip_special_tokens: - token_ids = [i for i in token_ids if i < self.eod_id] - - return self.tokenizer.decode(token_ids, errors=errors or self.errors) - - def bytes_to_str(self, byte_tokens: dict) -> str: - """Convert byte tokens to string representation. - - Args: - byte_tokens: A dictionary where keys are byte objects and values are integers, - or a single byte object. - - Returns: - If input is a dictionary, returns a new dictionary with byte keys converted to strings. - If input is a single byte object, returns the string representation. - """ - if isinstance(byte_tokens, dict): - return {k.decode("utf-8", errors=self.errors): v for k, v in byte_tokens.items() if isinstance(k, bytes)} - if isinstance(byte_tokens, bytes): - return byte_tokens.decode("utf-8", errors=self.errors) - return byte_tokens diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 42d6d8d45..2343ecd22 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -11,10 +11,6 @@ import importlib _CLASS_TO_MODULE: dict[str, str] = { "HiggsAudioV3Config": "vllm_omni.transformers_utils.configs.higgs_audio_v3", - "Mammothmoda2Config": "vllm_omni.transformers_utils.configs.mammoth_moda2", - "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", - "Mammothmoda2Qwen2_5_VLTextConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", - "Mammothmoda2Qwen2_5_VLVisionConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", "FishSpeechConfig": "vllm_omni.transformers_utils.configs.fish_speech", "FishSpeechSlowARConfig": "vllm_omni.transformers_utils.configs.fish_speech", "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", @@ -36,10 +32,6 @@ _CLASS_TO_MODULE: dict[str, str] = { __all__ = [ "HiggsAudioV3Config", - "Mammothmoda2Config", - "Mammothmoda2Qwen2_5_VLConfig", - "Mammothmoda2Qwen2_5_VLTextConfig", - "Mammothmoda2Qwen2_5_VLVisionConfig", "FishSpeechConfig", "FishSpeechSlowARConfig", "FishSpeechFastARConfig", @@ -80,7 +72,6 @@ from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3 # noqa: F401, E402 -from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 from vllm_omni.transformers_utils.configs import minimax_music3 as _minimax_music3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/mammoth_moda2.py b/vllm_omni/transformers_utils/configs/mammoth_moda2.py deleted file mode 100644 index d0848d1ca..000000000 --- a/vllm_omni/transformers_utils/configs/mammoth_moda2.py +++ /dev/null @@ -1,377 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import ClassVar, Literal - -from transformers import AutoConfig, AutoTokenizer, PretrainedConfig -from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( - Qwen2_5_VLConfig, - Qwen2_5_VLTextConfig, - Qwen2_5_VLVisionConfig, -) -from transformers.models.qwen3_vl.configuration_qwen3_vl import ( - Qwen3VLConfig, - Qwen3VLTextConfig, - Qwen3VLVisionConfig, -) - -from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer - -__all__ = [ - "Mammothmoda2Config", - "Mammothmoda2Qwen2_5_VLConfig", - "Mammothmoda2Qwen2_5_VLTextConfig", - "Mammothmoda2Qwen2_5_VLVisionConfig", - "Mammothmoda2Qwen3VLConfig", - "Mammothmoda2Qwen3VLTextConfig", - "Mammothmoda2Qwen3VLVisionConfig", -] - - -class Mammothmoda2Qwen2_5_VLVisionConfig(Qwen2_5_VLVisionConfig): - model_type = "mammothmoda2_qwen2_5_vl_vision" - base_config_key = "vision_config" - - def __init__( - self, - depth: int = 32, - hidden_size: int = 3584, - hidden_act: str = "silu", - intermediate_size: int = 3420, - num_heads: int = 16, - in_channels: int = 3, - patch_size: int = 14, - spatial_merge_size: int = 2, - temporal_patch_size: int = 2, - tokens_per_second: int = 4, - window_size: int = 112, - out_hidden_size: int = 3584, - fullatt_block_indexes: list[int] | None = None, - initializer_range: float = 0.02, - **kwargs, - ) -> None: - super().__init__( - depth=depth, - hidden_size=hidden_size, - hidden_act=hidden_act, - intermediate_size=intermediate_size, - num_heads=num_heads, - in_channels=in_channels, - patch_size=patch_size, - spatial_merge_size=spatial_merge_size, - temporal_patch_size=temporal_patch_size, - tokens_per_second=tokens_per_second, - window_size=window_size, - out_hidden_size=out_hidden_size, - fullatt_block_indexes=fullatt_block_indexes or [7, 15, 23, 31], - initializer_range=initializer_range, - **kwargs, - ) - - -class Mammothmoda2Qwen2_5_VLTextConfig(Qwen2_5_VLTextConfig): - model_type = "mammothmoda2_qwen2_5_vl_text" - base_config_key = "text_config" - keys_to_ignore_at_inference = ["past_key_values"] - base_model_tp_plan = Qwen2_5_VLTextConfig.base_model_tp_plan - base_model_pp_plan = Qwen2_5_VLTextConfig.base_model_pp_plan - - def __init__( - self, - vocab_size: int = 152064, - hidden_size: int = 8192, - intermediate_size: int = 29568, - num_hidden_layers: int = 80, - num_attention_heads: int = 64, - num_key_value_heads: int | None = 8, - hidden_act: str = "silu", - max_position_embeddings: int = 32768, - initializer_range: float = 0.02, - rms_norm_eps: float = 1e-05, - use_cache: bool = True, - tie_word_embeddings: bool = False, - rope_theta: float = 1000000.0, - use_sliding_window: bool = False, - sliding_window: int = 4096, - max_window_layers: int = 80, - layer_types: list[str] | None = None, - attention_dropout: float = 0.0, - rope_scaling: dict | None = None, - image_token_id: int | None = None, - video_token_id: int | None = None, - extra_gen_vocab: bool = True, - gen_vocab_size: int = 32800, - gen_vocab_start_index: int | None = None, - moe_type: str = "ffn", - **kwargs, - ) -> None: - super().__init__( - vocab_size=vocab_size, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - num_hidden_layers=num_hidden_layers, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - hidden_act=hidden_act, - max_position_embeddings=max_position_embeddings, - initializer_range=initializer_range, - rms_norm_eps=rms_norm_eps, - use_cache=use_cache, - tie_word_embeddings=tie_word_embeddings, - rope_theta=rope_theta, - use_sliding_window=use_sliding_window, - sliding_window=sliding_window, - max_window_layers=max_window_layers, - layer_types=layer_types, - attention_dropout=attention_dropout, - rope_scaling=rope_scaling, - **kwargs, - ) - - self.extra_gen_vocab = extra_gen_vocab - self.gen_vocab_size = gen_vocab_size - self.moe_type = moe_type - if gen_vocab_start_index is None: - self.gen_vocab_start_index = ( - self.vocab_size if self.extra_gen_vocab else self.vocab_size - self.gen_vocab_size - ) - else: - self.gen_vocab_start_index = gen_vocab_start_index - - # NOTE: vLLM uses `hf_text_config.vocab_size` for sampling parameter validation - # (e.g., allowed_token_ids). Although MammothModa2's gen vocab is implemented via - # independent gen_embed/gen_head, the overall vocab size should still cover the - # gen vocab token ID range from the perspective of "output logits dimension". - if self.extra_gen_vocab: - self.vocab_size = int(self.gen_vocab_start_index) + int(self.gen_vocab_size) - - # Extra token IDs for multi-modal placeholders. - self.image_token_id = image_token_id - self.video_token_id = video_token_id - - -class Mammothmoda2Qwen2_5_VLConfig(Qwen2_5_VLConfig): - """Combined configuration: text_config + vision_config.""" - - model_type = "mammothmoda2_qwen2_5_vl" - sub_configs = { - "vision_config": Mammothmoda2Qwen2_5_VLVisionConfig, - "text_config": Mammothmoda2Qwen2_5_VLTextConfig, - } - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - text_config: dict | PretrainedConfig | None = None, - vision_config: dict | PretrainedConfig | None = None, - image_token_id: int = 151655, - video_token_id: int = 151656, - vision_start_token_id: int = 151652, - vision_end_token_id: int = 151653, - extra_gen_vocab: bool = True, - gen_vocab_size: int = 32800, - gen_vocab_start_index: int | None = None, - moe_type: str = "ffn", - **kwargs, - ) -> None: - if isinstance(vision_config, dict): - vision_config = self.sub_configs["vision_config"](**vision_config) - elif vision_config is None: - vision_config = self.sub_configs["vision_config"]() - - text_extra_kwargs = { - "extra_gen_vocab": extra_gen_vocab, - "gen_vocab_size": gen_vocab_size, - "moe_type": moe_type, - "gen_vocab_start_index": gen_vocab_start_index, - } - if isinstance(text_config, dict): - for key, val in text_extra_kwargs.items(): - text_config.setdefault(key, val) - text_config = self.sub_configs["text_config"](**text_config) - elif text_config is None: - text_config = self.sub_configs["text_config"](**{**text_extra_kwargs, **kwargs}) - elif isinstance(text_config, PretrainedConfig): - text_config = text_config - - super().__init__( - text_config=text_config, - vision_config=vision_config, - image_token_id=image_token_id, - video_token_id=video_token_id, - vision_start_token_id=vision_start_token_id, - vision_end_token_id=vision_end_token_id, - **kwargs, - ) - - if not hasattr(self, "text_config"): - self.text_config = text_config - if not hasattr(self, "vision_config"): - self.vision_config = vision_config - - self.image_token_id = image_token_id - self.video_token_id = video_token_id - self.extra_gen_vocab = getattr(self.text_config, "extra_gen_vocab", extra_gen_vocab) - self.gen_vocab_size = getattr(self.text_config, "gen_vocab_size", gen_vocab_size) - self.moe_type = getattr(self.text_config, "moe_type", moe_type) - self.gen_vocab_start_index = getattr(self.text_config, "gen_vocab_start_index", gen_vocab_start_index) - self.tokenizer_class = "MammothUTokenizer" - - -class Mammothmoda2Qwen3VLVisionConfig(Qwen3VLVisionConfig): - model_type = "mammothmoda2_qwen3_vl_vision" - base_config_key = "vision_config" - - -class Mammothmoda2Qwen3VLTextConfig(Qwen3VLTextConfig): - model_type = "mammothmoda2_qwen3_vl_text" - base_config_key = "text_config" - - def __init__( - self, - extra_gen_vocab: bool = True, - gen_vocab_size: int = 32800, - gen_vocab_start_index: int | None = None, - moe_type: str = "ffn", - **kwargs, - ) -> None: - super().__init__(**kwargs) - self.base_vocab_size = int(self.vocab_size) - self.extra_gen_vocab = extra_gen_vocab - self.gen_vocab_size = gen_vocab_size - self.gen_vocab_start_index = self.vocab_size if gen_vocab_start_index is None else gen_vocab_start_index - self.moe_type = moe_type - if self.extra_gen_vocab: - self.vocab_size = int(self.gen_vocab_start_index) + int(self.gen_vocab_size) - - -class Mammothmoda2Qwen3VLConfig(Qwen3VLConfig): - """MammothModa2-Dev AR config backed by Qwen3-VL.""" - - model_type = "mammothmoda2_qwen3_vl" - sub_configs = { - "vision_config": Mammothmoda2Qwen3VLVisionConfig, - "text_config": Mammothmoda2Qwen3VLTextConfig, - } - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - text_config: dict | PretrainedConfig | None = None, - vision_config: dict | PretrainedConfig | None = None, - extra_gen_vocab: bool = True, - gen_vocab_size: int = 32800, - gen_vocab_start_index: int | None = None, - moe_type: str = "ffn", - **kwargs, - ) -> None: - if isinstance(vision_config, dict): - vision_config = self.sub_configs["vision_config"](**vision_config) - elif vision_config is None: - vision_config = self.sub_configs["vision_config"]() - extras = { - "extra_gen_vocab": extra_gen_vocab, - "gen_vocab_size": gen_vocab_size, - "gen_vocab_start_index": gen_vocab_start_index, - "moe_type": moe_type, - } - if isinstance(text_config, dict): - for key, value in extras.items(): - text_config.setdefault(key, value) - text_config = self.sub_configs["text_config"](**text_config) - elif text_config is None: - text_config = self.sub_configs["text_config"](**extras) - super().__init__(text_config=text_config, vision_config=vision_config, **kwargs) - self.extra_gen_vocab = getattr(self.text_config, "extra_gen_vocab", extra_gen_vocab) - self.gen_vocab_size = getattr(self.text_config, "gen_vocab_size", gen_vocab_size) - self.gen_vocab_start_index = getattr(self.text_config, "gen_vocab_start_index", gen_vocab_start_index) - self.moe_type = getattr(self.text_config, "moe_type", moe_type) - self.tokenizer_class = "MammothUTokenizer" - - -class Mammothmoda2Config(PretrainedConfig): - """Top-level MammothModa2 composition configuration""" - - model_type = "mammothmoda2" - is_composition = True - sub_configs: ClassVar = {"llm_config": AutoConfig} - - def __init__( - self, - *, - llm_config: dict | None = None, - gen_vae_config: dict | None = None, - gen_dit_config: dict | None = None, - gen_condition_mode: Literal["text", "image", "text_image"] = "image", - gen_image_condition_refiner_config: dict | None = None, - gen_axes_dim_rope: list[int] | None = None, - gen_axes_lens: list[int] | None = None, - gen_transport_config: dict | None = None, - initializer_range: float = 0.02, - architectures: list[str] | None = None, - **kwargs, - ) -> None: - # Must set llm_config before super().__init__() because the parent's - # validation calls get_text_config() which accesses self.llm_config - self.llm_config = AutoConfig.for_model(**llm_config) if llm_config is not None else None - self.gen_vae_config = gen_vae_config - self.gen_dit_config = gen_dit_config - - self.gen_condition_mode = gen_condition_mode - self.gen_image_condition_refiner_config = gen_image_condition_refiner_config - self.gen_axes_dim_rope = gen_axes_dim_rope or [40, 40, 40] - self.gen_axes_lens = gen_axes_lens or [10000, 10000, 10000] - self.gen_transport_config = gen_transport_config or {} - self.initializer_range = initializer_range - self.tokenizer_class = "MammothUTokenizer" - self.architectures = ["Mammothmoda2Model"] - - super().__init__(**kwargs) - - def get_text_config(self, decoder: bool = False) -> PretrainedConfig: # noqa: ARG002 - if self.llm_config is None: - return None - # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config - return self.llm_config.text_config - - def _require_llm_config(self) -> PretrainedConfig: - if self.llm_config is None: - raise AttributeError("Mammothmoda2Config.llm_config is None") - return self.llm_config - - # ---- Proxy attrs for vLLM multimodal/mrope helpers ---- - # vllm_omni/model_executor/layers/mrope.py expects these fields on `hf_config`. - # In MammothModa2, they live in the nested `llm_config` (VL config), so we - # expose them here to make the top-level composition config compatible. - @property - def vision_config(self): - return self._require_llm_config().vision_config - - @property - def image_token_id(self) -> int: - return int(self._require_llm_config().image_token_id) - - @property - def video_token_id(self) -> int: - return int(self._require_llm_config().video_token_id) - - @property - def vision_start_token_id(self) -> int: - return int(self._require_llm_config().vision_start_token_id) - - @property - def vision_end_token_id(self) -> int: - return int(self._require_llm_config().vision_end_token_id) - - -# Register model_type -> config class for AutoConfig -AutoConfig.register(Mammothmoda2Config.model_type, Mammothmoda2Config) -AutoConfig.register(Mammothmoda2Qwen2_5_VLConfig.model_type, Mammothmoda2Qwen2_5_VLConfig) -AutoConfig.register(Mammothmoda2Qwen2_5_VLTextConfig.model_type, Mammothmoda2Qwen2_5_VLTextConfig) -AutoConfig.register(Mammothmoda2Qwen2_5_VLVisionConfig.model_type, Mammothmoda2Qwen2_5_VLVisionConfig) -AutoConfig.register(Mammothmoda2Qwen3VLConfig.model_type, Mammothmoda2Qwen3VLConfig) -AutoConfig.register(Mammothmoda2Qwen3VLTextConfig.model_type, Mammothmoda2Qwen3VLTextConfig) -AutoConfig.register(Mammothmoda2Qwen3VLVisionConfig.model_type, Mammothmoda2Qwen3VLVisionConfig) -# Register tokenizer_type -> the configs & AutoTokenizer -AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer) -AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer) -AutoTokenizer.register(config_class=Mammothmoda2Qwen3VLConfig, slow_tokenizer_class=MammothUTokenizer) diff --git a/vllm_omni/transformers_utils/processors/audiox.py b/vllm_omni/transformers_utils/processors/audiox.py deleted file mode 100644 index b0e53cfc7..000000000 --- a/vllm_omni/transformers_utils/processors/audiox.py +++ /dev/null @@ -1,240 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Input transform utilities for the AudioX diffusion pipeline. - -Loads and normalizes the raw audio/video conditioning signals (file path / URL / -``data:`` URI / ``np.ndarray`` / ``torch.Tensor``) into the (channels, samples) and -[T, C, H, W] tensors the pipeline needs. The pipeline itself stays focused on model -forward + sampling logic. -""" - -from __future__ import annotations - -import base64 -import os -import tempfile -from typing import Any -from urllib.parse import urlparse -from urllib.request import urlopen - -import av -import numpy as np -import soundfile -import torch -import torch.nn.functional as F -import torchaudio.functional as taF -from einops import rearrange -from torchvision.io import read_image - -# AudioX task taxonomy. Tasks beginning with "v" require a video input; tasks containing -# "t" carry a text prompt. tv2*/v2* share the same conditioning pathways. -VIDEO_ONLY_TASKS = frozenset({"v2a", "v2m"}) -TEXT_VIDEO_TASKS = frozenset({"tv2a", "tv2m"}) -VIDEO_CONDITIONED_TASKS = VIDEO_ONLY_TASKS | TEXT_VIDEO_TASKS - -_IMAGE_EXTS = frozenset({".jpg", ".jpeg", ".png"}) - - -def normalize_prompts(prompts: list[Any]) -> list[dict[str, Any]]: - """Coerce raw prompt entries into ``{"prompt": str, ...}`` dicts (preserves extras).""" - out: list[dict[str, Any]] = [] - for i, raw in enumerate(prompts): - if isinstance(raw, str): - out.append({"prompt": raw.strip()}) - elif isinstance(raw, dict): - p = dict(raw) - p["prompt"] = str(p.get("prompt") or "").strip() - out.append(p) - else: - raise TypeError(f"AudioX prompt {i} must be str or dict, got {type(raw)!r}") - return out - - -def materialize_media_source(source: str) -> str: - """Return a local filesystem path for ``source``. - - Accepts a local path, a ``data:;base64,...`` URI, or an ``http(s)://`` URL. - Anything non-local is fetched into a NamedTemporaryFile and that path is returned; - callers don't need to clean the tempfile up (the OS does on exit). - """ - if source.startswith("data:"): - _, _, payload = source.partition(",") - raw = base64.b64decode(payload) - f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False) - f.write(raw) - f.close() - return f.name - parsed = urlparse(source) - if parsed.scheme in ("http", "https"): - with urlopen(source) as resp: - data = resp.read() - f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False) - f.write(data) - f.close() - return f.name - return source - - -def _load_video_path_pyav( - path: str, - *, - target_fps: int, - duration: float, - seek_time: float, -) -> torch.Tensor: - path = materialize_media_source(path) - seek_time = float(seek_time) - duration = float(duration) - end_time = seek_time + duration if duration > 0 else None - - with av.open(path) as container: - stream = container.streams.video[0] - time_base = float(stream.time_base) - src_fps = float(stream.average_rate or target_fps) - - if seek_time > 0: - container.seek(int(seek_time / time_base), stream=stream) - - frames: list[np.ndarray] = [] - for frame in container.decode(stream): - t = float(frame.pts) * time_base - if t < seek_time: - continue - if end_time is not None and t >= end_time: - break - frames.append(frame.to_ndarray(format="rgb24")) - - if not frames: - raise ValueError(f"No frames in range seek_time={seek_time!r}, duration={duration!r} for {path!r}") - - # PyAV gives [H, W, C] uint8 RGB per frame; AudioX expects [T, C, H, W]. - video = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).contiguous() - n_target = max(1, int(round(video.shape[0] * float(target_fps) / src_fps))) - if n_target >= video.shape[0]: - return video - indices = torch.linspace(0, video.shape[0] - 1, n_target).round().long() - return video[indices] - - -def load_video_source( - source: Any, - *, - target_fps: int, - duration: float, - seek_time: float = 0.0, -) -> torch.Tensor: - if isinstance(source, str): - ext = os.path.splitext(source)[1].lower() - if ext in _IMAGE_EXTS: - return read_image(materialize_media_source(source)).unsqueeze(0) - return _load_video_path_pyav( - source, - target_fps=target_fps, - duration=duration, - seek_time=seek_time, - ) - - if isinstance(source, torch.Tensor): - return source - if isinstance(source, np.ndarray): - return torch.from_numpy(source) - raise TypeError(f"Unsupported video source type: {type(source)!r}") - - -def normalize_video_tensor(frames: torch.Tensor, size: int = 224) -> torch.Tensor: - if frames.dim() != 4: - raise ValueError(f"Expected [T, C, H, W], got {tuple(frames.shape)}") - - frames = frames.float() - if frames.max() > 1.5: - frames = frames / 255.0 - - if frames.shape[-2:] != (size, size): - frames = F.interpolate(frames, size=(size, size), mode="bicubic", align_corners=False) - return frames - - -def adjust_video_duration(frames: torch.Tensor, duration: float, target_fps: int) -> torch.Tensor: - target_t = int(duration * target_fps) - cur_t = frames.shape[0] - - if cur_t > target_t: - return frames[:target_t] - if cur_t < target_t: - last = frames[-1:].repeat(target_t - cur_t, 1, 1, 1) - return torch.cat([frames, last], dim=0) - return frames - - -def prepare_video_reference( - source: Any, - *, - duration: float, - target_fps: int, - seek_time: float = 0.0, -) -> torch.Tensor: - """Decode a video clip (or single image) into the AudioX [T, 3, 224, 224] form.""" - frames = load_video_source( - source, - target_fps=target_fps, - duration=duration, - seek_time=seek_time, - ) - - if frames.dim() == 4 and frames.shape[-1] == 3: - frames = rearrange(frames, "t h w c -> t c h w") - - frames = normalize_video_tensor(frames, size=224) - if duration > 0: - frames = adjust_video_duration(frames, duration, target_fps) - return frames - - -def prepare_audio_reference( - source: Any, - *, - model_sample_rate: int, - seconds_start: float, - seconds_total: float, - device: torch.device, -) -> torch.Tensor: - """Decode an audio source into a stereo (2, samples) tensor at the model's rate.""" - target_len = int(model_sample_rate * seconds_total) - start = int(model_sample_rate * seconds_start) - if isinstance(source, str): - data, sr = soundfile.read(materialize_media_source(source), dtype="float32", always_2d=True) - # soundfile returns channels-last (T, C); project convention is (C, T). - wav = torch.from_numpy(data).transpose(0, 1).contiguous() - if sr != model_sample_rate: - wav = taF.resample(wav, sr, model_sample_rate) - elif isinstance(source, torch.Tensor): - wav = source.float() - elif isinstance(source, np.ndarray): - wav = torch.from_numpy(source).float() - else: - raise TypeError(f"Unsupported audio source type: {type(source)!r}") - if wav.dim() == 1: - wav = wav.unsqueeze(0) - if wav.shape[0] == 1: - wav = wav.repeat(2, 1) - elif wav.shape[0] > 2: - wav = wav[:2] - wav = wav[:, start : start + target_len] - if wav.shape[1] < target_len: - wav = F.pad(wav, (target_len - wav.shape[1], 0)) - return wav.to(device=device, dtype=torch.float32) - - -__all__ = [ - "VIDEO_ONLY_TASKS", - "TEXT_VIDEO_TASKS", - "VIDEO_CONDITIONED_TASKS", - "normalize_prompts", - "materialize_media_source", - "load_video_source", - "normalize_video_tensor", - "adjust_video_duration", - "prepare_video_reference", - "prepare_audio_reference", -] diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index a397c2e9d..19a6bdb93 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -815,9 +815,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None - # Some generation-stage models (e.g. MammothModa2DiTPipeline) require - # model-specific runtime information (such as image size and conditioning - # embeddings) even during the dummy profiling run that vLLM uses to + # Some generation-stage models require model-specific runtime + # information (such as image size and conditioning embeddings) + # even during the dummy profiling run that vLLM uses to # estimate KV-cache capacity. get_dummy_runtime_additional_information # provides placeholder values of the correct shape so that the profiling # run does not raise an error due to missing inputs. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index a4716d682..9a46ab4e2 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1296,13 +1296,9 @@ class OmniGPUModelRunner(GPUModelRunner): per_req_runtime_info = [] for req_id in self.input_batch.req_ids: req_state = self.requests.get(req_id) - # MammothModa2 AR grid constraint: the model must emit a special - # end-of-line (EOL) token at the end of each image row. To determine - # whether the current decoding step falls on a row boundary, the - # constraint logic (see MammothModa2ARForConditionalGeneration. - # _apply_t2i_token_constraints) computes: - # column_id = generated_len % (ar_width + 1) - # and forces the EOL token when column_id == ar_width. + # Per-request runtime metadata is stamped with the number of tokens + # already generated so models can apply step-dependent constraints + # (for example end-of-row tokens during image-grid decoding). generated_len = len(req_state.output_token_ids) if req_state is not None else 0 info = self.model_intermediate_buffer.get(req_id, {}) if info: