anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[CI] Remove AudioX and MammothModa2 support (#6353)

Browse Source 
Signed-off-by: wangyu <410167048@qq.com>
This commit is contained in:
wangyu authored and GitHub committed 2026-08-23 04:22:05 +00:00
1 parent 39c16d75a9
commit 1e2a6446f7
59 files changed
+42 -7029
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/api/README.md
-4
View File
@@ -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][]
docs/models/supported_models.md
-3
View File
@@ -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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
docs/user_guide/diffusion_features.md
-1
View File
@@ -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** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
examples/offline_inference/audex/README.md
+1 -1
View File
@@ -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
examples/offline_inference/audex/text_to_audio.py
+1 -1
View File
@@ -7,7 +7,7 @@ thinker generates interleaved 4-codebook <audiocodec_N> 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.)
examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py
-117
View File
@@ -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()
examples/offline_inference/text_to_audio/README.md
+4 -35
View File
@@ -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.
examples/offline_inference/text_to_audio/text_to_audio.py
+19 -91
View File
@@ -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"
examples/offline_inference/text_to_image/text_to_image.py
-11
View File
@@ -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]
examples/offline_inference/x_to_text/README.md
+3 -31
View File
@@ -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.
examples/offline_inference/x_to_text/x_to_text.py
+1 -4
View File
@@ -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:
examples/online_serving/audiox/README.md
-64
View File
@@ -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
```
examples/online_serving/audiox/openai_chat_client.py
-119
View File
@@ -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())
examples/online_serving/audiox/run_server.sh
-16
View File
@@ -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"
recipes/MammothModa2/MammothModa2.md
-212
View File
@@ -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.
recipes/README.md
-2
View File
@@ -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) |
recipes/audiox/AudioX.md
-130
View File
@@ -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: <https://zeyuet.github.io/AudioX/>
- vLLM-Omni weight bundle: <https://huggingface.co/zhangj1an/AudioX>
- 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.
tests/config/test_config_factory.py
-2
View File
@@ -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),
tests/diffusion/test_diffusion_output_metadata.py
-2
View File
@@ -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),
tests/e2e/offline_inference/test_audiox_model_expansion.py
-71
View File
@@ -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=<model-or-local-path>
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
tests/e2e/offline_inference/test_mammoth_moda2_expansion.py
-204
View File
@@ -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"
tests/e2e/offline_inference/test_stable_audio_expansion.py
+1 -2
View File
@@ -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",
tests/e2e/online_serving/test_audiox_expansion.py
-63
View File
@@ -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)
tests/helpers/runtime.py
+1 -1
View File
@@ -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:
tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py
-119
View File
@@ -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]
tests/model_executor/models/test_kv_cache_scale_mapper.py
-5
View File
@@ -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",
tests/model_extras/test_model_extras.py
-119
View File
@@ -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],
},
}
tests/model_tests/diffusion/test_alignment.py
-1
View File
@@ -59,7 +59,6 @@ EXCLUDED_MODELS = [
"HeliosPyramidPipeline",
"DreamIDOmniPipeline",
"SenseNovaU1Pipeline",
"AudioXPipeline",
"HunyuanVideo15Pipeline",
"HunyuanVideo15ImageToVideoPipeline",
"LingBotVideoPipeline",
tests/worker/test_omni_connector_mixin.py
-1
View File
@@ -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",
]
vllm_omni/config/pipeline_registry.py
-6
View File
@@ -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,
vllm_omni/deploy/mammoth_moda2.yaml
-27
View File
@@ -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
vllm_omni/deploy/mammoth_moda2_ar.yaml
-14
View File
@@ -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
vllm_omni/diffusion/models/audiox/__init__.py
-11
View File
@@ -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",
]
vllm_omni/diffusion/models/audiox/audiox_transformer.py
-487
View File
@@ -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}")
vllm_omni/diffusion/models/audiox/pipeline_audiox.py
-933
View File
@@ -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,
)
vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py
-839
View File
@@ -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
vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py
-405
View File
@@ -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)
vllm_omni/diffusion/models/mammoth_moda2/rope_real.py
-252
View File
@@ -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,
)
vllm_omni/diffusion/models/mammoth_moda2/schedulers.py
-137
View File
@@ -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
vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py
+1 -2
View File
@@ -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
vllm_omni/diffusion/registry.py
-7
View File
@@ -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",
vllm_omni/entrypoints/openai/serving_chat.py
+1 -1
View File
@@ -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)
vllm_omni/model_executor/models/mammoth_moda2/__init__.py
-8
View File
@@ -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.
vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py
-1032
View File
File diff suppressed because it is too large. Load diff
vllm_omni/model_executor/models/mammoth_moda2/pipeline.py
-68
View File
@@ -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",
),
),
)
vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py
-15
View File
@@ -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",
]
vllm_omni/model_executor/models/registry.py
-35
View File
@@ -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",
vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py
-115
View File
@@ -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
vllm_omni/model_extras/audiox.py
-27
View File
@@ -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",
}
)
vllm_omni/model_extras/mammothmodal2_preview.py
-101
View File
@@ -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
vllm_omni/model_extras/registry.py
-34
View File
@@ -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
vllm_omni/outputs/output_metadata.py
-1
View File
@@ -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:
vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+3 -3
View File
@@ -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.
vllm_omni/tokenizers/mammoth_moda2_tokenizer.py
-401
View File
@@ -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|>",
"<tool_call>",
"</tool_call>",
"<|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
vllm_omni/transformers_utils/configs/__init__.py
-9
View File
@@ -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
vllm_omni/transformers_utils/configs/mammoth_moda2.py
-377
View File
@@ -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)
vllm_omni/transformers_utils/processors/audiox.py
-240
View File
@@ -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:<mime>;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",
]
vllm_omni/worker/gpu_generation_model_runner.py
+3 -3
View File
@@ -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.
vllm_omni/worker/gpu_model_runner.py
+3 -7
View File
@@ -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: