diff --git a/README.md b/README.md index 5063f9a3f..2cd6a574a 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ vLLM-Omni is flexible and easy to use with: vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including: - **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL) -- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3) +- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, CosyVoice3) - **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2) - **Robot-policy and action models** (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1) diff --git a/docs/README.md b/docs/README.md index 90a8b6304..3733a9d2d 100644 --- a/docs/README.md +++ b/docs/README.md @@ -56,7 +56,7 @@ vLLM-Omni is flexible and easy to use with: vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including: - Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL) -- TTS models (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3) +- TTS models (e.g. Qwen3-TTS, IndexTTS 2.5, CosyVoice3) - Diffusion models — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2) - Robot-policy and action models (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index ec895f794..1f543a6d2 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -73,7 +73,6 @@ th { | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — | | `BreezeForConditionalGeneration` | Breeze-TTS-2 (TTS, voice clone, voice direction) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) | -| `DotsTTSForConditionalGeneration` | dots.tts | `dots-studio/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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | @@ -107,7 +106,6 @@ th { | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — | | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — | | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — | -| `DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — | | `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenBMB/MiniCPM-o-4_5.md) | | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — | diff --git a/examples/model_display_names.yml b/examples/model_display_names.yml index 7a04ff8fa..be4ba2f80 100644 --- a/examples/model_display_names.yml +++ b/examples/model_display_names.yml @@ -2,7 +2,6 @@ # The docs generator combines these names with the serving mode for each page H1. aura_omni: AURA Omni covo_audio: Covo-Audio-Chat -dynin_omni: Dynin-Omni lance: Lance mimo_audio: MiMo-Audio ming_flash_omni: Ming-flash-omni 2.0 diff --git a/examples/offline_inference/dynin_omni/README.md b/examples/offline_inference/dynin_omni/README.md deleted file mode 100644 index 6cf834142..000000000 --- a/examples/offline_inference/dynin_omni/README.md +++ /dev/null @@ -1,110 +0,0 @@ -# Dynin-Omni: Offline inference - -This folder contains a unified offline inference entrypoint: - -- `end2end.py` - -## 1. Environment Setup - -Run from repository root: - -```bash -cd -``` - -If needed, install this repo in editable mode: - -```bash -pip install -e . -``` - -## 2. Extra Dependencies (EMOVA) - -Install the following packages for EMOVA-related components: - -```bash -pip install \ - "phonemizer==3.3.0" \ - "Unidecode==1.4.0" \ - "hydra-core==1.3.2" \ - "pytorch-lightning==1.1.0" \ - "wget==3.2" \ - "wrapt==2.1.1" \ - "onnx==1.20.1" \ - "frozendict==2.4.7" \ - "inflect==7.5.0" \ - "braceexpand==0.1.7" \ - "webdataset==1.0.2" \ - "torch-stft==0.1.4" \ - "editdistance==0.8.1" -``` - -## 3. Hardware and VRAM Requirements - -This example uses a 3-stage pipeline on one GPU by default -([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)): - -- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5` -- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1` -- Stage-2 (`token2audio`): `gpu_memory_utilization: 0.1` - -### Requested GPU Memory Budget from `gpu_memory_utilization` - -| Stage | Utilization | A100 80GB | H200 141GB | -| :-- | :-- | :-- | :-- | -| Stage-0 (token2text) | 0.5 | ~40.0 GB | ~70.5 GB | -| Stage-1 (token2image) | 0.1 | ~8.0 GB | ~14.1 GB | -| Stage-2 (token2audio) | 0.1 | ~8.0 GB | ~14.1 GB | -| Total requested budget | 0.7 | ~56.0 GB | ~98.7 GB | - -### Observed Runtime Signal (from your log) - -- Stage-0 reported: `Model loading took 15.12 GiB memory` (weights footprint signal). -- Stages 1/2 can still add runtime memory depending on task path and backend allocations. -- Keep extra headroom for CUDA/PyTorch overhead and temporary allocations. - -### GPU Compatibility - -- Confirmed target GPUs for this setup: **NVIDIA H200**, **NVIDIA A100**. -- CI/e2e coverage in this repo also includes CUDA **L4** markers for Dynin tests. - -## 4. End2End Run Examples - -```bash -# t2t -python /examples/offline_inference/dynin_omni/end2end.py \ - --task t2t --model snu-aidas/Dynin-Omni --text - -# i2t -python /examples/offline_inference/dynin_omni/end2end.py \ - --task i2t --model snu-aidas/Dynin-Omni --image --text "Please describe this image in detail." - -# s2t -python /examples/offline_inference/dynin_omni/end2end.py \ - --task s2t --model snu-aidas/Dynin-Omni --audio --text "Transcribe the given audio." - -# t2i -python /examples/offline_inference/dynin_omni/end2end.py \ - --task t2i --model snu-aidas/Dynin-Omni --text - -# v2t -python /examples/offline_inference/dynin_omni/end2end.py \ - --task v2t --model snu-aidas/Dynin-Omni --video --text "Describe this video in detail." - -# i2i -python /examples/offline_inference/dynin_omni/end2end.py \ - --task i2i --model snu-aidas/Dynin-Omni --image --text - -# t2s -python /examples/offline_inference/dynin_omni/end2end.py \ - --task t2s --model snu-aidas/Dynin-Omni --text -``` - -## 5. Notes - -- Outputs are saved under task-specific directories in `/tmp` by default. -- You can override output path with `--output-dir`. -- If you want to force local config resolution, pass `--dynin-config-path `. -- If you see the warning - `max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`, - reduce `max_num_batched_tokens` in the deploy config (for example, `4096` in the CI config). diff --git a/examples/offline_inference/dynin_omni/end2end.py b/examples/offline_inference/dynin_omni/end2end.py deleted file mode 100644 index 74d33fe01..000000000 --- a/examples/offline_inference/dynin_omni/end2end.py +++ /dev/null @@ -1,1448 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import argparse -import json -import os -import re -import sys -import time -import types -from importlib.machinery import ModuleSpec -from pathlib import Path -from typing import Any - -import numpy as np -import torch -from PIL import Image - -TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t") - -TASK_DEFAULT_RUNTIME = { - "t2t": ("mmu", "mmu", 0, "text"), - "t2i": ("t2i", "t2i_gen", 2, "image"), - "t2s": ("t2s_mmu_like", "t2s_gen", 1, "audio"), - "i2i": ("i2i", "i2i", 2, "image"), - "i2t": ("mmu", "mmu", 0, "text"), - "s2t": ("s2t", "s2t", 0, "text"), - "v2t": ("v2t", "v2t", 0, "text"), -} - -TASK_RUNTIME_FALLBACKS: dict[str, dict[str, Any]] = { - "t2t": { - "output_dir": "/tmp/dynin_end2end_outputs", - "prompt_max_text_len": 1024, - "max_new_tokens": 1024, - "steps": 1024, - "block_length": 16, - "temperature": 0.0, - "cfg_scale": 0.0, - }, - "t2i": { - "output_dir": "/tmp/dynin_t2i_outputs", - "prompt_max_text_len": 128, - "image_token_count": 1024, - "mask_token_id": 126336, - "codebook_size": 8192, - "timesteps": 20, - "guidance_scale": 3.5, - "temperature": 1.0, - }, - "i2i": { - "output_dir": "/tmp/dynin_i2i_outputs", - "prompt_max_text_len": 128, - "mask_token_id": 126336, - "codebook_size": 8192, - "timesteps": 64, - "guidance_scale": 3.5, - "temperature": 1.0, - "image_resolution": 336, - "use_train_i2i_prompt": True, - }, - "i2t": { - "output_dir": "/tmp/dynin_i2t_outputs", - "prompt_max_text_len": 128, - "max_new_tokens": 128, - "steps": 128, - "block_length": 2, - "temperature": 0.0, - "cfg_scale": 0.0, - "mask_token_id": 126336, - "codebook_size": 8192, - "image_resolution": 480, - "remasking": "low_confidence", - }, - "s2t": { - "output_dir": "/tmp/dynin_s2t_outputs", - "prompt_max_text_len": 1024, - "max_new_tokens": 128, - "steps": 128, - "block_length": 2, - "temperature": 0.0, - "cfg_scale": 0.0, - "mask_token_id": 126336, - "codebook_size": 8192, - "remasking": "low_confidence", - }, - "t2s": { - "output_dir": "/tmp/dynin_t2s_outputs", - "runtime_task": "t2s_mmu_like", - "prompting_task": "t2s_gen", - "prompt_max_text_len": 1024, - "t2s_token_length": 512, - "mask_token_id": 126336, - "codebook_size": 8192, - "audio_codebook_size": 4096, - "steps": 512, - "block_length": 128, - "temperature": 1.0, - "cfg_scale": 2.5, - "t2s_condition": "gender-female_emotion-neutral_speed-normal_pitch-normal", - }, - "v2t": { - "output_dir": "/tmp/dynin_v2t_outputs", - "prompt_max_text_len": 1024, - "max_new_tokens": 128, - "steps": 128, - "block_length": 2, - "temperature": 0.0, - "cfg_scale": 0.0, - "mask_token_id": 126336, - "codebook_size": 8192, - "image_resolution": 224, - "num_frames": 5, - "remasking": "low_confidence", - }, -} - -DEFAULT_I2T_QUESTION = "Please describe this image in detail." -DEFAULT_S2T_INSTRUCTION = "Transcribe the given audio." -DEFAULT_V2T_QUESTION = "Please provide a detailed description of the video." -DEFAULT_T2T_PROMPT = "Explain multimodal LLM inference in 3 sentences." -DEFAULT_T2S_INSTRUCTION = "Convert the given text into spoken audio." -DEFAULT_T2S_PROMPT = "Hello. This is a default text-to-speech sample." - -DYNIN_SPECIAL_TOKENS = ( - "<|soi|>", - "<|eoi|>", - "<|sov|>", - "<|eov|>", - "<|t2i|>", - "<|mmu|>", - "<|t2v|>", - "<|v2v|>", - "<|lvg|>", - "<|i2i|>", - "<|ti2ti|>", - "<|v2t|>", - "<|v2s|>", - "<|s2t|>", - "<|t2s|>", - "<|s2s|>", - "<|soa|>", - "<|eoa|>", -) - - -def bootstrap_repo_path() -> Path: - repo_root = Path(__file__).resolve().parents[3] - repo_root_str = str(repo_root) - if repo_root_str not in sys.path: - sys.path.insert(0, repo_root_str) - return repo_root - - -def ensure_safe_import_for_vllm() -> None: - os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1") - try: - import torchvision # noqa: F401 - - return - except Exception: - pass - - import enum - - class _InterpolationMode(enum.Enum): - NEAREST = 0 - BILINEAR = 2 - BICUBIC = 3 - LANCZOS = 1 - HAMMING = 4 - BOX = 5 - - tv_mod = types.ModuleType("torchvision") - tv_mod.__dict__["__version__"] = "0.0-stub" - tv_mod.__spec__ = ModuleSpec(name="torchvision", loader=None) - transforms_mod = types.ModuleType("torchvision.transforms") - transforms_mod.__spec__ = ModuleSpec(name="torchvision.transforms", loader=None) - transforms_mod.InterpolationMode = _InterpolationMode - tv_mod.transforms = transforms_mod - sys.modules["torchvision"] = tv_mod - sys.modules["torchvision.transforms"] = transforms_mod - - -def sanitize_repo_id(repo_id: str) -> str: - return re.sub(r"[^a-zA-Z0-9._-]+", "_", repo_id) - - -def is_hf_repo_id(value: str) -> bool: - return isinstance(value, str) and value.count("/") == 1 and all(value.split("/", 1)) - - -def ensure_local_model_dir(model: str, cache_dir: Path, localize: bool) -> Path: - model_path = Path(model).expanduser() - if model_path.is_dir(): - return model_path.resolve() - if not localize: - return Path(model) - - from huggingface_hub import snapshot_download - - cache_dir.mkdir(parents=True, exist_ok=True) - os.environ.setdefault("HF_HOME", str(cache_dir / ".hf_home")) - local_dir = cache_dir / sanitize_repo_id(model) - if not local_dir.exists(): - print(f"[end2end] Downloading model into local cache: {local_dir}") - snapshot_download( - repo_id=model, - local_dir=str(local_dir), - local_dir_use_symlinks=True, - resume_download=True, - ) - return local_dir.resolve() - - -def resolve_local_only( - override: bool | None, - source: str, - default: bool, -) -> bool: - if override is not None: - return bool(override) - return default or Path(source).expanduser().is_dir() - - -def load_text_tokenizer(tokenizer_source: str, local_files_only: bool): - from transformers import AutoTokenizer - - kwargs = { - "trust_remote_code": True, - "padding_side": "left", - "local_files_only": bool(local_files_only), - } - try: - tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, **kwargs) - except TypeError: - kwargs.pop("local_files_only", None) - tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, **kwargs) - return tokenizer - - -def preprocess_image(image: Image.Image, resolution: int) -> torch.Tensor: - w, h = image.size - short_side = min(w, h) - scale = resolution / short_side - new_w, new_h = round(w * scale), round(h * scale) - image = image.resize((new_w, new_h), Image.BICUBIC) - left = (new_w - resolution) // 2 - top = (new_h - resolution) // 2 - image = image.crop((left, top, left + resolution, top + resolution)) - arr = np.array(image, dtype=np.float32) / 255.0 - tensor = torch.from_numpy(arr).permute(2, 0, 1) - return (tensor - 0.5) / 0.5 - - -def load_vq_image_encoder(source: str, local_files_only: bool, device: torch.device) -> Any: - from vllm_omni.model_executor.models.dynin_omni.dynin_omni_common import get_dynin_magvit_attr - - MAGVITv2 = get_dynin_magvit_attr("MAGVITv2", source=source, local_files_only=local_files_only) - vq_model = MAGVITv2.from_pretrained(source, local_files_only=local_files_only).to(device) - vq_model.requires_grad_(False) - vq_model.eval() - return vq_model - - -def encode_image_tokens( - image_path: Path, - vq_model: Any, - device: torch.device, - resolution: int, -) -> torch.Tensor: - image = Image.open(image_path).convert("RGB") - image_tensor = preprocess_image(image, resolution=resolution).unsqueeze(0).to(device) - with torch.no_grad(): - token_ids = vq_model.get_code(image_tensor) - token_ids = torch.as_tensor(token_ids, dtype=torch.long).detach().cpu() - if token_ids.ndim == 2 and token_ids.shape[0] == 1: - token_ids = token_ids[0] - return token_ids.contiguous() - - -def encode_video_tokens( - video_path: Path, - vq_model: Any, - device: torch.device, - resolution: int, - num_frames: int, -) -> torch.Tensor: - import cv2 - - cap = cv2.VideoCapture(str(video_path)) - frames: list[np.ndarray] = [] - while True: - ok, frame = cap.read() - if not ok: - break - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - frames.append(frame) - cap.release() - if not frames: - raise ValueError(f"Video has no readable frames: {video_path}") - if len(frames) < num_frames: - raise ValueError(f"Video has {len(frames)} frames, requires >= {num_frames}: {video_path}") - - indices = np.linspace(0, len(frames) - 1, num_frames).astype(int) - token_list: list[torch.Tensor] = [] - for idx in indices: - pil = Image.fromarray(frames[int(idx)]) - frame_tensor = preprocess_image(pil, resolution=resolution).unsqueeze(0).to(device) - with torch.no_grad(): - token_list.append(torch.as_tensor(vq_model.get_code(frame_tensor), dtype=torch.long)) - merged = torch.cat(token_list, dim=1).detach().cpu() - if merged.ndim == 2 and merged.shape[0] == 1: - merged = merged[0] - return merged.contiguous() - - -def load_vq_audio_encoder(source: str, local_files_only: bool, device: torch.device) -> Any: - from transformers import AutoModel - - kwargs = { - "trust_remote_code": True, - "local_files_only": bool(local_files_only), - "low_cpu_mem_usage": False, - } - try: - model = AutoModel.from_pretrained(source, **kwargs) - except TypeError: - kwargs.pop("low_cpu_mem_usage", None) - try: - model = AutoModel.from_pretrained(source, **kwargs) - except TypeError: - kwargs.pop("local_files_only", None) - model = AutoModel.from_pretrained(source, **kwargs) - model.requires_grad_(False) - model.eval() - if hasattr(model, "to"): - model = model.to(device) - return model - - -def encode_audio_tokens(audio_path: Path, vq_audio_model: Any) -> torch.Tensor: - encoded = vq_audio_model.encode(str(audio_path)) - if isinstance(encoded, dict): - for key in ("input_ids", "token_ids", "codes", "tokens"): - if key in encoded: - encoded = encoded[key] - break - encoded = torch.as_tensor(encoded, dtype=torch.long).detach().cpu() - if encoded.ndim == 1: - encoded = encoded.unsqueeze(0) - elif encoded.ndim > 2: - encoded = encoded.view(encoded.shape[0], -1) - return encoded.contiguous() - - -def build_chat_prompt(content: str) -> str: - return ( - f"<|start_header_id|>user<|end_header_id|>\n{content}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" - ) - - -def resolve_task_text( - *, - task_name: str, - text: str, - instruction: str = "", - raw_prompt: bool = False, -) -> str: - text = str(text or "").strip() - - if task_name == "t2t" and not text: - return DEFAULT_T2T_PROMPT - if task_name == "i2t" and not text: - return DEFAULT_I2T_QUESTION - if task_name == "s2t" and not text: - return DEFAULT_S2T_INSTRUCTION - if task_name == "v2t" and not text: - return DEFAULT_V2T_QUESTION - if task_name in {"t2i", "i2i"} and not text: - return "A high quality detailed image." - - if task_name != "t2s": - return text - - if not text: - text = DEFAULT_T2S_PROMPT - - if raw_prompt: - return text - - instruction = str(instruction or "").strip() or DEFAULT_T2S_INSTRUCTION - return build_chat_prompt(f"{instruction}\n{text}") - - -def load_universal_prompting( - *, - tokenizer: Any, - tokenizer_source: str, - max_text_len: int, - cond_dropout_prob: float, - local_files_only: bool, - max_audio_len: int = 512, - max_audio_len_short: int = 256, -) -> Any: - from vllm_omni.model_executor.models.dynin_omni.dynin_omni_common import ( - DYNIN_REMOTE_SETTINGS, - resolve_remote_attr, - ) - - UniversalPrompting = resolve_remote_attr( - "UniversalPrompting", - module_name="prompting_utils", - settings=DYNIN_REMOTE_SETTINGS, - source=tokenizer_source, - local_files_only=bool(local_files_only), - fallback_module_names=("modeling_dynin_omni",), - ) - init_kwargs: dict[str, Any] = { - "max_text_len": int(max_text_len), - "special_tokens": DYNIN_SPECIAL_TOKENS, - "ignore_id": -100, - "cond_dropout_prob": float(cond_dropout_prob), - "use_reserved_token": True, - "max_audio_len": int(max_audio_len), - "max_audio_len_short": int(max_audio_len_short), - } - try: - return UniversalPrompting(tokenizer, **init_kwargs) - except TypeError: - init_kwargs.pop("max_audio_len", None) - init_kwargs.pop("max_audio_len_short", None) - return UniversalPrompting(tokenizer, **init_kwargs) - - -def _runtime_fallback(task: str, key: str, value: Any) -> Any: - if isinstance(value, str): - if value.strip() != "": - return value - elif value is not None: - return value - return TASK_RUNTIME_FALLBACKS.get(task, {}).get(key) - - -def _validate_generation_args(*, task: str, max_new_tokens: int, steps: int, block_length: int) -> None: - # Keep i2t/v2t generation constraints aligned with i2t.py/v2t.py. - if task not in {"i2t", "v2t"}: - return - if max_new_tokens <= 0: - raise ValueError(f"{task} requires max_new_tokens > 0.") - if block_length <= 0: - raise ValueError(f"{task} requires block_length > 0.") - if steps <= 0: - raise ValueError(f"{task} requires steps > 0.") - if max_new_tokens % block_length != 0: - raise ValueError(f"{task} requires max_new_tokens % block_length == 0, got {max_new_tokens} % {block_length}") - num_blocks = max_new_tokens // block_length - if num_blocks <= 0: - raise ValueError(f"{task} has invalid num_blocks.") - if steps % num_blocks != 0: - raise ValueError( - f"{task} requires steps % (max_new_tokens // block_length) == 0, " - f"got steps={steps}, max_new_tokens={max_new_tokens}, block_length={block_length}" - ) - - -def make_prompt_payload( - *, - task: str, - text: str, - image_tokens: torch.Tensor | None, - audio_tokens: torch.Tensor | None, - video_tokens: torch.Tensor | None, - image_placeholder_tokens: int, - audio_placeholder_tokens: int, - image_token_offset: int, - speech_token_offset: int, - mask_token_id: int, - use_train_i2i_prompt: bool, -) -> tuple[Any, str]: - runtime_task, prompting_task, _, _ = TASK_DEFAULT_RUNTIME[task] - del runtime_task - - if task == "t2t": - payload = ([[]], [build_chat_prompt(text)]) - return payload, prompting_task - - if task == "i2t": - if image_tokens is None: - raise ValueError("i2t requires image tokens") - img = image_tokens.view(-1).long() + int(image_token_offset) - payload = ([[img]], [build_chat_prompt(text)]) - return payload, prompting_task - - if task == "s2t": - if audio_tokens is None: - raise ValueError("s2t requires audio tokens") - aud = audio_tokens.long() + int(speech_token_offset) - if aud.ndim == 1: - aud = aud.unsqueeze(0) - payload = ([aud], [build_chat_prompt(text)]) - return payload, prompting_task - - if task == "v2t": - if video_tokens is None: - raise ValueError("v2t requires video tokens") - vid = video_tokens.view(-1).long() + int(image_token_offset) - payload = (vid.unsqueeze(0), [build_chat_prompt(text)]) - return payload, prompting_task - - if task == "t2i": - image_placeholder = torch.full( - (1, int(image_placeholder_tokens)), - fill_value=int(mask_token_id), - dtype=torch.long, - ) - payload = ([text], image_placeholder) - return payload, prompting_task - - if task == "i2i": - if image_tokens is None: - raise ValueError("i2i requires image tokens") - src = image_tokens.view(1, -1).long() + int(image_token_offset) - target_len = int(image_placeholder_tokens) if image_placeholder_tokens > 0 else int(src.shape[1]) - image_placeholder = torch.full( - (1, target_len), - fill_value=int(mask_token_id), - dtype=torch.long, - ) - if use_train_i2i_prompt: - labels_placeholder = torch.full( - (1, target_len), - fill_value=-100, - dtype=torch.long, - ) - payload = ([text], src, image_placeholder, labels_placeholder) - return payload, "i2i" - payload = ([text], src, image_placeholder) - return payload, "i2i_gen" - - if task == "t2s": - audio_placeholder = torch.full( - (1, int(audio_placeholder_tokens)), - fill_value=int(mask_token_id), - dtype=torch.long, - ) - payload = ([text], audio_placeholder) - return payload, prompting_task - - raise ValueError(f"Unsupported task: {task}") - - -def _to_1d_int_list(value: Any) -> list[int]: - if value is None: - return [] - if isinstance(value, torch.Tensor): - tensor = value.detach().to(device="cpu", dtype=torch.long) - else: - tensor = torch.as_tensor(value, dtype=torch.long) - if tensor.ndim == 0: - tensor = tensor.view(1) - elif tensor.ndim >= 2: - tensor = tensor.view(tensor.shape[0], -1)[0] - return [int(v) for v in tensor.tolist()] - - -def _run_uni_prompting(uni_prompting: Any, payload: Any, prompting_task: str) -> tuple[list[int], list[int]]: - prepared = uni_prompting(payload, prompting_task) - if isinstance(prepared, tuple): - prepared_input_ids = prepared[0] if len(prepared) > 0 else None - prepared_attention_mask = prepared[1] if len(prepared) > 1 else None - else: - prepared_input_ids = prepared - prepared_attention_mask = None - - input_ids = _to_1d_int_list(prepared_input_ids) - attention_mask = _to_1d_int_list(prepared_attention_mask) - if not input_ids: - raise RuntimeError(f"UniversalPrompting returned empty input_ids for task={prompting_task}") - return input_ids, attention_mask - - -def _get_special_token_id(uni_prompting: Any, token: str) -> int: - sptids = getattr(uni_prompting, "sptids_dict", None) or {} - if token not in sptids: - raise KeyError(f"Special token not found in UniversalPrompting.sptids_dict: {token}") - token_ids = _to_1d_int_list(sptids[token]) - if not token_ids: - raise ValueError(f"Special token id is empty for token: {token}") - return int(token_ids[0]) - - -def _tokenize_chat_query(tokenizer: Any, text: str) -> list[int]: - encoded = tokenizer(build_chat_prompt(text), return_tensors="pt").input_ids[0] - token_ids = _to_1d_int_list(encoded) - if not token_ids: - raise RuntimeError("Failed to tokenize chat query text.") - return token_ids - - -def _flatten_media_token_ids_with_offset(token_ids: Any, token_offset: int) -> list[int]: - media_ids = token_ids - if isinstance(media_ids, torch.Tensor): - media_ids = media_ids.detach().cpu().reshape(-1).tolist() - else: - media_ids = np.asarray(media_ids).reshape(-1).tolist() - return [int(x) + int(token_offset) for x in media_ids] - - -def _scalar_token_id(value: Any) -> int: - if isinstance(value, torch.Tensor): - if value.numel() == 0: - raise ValueError("Empty special-token tensor.") - return int(value.view(-1)[0].item()) - if isinstance(value, (list, tuple)): - if not value: - raise ValueError("Empty special-token list.") - return int(value[0]) - return int(value) - - -def build_v2t_input_ids( - *, - video_token_ids: Any, - tokenizer: Any, - uni_prompting: Any, - question: str, - image_token_offset: int, -) -> tuple[list[int], str]: - media_ids = video_token_ids - if isinstance(media_ids, torch.Tensor): - media_ids = media_ids.detach().cpu().reshape(-1).tolist() - else: - media_ids = np.asarray(media_ids).reshape(-1).tolist() - media_ids = [int(x) + int(image_token_offset) for x in media_ids] - - sptids = uni_prompting.sptids_dict - task_id = _scalar_token_id(sptids["<|v2t|>"]) - soi_id = _scalar_token_id(sptids["<|soi|>"]) - eoi_id = _scalar_token_id(sptids["<|eoi|>"]) - sot_id = _scalar_token_id(sptids["<|sot|>"]) - - prompt_text = build_v2t_chat_prompt(question) - query_ids = tokenizer(prompt_text, return_tensors="pt").input_ids[0].detach().cpu().tolist() - input_ids = [task_id, soi_id] + media_ids + [eoi_id, sot_id] + [int(v) for v in query_ids] - return input_ids, prompt_text - - -def build_i2t_input_ids( - *, - image_token_ids: Any, - tokenizer: Any, - uni_prompting: Any, - question: str, - image_token_offset: int, -) -> tuple[list[int], str]: - image_ids = image_token_ids - if isinstance(image_ids, torch.Tensor): - image_ids = image_ids.detach().cpu().reshape(-1).tolist() - else: - image_ids = np.asarray(image_ids).reshape(-1).tolist() - image_ids = [int(x) + int(image_token_offset) for x in image_ids] - - sptids = uni_prompting.sptids_dict - task_id = _scalar_token_id(sptids["<|mmu|>"]) - soi_id = _scalar_token_id(sptids["<|soi|>"]) - eoi_id = _scalar_token_id(sptids["<|eoi|>"]) - sot_id = _scalar_token_id(sptids["<|sot|>"]) - - prompt_text = build_i2t_chat_prompt(question) - query_ids = tokenizer(prompt_text, return_tensors="pt").input_ids[0].detach().cpu().tolist() - input_ids = [task_id, soi_id] + image_ids + [eoi_id, sot_id] + [int(v) for v in query_ids] - return input_ids, prompt_text - - -def build_v2t_chat_prompt(question: str) -> str: - return ( - f"<|start_header_id|>user<|end_header_id|>\n{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" - ) - - -def build_i2t_chat_prompt(question: str) -> str: - return ( - f"<|start_header_id|>user<|end_header_id|>\n{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" - ) - - -def make_mmu_prompt( - *, - task: str, - text: str, - tokenizer: Any, - uni_prompting: Any, - image_tokens: torch.Tensor | None, - audio_tokens: torch.Tensor | None, - video_tokens: torch.Tensor | None, - image_token_offset: int, - speech_token_offset: int, -) -> tuple[list[int], list[int]]: - query_ids = _tokenize_chat_query(tokenizer, text) - - if task == "i2t": - token_ids, _ = build_i2t_input_ids( - image_token_ids=image_tokens, - tokenizer=tokenizer, - uni_prompting=uni_prompting, - question=text, - image_token_offset=int(image_token_offset), - ) - token_ids = [int(v) for v in token_ids] - return token_ids, [1] * len(token_ids) - - if task == "v2t": - token_ids, _ = build_v2t_input_ids( - video_token_ids=video_tokens, - tokenizer=tokenizer, - uni_prompting=uni_prompting, - question=text, - image_token_offset=int(image_token_offset), - ) - token_ids = [int(v) for v in token_ids] - return token_ids, [1] * len(token_ids) - - if task == "s2t": - if audio_tokens is None: - raise ValueError("s2t requires audio tokens") - audio_ids = _to_1d_int_list(audio_tokens.long() + int(speech_token_offset)) - token_ids = [ - _get_special_token_id(uni_prompting, "<|s2t|>"), - _get_special_token_id(uni_prompting, "<|soa|>"), - *audio_ids, - _get_special_token_id(uni_prompting, "<|eoa|>"), - *query_ids, - ] - return token_ids, [1] * len(token_ids) - - raise ValueError(f"Unsupported task for validation-style MMU prompt: {task}") - - -def iter_mm_outputs(outputs: list[Any]): - for omni_out in outputs: - req_out = omni_out - req_list = req_out if isinstance(req_out, list) else [req_out] - for item in req_list: - if item is None: - continue - mm_out = getattr(item, "multimodal_output", None) or {} - if mm_out: - yield mm_out - completions = getattr(item, "outputs", None) or [] - for completion in completions: - c_mm_out = getattr(completion, "multimodal_output", None) or {} - if c_mm_out: - yield c_mm_out - omni_mm = getattr(omni_out, "multimodal_output", None) or {} - if omni_mm: - yield omni_mm - - -def _to_token_list(value: Any) -> list[int]: - if value is None: - return [] - if hasattr(value, "detach"): - value = value.detach() - if hasattr(value, "cpu"): - value = value.cpu() - if hasattr(value, "flatten"): - value = value.flatten().tolist() - if isinstance(value, tuple): - value = list(value) - if not isinstance(value, list): - return [] - out: list[int] = [] - for token in value: - if isinstance(token, bool): - continue - try: - out.append(int(token)) - except Exception: - continue - return out - - -def extract_text_output(outputs: list[Any], tokenizer: Any) -> str: - for mm_out in iter_mm_outputs(outputs): - text = mm_out.get("text") - if isinstance(text, list) and text: - text = text[-1] - if isinstance(text, str) and text.strip(): - return text.strip() - for key in ("text_tokens", "token_ids"): - token_ids = _to_token_list(mm_out.get(key)) - if not token_ids: - continue - decoded = tokenizer.decode(token_ids, skip_special_tokens=True) - if isinstance(decoded, str) and decoded.strip(): - return decoded.strip() - return "" - - -def extract_image_output(outputs: list[Any]) -> torch.Tensor | None: - for mm_out in iter_mm_outputs(outputs): - image = mm_out.get("image") - if isinstance(image, list) and image: - image = image[-1] - if isinstance(image, torch.Tensor): - return image - return None - - -def tensor_to_pil_image(image: torch.Tensor) -> Image.Image: - arr = image.detach().cpu().numpy() - if arr.ndim == 4: - arr = arr[0] - if arr.ndim == 3 and arr.shape[0] in (1, 3, 4): - arr = np.transpose(arr, (1, 2, 0)) - if arr.dtype != np.uint8: - arr = arr.astype(np.float32) - if arr.max() <= 1.0: - arr = arr * 255.0 - arr = np.clip(arr, 0.0, 255.0).astype(np.uint8) - if arr.ndim == 3 and arr.shape[-1] == 1: - arr = arr[..., 0] - return Image.fromarray(arr) - - -def extract_audio_output(outputs: list[Any]) -> tuple[np.ndarray, int] | None: - for mm_out in iter_mm_outputs(outputs): - audio = mm_out.get("audio") - if audio is None: - audio = mm_out.get("speech") - if audio is None: - continue - - def _to_wav_array(value: Any) -> np.ndarray: - if isinstance(value, torch.Tensor): - return value.detach().cpu().numpy().reshape(-1).astype(np.float32) - return np.asarray(value).reshape(-1).astype(np.float32) - - if isinstance(audio, list): - chunks = [_to_wav_array(chunk) for chunk in audio] - wav = np.concatenate(chunks, axis=0) if chunks else np.zeros((0,), dtype=np.float32) - else: - wav = _to_wav_array(audio) - sr = mm_out.get("sr", 24000) - if hasattr(sr, "item"): - try: - sr = int(sr.item()) - except Exception: - sr = 24000 - elif isinstance(sr, list): - sr = int(sr[0]) if sr else 24000 - else: - sr = int(sr) - return wav, sr - return None - - -def save_audio_wav(path: Path, wav: np.ndarray, sr: int) -> None: - try: - import soundfile as sf - - sf.write(str(path), wav, int(sr), format="WAV") - except Exception: - from scipy.io import wavfile - - wav_i16 = np.clip(wav, -1.0, 1.0) - wav_i16 = (wav_i16 * 32767.0).astype(np.int16) - wavfile.write(str(path), int(sr), wav_i16) - - -def parse_args(repo_root: Path) -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Dynin-Omni unified offline end2end example.") - parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES) - parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.") - parser.add_argument( - "--deploy-config", - type=str, - default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"), - help="Path to deploy config yaml.", - ) - parser.add_argument( - "--dynin-config-path", - type=str, - default="", - help="Path to DYNIN config yaml (passed through additional_information).", - ) - parser.add_argument( - "--model-cache-dir", - type=str, - default="/tmp/dynin_localized_models", - help="Cache directory used when --model is HF repo id.", - ) - parser.add_argument( - "--localize-model", - action=argparse.BooleanOptionalAction, - default=True, - help="If true and --model is HF repo id, snapshot it under --model-cache-dir.", - ) - parser.add_argument("--text", type=str, default="", help="Prompt/edit/question text.") - parser.add_argument("--instruction", type=str, default="", help="Optional extra instruction.") - parser.add_argument("--raw-prompt", action=argparse.BooleanOptionalAction, default=False) - parser.add_argument("--image", type=str, default="", help="Input image path for i2i/i2t.") - parser.add_argument("--audio", type=str, default="", help="Input audio path for s2t.") - parser.add_argument("--video", type=str, default="", help="Input video path for v2t.") - parser.add_argument("--image-resolution", type=int, default=None) - parser.add_argument("--num-frames", type=int, default=None) - parser.add_argument( - "--output-dir", - type=str, - default="", - help="Directory for generated outputs.", - ) - parser.add_argument("--output-prefix", type=str, default="") - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--dtype", type=str, default="auto") - parser.add_argument("--max-tokens-per-stage", type=int, default=1) - - parser.add_argument("--runtime-task", type=str, default="", help="Override runtime task key.") - parser.add_argument("--prompting-task", type=str, default="", help="Override prompting task key.") - parser.add_argument("--detok-id", type=int, default=None, help="Override detok id.") - - parser.add_argument("--prompt-max-text-len", type=int, default=None) - parser.add_argument("--cond-dropout-prob", type=float, default=0.0) - parser.add_argument("--max-new-tokens", type=int, default=None) - parser.add_argument("--steps", type=int, default=None) - parser.add_argument("--block-length", type=int, default=None) - parser.add_argument("--temperature", type=float, default=None) - parser.add_argument("--cfg-scale", type=float, default=None) - parser.add_argument("--remasking", type=str, default="low_confidence") - - parser.add_argument("--timesteps", type=int, default=None) - parser.add_argument("--guidance-scale", type=float, default=None) - parser.add_argument("--noise-type", type=str, default="mask") - parser.add_argument("--noise-schedule-name", type=str, default="cosine") - parser.add_argument("--noise-schedule-params", type=str, default="{}") - - parser.add_argument("--mask-token-id", type=int, default=None) - parser.add_argument("--codebook-size", type=int, default=None) - parser.add_argument("--audio-codebook-size", type=int, default=None) - parser.add_argument("--image-token-count", type=int, default=None) - parser.add_argument("--t2s-token-length", type=int, default=None) - parser.add_argument( - "--t2s-condition", - type=str, - default="", - ) - parser.add_argument( - "--use-train-i2i-prompt", - action="store_true", - help="Use i2i training prompt template (default behavior of i2i.py).", - ) - parser.add_argument( - "--no-use-train-i2i-prompt", - dest="use_train_i2i_prompt", - action="store_false", - help="Use i2i_gen prompt template.", - ) - parser.set_defaults(use_train_i2i_prompt=None) - - parser.add_argument("--tokenizer-path", type=str, default="") - parser.add_argument("--model-local-files-only", action=argparse.BooleanOptionalAction, default=None) - parser.add_argument("--tokenizer-local-files-only", action=argparse.BooleanOptionalAction, default=None) - - parser.add_argument("--vq-model-image-path", type=str, default="") - parser.add_argument("--vq-model-image-local-files-only", action=argparse.BooleanOptionalAction, default=None) - parser.add_argument("--vq-model-audio-path", type=str, default="") - parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None) - - parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True) - return parser.parse_args() - - -def main() -> None: - repo_root = bootstrap_repo_path() - ensure_safe_import_for_vllm() - from vllm_omni.model_executor.models.dynin_omni.dynin_omni_common import ( - DYNIN_PROMPT_SOURCE_KEY, - DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT, - ) - - args = parse_args(repo_root) - - if args.disable_hf_xet: - os.environ.setdefault("HF_HUB_DISABLE_XET", "1") - - np.random.seed(args.seed) - torch.manual_seed(args.seed) - - model_dir = ensure_local_model_dir( - model=args.model, - cache_dir=Path(args.model_cache_dir).expanduser(), - localize=bool(args.localize_model), - ) - model_source = str(model_dir) - - task_name = str(args.task) - dynin_config_path = str(Path(args.dynin_config_path).expanduser()) - os.environ["DYNIN_CONFIG_PATH"] = dynin_config_path - default_runtime_task, default_prompting_task, default_detok_id, final_modality = TASK_DEFAULT_RUNTIME[task_name] - runtime_task = args.runtime_task.strip() or str( - _runtime_fallback(task_name, "runtime_task", None) or default_runtime_task - ) - prompting_task = args.prompting_task.strip() or str( - _runtime_fallback(task_name, "prompting_task", None) or default_prompting_task - ) - detok_id_default = _runtime_fallback(task_name, "detok_id", None) - if detok_id_default is None: - detok_id_default = default_detok_id - detok_id = int(detok_id_default if args.detok_id is None else args.detok_id) - - output_dir_default = _runtime_fallback(task_name, "output_dir", args.output_dir) - resolved_output_dir = str(output_dir_default or "/tmp/dynin_end2end_outputs") - - image_resolution_value = _runtime_fallback( - task_name, - "image_resolution", - args.image_resolution, - ) - if image_resolution_value is None: - image_resolution_value = 336 - image_resolution = int(image_resolution_value) - - num_frames_value = _runtime_fallback( - task_name, - "num_frames", - args.num_frames, - ) - if num_frames_value is None: - num_frames_value = 8 - num_frames = int(num_frames_value) - - prompt_max_text_len_value = _runtime_fallback( - task_name, - "prompt_max_text_len", - args.prompt_max_text_len, - ) - if prompt_max_text_len_value is None: - prompt_max_text_len_value = 1024 - prompt_max_text_len = int(prompt_max_text_len_value) - - max_new_tokens_value = _runtime_fallback( - task_name, - "max_new_tokens", - args.max_new_tokens, - ) - if max_new_tokens_value is None: - max_new_tokens_value = 256 - max_new_tokens = int(max_new_tokens_value) - - steps_value = _runtime_fallback( - task_name, - "steps", - args.steps, - ) - if steps_value is None: - steps_value = 256 - steps = int(steps_value) - - block_length_value = _runtime_fallback( - task_name, - "block_length", - args.block_length, - ) - if block_length_value is None: - block_length_value = 2 - block_length = int(block_length_value) - - temperature_value = _runtime_fallback( - task_name, - "temperature", - args.temperature, - ) - if temperature_value is None: - temperature_value = 0.0 - temperature = float(temperature_value) - - cfg_scale_value = _runtime_fallback( - task_name, - "cfg_scale", - args.cfg_scale, - ) - if cfg_scale_value is None: - cfg_scale_value = 0.0 - cfg_scale = float(cfg_scale_value) - - remasking = str(_runtime_fallback(task_name, "remasking", args.remasking) or "low_confidence") - - timesteps_value = _runtime_fallback( - task_name, - "timesteps", - args.timesteps, - ) - if timesteps_value is None: - timesteps_value = 20 - timesteps = int(timesteps_value) - - guidance_scale_value = _runtime_fallback( - task_name, - "guidance_scale", - args.guidance_scale, - ) - if guidance_scale_value is None: - guidance_scale_value = 0.0 - guidance_scale = float(guidance_scale_value) - - mask_token_id_value = _runtime_fallback( - task_name, - "mask_token_id", - args.mask_token_id, - ) - if mask_token_id_value is None: - mask_token_id_value = 126336 - mask_token_id = int(mask_token_id_value) - - codebook_size_value = _runtime_fallback( - task_name, - "codebook_size", - args.codebook_size, - ) - if codebook_size_value is None: - codebook_size_value = 8192 - codebook_size = int(codebook_size_value) - - audio_codebook_size_value = _runtime_fallback( - task_name, - "audio_codebook_size", - args.audio_codebook_size, - ) - if audio_codebook_size_value is None: - audio_codebook_size_value = 4096 - audio_codebook_size = int(audio_codebook_size_value) - - image_token_count_value = _runtime_fallback( - task_name, - "image_token_count", - args.image_token_count, - ) - image_token_count = int(image_token_count_value) if image_token_count_value is not None else 0 - - t2s_token_length_value = _runtime_fallback( - task_name, - "t2s_token_length", - args.t2s_token_length, - ) - if t2s_token_length_value is None: - t2s_token_length_value = 383 - t2s_token_length = int(t2s_token_length_value) - - t2s_condition = str( - _runtime_fallback(task_name, "t2s_condition", args.t2s_condition) - or "gender-female_emotion-neutral_speed-normal_pitch-normal" - ) - - _validate_generation_args( - task=task_name, - max_new_tokens=max_new_tokens, - steps=steps, - block_length=block_length, - ) - - use_train_i2i_prompt = _runtime_fallback(task_name, "use_train_i2i_prompt", args.use_train_i2i_prompt) - if use_train_i2i_prompt is None: - use_train_i2i_prompt = bool(task_name == "i2i") - use_train_i2i_prompt = bool(use_train_i2i_prompt) - - if task_name in {"i2i", "i2t"} and not args.image: - raise ValueError(f"--task {task_name} requires --image") - if task_name == "s2t" and not args.audio: - raise ValueError("--task s2t requires --audio") - if task_name == "v2t" and not args.video: - raise ValueError("--task v2t requires --video") - - text = resolve_task_text( - task_name=task_name, - text=args.text, - instruction=args.instruction, - raw_prompt=bool(args.raw_prompt), - ) - - tokenizer_source = args.tokenizer_path.strip() or model_source - model_local_only = resolve_local_only( - args.model_local_files_only, model_source, default=Path(model_source).is_dir() - ) - tokenizer_local_only = resolve_local_only( - args.tokenizer_local_files_only, - tokenizer_source, - default=model_local_only, - ) - tokenizer = load_text_tokenizer(tokenizer_source, local_files_only=tokenizer_local_only) - text_vocab_size = int(len(tokenizer)) - - image_tokens: torch.Tensor | None = None - audio_tokens: torch.Tensor | None = None - video_tokens: torch.Tensor | None = None - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - vq_image_source = args.vq_model_image_path.strip() or "snu-aidas/magvitv2" - vq_audio_source = args.vq_model_audio_path.strip() or "snu-aidas/emova_speech_tokenizer_vllm" - vq_image_local_only = resolve_local_only(args.vq_model_image_local_files_only, vq_image_source, default=False) - vq_audio_local_only = resolve_local_only(args.vq_model_audio_local_files_only, vq_audio_source, default=False) - - if task_name in {"i2i", "i2t", "v2t"}: - vq_image = load_vq_image_encoder(vq_image_source, vq_image_local_only, device) - if task_name in {"i2i", "i2t"}: - image_tokens = encode_image_tokens( - Path(args.image).expanduser().resolve(), - vq_model=vq_image, - device=device, - resolution=int(image_resolution), - ) - if task_name == "v2t": - video_tokens = encode_video_tokens( - Path(args.video).expanduser().resolve(), - vq_model=vq_image, - device=device, - resolution=int(image_resolution), - num_frames=int(num_frames), - ) - if hasattr(vq_image, "cpu"): - vq_image = vq_image.cpu() - - if task_name == "s2t": - vq_audio = load_vq_audio_encoder(vq_audio_source, vq_audio_local_only, device) - audio_tokens = encode_audio_tokens(Path(args.audio).expanduser().resolve(), vq_audio) - if hasattr(vq_audio, "cpu"): - vq_audio = vq_audio.cpu() - - noise_schedule_params: dict[str, Any] = {} - try: - parsed = json.loads(args.noise_schedule_params) - if isinstance(parsed, dict): - noise_schedule_params = {str(k): v for k, v in parsed.items()} - except Exception: - noise_schedule_params = {} - - image_token_count = int(image_token_count) - if image_token_count <= 0: - if image_tokens is not None: - image_token_count = int(image_tokens.numel()) - else: - base_res = int(image_resolution) - image_token_count = max(1, (base_res // 16) ** 2) - - uncond_input_ids: list[int] | None = None - uncond_attention_mask: list[int] | None = None - if task_name == "t2t": - messages = [{"role": "user", "content": text}] - if getattr(tokenizer, "chat_template", None): - prompt_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) - encoded = tokenizer(prompt_text, return_tensors="pt", add_special_tokens=False) - else: - encoded = tokenizer(text, return_tensors="pt", add_special_tokens=True) - prompt_token_ids = _to_1d_int_list(encoded["input_ids"]) - prompt_attention_mask = _to_1d_int_list(encoded.get("attention_mask")) - if not prompt_attention_mask: - prompt_attention_mask = [1] * len(prompt_token_ids) - else: - max_audio_len_for_prompt = int(max(t2s_token_length, 512)) - if audio_tokens is not None: - max_audio_len_for_prompt = max(max_audio_len_for_prompt, int(audio_tokens.numel())) - max_audio_len_short_for_prompt = max(256, max_audio_len_for_prompt // 2) - - uni_prompting = load_universal_prompting( - tokenizer=tokenizer, - tokenizer_source=tokenizer_source, - max_text_len=int(prompt_max_text_len), - cond_dropout_prob=float(args.cond_dropout_prob), - local_files_only=bool(tokenizer_local_only), - max_audio_len=int(max_audio_len_for_prompt), - max_audio_len_short=int(max_audio_len_short_for_prompt), - ) - prompting_text_vocab_size = int(len(uni_prompting.text_tokenizer)) - - is_mmu_task = task_name in {"i2t", "s2t", "v2t"} and not args.prompting_task.strip() - if is_mmu_task: - prompt_token_ids, prompt_attention_mask = make_mmu_prompt( - task=task_name, - text=text, - tokenizer=uni_prompting.text_tokenizer, - uni_prompting=uni_prompting, - image_tokens=image_tokens, - audio_tokens=audio_tokens, - video_tokens=video_tokens, - image_token_offset=prompting_text_vocab_size, - speech_token_offset=prompting_text_vocab_size + int(codebook_size), - ) - else: - prompt_payload, prompting_task = make_prompt_payload( - task=task_name, - text=text, - image_tokens=image_tokens, - audio_tokens=audio_tokens, - video_tokens=video_tokens, - image_placeholder_tokens=image_token_count, - audio_placeholder_tokens=int(t2s_token_length), - image_token_offset=text_vocab_size, - speech_token_offset=text_vocab_size + int(codebook_size), - mask_token_id=int(mask_token_id), - use_train_i2i_prompt=use_train_i2i_prompt, - ) - if args.prompting_task.strip(): - prompting_task = args.prompting_task.strip() - - prompt_token_ids, prompt_attention_mask = _run_uni_prompting( - uni_prompting, - prompt_payload, - prompting_task, - ) - - if task_name in {"i2t", "s2t", "v2t"}: - prompt_attention_mask = [1] * len(prompt_token_ids) - if not prompt_attention_mask: - prompt_attention_mask = [1] * len(prompt_token_ids) - - if task_name in {"t2i", "i2i"} and guidance_scale > 0: - uncond_payload, uncond_prompting_task = make_prompt_payload( - task=task_name, - text="", - image_tokens=image_tokens, - audio_tokens=audio_tokens, - video_tokens=video_tokens, - image_placeholder_tokens=image_token_count, - audio_placeholder_tokens=int(t2s_token_length), - image_token_offset=text_vocab_size, - speech_token_offset=text_vocab_size + int(codebook_size), - mask_token_id=int(mask_token_id), - use_train_i2i_prompt=use_train_i2i_prompt, - ) - uncond_input_ids, uncond_attention_mask = _run_uni_prompting( - uni_prompting, - uncond_payload, - args.prompting_task.strip() or uncond_prompting_task, - ) - if not uncond_attention_mask: - uncond_attention_mask = [1] * len(uncond_input_ids) - - runtime_info: dict[str, Any] = { - "task": [runtime_task], - "detok_id": [int(detok_id)], - DYNIN_PROMPT_SOURCE_KEY: [DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT], - "dynin_config_path": [str(dynin_config_path)], - "attention_mask": [prompt_attention_mask], - "prompt_max_text_len": [int(prompt_max_text_len)], - "prompting_max_text_len": [int(prompt_max_text_len)], - "cond_dropout_prob": [float(args.cond_dropout_prob)], - "prompting_cond_dropout_prob": [float(args.cond_dropout_prob)], - "tokenizer_path": [str(tokenizer_source)], - "text_vocab_size": [int(text_vocab_size)], - "model_local_files_only": [bool(model_local_only)], - "max_new_tokens": [int(max_new_tokens)], - "steps": [int(steps)], - "block_length": [int(block_length)], - "temperature": [float(temperature)], - "cfg_scale": [float(cfg_scale)], - "remasking": [str(remasking)], - "mask_id": [int(mask_token_id)], - "mask_token_id": [int(mask_token_id)], - "codebook_size": [int(codebook_size)], - "audio_codebook_size": [int(audio_codebook_size)], - "timesteps": [int(timesteps)], - "guidance_scale": [float(guidance_scale)], - "noise_type": [str(args.noise_type)], - "noise_schedule_name": [str(args.noise_schedule_name)], - "noise_schedule_params": [noise_schedule_params], - "seq_len": [int(image_token_count)], - "condition": [str(t2s_condition)], - "vq_model_image_path": [str(vq_image_source)], - "vq_model_image_local_files_only": [bool(vq_image_local_only)], - "vq_model_audio_path": [str(vq_audio_source)], - "vq_model_audio_local_files_only": [bool(vq_audio_local_only)], - } - - if task_name in {"t2t", "i2t", "s2t", "v2t"}: - runtime_info["prompt_length"] = [int(len(prompt_token_ids))] - if uncond_input_ids is not None: - runtime_info["uncond_input_ids"] = [uncond_input_ids] - if uncond_attention_mask is not None: - runtime_info["uncond_attention_mask"] = [uncond_attention_mask] - - if task_name == "t2s": - runtime_info["max_new_tokens"] = [int(t2s_token_length)] - - prompt = { - "prompt_token_ids": [int(v) for v in prompt_token_ids], - "additional_information": runtime_info, - "modalities": [final_modality], - } - - from vllm import SamplingParams - - from vllm_omni.entrypoints.omni import Omni - - deploy_config_path = str(Path(args.deploy_config).expanduser()) - omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype) - sampling_params_list = [ - SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False) - for _ in range(omni.num_stages) - ] - - try: - outputs = list(omni.generate(prompt, sampling_params_list)) - finally: - omni.close() - - out_dir = Path(resolved_output_dir).expanduser() - out_dir.mkdir(parents=True, exist_ok=True) - stamp = time.strftime("%Y%m%d_%H%M%S") - prefix = args.output_prefix.strip() or f"{task_name}_{stamp}" - - if final_modality == "text": - text_out = extract_text_output(outputs, tokenizer=tokenizer) - if not text_out: - raise RuntimeError("No text output found.") - out_path = out_dir / f"{prefix}.txt" - out_path.write_text(text_out + "\n", encoding="utf-8") - print(f"[end2end] text saved: {out_path}") - print(text_out) - return - - if final_modality == "image": - image_out = extract_image_output(outputs) - if image_out is None: - raise RuntimeError("No image output found.") - pil = tensor_to_pil_image(image_out) - out_path = out_dir / f"{prefix}.png" - pil.save(out_path) - print(f"[end2end] image saved: {out_path}") - return - - if final_modality == "audio": - audio_out = extract_audio_output(outputs) - if audio_out is None: - raise RuntimeError("No audio output found.") - wav, sr = audio_out - out_path = out_dir / f"{prefix}.wav" - save_audio_wav(out_path, wav, sr) - print(f"[end2end] audio saved: {out_path} (sr={sr}, samples={wav.shape[0]})") - return - - raise RuntimeError(f"Unsupported final modality: {final_modality}") - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 04ff1a97d..36d1b44ee 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -13,7 +13,7 @@ list of supported architectures across all modalities, see ## Supported Models | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | -|---|---|---|---|---|---|---| +| --- | --- | --- | --- | --- | --- | --- | | AuK | `tencent/AuK`, `tencent/AuK-Flash` (assembled with `tools/prepare_auk_checkpoint.py`) | 2 (thinker encoder + DiT) | ✓ | — | instruction-driven editing, enhancement, separation; no example script (use `vllm_omni.model_extras.auk`) | 24 kHz | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz | | Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | 2 (talker + codec) | ✓ | ✓ (async chunk) | voice design / voice direction (`--instruction`) | 24 kHz | @@ -27,7 +27,6 @@ list of supported architectures across all modalities, see | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz | | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz | -| dots.tts | `dots-studio/dots.tts-soar` | single (native AR) | — (text-only) | ✓ (online) | no-reference synthesis | 48 kHz | | IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz | | IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz | @@ -54,32 +53,38 @@ codec: a slow AR transformer predicts one semantic token per codec frame, and a 4-layer fast AR predicts that frame's remaining 9 codebooks. ### Prerequisites + No extra packages: the codec architecture ships in tree and its weights (`codec.pth`) come with the checkpoint. ### Quick start + ```bash python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ --text "Welcome to Audio8 TTS, a compact multilingual text to speech model." ``` ### Voice cloning + ```bash python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ --text "Welcome to Audio8 TTS." \ --ref-audio /path/to/reference.wav \ --ref-text "The exact transcript of the reference recording." ``` + The reference transcript must match what is actually spoken in the clip; a mismatch degrades both stability and speaker similarity. ### Streaming + ```bash python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ --text "Welcome to Audio8 TTS." --streaming ``` ### Notes + - Output: 44.1 kHz mono WAV; the codec emits ~21.5 frames/s (2048 samples per frame). - Context is 2048 packed text+audio positions, so `max_model_len` is 2048. - Greedy decoding (`temperature: 0`) is degenerate for this checkpoint — it @@ -92,6 +97,7 @@ python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ - `--num-prompts 4` smoke-tests concurrency (stage 0 runs `max_num_seqs: 4`). ### Measured behaviour (1x H20, shared GPU) + `--streaming` on a 58-character English sentence: TTFA ~0.74 s (4-frame first chunk), 7 chunks, 1.55 s wall for 3.16 s of audio (RTF ~0.49). Treat these as smoke-test magnitudes, not a benchmark. @@ -103,12 +109,14 @@ smoke-test magnitudes, not a benchmark. 2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz. ### Prerequisites + ```bash uv pip install -e . # Includes soundfile, onnxruntime, x-transformers, einops via requirements. ``` Download the model snapshot: + ```python from huggingface_hub import snapshot_download snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', @@ -116,15 +124,18 @@ snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', ``` If your downloaded checkpoint lacks `config.json`, add it: + ```json { "model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"] } ``` + This is required because `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers the class mapping; the loader still reads `model_type` from `config.json` to select the class. ### Quick start + ```bash python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ --model pretrained_models/Fun-CosyVoice3-0.5B \ @@ -132,9 +143,11 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ ``` ### Voice cloning + If `--ref-audio` is omitted, the script downloads the upstream [`zero_shot_prompt.wav`](https://github.com/FunAudioLLM/CosyVoice/blob/main/asset/zero_shot_prompt.wav) from the CosyVoice repo into the current directory. To use your own clip, pass `--ref-audio /path/to/reference.wav`, and modify `--prompt-text` correspondingly. + ```bash python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ --model pretrained_models/Fun-CosyVoice3-0.5B \ @@ -144,9 +157,11 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \ ``` ### Streaming + Streaming is enabled by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` on `vllm serve` to switch to the legacy synchronous path. ### Notes + - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform. - Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config ` to override. @@ -157,6 +172,7 @@ Streaming is enabled by default via `async_chunk: true` in `vllm_omni/deploy/cos 2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning. ### Quick start + ```bash python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ --model zai-org/GLM-TTS \ @@ -167,12 +183,14 @@ python examples/offline_inference/text_to_speech/glm_tts/end2end.py \ ``` ### Architecture + ```text Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz) (Llama-based) (32k vocab) (Flow Matching) ``` ### Notes + - `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis. - Reference audio should be 3-10 seconds. - First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder. @@ -187,12 +205,14 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 No extra packages are required; the DAC codec is vendored in vLLM-Omni. ### Quick start + ```bash python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ --text "Hello, this is a test of the Fish Speech text to speech system." ``` ### Voice cloning + ```bash python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ --text "Hello, this is a cloned voice." \ @@ -201,14 +221,17 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ ``` ### Streaming + ```bash python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ --text "Hello, this is a streaming test." \ --streaming ``` + Streaming requires `async_chunk: true` in the deploy config. ### Notes + - Output: 44.1 kHz mono WAV. - DAC codec weights (`codec.pth`) are loaded lazily from the model directory. @@ -219,6 +242,7 @@ Streaming requires `async_chunk: true` in the deploy config. Dense 0.5B two-stage TTS pipeline (`AR + flow` + audio VAE) at 44.1 kHz. The example covers style, IP voice, music-only generation, text-to-audio events, emotion, dialect, zero-shot cloning, podcast, speech+BGM, and speech+environment-sound cases. ### Quick start + ```bash python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ --case style \ @@ -227,6 +251,7 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ ``` ### Voice cloning + ```bash python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ --case zero_shot \ @@ -237,6 +262,7 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ ``` ### Streaming + ```bash python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ --case basic \ @@ -247,6 +273,7 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ ``` ### Notes + - `style`, `ip`, `bgm`, and `tta` do not require reference audio. - Reference-audio cases use `--ref-audio`; `zero_shot` also requires `--ref-text`. - `podcast` uses multiple references via `--ref-audio-paths`. @@ -259,14 +286,17 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ Standalone talker-only deployment of Ming-flash-omni-2.0 at 44.1 kHz. Voice is controlled through caption fields (`风格` / `IP` / `语速`/`基频`/`音量`) rather than reference audio. ### Prerequisites + The example calls into `vllm_omni.model_executor.models.ming_flash_omni.prompt_utils` for the default prompt and instruction builder; no extra pip install on top of the base vLLM-Omni install. ### Quick start + ```bash python examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py --case style ``` ### Cases + ```bash # ASMR-style whisper (caption-driven) python examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py --case style @@ -281,6 +311,7 @@ python examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py Override the default text per case with `--text`, write to a custom path with `--output`. ### Notes + - Talker-only deployment — for the multimodal Ming-flash-omni example, see [`examples/offline_inference/ming_flash_omni/`](../../ming_flash_omni/). - Deploy config: `vllm_omni/deploy/ming_flash_omni_tts.yaml` (single GPU, `enforce_eager`, `max_num_seqs: 1`). - Decode defaults from the Ming cookbook: `max_decode_steps=200`, `cfg=2.0`, `sigma=0.25`, `temperature=0.0`, `use_zero_spk_emb=True`. @@ -294,6 +325,7 @@ Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono (mixed > **No built-in speaker presets.** `--ref-audio` is required on every call. Default `--mode voice_clone` matches upstream's recommended workflow; `--mode continuation` is exposed for completeness but upstream's continuation-with-prompt path emits very short / near-silent output, so it is rarely useful in practice. Sample reference clips ship in the upstream repo under [`assets/audio/`](https://github.com/OpenMOSS/MOSS-TTS-Nano/tree/main/assets/audio) (e.g. `zh_1.wav`, `en_2.wav`, `jp_2.wav`). ### Quick start + ```bash # Fetch a sample reference clip (one-off, user-scoped cache). REF_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/moss-tts-nano" @@ -305,9 +337,11 @@ python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \ --text "你好，这是MOSS-TTS-Nano的语音合成演示。" \ --ref-audio "$REF_DIR/zh_1.wav" ``` + The first run downloads `OpenMOSS-Team/MOSS-TTS-Nano` and `OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano` from Hugging Face. ### Reproducible runs + ```bash python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \ --text "Deterministic test." \ @@ -316,6 +350,7 @@ python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \ ``` ### Notes + - Output: 48 kHz mono WAV (the tokenizer is internally stereo at 48 kHz; the wrapper averages to mono before reaching the engine). - Deploy config: `vllm_omni/deploy/moss_tts_nano.yaml` (auto-loaded; override with `--deploy-config`). - Default `--max-new-frames 375` ≈ 14 s of audio; raise for longer outputs. @@ -329,12 +364,15 @@ python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \ Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design). ### Prerequisites + ```bash huggingface-cli download k2-fsa/OmniVoice ``` + Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`. ### Quick start (auto voice) + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -342,6 +380,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Voice cloning + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -351,6 +390,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Voice design + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -359,6 +399,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Language hint + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -367,6 +408,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Seed for Reproducibility + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -375,6 +417,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ``` ### Notes + - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking. - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz. @@ -385,13 +428,16 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ 3-task-variant TTS with 24 kHz output. Has its own argparse surface (this script does not follow the common `--text` / `--ref-audio` shape). ### Prerequisites + For ROCm builds, replace `onnxruntime` with `onnxruntime-rocm`: + ```bash pip uninstall onnxruntime pip install onnxruntime-rocm ``` ### Task variants + - `CustomVoice`: predefined speaker (speaker ID) with optional style instruction. - `VoiceDesign`: text + descriptive instruction designs a new voice. - `Base`: voice cloning from reference audio + transcript. @@ -415,22 +461,27 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty ``` ### Streaming + ```bash python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --query-type CustomVoice \ --streaming \ --output-dir /tmp/out_stream ``` + Streaming requires `async_chunk: true` in the deploy config. ### Word Timestamps + Word-level timestamps are currently a serving-path feature: launch `vllm-omni serve` with `--forced-aligner` and request `word_timestamps` (see `examples/online_serving/text_to_speech/README.md`). An offline example is not provided in this release. ### Batched decoding + The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, create an overlay based on `vllm_omni/deploy/qwen3_tts.yaml` and pass it with `--deploy-config`: + ```bash python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --query-type CustomVoice \ @@ -438,9 +489,11 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --batch-size 4 \ --deploy-config /path/to/qwen3_tts_batched.yaml ``` + `--batch-size` must match a CUDA-graph capture size (1, 2, 4, 8, 16…). ### Notes + - Run `--help` for the full argument surface. - See `qwen3_tts/end2end.py` for the prompt-length-estimation logic the Talker uses. @@ -451,6 +504,7 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up). ### Prerequisites + The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately: ```bash @@ -480,15 +534,18 @@ failure rather than an environment one. The PyTorch sampler it falls back to cos here. ### Quick start (zero-shot, default voice) + ```bash python examples/offline_inference/text_to_speech/gepard/end2end.py \ --text "Hello, this is a Gepard demo." ``` ### Voice cloning + Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR. ### Notes + - Output: 22.05 kHz mono WAV. - First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower. - Offline and online: `examples/offline_inference/text_to_speech/gepard/end2end.py` and `/v1/audio/speech` (see the [online Gepard section](../../online_serving/text_to_speech/README.md#gepard-10)). @@ -502,6 +559,7 @@ Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Ref Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`. ### Prerequisites + ```bash pip install voxcpm # or, for a local source checkout: @@ -509,6 +567,7 @@ export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm ``` ### Quick start + ```bash python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ --model openbmb/VoxCPM2 \ @@ -516,7 +575,9 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ ``` ### Voice cloning + Pass a reference audio for isolated cloning, or both `--ref-audio` + `--ref-text` for prompt continuation: + ```bash python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ --text "Hello, this is a voice clone demo." \ @@ -525,39 +586,16 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ ``` ### Streaming + Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [`examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py`](../../online_serving/text_to_speech/voxcpm2/gradio_demo.py) for an AudioWorklet-based gapless streaming player; the offline `end2end.py` script does not expose a streaming path. ### Notes + - Output: 48 kHz mono WAV. - Deploy config: `vllm_omni/deploy/voxcpm2.yaml` (auto-loaded by HF `model_type`). --- -## dots.tts - -Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage. - -### Quick start -```bash -python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ - --model dots-studio/dots.tts-soar \ - --text "Hello, this is a test of dots TTS running on vLLM Omni." -``` - -### Voice cloning -Not wired in this release. Generation is text-only and does not consume reference audio, a named voice, or a precomputed speaker embedding. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet. - -### Streaming -The AudioVAE decoder emits incremental audio through its internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`). Online streaming is exposed through the OpenAI-compatible `/v1/audio/speech` endpoint with `stream=true`; see the [online serving guide](../../online_serving/text_to_speech/README.md#dotstts). The offline `end2end.py` script returns the consolidated waveform. - -### Notes -- Output: 48 kHz mono WAV. -- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`). -- Checkpoints: `dots-studio/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet. -- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap. - ---- - ## IndexTTS-2 and IndexTTS-2.5 Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning. @@ -592,6 +630,7 @@ python examples/offline_inference/text_to_speech/indextts2/end2end.py \ ``` ### Emotion control + ```bash # Emotion from reference audio python examples/offline_inference/text_to_speech/indextts2/end2end.py \ @@ -616,6 +655,7 @@ python examples/offline_inference/text_to_speech/indextts2/end2end.py \ ``` ### Notes + - `--ref-audio` is **required** — neither version provides text-only synthesis. - IndexTTS-2.5 Stage 0 (AR Talker) generates mel codes from text + reference audio using plain vLLM sampling. It does not reproduce the official default `num_beams=3` beam search. For parity comparisons, run upstream with `num_beams=1`; output quality can differ from the official beam-search result. - IndexTTS-2.5 accepts language codes such as `zh`, `en`, `zhen` (mixed Chinese/English), `ja`, and `yue`. `Mandarin` is a vLLM-Omni convenience alias for `zh`. Japanese (`ja`) uses `fugashi` tokenization and produces audio, but it does not automatically expand numbers, dates, or percentages; write those inputs as readable Japanese text before calling the model. Use `--no-text-normalization` only when the input is already normalized. @@ -631,12 +671,15 @@ python examples/offline_inference/text_to_speech/indextts2/end2end.py \ Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol. ### Prerequisites + Latest `mistral_common` with `SpeechRequest` support: + ```bash pip install -e /path/to/mistral-common # or upgrade from PyPI when available ``` ### Quick start (voice preset) + ```bash python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ --write-audio --voice cheerful_female \ @@ -645,6 +688,7 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ ``` ### Voice cloning (capability gated upstream) + ```bash python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ --write-audio \ @@ -654,12 +698,14 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ ``` ### Streaming + concurrency + ```bash python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ --num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \ --model mistralai/Voxtral-4B-TTS-2603 \ --text "..." ``` + Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). ### Notes diff --git a/examples/offline_inference/text_to_speech/dots_tts/end2end.py b/examples/offline_inference/text_to_speech/dots_tts/end2end.py deleted file mode 100644 index aa710db4a..000000000 --- a/examples/offline_inference/text_to_speech/dots_tts/end2end.py +++ /dev/null @@ -1,125 +0,0 @@ -"""Offline dots.tts inference example (single-stage native AR pipeline). - -Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml). -Reference checkpoint: dots-studio/dots.tts-soar. -""" - -from __future__ import annotations - -import os -import time -from pathlib import Path - -import soundfile as sf -import torch - -from vllm_omni import Omni -from vllm_omni.utils.tracking_parser import TrackingArgumentParser - -REPO_ROOT = Path(__file__).resolve().parents[4] -SAMPLE_RATE = 48_000 - - -def parse_args(): - parser = TrackingArgumentParser(description="Offline dots.tts native AR inference") - parser.add_argument( - "--model", - type=str, - default="dots-studio/dots.tts-soar", - help="dots.tts model path or HuggingFace repo ID.", - ) - parser.add_argument( - "--text", - type=str, - default="Hello, this is a test of dots TTS running on vLLM Omni.", - help="Text to synthesize.", - ) - parser.add_argument( - "--output-dir", - type=str, - default="output_audio", - help="Directory for output WAV files.", - ) - parser.add_argument( - "--deploy-config", - type=str, - default=None, - help="Override the deploy config path. If unset, auto-loads " - "vllm_omni/deploy/dots_tts.yaml based on the HF model_type.", - ) - return parser.parse_args() - - -def extract_audio(multimodal_output: dict) -> torch.Tensor: - """Extract the final complete audio tensor from multimodal output. - - The output processor concatenates per-step delta tensors under - ``model_outputs``. Falls back to ``audio`` for backwards compat. - """ - audio = multimodal_output.get("model_outputs") - if audio is None: - audio = multimodal_output.get("audio") - if audio is None: - raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}") - - if isinstance(audio, list): - valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] - if not valid: - raise ValueError("Audio list is empty or all elements are None.") - return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] - - return torch.as_tensor(audio).float().cpu().reshape(-1) - - -def main(): - args = parse_args() - - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - engine = Omni( - model=args.model, - deploy_config=args.deploy_config, - ) - - from transformers import AutoTokenizer - - from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import ( - build_dots_tts_prompt, - ) - - tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) - - prompt = build_dots_tts_prompt( - tokenizer=tokenizer, - text=args.text, - ) - - print(f"Model : {args.model}") - print(f"Text : {args.text}") - print(f"Output dir : {output_dir}") - print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens") - - t_start = time.perf_counter() - outputs = engine.generate([prompt]) - elapsed = time.perf_counter() - t_start - - request_output = outputs[0] - mm = request_output.outputs[0].multimodal_output - audio = extract_audio(mm) - - duration = audio.numel() / SAMPLE_RATE - rtf = elapsed / duration if duration > 0 else float("inf") - - output_path = output_dir / "output.wav" - sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV") - - print(f"Saved : {output_path}") - print(f"Duration : {duration:.2f}s") - print(f"Inference : {elapsed:.2f}s") - print(f"RTF : {rtf:.3f}") - - -if __name__ == "__main__": - os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - main() diff --git a/examples/online_serving/dynin_omni/README.md b/examples/online_serving/dynin_omni/README.md deleted file mode 100644 index 6e78a5c1f..000000000 --- a/examples/online_serving/dynin_omni/README.md +++ /dev/null @@ -1,97 +0,0 @@ -# Dynin-Omni: Online serving - -## Installation - -Please refer to [README.md](../../../README.md). - -## Launch the Server - -First, find the `transformers_modules` path: - -```bash -python - <<'PY' -from transformers.utils.hub import HF_MODULES_CACHE -print(HF_MODULES_CACHE) -PY -``` - -Then export it for both `PYTHONPATH` and `HF_MODULES_CACHE`: - -```bash -export PYTHONPATH=:$PYTHONPATH -export HF_MODULES_CACHE= -``` - -Run from repository root. The 3-stage pipeline is resolved automatically from -the model type (`vllm_omni/deploy/dynin_omni.yaml`), so no config path is -required. Pass `--deploy-config` only to override the defaults: - -```bash -vllm-omni serve snu-aidas/Dynin-Omni \ - --omni \ - --port 8091 -``` - -If `vllm-omni` is not in PATH, run: - -```bash -PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \ - --omni \ - --port 8091 -``` - -Wait until the server logs show both `All stages initialized successfully` and -`Application startup complete.` before sending requests. - -## Send Requests via Python Client - -Move to the example directory: - -```bash -cd examples/online_serving/dynin_omni -``` - -### Text -> Image - -```bash -python openai_chat_completion_client_for_multimodal_generation.py \ - --query-type t2i \ - --prompt "A realistic indoor living room with natural daylight." -``` - -### Image -> Image - -```bash -python openai_chat_completion_client_for_multimodal_generation.py \ - --query-type i2i \ - --image-path ../../offline_inference/dynin_omni/data/image/sofa_under_water.jpg \ - --prompt "Transform this surreal underwater setting into a realistic indoor living room while preserving the sofa layout." -``` - -### Text -> Speech - -```bash -python openai_chat_completion_client_for_multimodal_generation.py \ - --query-type t2s \ - --prompt "Hello. This is Dynin-omni." -``` - -## CLI Arguments - -- `--query-type` (`t2i|t2s|i2i`) -- `--model` (default: `snu-aidas/Dynin-Omni`) -- `--host` / `--port` (OpenAI-compatible vLLM endpoint) -- `--prompt` (custom text) -- `--image-path` (required for `i2i`) -- `--modalities` (optional output modalities override) -- `--output-dir` (default: `/tmp/dynin_online_outputs`) - -## Notes - -- This client currently supports only `t2i`, `t2s`, and `i2i`. -- `t2t` is intentionally not exposed in this online example. -- This example intentionally uses the OpenAI-compatible chat completion endpoint. -- Task routing for non-text outputs relies on Dynin task trigger tokens (`<|t2i|>`, `<|i2i|>`, `<|t2s|>`) injected by the client. -- Outputs are saved under `/tmp/dynin_online_outputs` by default. -- Dynin stage-0 warmup can take a while on first startup; do not send requests before startup completes. -- Dynin itself can execute text-returning tasks such as `t2t`, `s2t`, `i2t`, and `v2t`, but this online serving example currently runs stage-0 in `generation` mode. In that path, the generation worker does not surface the final text as `output.text`, so OpenAI chat responses for those text-output tasks may complete internally but still return empty text. diff --git a/examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py deleted file mode 100644 index 972855543..000000000 --- a/examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py +++ /dev/null @@ -1,342 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import argparse -import base64 -import json -import mimetypes -import os -import time -from pathlib import Path -from typing import Any - -DEFAULT_MODEL = "snu-aidas/Dynin-Omni" -DEFAULT_OUTPUT_DIR = "/tmp/dynin_online_outputs" - -QUERY_CHOICES = ("t2i", "t2s", "i2i") -DEFAULT_PROMPT_BY_QUERY = { - "t2i": "A high quality detailed living room interior photo.", - "t2s": "Please read this sentence naturally: Hello from Dynin-Omni online serving.", - "i2i": "Transform this image into a realistic indoor living room while preserving layout.", -} -DEFAULT_MODALITIES_BY_QUERY = { - "t2i": ["image"], - "t2s": ["audio"], - "i2i": ["image"], -} -OFFLINE_PARITY_STAGE_COUNT = 3 -OFFLINE_PARITY_STAGE_SAMPLING = { - "max_tokens": 1, - "temperature": 0.0, - "top_p": 1.0, - "detokenize": False, -} - - -def _infer_mime_type(path: Path) -> str: - mime_type, _ = mimetypes.guess_type(str(path)) - return mime_type or "application/octet-stream" - - -def _encode_file_as_data_url(path: Path) -> str: - mime_type = _infer_mime_type(path) - raw = path.read_bytes() - encoded = base64.b64encode(raw).decode("utf-8") - return f"data:{mime_type};base64,{encoded}" - - -def _to_image_url(path_or_url: str) -> str: - value = str(path_or_url) - if value.startswith(("http://", "https://", "data:image/")): - return value - path = Path(value).expanduser().resolve() - if not path.exists(): - raise FileNotFoundError(f"Image file not found: {path}") - return _encode_file_as_data_url(path) - - -def _build_user_content(query_type: str, prompt: str, image_path: str | None) -> list[dict[str, Any]]: - if query_type == "t2i": - return [{"type": "text", "text": f"<|t2i|> {prompt}"}] - - if query_type == "t2s": - return [{"type": "text", "text": f"<|t2s|> {prompt}"}] - - if query_type == "i2i": - if not image_path: - raise ValueError("--image-path is required for query type i2i") - return [ - {"type": "text", "text": f"<|i2i|> {prompt}"}, - {"type": "image_url", "image_url": {"url": _to_image_url(image_path)}}, - ] - - raise ValueError(f"Unsupported query_type: {query_type}") - - -def _collect_text_from_content(content: Any) -> list[str]: - texts: list[str] = [] - if isinstance(content, str): - stripped = content.strip() - if stripped: - texts.append(stripped) - return texts - - if isinstance(content, dict): - for key in ("text", "content", "value", "output_text"): - text_value = content.get(key) - if isinstance(text_value, str) and text_value.strip(): - texts.append(text_value.strip()) - return texts - - if isinstance(content, list): - for item in content: - texts.extend(_collect_text_from_content(item)) - return texts - - content_text = getattr(content, "text", None) - if isinstance(content_text, str) and content_text.strip(): - texts.append(content_text.strip()) - content_value = getattr(content, "content", None) - if isinstance(content_value, str) and content_value.strip(): - texts.append(content_value.strip()) - output_text = getattr(content, "output_text", None) - if isinstance(output_text, str) and output_text.strip(): - texts.append(output_text.strip()) - return texts - - -def _extract_text_outputs(chat_completion: Any) -> list[str]: - texts: list[str] = [] - for choice in getattr(chat_completion, "choices", []) or []: - message = getattr(choice, "message", None) - if message is None: - continue - content = getattr(message, "content", None) - texts.extend(_collect_text_from_content(content)) - reasoning_content = getattr(message, "reasoning_content", None) - if isinstance(reasoning_content, str) and reasoning_content.strip(): - texts.append(reasoning_content.strip()) - choice_text = getattr(choice, "text", None) - if isinstance(choice_text, str) and choice_text.strip(): - texts.append(choice_text.strip()) - top_level_output_text = getattr(chat_completion, "output_text", None) - if isinstance(top_level_output_text, str) and top_level_output_text.strip(): - texts.append(top_level_output_text.strip()) - return texts - - -def _extract_image_data_urls(chat_completion: Any) -> list[str]: - urls: list[str] = [] - for choice in getattr(chat_completion, "choices", []) or []: - message = getattr(choice, "message", None) - if message is None: - continue - content = getattr(message, "content", None) - if not isinstance(content, list): - continue - for item in content: - if not isinstance(item, dict): - continue - if item.get("type") != "image_url": - continue - image_url = (item.get("image_url") or {}).get("url") - if isinstance(image_url, str) and image_url.startswith("data:image"): - urls.append(image_url) - return urls - - -def _extract_audio_payloads(chat_completion: Any) -> list[bytes]: - payloads: list[bytes] = [] - for choice in getattr(chat_completion, "choices", []) or []: - message = getattr(choice, "message", None) - if message is None: - continue - message_audio = getattr(message, "audio", None) - if message_audio is None: - continue - data_b64 = getattr(message_audio, "data", None) - if isinstance(data_b64, str) and data_b64: - try: - payloads.append(base64.b64decode(data_b64)) - except Exception: - continue - return payloads - - -def _decode_data_url(data_url: str) -> tuple[bytes, str]: - header, data = data_url.split(",", 1) - mime_type = "image/png" - if ";" in header and ":" in header: - mime_type = header.split(":", 1)[1].split(";", 1)[0] - return base64.b64decode(data), mime_type - - -def _image_extension_from_mime(mime_type: str) -> str: - if mime_type == "image/jpeg": - return ".jpg" - if mime_type == "image/webp": - return ".webp" - if mime_type == "image/gif": - return ".gif" - return ".png" - - -def _save_outputs( - *, - query_type: str, - chat_completion: Any, - output_dir: Path, -) -> None: - output_dir.mkdir(parents=True, exist_ok=True) - stamp = time.strftime("%Y%m%d_%H%M%S") - - text_outputs = _extract_text_outputs(chat_completion) - image_data_urls = _extract_image_data_urls(chat_completion) - audio_payloads = _extract_audio_payloads(chat_completion) - - if text_outputs: - text_path = output_dir / f"{query_type}_{stamp}.txt" - text_path.write_text("\n\n".join(text_outputs) + "\n", encoding="utf-8") - print(f"[dynin-online] text saved: {text_path}") - print(text_outputs[0]) - - for idx, image_url in enumerate(image_data_urls): - image_bytes, mime_type = _decode_data_url(image_url) - ext = _image_extension_from_mime(mime_type) - image_path = output_dir / f"{query_type}_{stamp}_{idx}{ext}" - image_path.write_bytes(image_bytes) - print(f"[dynin-online] image saved: {image_path}") - - for idx, audio_bytes in enumerate(audio_payloads): - audio_path = output_dir / f"{query_type}_{stamp}_{idx}.wav" - audio_path.write_bytes(audio_bytes) - print(f"[dynin-online] audio saved: {audio_path}") - - if not text_outputs and not image_data_urls and not audio_payloads: - print("[dynin-online] no output extracted from response") - raw_path = output_dir / f"{query_type}_{stamp}_raw_response.json" - try: - if hasattr(chat_completion, "model_dump_json"): - serialized = chat_completion.model_dump_json(indent=2) - else: - if hasattr(chat_completion, "model_dump"): - raw_payload: Any = chat_completion.model_dump(mode="json") - else: - raw_payload = chat_completion - try: - serialized = json.dumps(raw_payload, ensure_ascii=False, indent=2) - except Exception: - serialized = json.dumps({"repr": repr(raw_payload)}, ensure_ascii=False, indent=2) - raw_path.write_text(serialized + "\n", encoding="utf-8") - print(f"[dynin-online] raw response saved: {raw_path}") - except Exception: - pass - - -def _build_offline_parity_sampling_params_list() -> list[dict[str, Any]]: - return [dict(OFFLINE_PARITY_STAGE_SAMPLING) for _ in range(OFFLINE_PARITY_STAGE_COUNT)] - - -def run_request(args: argparse.Namespace) -> None: - from openai import OpenAI - - client = OpenAI( - api_key="EMPTY", - base_url=f"http://{args.host}:{args.port}/v1", - ) - prompt = args.prompt.strip() if args.prompt else DEFAULT_PROMPT_BY_QUERY[args.query_type] - user_content = _build_user_content( - query_type=args.query_type, - prompt=prompt, - image_path=args.image_path, - ) - if args.modalities: - modalities = [item.strip() for item in args.modalities.split(",") if item.strip()] - else: - modalities = DEFAULT_MODALITIES_BY_QUERY[args.query_type] - - extra_body = { - "sampling_params_list": _build_offline_parity_sampling_params_list(), - } - chat_completion = client.chat.completions.create( - model=args.model, - messages=[{"role": "user", "content": user_content}], - modalities=modalities, - extra_body=extra_body, - ) - _save_outputs( - query_type=args.query_type, - chat_completion=chat_completion, - output_dir=Path(args.output_dir).expanduser(), - ) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Dynin-Omni online chat completion client") - parser.add_argument( - "--query-type", - "-q", - type=str, - default="t2i", - choices=QUERY_CHOICES, - help="Dynin query type", - ) - parser.add_argument( - "--model", - "-m", - type=str, - default=DEFAULT_MODEL, - help="Model name/path", - ) - parser.add_argument( - "--host", - type=str, - default="localhost", - help="Host/IP of the vLLM Omni API server", - ) - parser.add_argument( - "--port", - type=int, - default=8091, - help="Port of the vLLM Omni API server", - ) - parser.add_argument( - "--prompt", - "-p", - type=str, - default="", - help="Custom prompt text", - ) - parser.add_argument( - "--image-path", - "-i", - type=str, - default=None, - help="Image path/URL for i2i", - ) - parser.add_argument( - "--modalities", - type=str, - default="", - help="Comma-separated output modalities override (e.g., text,image,audio)", - ) - parser.add_argument( - "--output-dir", - "-o", - type=str, - default=DEFAULT_OUTPUT_DIR, - help="Directory to save outputs", - ) - return parser.parse_args() - - -def main() -> None: - args = parse_args() - os.environ.setdefault("HF_HUB_DISABLE_XET", "1") - run_request(args) - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 77e9ec4f1..711fdec6b 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -27,7 +27,6 @@ For the full list of supported architectures across all modalities, see | OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) | | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ | -| dots.tts | `dots-studio/dots.tts-soar` | — (text-only) | ✓ (PCM stream) | — (default placeholder only) | — | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ | CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead. @@ -143,42 +142,6 @@ python examples/online_serving/text_to_speech/gepard/speech_client.py \ - Concurrent requests at `max_num_seqs: 4` are supported. Native-AR recompute preemption is a known limitation of this architecture (a request that is preempted mid-generation can resume incorrectly); keep concurrency at or below `max_num_seqs` and treat preemption as out of scope until the platform fix lands. - Optional comparison against the upstream Gepard reference server needs Blackwell/Hopper + CUDA 13 + Postgres and is not part of CI. -## dots.tts - -Single-stage, text-only TTS at 48 kHz. The current integration supports -no-reference synthesis; voice cloning, reference audio, named voices, and -precomputed speaker embeddings are not supported yet. - -### Launch - -```bash -vllm serve dots-studio/dots.tts-soar --omni --trust-remote-code --port 8091 -``` - -### Sending requests - -```bash -# Non-streaming WAV -curl -X POST http://localhost:8091/v1/audio/speech \ - -H "Content-Type: application/json" \ - -d '{ - "input": "Hello, this is a dots TTS online-serving test.", - "voice": "default", - "response_format": "wav" - }' --output output.wav - -# Streaming 48 kHz mono PCM -curl -X POST http://localhost:8091/v1/audio/speech \ - -H "Content-Type: application/json" \ - -d '{ - "input": "Hello, this is a streaming dots TTS test.", - "voice": "default", - "stream": true, - "stream_format": "audio", - "response_format": "pcm" - }' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 1 - -``` - --- ## GLM-TTS diff --git a/recipes/rednote-hilab/dots.tts.md b/recipes/rednote-hilab/dots.tts.md deleted file mode 100644 index cc180020e..000000000 --- a/recipes/rednote-hilab/dots.tts.md +++ /dev/null @@ -1,166 +0,0 @@ -# dots.tts - -> Continuous-AR TTS at 48 kHz (rednote-hilab) - -## Summary - -- Vendor: rednote-hilab -- Model: `dots-studio/dots.tts-soar` -- Task: Text-to-speech, zero-shot synthesis only -- Mode: Offline inference and OpenAI-compatible online serving -- Maintainer: Community - -## When to use this recipe - -Use this recipe as a known-good starting point for running -`dots-studio/dots.tts-soar` on vLLM-Omni on consumer-class GPUs. -dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM -+ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio. -It follows the same "vLLM-native base LM + side-path computation" pattern as -VoxCPM2 — single-stage pipeline -`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR -loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone -instead of MiniCPM4, and no FSQ / residual-LM stage. - -This is an early integration; review the [Known limitations](#known-limitations) -before deploying it in production. - -## References - -- Offline end-to-end script: - [`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py) -- Example guide: - [`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts) -- Online serving guide: - [`examples/online_serving/text_to_speech/README.md`](../../examples/online_serving/text_to_speech/README.md#dotstts) -- Default deploy config: - [`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) -- Talker / pipeline source: - [`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/) -- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts) - -## Hardware Support - -This recipe documents one tested 16 GB consumer-GPU configuration. Other -vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as -community validation lands. - -## GPU - -### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended) - -dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++ -speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The -bundled default config at -[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) -(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`, -`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model -weights and ~0.3 GiB peak activation; the remainder of the configured -budget is available for KV cache. Total resident footprint at idle is -roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph -roadmap would be step 8's graph capture (not implemented yet; this -release runs `enforce_eager: true`, so it doesn't apply today). - -#### Environment - -- OS: Linux (WSL2) -- Python: 3.12 -- Driver / runtime: NVIDIA driver 595.95 -- torch: 2.11.0+cu130 -- vLLM: 0.26.0 -- vLLM-Omni: 0.22.1.dev (current `main`) - -#### Command - -```bash -python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ - --model dots-studio/dots.tts-soar \ - --text "Hello, this is a test of dots TTS running on vLLM Omni." -``` - -The deploy config at -[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) -is loaded automatically by the model registry (HF `model_type=dots_tts`). -Pass `--deploy-config ` to override. - -#### Verification - -**T1 — offline zero-shot synthesis**: - -```bash -python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ - --model dots-studio/dots.tts-soar \ - --text "Hello, this is a test of dots TTS running on vLLM Omni." -``` - -Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single -one-shot process (init → one `generate()` → exit), so the reported -numbers include engine init and first-request warmup, not just steady-state -per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`, -`RTF: 1.569`. - -Weight-loading breakdown from the same run (all tensors matched, no -missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270 -patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder. - -Whisper transcription of the output matched the input text with no -dropped leading word (confirms the streaming-vocoder patch-boundary fix -described in [Known limitations](#known-limitations)). - -**T2 — online text-only synthesis**: - -```bash -vllm serve dots-studio/dots.tts-soar --omni --trust-remote-code --port 8091 -``` - -```bash -curl -X POST http://localhost:8091/v1/audio/speech \ - -H "Content-Type: application/json" \ - -d '{ - "input": "Hello, this is a test of dots TTS online serving.", - "voice": "default", - "response_format": "wav" - }' --output output.wav -``` - -The endpoint also supports raw streaming audio with `stream=true`, -`stream_format="audio"`, and `response_format="pcm"`. - -#### Notes - -- Output: 48 kHz mono WAV. -- Checkpoints: `dots-studio/dots.tts-soar` is the validated default - used throughout this recipe. `rednote-hilab/dots.tts-base` shares the - same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf` - (MeanFlow, 2-4 step) is not supported — see below. -- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy - config are load-bearing, not just conservative defaults: with prefix - caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path - does not preserve the `sparse_audio` marker this model relies on to - route audio output correctly, and generation silently truncates to a - single ~160 ms patch. Do not override `enable_prefix_caching` for this - model until that framework-level gap is fixed. - -## Known limitations - -- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder - weights load and are exercised by `load_weights()`, but nothing in the - prompt builder consumes reference audio or a precomputed speaker embedding - yet. Both offline and online generation are text-only; `ref_audio`, - `ref_text`, named voices, and speaker embeddings are not supported. -- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only - the fixed 10-step Euler DiT sampler used by `dots.tts-soar` / - `dots.tts-base` is implemented. -- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three - captured graphs (base LM decode, CFM solver, VAE decode) have no - dots.tts equivalent yet. -- **Concurrent requests do not scale.** Each request's 10-step DiT Euler - integration runs serially in the side path (no cross-request batching, - unlike voxcpm2's `enable_batched_cfm`). A community review of this - integration measured no throughput gain at `c=4` concurrent requests - versus `c=1`. -- **`SamplingParams.seed` does not control audio-generation randomness.** - The DiT's flow-matching noise is deterministically derived per-request - from a fixed internal seed (reproducible run-to-run, matching - voxcpm2's `deterministic_cfm_seed` convention) rather than from the - caller-supplied `seed` field — the same limitation voxcpm2 has today. diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index ec3db0501..039073303 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 28, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 67, + EnvironmentVariableCategory.MODEL_SPECIFIC: 57, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -178,10 +178,10 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 38, - ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, + ModelEnvironmentVariableDisposition.PROMOTE: 31, + ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 16, + ModelEnvironmentVariableDisposition.INTERNALIZE: 14, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7, } diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 1f6111d64..76bf3c319 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -272,7 +272,7 @@ def test_from_pipeline_config_normalizes_stage_engine_extras_without_expanding_s def test_frontend_log_stats_flag_is_not_an_unowned_stage_argument(disabled): from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config - config = _from_pipeline_key("dots_tts", cli_overrides={"disable_log_stats": disabled}) + config = _from_pipeline_key("voxcpm2", cli_overrides={"disable_log_stats": disabled}) assert config.stage_configs engine_args = build_engine_args_dict_from_omni_stage_config(config.stage_by_id(0), model="test-model") assert "disable_log_stats" not in engine_args diff --git a/tests/e2e/offline_inference/test_dots_tts_expansion.py b/tests/e2e/offline_inference/test_dots_tts_expansion.py deleted file mode 100644 index 4d13f03fd..000000000 --- a/tests/e2e/offline_inference/test_dots_tts_expansion.py +++ /dev/null @@ -1,129 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""E2E tests for dots.tts native AR offline inference.""" - -from collections.abc import Mapping - -import pytest -import torch - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import get_deploy_config_path - -DOTS_TTS_MODEL = "rednote-hilab/dots.tts-soar" -DEPLOY_CONFIG = get_deploy_config_path("dots_tts.yaml") -SAMPLE_RATE = 48000 - -# dots.tts ships a custom tokenizer carrying the added ``<|audio_gen_start|>`` -# token, so remote code must be explicitly enabled. -_OMNI_RUNNER_PARAM = (DOTS_TTS_MODEL, DEPLOY_CONFIG, {"trust_remote_code": True}) - -pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True) - - -def _build_prompt(omni_runner: OmniRunner, text: str) -> dict: - """Wrap ``text`` in the prefix scaffold dots.tts prefill expects. - - Unlike most TTS models here, a bare ``{"prompt": ...}`` string does not - work: prefill needs ``[文本][文本对应语音]<|audio_gen_start|>`` as - real Qwen2 token ids. - """ - from transformers import AutoTokenizer - - from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import ( - build_dots_tts_prompt, - ) - - tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True) - return build_dots_tts_prompt(tokenizer=tokenizer, text=text) - - -def _extract_audio(multimodal_output: dict) -> torch.Tensor: - """Extract the final complete audio tensor from multimodal output.""" - assert isinstance(multimodal_output, (dict, Mapping)), f"Expected dict/Mapping, got {type(multimodal_output)}" - - # The talker emits delta chunks; the output processor accumulates them - # under "model_outputs" ("audio" kept for backwards compat). - audio = multimodal_output.get("model_outputs") - if audio is None: - audio = multimodal_output.get("audio") - assert audio is not None, f"No audio key, got {list(multimodal_output.keys())}" - - if isinstance(audio, list): - valid = [torch.as_tensor(x).float().cpu().reshape(-1) for x in audio if x is not None] - assert valid, "No valid audio tensors in output list" - audio = torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] - - assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}" - return audio.float().cpu().reshape(-1) - - -@pytest.mark.slow -@pytest.mark.tts -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_dots_tts_zero_shot_001(omni_runner: OmniRunner) -> None: - """Zero-shot synthesis produces audio of a plausible duration.""" - prompt = _build_prompt(omni_runner, "Hello, this is a test of dots TTS running on vLLM Omni.") - - outputs = omni_runner.omni.generate([prompt]) - assert len(outputs) == 1 - - audio = _extract_audio(outputs[0].outputs[0].multimodal_output) - duration_s = audio.shape[0] / SAMPLE_RATE - assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s" - assert torch.isfinite(audio).all(), "Audio contains NaN/Inf" - assert audio.abs().max() > 1e-3, "Audio is silent" - - -@pytest.mark.slow -@pytest.mark.tts -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_dots_tts_prefill_decode_mixed_batch_002(omni_runner: OmniRunner) -> None: - """Mixed prefill+decode batches must keep per-request rows aligned. - - ``compute_logits`` fills ``logits[i]`` positionally from - ``_results_queue``, so every request in the batch has to push exactly one - entry per step, prefills included. When the prefill branch pushed nothing, - a shorter request's stop signal could be read into another request's row. - Regression guard for that pairing (PR #4765 review). - """ - long_text = ( - "This is a deliberately long sentence that stays in the decode phase " - "for many steps, so the shorter requests behind it keep entering " - "prefill alongside it and reproduce the mixed batch pattern." - ) - short_texts = ["Hello one.", "Hello two.", "Hello three.", "Hello four."] - prompts = [_build_prompt(omni_runner, long_text)] + [_build_prompt(omni_runner, t) for t in short_texts] - - outputs = omni_runner.omni.generate(prompts) - assert len(outputs) == len(prompts) - - for i, out in enumerate(outputs): - audio = _extract_audio(out.outputs[0].multimodal_output) - duration_s = audio.shape[0] / SAMPLE_RATE - assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s" - - -@pytest.mark.slow -@pytest.mark.tts -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_dots_tts_deterministic_flow_matching_noise_003(omni_runner: OmniRunner) -> None: - """The same text twice must yield the same audio. - - The DiT flow-matching noise is drawn from a per-request generator seeded - from ``blake2b(seed:request_key:noise_step)`` rather than the global CUDA - RNG. Drawing from the global RNG made output non-reproducible and let - concurrent requests perturb each other's noise streams (PR #4765 review). - """ - text = "Reading aloud every evening helps children build a lasting love of language." - - first = _extract_audio( - omni_runner.omni.generate([_build_prompt(omni_runner, text)])[0].outputs[0].multimodal_output - ) - second = _extract_audio( - omni_runner.omni.generate([_build_prompt(omni_runner, text)])[0].outputs[0].multimodal_output - ) - - assert first.shape == second.shape, f"Length differs across runs: {first.shape} vs {second.shape}" - assert torch.allclose(first, second, atol=1e-4), "Audio differs across runs; flow-matching noise is not seeded" diff --git a/tests/e2e/offline_inference/test_dynin_omni_expansion.py b/tests/e2e/offline_inference/test_dynin_omni_expansion.py deleted file mode 100644 index 922518557..000000000 --- a/tests/e2e/offline_inference/test_dynin_omni_expansion.py +++ /dev/null @@ -1,374 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -E2E offline smoke tests for Dynin-Omni. - -- model: "snu-aidas/Dynin-Omni" -- stage config: vllm_omni/deploy/dynin_omni_ci.yaml -""" - -from __future__ import annotations - -import os -from pathlib import Path -from typing import Any - -import numpy as np -import pytest -import torch -from transformers import AutoTokenizer - -from tests.helpers.mark import hardware_test -from tests.helpers.stage_config import get_deploy_config_path - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - -_REPO_ROOT = Path(__file__).resolve().parents[3] -_DEFAULT_DYNIN_CONFIG_PATH: Path | None = None -_DEFAULT_STAGE_CONFIG_PATH = Path(get_deploy_config_path("dynin_omni_ci.yaml")) - -models = ["snu-aidas/Dynin-Omni"] -stage_configs = [str(_DEFAULT_STAGE_CONFIG_PATH)] -test_params = [(model, stage_config, {"trust_remote_code": True}) for model in models for stage_config in stage_configs] - -DYNIN_CONFIG_PATH = str(_DEFAULT_DYNIN_CONFIG_PATH) if _DEFAULT_DYNIN_CONFIG_PATH is not None else None - -pytestmark = [ - pytest.mark.slow, - pytest.mark.omni, - pytest.mark.parametrize("omni_runner", test_params, indirect=True), -] - - -# prompting util -def _build_mmu_prompt(tokenizer: Any, question: str, dynin_config_path: str | None) -> dict[str, Any]: - encoded = tokenizer(question, return_tensors="pt", add_special_tokens=True) - token_ids = [int(v) for v in encoded["input_ids"][0].tolist()] - attention_mask = [int(v) for v in encoded["attention_mask"][0].tolist()] - additional_information: dict[str, Any] = { - "task": ["mmu"], - "detok_id": [0], - "prompt_length": [len(token_ids)], - "attention_mask": [attention_mask], - "max_new_tokens": [64], - "steps": [64], - "block_length": [16], - "temperature": [0.0], - } - if dynin_config_path: - additional_information["dynin_config_path"] = [str(dynin_config_path)] - return { - "prompt_token_ids": token_ids, - "additional_information": additional_information, - "modalities": ["text"], - } - - -def _build_mmu_multimodal_prompt( - tokenizer: Any, - question: str, - dynin_config_path: str | None, - *, - image: Any | None = None, - audio: tuple[np.ndarray, int] | None = None, -) -> dict[str, Any]: - if image is None and audio is None: - raise ValueError("At least one multimodal input (image or audio) must be provided.") - - prefix_chunks: list[str] = [] - mm_data: dict[str, Any] = {} - if image is not None: - prefix_chunks.append("<|soi|><|image|><|eoi|>") - mm_data["image"] = image - if audio is not None: - prefix_chunks.append("<|soa|><|audio|><|eoa|>") - mm_data["audio"] = audio - - prefixed_question = " ".join(prefix_chunks + [question]).strip() - prompt = _build_mmu_prompt( - tokenizer=tokenizer, - question=prefixed_question, - dynin_config_path=dynin_config_path, - ) - prompt["multi_modal_data"] = mm_data - prompt["modalities"] = ["text"] - return prompt - - -def _generate_synthetic_image(width: int = 224, height: int = 224) -> np.ndarray: - x = np.linspace(0, 255, width, dtype=np.uint8) - y = np.linspace(0, 255, height, dtype=np.uint8)[:, None] - red = np.tile(x, (height, 1)) - green = np.tile(y, (1, width)) - blue = ((red.astype(np.uint16) + green.astype(np.uint16)) // 2).astype(np.uint8) - return np.stack([red, green, blue], axis=-1) - - -def _generate_synthetic_audio(duration_s: int = 5, sample_rate: int = 48_000) -> tuple[np.ndarray, int]: - t = np.linspace(0, duration_s, int(sample_rate * duration_s), endpoint=False, dtype=np.float32) - waveform = 0.1 * np.sin(2.0 * np.pi * 440.0 * t) - return waveform.astype(np.float32), sample_rate - - -# prompting util -def _build_t2s_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]: - # Bypass stage-0 generation and directly validate token->audio decode path. - generated_audio_token_ids = [int(v) for v in ([10, 11, 12, 13, 14] * 32)] - additional_information: dict[str, Any] = { - "task": ["t2s"], - "detok_id": [1], - "generated_token_ids": [generated_audio_token_ids], - "audio_codebook_size": [4096], - } - if dynin_config_path: - additional_information["dynin_config_path"] = [str(dynin_config_path)] - return { - "prompt_token_ids": [0], - "additional_information": additional_information, - "modalities": ["audio"], - } - - -# prompting util -def _build_t2i_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]: - # Bypass stage-0 generation and directly validate token->image decode path. - # MAGVIT decode path expects a square token grid; 1024 tokens -> 32x32. - generated_image_token_ids = [int(v) for v in ([10, 11, 12, 13, 14, 15, 16, 17] * 128)] - additional_information: dict[str, Any] = { - "task": ["t2i"], - "detok_id": [2], - "generated_token_ids": [generated_image_token_ids], - "codebook_size": [8192], - } - if dynin_config_path: - additional_information["dynin_config_path"] = [str(dynin_config_path)] - return { - "prompt_token_ids": [0], - "additional_information": additional_information, - "modalities": ["image"], - } - - -def _configure_dynin_config_env() -> None: - if DYNIN_CONFIG_PATH: - os.environ["DYNIN_CONFIG_PATH"] = str(DYNIN_CONFIG_PATH) - else: - os.environ.pop("DYNIN_CONFIG_PATH", None) - - -def _is_finished_request_output(request_output: Any) -> bool: - if request_output is None: - return False - req_list = request_output if isinstance(request_output, list) else [request_output] - for req in req_list: - if req is not None and bool(getattr(req, "finished", False)): - return True - return False - - -def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None: - matched = [ - stage_output for stage_output in outputs if getattr(stage_output, "final_output_type", None) == output_type - ] - if not matched: - return None - - # Prefer the latest finished chunk to avoid picking an intermediate stream output. - for stage_output in reversed(matched): - if _is_finished_request_output(stage_output): - return stage_output - return matched[-1] - - -def _to_token_list(value: Any) -> list[int]: - if value is None: - return [] - if hasattr(value, "detach"): - value = value.detach() - if hasattr(value, "cpu"): - value = value.cpu() - if hasattr(value, "flatten"): - value = value.flatten().tolist() - if isinstance(value, tuple): - value = list(value) - if not isinstance(value, list): - return [] - out: list[int] = [] - for token in value: - if isinstance(token, bool): - continue - try: - out.append(int(token)) - except Exception: - continue - return out - - -def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str: - request_output = stage_output - if request_output is None: - return "" - req_list = request_output if isinstance(request_output, list) else [request_output] - for req in req_list: - completions = getattr(req, "outputs", None) or [] - if not completions: - continue - completion = completions[0] - mm_out = ( - getattr(completion, "multimodal_output", None) - or getattr(req, "multimodal_output", None) - or getattr(stage_output, "multimodal_output", None) - or {} - ) - text = mm_out.get("text") - if isinstance(text, list) and text: - text = text[-1] - if isinstance(text, str) and text.strip(): - return text.strip() - if tokenizer is not None: - for key in ("text_tokens", "token_ids"): - token_ids = _to_token_list(mm_out.get(key)) - if not token_ids: - continue - decoded = tokenizer.decode(token_ids, skip_special_tokens=True) - if isinstance(decoded, str) and decoded.strip(): - return decoded.strip() - fallback = getattr(completion, "text", None) - if isinstance(fallback, str) and fallback.strip(): - return fallback.strip() - return "" - - -def _extract_audio(stage_output: Any) -> Any | None: - request_output = stage_output - if request_output is None: - return None - req_list = request_output if isinstance(request_output, list) else [request_output] - for req in req_list: - completions = getattr(req, "outputs", None) or [] - if not completions: - continue - completion = completions[0] - mm_out = getattr(completion, "multimodal_output", None) or {} - if "audio" in mm_out: - return mm_out["audio"] - return None - - -def _extract_image(stage_output: Any) -> Any | None: - request_output = stage_output - if request_output is None: - return None - req_list = request_output if isinstance(request_output, list) else [request_output] - for req in req_list: - completions = getattr(req, "outputs", None) or [] - if not completions: - continue - completion = completions[0] - mm_out = getattr(completion, "multimodal_output", None) or {} - if "image" in mm_out: - return mm_out["image"] - return None - - -def _numel(value: Any) -> int: - if value is None: - return 0 - if isinstance(value, torch.Tensor): - return int(value.numel()) - shape = getattr(value, "shape", None) - if shape is not None: - try: - total = 1 - for dim in shape: - total *= int(dim) - return int(total) - except Exception: - pass - if isinstance(value, (list, tuple)): - return len(value) - return 0 - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) -def test_dynin_t2i_decode_to_image(omni_runner) -> None: - _configure_dynin_config_env() - prompt = _build_t2i_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH) - - outputs = omni_runner.generate([prompt]) - - image_output = _find_stage_output(outputs, "image") - assert image_output is not None - image_value = _extract_image(image_output) - assert image_value is not None - assert _numel(image_value) > 0 - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) -def test_dynin_mmu_to_text(omni_runner) -> None: - _configure_dynin_config_env() - tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True) - prompt = _build_mmu_prompt( - tokenizer=tokenizer, - question="What is 2 + 2? Answer in one short sentence.", - dynin_config_path=DYNIN_CONFIG_PATH, - ) - - outputs = omni_runner.generate([prompt]) - - text_output = _find_stage_output(outputs, "text") - assert text_output is not None - text_content = _extract_text(text_output, tokenizer=tokenizer) - assert text_content - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) -def test_dynin_image_to_text(omni_runner) -> None: - _configure_dynin_config_env() - tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True) - prompt = _build_mmu_multimodal_prompt( - tokenizer=tokenizer, - question="Describe the image briefly in one sentence.", - dynin_config_path=DYNIN_CONFIG_PATH, - image=_generate_synthetic_image(), - ) - - outputs = omni_runner.generate([prompt]) - - text_output = _find_stage_output(outputs, "text") - assert text_output is not None - text_content = _extract_text(text_output, tokenizer=tokenizer) - assert text_content - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) -def test_dynin_speech_to_text(omni_runner) -> None: - _configure_dynin_config_env() - tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True) - prompt = _build_mmu_multimodal_prompt( - tokenizer=tokenizer, - question="Transcribe the audio briefly in one sentence.", - dynin_config_path=DYNIN_CONFIG_PATH, - audio=_generate_synthetic_audio(), - ) - - outputs = omni_runner.generate([prompt]) - - text_output = _find_stage_output(outputs, "text") - assert text_output is not None - text_content = _extract_text(text_output, tokenizer=tokenizer) - assert text_content - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) -def test_dynin_t2s_decode_to_audio(omni_runner) -> None: - _configure_dynin_config_env() - prompt = _build_t2s_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH) - - outputs = omni_runner.generate([prompt]) - - audio_output = _find_stage_output(outputs, "audio") - assert audio_output is not None - audio_value = _extract_audio(audio_output) - assert audio_value is not None - assert _numel(audio_value) > 0 diff --git a/tests/e2e/online_serving/test_dots_tts_expansion.py b/tests/e2e/online_serving/test_dots_tts_expansion.py deleted file mode 100644 index e5df55a33..000000000 --- a/tests/e2e/online_serving/test_dots_tts_expansion.py +++ /dev/null @@ -1,91 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -E2E Online expansion tests for dots.tts(text-only). -""" - -import os - -import pytest - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path - -pytestmark = [pytest.mark.full_model, pytest.mark.tts] - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - -MODEL = "dots-studio/dots.tts-soar" -DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0 -_MIN_AUDIO_BYTES = 40_000 -MAX_CONCURRENT = 4 - - -def get_prompt(prompt_type="text"): - prompts = { - "text": "The weather is nice today, perfect for a walk in the park.", - } - return prompts.get(prompt_type, prompts["text"]) - - -tts_server_params = [ - pytest.param( - OmniServerParams( - model=MODEL, - stage_config_path=get_deploy_config_path("dots_tts.yaml"), - server_args=["--trust-remote-code", "--disable-log-stats"], - ), - id="dots_tts", - ) -] - - -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_only_streaming_001(omni_server, openai_client) -> None: - """ - Test text-only text-to-audio via OpenAI API. - Deploy Setting: dots_tts.yaml - Input Modal: text - Output Modal: audio - Input Setting: stream=True - Datasets: few requests (max_num_seqs=4) - """ - request_config = { - "model": omni_server.model, - "input": get_prompt(), - "stream": True, - "stream_format": "audio", - "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, - "response_format": "wav", - "voice": "default", - "min_audio_bytes": _MIN_AUDIO_BYTES, - } - openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) - - -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_only_nonstreaming_pcm_001(omni_server, openai_client) -> None: - """ - Test text-only non-stream PCM output via OpenAI API. - Deploy Setting: dots_tts.yaml - Input Modal: text - Output Modal: audio (pcm) - Input Setting: stream=False - Datasets: single request - """ - request_config = { - "model": omni_server.model, - "input": get_prompt(), - "response_format": "pcm", - "stream": False, - "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, - "voice": "default", - "min_audio_bytes": _MIN_AUDIO_BYTES, - "min_hnr_db": -2.0, - # dots.tts-soar emits native 48 kHz PCM, which has no rate header. - "expected_sample_rate": 48_000, - } - openai_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_dynin_omni_expansion.py b/tests/e2e/online_serving/test_dynin_omni_expansion.py deleted file mode 100644 index c4cd40ff8..000000000 --- a/tests/e2e/online_serving/test_dynin_omni_expansion.py +++ /dev/null @@ -1,99 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -""" -Example online tests for Dynin-Omni model. -""" - -import base64 -import os -from io import BytesIO - -import pytest -from vllm.assets.image import ImageAsset - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path - -pytestmark = [pytest.mark.slow, pytest.mark.omni] - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - -MODEL = "snu-aidas/Dynin-Omni" -STAGE_CONFIG = get_deploy_config_path("dynin_omni_ci.yaml") - -T2I_PROMPT = "A high quality detailed living room interior photo." -T2S_PROMPT = "Please read this sentence naturally: Hello from online serving." -I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors." - -TEST_PARAMS = [ - OmniServerParams( - model=MODEL, - stage_config_path=STAGE_CONFIG, - stage_init_timeout=600, - server_args=["--trust-remote-code"], - ) -] -_STAGE_COUNT = 3 -_I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False} - - -def _build_t2i_messages(prompt: str) -> list[dict]: - return [{"role": "user", "content": [{"type": "text", "text": f"<|t2i|> {prompt}"}]}] - - -def _build_t2s_messages(prompt: str) -> list[dict]: - return [{"role": "user", "content": [{"type": "text", "text": f"<|t2s|> {prompt}"}]}] - - -def _build_i2i_messages(prompt: str) -> list[dict]: - input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") - buffer = BytesIO() - input_image.save(buffer, format="JPEG") - image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8") - return [ - { - "role": "user", - "content": [ - {"type": "text", "text": f"<|i2i|> {prompt}"}, - {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}}, - ], - } - ] - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) -@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -def test_send_i2i_request_001(omni_server, online_client) -> None: - request_config = { - "model": omni_server.model, - "messages": _build_i2i_messages(I2I_PROMPT), - "modalities": ["image"], - "extra_body": { - "sampling_params_list": [dict(_I2I_STAGE_SAMPLING) for _ in range(_STAGE_COUNT)], - }, - } - online_client.send_diffusion_request(request_config) - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) -@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -def test_send_t2i_request_001(omni_server, online_client) -> None: - request_config = { - "model": omni_server.model, - "messages": _build_t2i_messages(T2I_PROMPT), - "modalities": ["image"], - } - online_client.send_diffusion_request(request_config) - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) -@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -def test_send_t2s_request_001(omni_server, online_client) -> None: - request_config = { - "model": omni_server.model, - "messages": _build_t2s_messages(T2S_PROMPT), - "modalities": ["audio"], - "audio_ref_text": T2S_PROMPT, - } - online_client.send_omni_request(request_config) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 0e88a00ef..7132d7512 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -539,7 +539,7 @@ def test_tensor_parallel_size_none_is_handled(): engine_args = OmegaConf.create({"stage_id": 0, "engine_args": {"tensor_parallel_size": None}}) args = build_engine_args_dict( engine_args, - model="snu-aidas/Dynin-Omni", + model="Qwen/Qwen2-VL-2B-Instruct", ) assert isinstance(args, dict) assert "tensor_parallel_size" not in args diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 77117ed0d..36967e706 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -5035,89 +5035,6 @@ class TestMingFlashOmniTTSServing: assert prompt["additional_information"]["voice_name"] == "test" -@pytest.fixture -def dots_tts_server(mocker: MockerFixture): - mocker.patch( - "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", - return_value=set(), - ) - mocker.patch( - "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", - return_value=None, - ) - - mock_engine_client = mocker.MagicMock() - mock_engine_client.errored = False - mock_engine_client.model_config = mocker.MagicMock( - model="dots-studio/dots.tts-soar", - ) - mock_engine_client.default_sampling_params_list = [ - SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None) - ] - mock_engine_client.tts_batch_max_items = 32 - mock_engine_client.generate = mocker.MagicMock(return_value="generator") - mock_engine_client.stage_configs = [ - SimpleNamespace( - engine_args=SimpleNamespace(model_stage="latent_generator", model_arch="DotsTTSForConditionalGeneration"), - tts_args={}, - ) - ] - - mock_models = mocker.MagicMock() - mock_models.is_base_model.return_value = True - - return OmniOpenAIServingSpeech( - engine_client=mock_engine_client, - models=mock_models, - request_logger=mocker.MagicMock(), - ) - - -class TestDotsTTSServing: - def test_dots_tts_prompt_validation(self, dots_tts_server): - request = OpenAICreateSpeechRequest(input="Hello", ref_text="Reference transcript") - error = dots_tts_server._validate_tts_request(request) - assert error is not None - assert "ref_text" in error - - request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc") - error = dots_tts_server._validate_tts_request(request) - assert error is not None - assert "ref_audio" in error - - request = OpenAICreateSpeechRequest(input="Hello", voice="test") - error = dots_tts_server._validate_tts_request(request) - assert error is not None - assert "voice" in error - - request = OpenAICreateSpeechRequest(input="Hello", speaker_embedding=[1, 2, 3]) - error = dots_tts_server._validate_tts_request(request) - assert error is not None - assert "speaker_embedding" in error - - request = OpenAICreateSpeechRequest(input="Hello", x_vector_only_mode=True) - error = dots_tts_server._validate_tts_request(request) - assert error is not None - assert "x_vector_only_mode" in error - - def test_dots_tts_adapter_awaits_async_prompt_builder(self, dots_tts_server, mocker: MockerFixture): - build_prompt_async = mocker.patch.object( - dots_tts_server._adapter, - "_build_prompt_async", - new=mocker.AsyncMock(return_value={"prompt_token_ids": [1, 2, 3]}), - ) - request = OpenAICreateSpeechRequest(input="Hello") - asyncio.run(dots_tts_server._prepare_speech_generation(request)) - build_prompt_async.assert_awaited_once_with("Hello") - - def test_dots_tts_adapter_apply_sampling_overrides(self, dots_tts_server, mocker: MockerFixture): - mocker.patch.object(dots_tts_server._adapter, "build", return_value=PreparedRequest(prompt="Hello")) - request = OpenAICreateSpeechRequest(input="Hello", max_new_tokens=10) - asyncio.run(dots_tts_server._prepare_speech_generation(request)) - sampling_params_list = dots_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"] - assert sampling_params_list[0].max_tokens == 10 - - class TestTTSAsyncOffloading: """Tests for event-loop-safe offloading of blocking TTS operations.""" diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index a2507a056..22aab89ac 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -71,7 +71,6 @@ EXPECTED_MODEL_TYPES = { "indextts2", "indextts2_5", "gepard", - "dots_tts", } diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 90e7e0fee..8ecb137a2 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -263,7 +263,6 @@ def test_arch_matching_is_a_fallback_not_an_override(): def test_shared_latent_generator_resolves_by_architecture_priority(): assert detect_tts_model_type("latent_generator", "VoxCPM2TalkerForConditionalGeneration") == "voxcpm2" - assert detect_tts_model_type("latent_generator", "DotsTTSForConditionalGeneration") == "dots_tts" def test_stage_keys_cover_legacy_stage_set(): diff --git a/tests/model_executor/models/dots_tts/__init__.py b/tests/model_executor/models/dots_tts/__init__.py deleted file mode 100644 index 208f01a7c..000000000 --- a/tests/model_executor/models/dots_tts/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/model_executor/models/dots_tts/test_talker.py b/tests/model_executor/models/dots_tts/test_talker.py deleted file mode 100644 index ed3c6bcaa..000000000 --- a/tests/model_executor/models/dots_tts/test_talker.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Regression tests for the dots.tts talker's stop-signal pairing and -per-request state lifecycle. - -Pure-CPU tests on a bare talker (``__new__`` + the handful of attributes -the methods under test read) — no weights, no GPU, no engine. Mirrors -tests/model_executor/models/voxcpm2/test_talker_state_eviction.py. -""" - -from __future__ import annotations - -import functools -from types import SimpleNamespace - -import pytest - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -torch = pytest.importorskip("torch") - -_VOCAB_SIZE = 151672 - - -@functools.lru_cache(maxsize=1) -def _dots_tts_talker_mod(): - """Defer talker import (pulls vLLM model_executor) until first use.""" - from vllm_omni.model_executor.models.dots_tts.dots_tts_talker import ( - DotsTTSForConditionalGeneration, - _RequestState, - ) - - return DotsTTSForConditionalGeneration, _RequestState - - -def _make_bare_talker(): - DotsTTSForConditionalGeneration, _ = _dots_tts_talker_mod() - talker = DotsTTSForConditionalGeneration.__new__(DotsTTSForConditionalGeneration) - talker.config = SimpleNamespace(vocab_size=_VOCAB_SIZE) - talker._active_states = {} - talker._pending_requests = [] - talker._results_queue = [] - talker._audio_queue = [] - talker._deferred_cleanup_ids = set() - talker._beta_trace = False - return talker - - -def _add_state(talker, req_id: str, *, is_stopping: bool = False): - _, _RequestState = _dots_tts_talker_mod() - state = _RequestState(request_id=req_id) - state.prefill_completed = True - state.is_stopping = is_stopping - talker._active_states[req_id] = state - return state - - -def _hidden(bsz: int) -> torch.Tensor: - return torch.zeros(bsz, 1536, dtype=torch.float32) - - -class TestComputeLogitsPairing: - """compute_logits pairs _results_queue entries with batch rows - positionally, so every request must push exactly one entry per step — - including prefills, which push a (req_id, None) placeholder.""" - - def test_mixed_batch_prefill_placeholder_keeps_rows_aligned(self) -> None: - talker = _make_bare_talker() - _add_state(talker, "prefill-req") - decode_state = _add_state(talker, "decode-req", is_stopping=True) - decode_state.precomputed_stop_logits = torch.tensor([[0.1, 0.9]]) - - # Row 0 is the prefill (None placeholder), row 1 the stopping decode. - talker._results_queue = [ - ("prefill-req", None), - ("decode-req", torch.tensor([[0.1, 0.9]])), - ] - logits = talker.compute_logits(_hidden(2)) - - assert logits.shape == (2, _VOCAB_SIZE) - # Prefill row: forced continue. - assert logits[0, 0].item() == 1.0 - assert logits[0, 1].item() == float("-inf") - # Stopping decode row: stop slot wins. - assert logits[1, 1].item() == 1.0 - assert logits[1, 0].item() == 0.0 - # Everything else stays masked. - assert logits[0, 2].item() == float("-inf") - assert logits[1, 2].item() == float("-inf") - # Queue drained. - assert talker._results_queue == [] - - def test_below_threshold_stop_logits_force_continue(self) -> None: - """Raw (continue, stop) softmax must not reach the sampler: a - prob_stop in (0.5, 0.8] would flip greedy argmax even though - is_stopping stayed False (the 0.8-threshold contract).""" - talker = _make_bare_talker() - _add_state(talker, "req", is_stopping=False) - talker._results_queue = [("req", torch.tensor([[0.3, 0.7]]))] - - logits = talker.compute_logits(_hidden(1)) - - assert logits[0, 0].item() == 1.0 - assert logits[0, 1].item() == float("-inf") - - def test_empty_queue_defaults_to_all_continue(self) -> None: - talker = _make_bare_talker() - logits = talker.compute_logits(_hidden(3)) - - assert torch.all(logits[:, 0] == 1.0) - assert torch.all(logits[:, 1] == float("-inf")) - - -class TestStateEviction: - """on_requests_finished fires BEFORE forward(), so eviction must be - deferred to _flush_deferred_cleanup at the end of the step — the - finishing request's audio still drains through the current forward.""" - - def test_finished_request_survives_until_flush(self) -> None: - talker = _make_bare_talker() - _add_state(talker, "req-a") - _add_state(talker, "req-b") - - talker.on_requests_finished(["req-a"]) - assert "req-a" in talker._active_states # not evicted yet - - talker._flush_deferred_cleanup() - assert "req-a" not in talker._active_states - assert "req-b" in talker._active_states - assert talker._deferred_cleanup_ids == set() - - def test_unknown_request_id_is_ignored(self) -> None: - talker = _make_bare_talker() - talker.on_requests_finished(["never-seen"]) - talker._flush_deferred_cleanup() - assert talker._active_states == {} diff --git a/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py b/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py deleted file mode 100644 index d2138541f..000000000 --- a/tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py +++ /dev/null @@ -1,40 +0,0 @@ -from __future__ import annotations - -import pytest -from transformers import PreTrainedModel - -from vllm_omni.model_executor.models.dynin_omni.dynin_omni_token2audio import ( - _ensure_transformers_tied_weights_compat, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_transformers_tied_weights_compat(monkeypatch): - """Compat shim must not permanently mutate ``PreTrainedModel``. - - ``monkeypatch.delattr(..., raising=False)`` is a no-op undo when the attr - was absent. ``_ensure_transformers_tied_weights_compat`` then assigns a - read-only ``property`` on the class that later PreTrainedModel subclasses - (e.g. qwen3 TTS mask-cache) can trip over with - ``AttributeError: property ... has no setter``. - """ - monkeypatch.delattr(PreTrainedModel, "all_tied_weights_keys", raising=False) - - _ensure_transformers_tied_weights_compat() - installed = PreTrainedModel.__dict__.get("all_tied_weights_keys") - assert installed is not None - - # Re-own the installed descriptor through monkeypatch so teardown always - # removes/restores it, even when the attr did not exist before this test. - delattr(PreTrainedModel, "all_tied_weights_keys") - monkeypatch.setattr(PreTrainedModel, "all_tied_weights_keys", installed, raising=False) - - model = PreTrainedModel.__new__(PreTrainedModel) - model._tied_weights_keys = ["decoder.weight"] - model._dynamic_tied_weights_keys = {"lm_head.weight": "decoder.weight"} - - assert model.all_tied_weights_keys == { - "decoder.weight": None, - "lm_head.weight": "decoder.weight", - } diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index c1a6041c7..208cc72c6 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for Qwen3-Omni streaming thinker→talker / talker→codec helpers (PR #2581).""" from __future__ import annotations @@ -356,7 +356,7 @@ def test_accumulator_concat_default_when_no_replace_keys() -> None: class _StubMixin(OmniConnectorModelRunnerMixin): def __init__(self): self._pending_full_payload_send = {} - self._full_payload_replace_keys_cached = frozenset() + self._full_payload_replace_keys_cached: frozenset[str] = frozenset() stub = _StubMixin() stub.accumulate_full_payload_output( @@ -415,51 +415,6 @@ def test_covo_audio_llm2code2wav_full_payload_smoke() -> None: assert payload["meta"]["finished"].item() is True -def test_dynin_omni_token_only_smoke() -> None: - """Smoke: dynin_omni token-only builders return placeholders.""" - from vllm_omni.model_executor.stage_input_processors.dynin_omni import ( - token2text_to_token2image_token_only, - ) - - class _Out: - def __init__(self, tids, mm=None): - self.token_ids = tids - self.multimodal_output = mm - - class _Wrapper: - def __init__(self, tids, mm=None): - self.outputs = [_Out(tids, mm)] - self.request_id = "r0" - - class _Stage: - def __init__(self, outs): - self.engine_outputs = outs - - src = [_Wrapper([10, 11, 12])] - out = token2text_to_token2image_token_only([_Stage(src)], [0]) - assert len(out) == 1 - assert len(out[0]["prompt_token_ids"]) == 3 - assert out[0]["additional_information"] is None - - -def test_dynin_omni_full_payload_smoke() -> None: - """Smoke: dynin_omni producer-side payload builder returns nested OmniPayload + carries metadata.""" - from types import SimpleNamespace - - from vllm_omni.model_executor.stage_input_processors.dynin_omni import ( - token2text_to_token2image_full_payload, - ) - - pooling = {"token_ids": [1, 2, 3]} - req = SimpleNamespace(output_token_ids=[], additional_information={"speaker": ["alice"]}) - payload = token2text_to_token2image_full_payload(None, pooling, req) - assert payload is not None - assert payload["codes"]["audio"] == [1, 2, 3] - assert payload["meta"]["finished"].item() is True - # additional_information is normalized + carried forward (speaker stays list-wrapped). - assert payload.get("speaker") == ["alice"] - - def test_qwen2_5_omni_talker2code2wav_token_only_smoke() -> None: """Smoke: qwen2_5_omni talker→code2wav token_only marker + boundary strip.""" from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import ( diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index f57d7ffef..4ba700e84 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -145,13 +145,6 @@ _MODEL_PROMOTE = ( "COSYVOICE3_ESTIMATOR_ONNX", "COSYVOICE3_TRT", "COSYVOICE3_TRT_CACHE", - "DYNIN_CONFIG_PATH", - "DYNIN_MAGVIT_REMOTE_CODE_LOCAL_FILES_ONLY", - "DYNIN_MAGVIT_REMOTE_CODE_REPO_ID", - "DYNIN_MAGVIT_REMOTE_CODE_REVISION", - "DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY", - "DYNIN_REMOTE_CODE_REPO_ID", - "DYNIN_REMOTE_CODE_REVISION", "FASTVIDEO_VSA_SM100A", "HIGGS_AUDIO_TOKENIZER_PATH", "INTERNVLA_A1_COSMOS_DECODER_PATH", @@ -179,7 +172,6 @@ _MODEL_PROMOTE = ( ) _MODEL_REQUEST_SCOPE = ( - "DOTS_TTS_DIT_NUM_STEPS", "MAGI2_NEGATIVE_PROMPT", "STEP_AUDIO2_DEFAULT_PROMPT_WAV", "VLLM_GEPARD_END_FADE_MS", @@ -194,8 +186,6 @@ _MODEL_EXTERNAL: tuple[str, ...] = () _MODEL_INTERNALIZE = ( "COSYVOICE3_BATCH_FLOW_DEBUG", - "DOTS_TTS_BETA_TRACE", - "DYNIN_S2U_VENDOR_ROOT", "DZ_PHASE_TIMING", "MAGI2_ALLOW_UNSUPPORTED_TOPOLOGY", "MAGI2_ROUTER_BIAS_SOURCE", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index eada1019d..4876bd0e1 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -52,9 +52,7 @@ from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_P from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE -from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE -from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE from vllm_omni.model_executor.models.gepard.pipeline import GEPARD_PIPELINE from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE @@ -168,7 +166,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE, "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, - "dots_tts": DOTS_TTS_PIPELINE, "auk": AUK_PIPELINE, "cosyvoice3": COSYVOICE3_PIPELINE, "audex_tts": AUDEX_TTS_PIPELINE, @@ -202,7 +199,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "minimax_music3": MINIMAX_MUSIC3_PIPELINE, "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE, "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE, - "dynin_omni": DYNIN_OMNI_PIPELINE, "indextts2": INDEXTTS2_PIPELINE, "indextts2_5": INDEXTTS25_PIPELINE, } diff --git a/vllm_omni/deploy/dots_tts.yaml b/vllm_omni/deploy/dots_tts.yaml deleted file mode 100644 index acd28de52..000000000 --- a/vllm_omni/deploy/dots_tts.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# dots.tts deployment config (single-stage AR talker). -# -# Conservative first-integration settings: eager execution, no CUDA -# Graphs, no batched side-path optimizations yet — see voxcpm2.yaml for -# the runtime knobs a mature deployment of this model family grows. -# -# Verified locally on RTX 5080 16GB. - -async_chunk: false -dtype: bfloat16 - -stages: - - stage_id: 0 - max_num_seqs: 4 - gpu_memory_utilization: 0.8 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - async_scheduling: true - max_num_batched_tokens: 4096 - max_model_len: 4096 - devices: "0" - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 4096 - repetition_penalty: 1.0 diff --git a/vllm_omni/deploy/dynin_omni.yaml b/vllm_omni/deploy/dynin_omni.yaml deleted file mode 100644 index f7a457695..000000000 --- a/vllm_omni/deploy/dynin_omni.yaml +++ /dev/null @@ -1,38 +0,0 @@ -# Dynin-Omni 3-stage deploy: token2text → token2image → token2audio. -# -# Topology (stages, processor functions, model_arch) is frozen in -# vllm_omni/model_executor/models/dynin_omni/pipeline.py. This file carries -# only the deploy-time knobs (devices, GPU memory, batched tokens). -# -# async_chunk=false: dynin hands the full token payload across stages through -# the worker-connector data plane (``*_full_payload`` producers), not chunked -# streaming. Pipeline-wide ``trust_remote_code`` is required by the custom -# Dynin HF modeling code. -async_chunk: false -trust_remote_code: true -enable_prefix_caching: false - -stages: - - stage_id: 0 - devices: "0" - max_num_seqs: 1 - gpu_memory_utilization: 0.5 - max_num_batched_tokens: 32768 - enforce_eager: true - async_scheduling: false - - - stage_id: 1 - devices: "0" - max_num_seqs: 1 - gpu_memory_utilization: 0.1 - max_num_batched_tokens: 32768 - enforce_eager: true - async_scheduling: false - - - stage_id: 2 - devices: "0" - max_num_seqs: 1 - gpu_memory_utilization: 0.1 - max_num_batched_tokens: 32768 - enforce_eager: true - async_scheduling: false diff --git a/vllm_omni/deploy/dynin_omni_ci.yaml b/vllm_omni/deploy/dynin_omni_ci.yaml deleted file mode 100644 index ce20c093b..000000000 --- a/vllm_omni/deploy/dynin_omni_ci.yaml +++ /dev/null @@ -1,35 +0,0 @@ -# Dynin-Omni deploy for e2e smoke tests (3-stage, one GPU). -# -# Topology is frozen in models/dynin_omni/pipeline.py; ``pipeline: dynin_omni`` -# pins it explicitly so this CI YAML resolves regardless of model_type -# auto-detection. Differs from the default deploy only in reduced GPU memory -# and a smaller max_num_batched_tokens so the pipeline fits L4-class CI GPUs. -pipeline: dynin_omni -async_chunk: false -trust_remote_code: true -enable_prefix_caching: false - -stages: - - stage_id: 0 - devices: "0" - max_num_seqs: 1 - gpu_memory_utilization: 0.5 - max_num_batched_tokens: 4096 - enforce_eager: true - async_scheduling: false - - - stage_id: 1 - devices: "0" - max_num_seqs: 1 - gpu_memory_utilization: 0.2 - max_num_batched_tokens: 4096 - enforce_eager: true - async_scheduling: false - - - stage_id: 2 - devices: "0" - max_num_seqs: 1 - gpu_memory_utilization: 0.2 - max_num_batched_tokens: 4096 - enforce_eager: true - async_scheduling: false diff --git a/vllm_omni/deploy/dynin_omni_multiconnector.yaml b/vllm_omni/deploy/dynin_omni_multiconnector.yaml deleted file mode 100644 index 5abb4c5a3..000000000 --- a/vllm_omni/deploy/dynin_omni_multiconnector.yaml +++ /dev/null @@ -1,63 +0,0 @@ -# Dynin-Omni 3-stage deploy with cross-node connectors (Mooncake demo). -# -# Topology is frozen in models/dynin_omni/pipeline.py (``pipeline: dynin_omni``). -# This variant wires a MooncakeConnector for each stage hand-off in place of -# the default SharedMemoryConnector. Update the ``extra`` host/master/metadata -# fields to match your Mooncake deployment before use. -pipeline: dynin_omni -async_chunk: false -trust_remote_code: false -enable_prefix_caching: false - -connectors: - # Mooncake connector for cross-node/intra-node communication. - mooncake_connector: - name: MooncakeConnector - extra: - host: "127.0.0.1" - metadata_server: "http://10.90.67.86:8080/metadata" - master: "10.90.67.86:50051" - segment: 512000000 # 512MB - localbuf: 64000000 # 64MB - proto: "tcp" - - # Yuanrong connector for cross-node/intra-node communication. - yuanrong_connector: - name: YuanrongConnector - extra: - host: "127.0.0.1" - port: "35000" - - # Alternative SHM connector with an explicit threshold. - shared_memory_connector: - name: SharedMemoryConnector - -stages: - - stage_id: 0 - devices: "0" - max_num_seqs: 1 - max_num_batched_tokens: 32768 - enforce_eager: true - async_scheduling: false - output_connectors: - to_stage_1: mooncake_connector - - - stage_id: 1 - devices: "0" - max_num_seqs: 1 - max_num_batched_tokens: 32768 - enforce_eager: true - async_scheduling: false - input_connectors: - from_stage_0: mooncake_connector - output_connectors: - to_stage_2: mooncake_connector - - - stage_id: 2 - devices: "0" - max_num_seqs: 1 - max_num_batched_tokens: 32768 - enforce_eager: true - async_scheduling: false - input_connectors: - from_stage_1: mooncake_connector diff --git a/vllm_omni/diffusion/model_loader/hub_prefetch.py b/vllm_omni/diffusion/model_loader/hub_prefetch.py index a1e893b3b..a99751d6f 100644 --- a/vllm_omni/diffusion/model_loader/hub_prefetch.py +++ b/vllm_omni/diffusion/model_loader/hub_prefetch.py @@ -30,8 +30,7 @@ Two environmental factors mask the race on main: ``cached_files`` (plural) which batch-resolves every shard listed in the index up-front via ``os.path.isfile`` and raises immediately if any shard is still sitting under its ``*.incomplete`` name. Same wave of v5 changes - that introduced ``tie_weights(missing_keys=..., recompute_mapping=...)`` - (see the Dynin shim in ``dynin_omni_token2text.py``). + that introduced ``tie_weights(missing_keys=..., recompute_mapping=...)``. * CI shares ``HF_HOME=/fsx/hf_cache`` across pipelines (both the ``vllm-omni`` and ``vllm-omni-rebase`` pipelines mount the same FS). That cache is normally warm for long-lived repos like ``Qwen-Image-Edit-2509``, @@ -195,19 +194,22 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]: flock_held = False # --- fcntl.flock path --- + fcntl_mod: Any = None try: - import fcntl # type: ignore[import-not-found] - except ImportError: # pragma: no cover - non-POSIX (Windows) - fcntl = None + import fcntl as _fcntl - if fcntl is not None: + fcntl_mod = _fcntl + except ImportError: # pragma: no cover - non-POSIX (Windows) + pass + + if fcntl_mod is not None: try: lock_dir = _node_lock_dir() except OSError as exc: logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc) - fcntl = None # force dotfile fallback + fcntl_mod = None # force dotfile fallback - if fcntl is not None and lock_dir is not None: + if fcntl_mod is not None and lock_dir is not None: lock_path = os.path.join(lock_dir, _safe_repo_filename(model)) try: fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644) @@ -217,7 +219,7 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]: if fd is not None: try: - fcntl.flock(fd, fcntl.LOCK_EX) + fcntl_mod.flock(fd, fcntl_mod.LOCK_EX) flock_held = True logger.info("Acquired flock prefetch lock for %s at %s", model, lock_path) except OSError as exc: @@ -249,9 +251,9 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]: try: yield finally: - if flock_held and fd is not None: + if flock_held and fd is not None and fcntl_mod is not None: with contextlib.suppress(OSError): - fcntl.flock(fd, fcntl.LOCK_UN) + fcntl_mod.flock(fd, fcntl_mod.LOCK_UN) if fd is not None: with contextlib.suppress(OSError): os.close(fd) diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 2343315b9..826ccec72 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -7,7 +7,7 @@ from __future__ import annotations import importlib import inspect from collections import deque -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import torch from vllm.distributed.parallel_state import get_tp_group @@ -29,21 +29,23 @@ if TYPE_CHECKING: class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): """Own payload caching, full/chunk transport, and connector I/O.""" + _full_payload_replace_keys_cached: frozenset[str] + # ------------------------------------------------------------------ # # Local payload cache (RFC §2.4 – Model Runner ownership) # ------------------------------------------------------------------ # def put_local_stage_payload(self, req_id: str, payload: OmniPayload) -> None: """Store a full stage payload in the local cache.""" - self._local_stage_payload_cache[req_id] = payload + self._local_stage_payload_cache[req_id] = cast(dict[str, Any], payload) def get_local_stage_payload(self, req_id: str) -> OmniPayload | None: """Read a stage payload without removing it.""" - return self._local_stage_payload_cache.get(req_id) + return cast(OmniPayload | None, self._local_stage_payload_cache.get(req_id)) def pop_local_stage_payload(self, req_id: str) -> OmniPayload | None: """Remove and return a stage payload (consume after use).""" - return self._local_stage_payload_cache.pop(req_id, None) + return cast(OmniPayload | None, self._local_stage_payload_cache.pop(req_id, None)) def put_local_request_metadata(self, req_id: str, metadata: dict[str, Any]) -> None: """Store lightweight scheduling metadata for a request.""" @@ -58,7 +60,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): # ------------------------------------------------------------------ # @classmethod - def _extract_scheduling_metadata(cls, payload: OmniPayload) -> dict[str, Any]: + def _extract_scheduling_metadata(cls, payload: dict[str, Any] | OmniPayload) -> dict[str, Any]: """Extract only the fields the scheduler needs from a full payload.""" extracted: dict[str, Any] = {} meta = payload.get("meta") if isinstance(payload, dict) else None @@ -66,11 +68,13 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): if "next_stage_prompt_len" in meta: extracted["next_stage_prompt_len"] = meta["next_stage_prompt_len"] - elif "next_stage_prompt_len" in payload: - logger.warning_once( - "legacy flat 'next_stage_prompt_len' key in payload; expected 'meta.next_stage_prompt_len'" - ) - extracted["next_stage_prompt_len"] = payload["next_stage_prompt_len"] + else: + legacy_prompt_len = cast(dict[str, Any], payload).get("next_stage_prompt_len") + if legacy_prompt_len is not None: + logger.warning_once( + "legacy flat 'next_stage_prompt_len' key in payload; expected 'meta.next_stage_prompt_len'" + ) + extracted["next_stage_prompt_len"] = legacy_prompt_len audio_codes = cls._payload_audio_codes(payload) if audio_codes is not None: @@ -129,7 +133,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): return None @classmethod - def _payload_is_consumable(cls, payload: OmniPayload | None) -> bool: + def _payload_is_consumable(cls, payload: dict[str, Any] | OmniPayload | None) -> bool: """Return True when an async payload can drive a real forward step. Metadata-only wake-ups should not transition WAITING_FOR_CHUNK requests @@ -502,8 +506,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): # Skip the per-step accumulator+build that would otherwise be # silently discarded. Defends against a terminal stage whose # custom_process_input_func has a *_full_payload derivative in - # the same module (e.g. dynin stage 2 token2image_to_token2audio - # in pipelines that don't configure any connector at all). + # the same module, even in pipelines that don't configure any + # connector at all. # # Known limitation: a *terminal-consumer* stage that has a # connector configured for receiving upstream input is NOT @@ -556,7 +560,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) return output, request - def _resolve_full_payload_replace_keys(self) -> frozenset: + def _resolve_full_payload_replace_keys(self) -> frozenset[str]: """Per-model REPLACE-key set for the full-payload accumulator. Looked up from the stage-input-processor module that ships the model's sync builder @@ -584,7 +588,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): import sys as _sys mod = _sys.modules.get(module_name) or importlib.import_module(module_name) - keys = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset()) + keys: Any = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset()) except ImportError: logger.debug( "Could not import stage input processor module %s while resolving " @@ -748,7 +752,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): code_len = len(audio_codes) else: code_len = None - meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {} + raw_meta = payload.get("meta") + meta = raw_meta if isinstance(raw_meta, dict) else {} logger.debug( "[Stage-%s] send_full_payload_outputs: req=%s payload_keys=%s code_len=%s left_context_size=%s", self._stage_id, @@ -863,6 +868,9 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): if not self.is_data_transfer_rank(): return True raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None) + if not isinstance(raw_req_id, str): + logger.warning("[Stage-%s] send_chunk: missing request_id", self._stage_id) + return False request_id = self._resolve_external_req_id(request, raw_req_id) # Cache the internal→external mapping so that finish sentinels can # resolve the external ID even after the request is freed. @@ -988,6 +996,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): """Re-enqueue a failed send task or drop it after max retries.""" retry_count = task.get("_retry_count", 0) + 1 req_id = task.get("request_id") + if not isinstance(req_id, str): + logger.error( + "[Stage-%s] Failed send task is missing request_id; dropping it", + getattr(self, "_stage_id", "?"), + ) + return if retry_count <= self._MAX_SEND_RETRIES: task["_retry_count"] = retry_count logger.warning( @@ -1268,7 +1282,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): if not success: return False - self._decrement_pending_save_count(request_id) + if isinstance(request_id, str): + self._decrement_pending_save_count(request_id) return True def _decrement_pending_save_count(self, request_id: str) -> None: @@ -1299,7 +1314,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): """Accumulate chunk payloads (concat tensors, extend lists).""" if req_id not in self._send_side_request_payload: self._send_side_request_payload[req_id] = dict(payload_data) - return dict(self._send_side_request_payload[req_id]) + return cast(OmniPayload, dict(self._send_side_request_payload[req_id])) origin = self._send_side_request_payload[req_id] merged = dict(origin) @@ -1338,4 +1353,4 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): merged[key] = value self._send_side_request_payload[req_id] = merged - return dict(merged) + return cast(OmniPayload, dict(merged)) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index 9f5387a7f..2b7f93a62 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -108,7 +108,6 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 breeze_tts_2, cosyvoice3, covo_audio, - dots_tts, fish_speech, gepard, glm_tts, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/dots_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/dots_tts.py deleted file mode 100644 index 75cfb8455..000000000 --- a/vllm_omni/entrypoints/openai/tts_adapters/dots_tts.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""dots.tts serving adapter.""" - -from typing import TYPE_CHECKING - -from transformers import AutoTokenizer -from vllm.utils.async_utils import make_async - -from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens - -if TYPE_CHECKING: - from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest - - -@register_tts_adapter -class DotsTTSAdapter(ARTTSAdapter): - """Adapter for dots.tts (AR ``engine_client`` backend). - - The current integration supports text-only, no-reference speech synthesis. - Voice cloning and reference-audio conditioning are not supported yet. - """ - - stage_keys = frozenset() - model_archs = frozenset({"DotsTTSForConditionalGeneration"}) - name = "dots_tts" - detect_priority = 5 - - def __init__(self, ctx): - super().__init__(ctx) - self.tokenizer = None - self._build_prompt_async = None - - def _build_prompt(self, text: str) -> dict: - from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import build_dots_tts_prompt - - if self.tokenizer is None: - self.tokenizer = AutoTokenizer.from_pretrained( - self.ctx.engine_client.model_config.model, - trust_remote_code=True, - ) - return build_dots_tts_prompt(self.tokenizer, text) - - def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - - if request.voice not in (None, "default"): - return "'voice' is not supported for dots.tts" - - if request.ref_audio is not None: - return "'ref_audio' is not supported for dots.tts" - - if request.ref_text is not None: - return "'ref_text' is not supported for dots.tts" - - if request.speaker_embedding is not None: - return "'speaker_embedding' is not supported for dots.tts" - - if request.x_vector_only_mode is not None: - return "'x_vector_only_mode' is not supported for dots.tts" - return None - - async def build( - self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool - ) -> PreparedRequest: - if self._build_prompt_async is None: - self._build_prompt_async = make_async( - self._build_prompt, - executor=self.ctx.server._tts_executor, - ) - prompt = await self._build_prompt_async(request.input) - return PreparedRequest(prompt=prompt, tts_params={}, model_type="dots_tts") - - def apply_sampling_overrides( - self, - sampling_params_list: list, - request: "OpenAICreateSpeechRequest", - prompt: dict | None = None, - request_id: str | None = None, - ) -> list: - return apply_max_new_tokens(sampling_params_list, request) diff --git a/vllm_omni/model_executor/models/dots_tts/__init__.py b/vllm_omni/model_executor/models/dots_tts/__init__.py deleted file mode 100644 index fb6738a29..000000000 --- a/vllm_omni/model_executor/models/dots_tts/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# NOTE: Do not import model classes in this file. Importing any -# submodule in this package triggers __init__.py execution, and -# both the model registry and pipeline registry import submodules -# directly — heavy imports here would be loaded as a side effect -# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py deleted file mode 100644 index 086ac109a..000000000 --- a/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py +++ /dev/null @@ -1,457 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright 2026 The vLLM-Omni team. -# Copyright (c) rednote-hilab. All rights reserved. -# Adapted from: -# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone -# -# Merged from two upstream files (preserved as-is, inference behaviour -# unchanged; only internal imports were folded inline): -# - modules/backbone/layers.py -> Dropout, Mlp, rotate_half, -# apply_rotary_pos_emb, RotaryEmbedding, -# MultiHeadAttention -# (Conv1d / ConvTranspose1d are NOT -# re-vendored here — they already live in -# dots_tts_vocoder.py) -# - modules/backbone/dit.py -> modulate, TimestepEmbedder, FinalLayer, -# DiTBlock, DiT -# -# Per-section upstream attributions are preserved at each section header. -# Nothing in this file diverges from upstream behaviour. - -from __future__ import annotations - -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange - -# ============================================================================ -# Building blocks (adapted from modules/backbone/layers.py) -# ---------------------------------------------------------------------------- -# Conv1d / ConvTranspose1d from layers.py are already vendored in -# dots_tts_vocoder.py (they are used by the BigVGAN-style AudioVAE). The -# pieces below are the ones consumed by the DiT flow-matching head. -# ============================================================================ - - -class Dropout(nn.Module): - def __init__(self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **_kwargs): - super().__init__() - if p < 0.0 or p > 1.0: - raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}") - self.p = p - self.inplace = inplace - self.force_drop = force_drop - - def forward(self, x, **_kwargs): - return F.dropout( - x, - p=self.p, - training=True if self.force_drop else self.training, - inplace=self.inplace, - ) - - -class Mlp(nn.Module): - def __init__( - self, - hidden_size, - ffn_hidden_size=4096, - act_layer=nn.GELU, - dropout=0.0, - **_kwargs, - ): - super().__init__() - self.fc1 = nn.Linear(hidden_size, ffn_hidden_size) - self.act = act_layer() - self.fc2 = nn.Linear(ffn_hidden_size, hidden_size) - self.drop = Dropout(dropout) - - def forward(self, x, _mask=None): - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - return self.drop(x) - - -def rotate_half(x): - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - - -@torch.autocast(enabled=False, device_type="cuda") -def apply_rotary_pos_emb(pos, t): - if pos.dim() == 3: - pos = pos.unsqueeze(1) - return t * pos.cos() + rotate_half(t) * pos.sin() - - -class RotaryEmbedding(nn.Module): - def __init__(self, dim, theta=50000): - super().__init__() - self.register_buffer( - "inv_freq", - 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)), - persistent=False, - ) - self._theta = float(theta) - - def _apply(self, fn): - inv_freq = self.inv_freq - super()._apply(fn) - self.inv_freq = inv_freq.to(device=self.inv_freq.device, dtype=torch.float32) - return self - - @torch.autocast(enabled=False, device_type="cuda") - def forward(self, t): - inv_freq = self.inv_freq - if inv_freq.device != t.device: - raise RuntimeError(f"RotaryEmbedding buffer device mismatch: inv_freq={inv_freq.device} input={t.device}.") - t = t.to(dtype=inv_freq.dtype) - if t.dim() == 1: - freqs = torch.einsum("i , j -> i j", t, inv_freq) - else: - freqs = torch.einsum("bi, j -> bij", t, inv_freq) - return torch.cat((freqs, freqs), dim=-1) - - -class MultiHeadAttention(nn.Module): - """Multi-head attention""" - - def __init__( - self, - hidden_size: int, - num_heads: int = 8, - qkv_bias: bool = False, - qk_norm: bool = False, - attn_drop: float = 0.0, - dropout: float = 0.0, - norm_layer: str = "LayerNorm", - rotary_bias: bool = False, - rotary_theta: float | None = 50000, - **_kwargs, - ): - super().__init__() - assert hidden_size % num_heads == 0, "hidden_size should be divisible by num_heads" - self.num_heads = num_heads - self.head_dim = hidden_size // num_heads - self.scale = self.head_dim**-0.5 - self.rotary_bias = rotary_bias - - self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) - self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) - self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) - - norm_layer = getattr(nn, norm_layer) - self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() - self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() - - self.attn_drop = Dropout(attn_drop) - self.o_proj = nn.Linear(hidden_size, hidden_size) - self.o_dropout = Dropout(dropout) - - if self.rotary_bias: - self.rotary = RotaryEmbedding(self.head_dim, theta=rotary_theta) - - def forward(self, q, k=None, v=None, mask=None, pos_ids=None, **_kwargs): - k = k or q - v = v or q - B, L, _ = q.shape - _, S, _ = v.shape - if mask is not None: - if mask.ndim == 2: # [B, L] - assert L == S - mask = rearrange(mask, "b j -> b 1 1 j") - mask = mask.expand(-1, self.num_heads, L, -1) - elif mask.ndim == 3: # [B, L, S] - assert mask.size(1) == L and mask.size(2) == S - mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1) - - q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v) - q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads) - k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads) - v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads) - q, k = self.q_norm(q), self.k_norm(k) - - # Apply rotary - if self.rotary_bias: - if L == S: - if pos_ids is None: - rotary_emb = self.rotary(torch.arange(L, device=q.device)) - else: - rotary_emb = self.rotary(pos_ids) - q, k = (apply_rotary_pos_emb(rotary_emb, tensor) for tensor in (q, k)) - else: - q_rotary_emb = self.rotary(torch.arange(L, device=q.device)) - k_rotary_emb = self.rotary(torch.arange(S, device=k.device)) - q = apply_rotary_pos_emb(q_rotary_emb, q) - k = apply_rotary_pos_emb(k_rotary_emb, k) - - attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device) - - if mask is not None: - attn_bias.masked_fill_(mask.logical_not(), float("-inf")) - - out = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=attn_bias, - dropout_p=self.attn_drop.p if self.training else 0.0, - ) - - out = rearrange(out, "b h n d -> b n (h d)") - return self.o_dropout(self.o_proj(out)) - - def decode_step(self, x, *, cache, positions: torch.Tensor): - if x.size(1) <= 0: - raise ValueError("MultiHeadAttention.decode_step expects a non-empty input.") - if positions.ndim != 1 or positions.size(0) != x.size(1): - raise ValueError("MultiHeadAttention.decode_step positions must match the decode block length.") - - q = self.q_proj(x) - k = self.k_proj(x) - v = self.v_proj(x) - - q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads) - k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads) - v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads) - q, k = self.q_norm(q), self.k_norm(k) - block_len = q.size(2) - - if self.rotary_bias: - rotary_emb = self.rotary(positions) - q = apply_rotary_pos_emb(rotary_emb, q) - k = apply_rotary_pos_emb(rotary_emb, k) - - cached_k, cached_v = cache - cached_k.index_copy_(2, positions, k) - cached_v.index_copy_(2, positions, v) - - cache_capacity = cached_k.size(2) - key_positions = torch.arange( - cache_capacity, - device=x.device, - dtype=torch.long, - ).unsqueeze(0) - query_positions = positions.unsqueeze(1) - causal_mask = key_positions <= query_positions - valid_mask = key_positions <= positions[-1] - attn_bias = torch.zeros( - q.size(0), - self.num_heads, - block_len, - cache_capacity, - dtype=q.dtype, - device=q.device, - ) - attn_bias.masked_fill_( - (causal_mask & valid_mask).unsqueeze(0).unsqueeze(0).logical_not(), - float("-inf"), - ) - - out = F.scaled_dot_product_attention( - q, - cached_k, - cached_v, - attn_mask=attn_bias, - dropout_p=self.attn_drop.p if self.training else 0.0, - ) - out = rearrange(out, "b h n d -> b n (h d)") - return self.o_dropout(self.o_proj(out)), cache - - -# ============================================================================ -# DiT flow-matching head (adapted from modules/backbone/dit.py) -# ---------------------------------------------------------------------------- -# DiT.__init__ takes a ``transformer_config`` object exposing ``.to_dict()``, -# ``.hidden_size`` and ``.num_layers``. Upstream constructs it from a pydantic -# config class higher up the stack; vLLM-Omni callers can pass anything that -# satisfies the same minimal interface (e.g. a stdlib @dataclass with a -# ``to_dict`` method or SimpleNamespace). See dots_tts_talker.py for the -# concrete soar-checkpoint config factory. -# ============================================================================ - - -def modulate(x, shift, scale, **_kwargs): - return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) - - -class TimestepEmbedder(nn.Module): - def __init__(self, hidden_size, frequency_embedding_size=256): - super().__init__() - self.mlp = nn.Sequential( - nn.Linear(frequency_embedding_size, hidden_size, bias=True), - nn.SiLU(), - nn.Linear(hidden_size, hidden_size, bias=True), - ) - self.frequency_embedding_size = frequency_embedding_size - - @staticmethod - def timestep_embedding(t, dim, max_period=10000): - half = dim // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( - device=t.device - ) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - - def forward(self, t): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) - return self.mlp(t_freq) - - -class FinalLayer(nn.Module): - def __init__(self, hidden_size, output_size): - super().__init__() - self.adaLN_modulation = nn.Sequential( - nn.SiLU(), - nn.Linear(hidden_size, 2 * hidden_size, bias=True), - ) - self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-5) - self.linear = nn.Linear(hidden_size, output_size, bias=True) - - def forward(self, x, c, **_kwargs): - shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) - x = modulate(self.norm(x), shift, scale) - return self.linear(x) - - -class DiTBlock(nn.Module): - def __init__( - self, - attention: nn.Module, - ffn: nn.Module, - hidden_size: int = 1024, - modulation: bool = False, - eps: float = 1e-5, - **_kwargs, - ): - super().__init__() - self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps) - self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps) - self.attn = attention - self.ffn = ffn - self.modulation = modulation - if modulation: - self.adaLN_modulation = nn.Sequential( - nn.SiLU(), - nn.Linear(hidden_size, 6 * hidden_size, bias=True), - ) - - def forward(self, x, condition=None, mask=None, **kwargs): - if condition is None: - assert not self.modulation, "Without global condition, must set modulation to False" - else: - assert self.modulation, "With global condition, must set modulation to True" - shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = self.adaLN_modulation(condition).chunk( - 6, dim=1 - ) - - if condition is not None: - pack_indices = kwargs.get("pack_indices") - if pack_indices is not None: - gate_attn = gate_attn[pack_indices] - gate_ffn = gate_ffn[pack_indices] - else: - gate_attn = gate_attn.unsqueeze(1) - gate_ffn = gate_ffn.unsqueeze(1) - - if condition is not None: - x = x + gate_attn * self.attn( - modulate(self.norm1(x), shift_attn, scale_attn, **kwargs), - mask=mask, - **kwargs, - ) - else: - x = x + self.attn(self.norm1(x), mask=mask, **kwargs) - - if condition is not None: - x = x + gate_ffn * self.ffn(modulate(self.norm2(x), shift_ffn, scale_ffn, **kwargs)) - else: - x = x + self.ffn(self.norm2(x), mask=mask) - return x - - -class DiT(nn.Module): - def __init__( - self, - in_dim, - out_dim, - transformer_config, - *, - mode: str = "flow_matching", - ): - super().__init__() - if mode not in {"flow_matching", "meanflow"}: - raise ValueError(f"DiT mode must be 'flow_matching' or 'meanflow', got {mode!r}.") - - transformer_kwargs = transformer_config.to_dict() - model_dim = transformer_config.hidden_size - self.mode = mode - self.num_layers = transformer_config.num_layers - - self.input_layer = nn.Linear(in_dim, model_dim) - self.time_embedder = TimestepEmbedder(model_dim) - if mode == "meanflow": - self.duration_embedder = TimestepEmbedder(model_dim) - - self.blocks = nn.ModuleList() - for i in range(self.num_layers): - attn_block = MultiHeadAttention(**transformer_kwargs, name=f"layer_{i}") - ffn_block = Mlp(act_layer=lambda: nn.GELU(approximate="tanh"), **transformer_kwargs) - self.blocks.append(DiTBlock(attention=attn_block, ffn=ffn_block, **transformer_kwargs)) - - self.output_layer = FinalLayer(model_dim, out_dim) - self.initialize_weights() - - def initialize_weights(self): - def _basic_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.xavier_uniform_(module.weight) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - - self.apply(_basic_init) - - nn.init.normal_(self.time_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.time_embedder.mlp[2].weight, std=0.02) - - for block in self.blocks: - if hasattr(block, "adaLN_modulation"): - nn.init.constant_(block.adaLN_modulation[-1].weight, 0) - nn.init.constant_(block.adaLN_modulation[-1].bias, 0) - - nn.init.constant_(self.output_layer.adaLN_modulation[-1].weight, 0) - nn.init.constant_(self.output_layer.adaLN_modulation[-1].bias, 0) - nn.init.constant_(self.output_layer.linear.weight, 0) - nn.init.constant_(self.output_layer.linear.bias, 0) - - def forward( - self, - x, - timesteps, - duration: torch.Tensor | None = None, - mask=None, - attn_mask=None, - g_cond: torch.Tensor | None = None, - **kwargs, - ): - t = self.time_embedder(timesteps) - c = t - duration_embedder = getattr(self, "duration_embedder", None) - if duration_embedder is not None and duration is not None: - c = c + duration_embedder(duration) - if g_cond is not None: - c = c + g_cond - - x = self.input_layer(x) - for block in self.blocks: - x = block(x, c, mask=attn_mask, **kwargs) - return self.output_layer(x, c, **kwargs) diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py deleted file mode 100644 index 39d1a4973..000000000 --- a/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py +++ /dev/null @@ -1,355 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright 2026 The vLLM-Omni team. -# Copyright (c) rednote-hilab. All rights reserved. -# Adapted from: -# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone -# -# Single upstream source (preserved as-is, inference behaviour unchanged; -# only the internal `from dots_tts.modules.backbone.layers import ...` line -# was folded to point at our sibling vendor files): -# - modules/backbone/semantic_encoder.py -> SemanticEncoderDecodeState, -# TransformerEncoderLayer, -# SuperviseEncoder, -# VAESemanticEncoder -# -# Building blocks (Conv1d / Mlp / MultiHeadAttention) are imported from the -# sibling vendor files where they already live: -# - Conv1d <- dots_tts_vocoder.py -# - Mlp <- dots_tts_dit.py -# - MultiHeadAttention <- dots_tts_dit.py -# -# Nothing else in this file diverges from upstream behaviour. - -from __future__ import annotations - -from dataclasses import dataclass - -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange - -from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import ( - Mlp, - MultiHeadAttention, -) -from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import Conv1d - - -@dataclass -class SemanticEncoderDecodeState: - conv_tail: torch.Tensor - layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...] - seq_len: int - - -class TransformerEncoderLayer(nn.Module): - def __init__( - self, - hidden_size, - num_heads=16, - ffn_hidden_size=4096, - attn_dropout=0.0, - ffn_dropout=0.0, - norm_layer="LayerNorm", - **kwargs, - ): - super().__init__() - self.attn = MultiHeadAttention( - hidden_size, - num_heads, - attn_drop=attn_dropout, - norm_layer=norm_layer, - **kwargs, - ) - norm_cls = getattr(nn, norm_layer) - self.attn_norm = norm_cls(hidden_size) - self.ffn = Mlp(hidden_size, ffn_hidden_size, dropout=ffn_dropout, act_layer=nn.SiLU) - self.ffn_norm = norm_cls(hidden_size) - self.hidden_size = hidden_size - - def _build_causal_mask(self, T: int, device): # noqa: N803 - return torch.tril(torch.ones(T, T, dtype=torch.bool, device=device)) - - def _build_padding_mask(self, x_lens, max_len: int, device): - B = x_lens.size(0) - positions = torch.arange(max_len, device=device).unsqueeze(0).expand(B, -1) - return positions < x_lens.unsqueeze(1) - - def _fuse_attn_mask(self, causal_mask, padding_mask): - if causal_mask is None and padding_mask is None: - return None - if causal_mask is None: - row = padding_mask.unsqueeze(2) - col = padding_mask.unsqueeze(1) - return row & col - if padding_mask is None: - return causal_mask.unsqueeze(0) - - _B, _T = padding_mask.shape - causal = causal_mask.unsqueeze(0) - row = padding_mask.unsqueeze(2) - col = padding_mask.unsqueeze(1) - pad_2d = row & col - return causal & pad_2d - - def forward( - self, - x, - x_lens=None, - causal=True, - ): - _B, T, C = x.shape - assert self.hidden_size == C - device = x.device - - causal_mask = self._build_causal_mask(T, device) if causal else None - if x_lens is not None: - padding_mask = self._build_padding_mask(x_lens, T, device) - else: - padding_mask = None - fused_mask = self._fuse_attn_mask(causal_mask, padding_mask) - - h = self.attn_norm(x) - h = self.attn( - q=h, - mask=fused_mask, - ) - x = x + h - - h = self.ffn_norm(x) - h = self.ffn(h) - return x + h - - def decode_step( - self, - x, - *, - cache: tuple[torch.Tensor, torch.Tensor], - positions: torch.Tensor, - ): - if x.size(1) <= 0: - raise ValueError("TransformerEncoderLayer.decode_step expects a non-empty input.") - - h = self.attn_norm(x) - h, cache = self.attn.decode_step(h, cache=cache, positions=positions) - x = x + h - - h = self.ffn_norm(x) - h = self.ffn(h) - return x + h, cache - - -class SuperviseEncoder(nn.Module): - def __init__(self, config): - super().__init__() - self.hidden_size = config.get("hidden_size", 1024) - self.layers = nn.ModuleList( - [ - TransformerEncoderLayer( - hidden_size=self.hidden_size, - num_heads=config.get("num_heads", 16), - ffn_hidden_size=config.get("ffn_hidden_size", 4096), - norm_layer=config.get("norm_layer", "LayerNorm"), - ) - for _ in range(config.get("num_layers", 6)) - ] - ) - self.causal = config.get("causal", False) - - def forward(self, x, x_lens=None): - batch_size, seq_len, _ = x.shape - if x_lens is None: - x_lens = torch.full((batch_size,), seq_len, device=x.device, dtype=torch.long) - for layer in self.layers: - x = layer(x, x_lens=x_lens, causal=self.causal) - return x - - def init_decode_state( - self, - *, - batch_size: int, - max_seq_len: int, - device: torch.device, - dtype: torch.dtype, - ): - layer_caches = [] - for layer in self.layers: - cache_shape = ( - batch_size, - layer.attn.num_heads, - max_seq_len, - layer.attn.head_dim, - ) - layer_caches.append( - ( - torch.zeros(cache_shape, dtype=dtype, device=device), - torch.zeros(cache_shape, dtype=dtype, device=device), - ) - ) - return tuple(layer_caches) - - def reset_decode_state( - self, - layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...], - ) -> None: - if len(layer_caches) != len(self.layers): - raise ValueError("Layer cache count does not match encoder depth.") - for key_cache, value_cache in layer_caches: - key_cache.zero_() - value_cache.zero_() - - def decode_step(self, x, *, layer_caches, positions: torch.Tensor): - if len(layer_caches) != len(self.layers): - raise ValueError("Layer cache count does not match encoder depth.") - - for layer, cache in zip(self.layers, layer_caches, strict=True): - x, _ = layer.decode_step(x, cache=cache, positions=positions) - return x - - -class VAESemanticEncoder(nn.Module): - def __init__(self, in_dim, out_dim, config): - super().__init__() - in_ds_rate = 2 - self.patch_size = int(config.patch_size) - self.in_ds_rate = in_ds_rate - self.ds_proj = Conv1d(in_dim, in_dim, kernel_size=in_ds_rate, stride=in_ds_rate, causal=True) - self.in_proj = nn.Linear(in_dim, config.PatchEncoder.hidden_size) - self.encoder = SuperviseEncoder(config.PatchEncoder) - self.out_ds_rate = self.patch_size // in_ds_rate - self.out_proj = nn.Linear(config.PatchEncoder.hidden_size * self.out_ds_rate, out_dim) - - def forward(self, x, x_lens=None): - x = self._downsample(x) - x = self.in_proj(x) - z = self.encoder(x, x_lens=x_lens) - return self._project_embeddings(z) - - def init_decode_state( - self, - *, - max_audio_patch_count: int, - batch_size: int, - device: torch.device, - dtype: torch.dtype, - ) -> SemanticEncoderDecodeState: - return SemanticEncoderDecodeState( - conv_tail=torch.zeros( - (batch_size, self.ds_proj.in_channels, self.ds_proj.left_padding), - dtype=dtype, - device=device, - ), - layer_caches=self.encoder.init_decode_state( - batch_size=batch_size, - max_seq_len=max_audio_patch_count * self.out_ds_rate, - device=device, - dtype=dtype, - ), - seq_len=0, - ) - - def reset_decode_state(self, state: SemanticEncoderDecodeState) -> None: - state.conv_tail.zero_() - self.encoder.reset_decode_state(state.layer_caches) - state.seq_len = 0 - - def prefill( - self, - x, - state: SemanticEncoderDecodeState, - ) -> tuple[torch.Tensor, SemanticEncoderDecodeState]: - if x.ndim != 3: - raise ValueError(f"VAESemanticEncoder.prefill expects rank-3 input, got {tuple(x.shape)}.") - if x.size(1) % self.patch_size != 0: - raise ValueError(f"Prompt latent length {x.size(1)} must be divisible by patch_size={self.patch_size}.") - - if x.size(1) == 0: - return ( - x.new_zeros((x.size(0), 0, self.out_proj.out_features)), - state, - ) - if state.conv_tail.size(0) != x.size(0): - raise ValueError("VAESemanticEncoder.prefill batch size does not match decode state.") - - step_inputs = self.in_proj(self._downsample(x)) - expected_token_count = (x.size(1) // self.patch_size) * self.out_ds_rate - if step_inputs.size(1) != expected_token_count: - raise RuntimeError( - "Patch encoder prefill produced an unexpected token count: " - f"expected={expected_token_count} actual={step_inputs.size(1)}." - ) - - current_seq_len = state.seq_len - next_seq_len = current_seq_len + step_inputs.size(1) - cache_capacity = state.layer_caches[0][0].size(2) - if next_seq_len > cache_capacity: - raise ValueError( - "Patch encoder prefill exceeds decode-state capacity: " - f"required={next_seq_len} capacity={cache_capacity}." - ) - - positions = torch.arange(step_inputs.size(1), device=x.device, dtype=torch.long) + current_seq_len - encoded = self.encoder.decode_step( - step_inputs, - layer_caches=state.layer_caches, - positions=positions, - ) - embedding = self._project_embeddings(encoded) - raw = x.transpose(1, 2) - state.conv_tail.copy_(raw[..., -self.ds_proj.left_padding :]) - state.seq_len = next_seq_len - return embedding, state - - def decode_patch( - self, - latent_patch, - conv_tail: torch.Tensor, - layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...], - positions: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - if latent_patch.ndim != 3: - raise ValueError(f"VAESemanticEncoder.decode_patch expects rank-3 input, got {tuple(latent_patch.shape)}.") - if latent_patch.size(1) != self.patch_size: - raise ValueError(f"decode_patch expects patch length {self.patch_size}, got {latent_patch.size(1)}.") - if positions.ndim != 1 or positions.size(0) != self.out_ds_rate: - raise ValueError("decode_patch positions must be a rank-1 tensor matching out_ds_rate.") - - step_inputs, conv_tail = self._downsample_step( - latent_patch, - conv_tail=conv_tail, - ) - if step_inputs.size(1) != self.out_ds_rate: - raise RuntimeError(f"Downsample step produced {step_inputs.size(1)} tokens, expected {self.out_ds_rate}.") - - encoded = self.encoder.decode_step( - step_inputs, - layer_caches=layer_caches, - positions=positions, - ) - embedding = self._project_embeddings(encoded) - return embedding, conv_tail - - def _downsample(self, x): - return self.ds_proj(x.transpose(1, 2)).transpose(1, 2) - - def _project_embeddings(self, z): - if self.out_ds_rate > 1: - z = rearrange(z, "b (s d) h -> b s (d h)", d=self.out_ds_rate) - return self.out_proj(z) - - def _downsample_step(self, latent_patch, *, conv_tail): - raw = latent_patch.transpose(1, 2) - conv_input = torch.cat([conv_tail, raw], dim=-1) - - projected = F.conv1d( - conv_input, - self.ds_proj.weight, - self.ds_proj.bias, - stride=self.ds_proj.stride[0], - padding=0, - dilation=self.ds_proj.dilation[0], - groups=self.ds_proj.groups, - ).transpose(1, 2) - new_conv_tail = raw[..., -self.ds_proj.left_padding :] - return self.in_proj(projected), new_conv_tail diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py deleted file mode 100644 index a2c4ac9f9..000000000 --- a/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py +++ /dev/null @@ -1,70 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""dots.tts prompt builder for vllm-omni's Omni engine + serving layer. - -Mirrors upstream ``build_generation_schedule`` (rednote-hilab/dots.tts @ -a393d2e data/pipelines/tokenizing.py:169) for the -no-prompt-audio + audio-placeholder path — text-only zero-shot TTS. - -V1 returns real Qwen2 token IDs (not placeholders ``[1] * N`` like -voxcpm2's builder). Reason: our talker's preprocess() runs -``self.model.embed_tokens(input_ids)`` directly, so real IDs flow -through naturally. voxcpm2 uses placeholders because its prefill mixes -reference-audio embeddings (no token IDs); dots.tts V1 zero-shot is -purely text + audio_gen_start, all real IDs. - -Voice clone (reference audio + speaker x-vector) is intentionally not -wired here — V1 is zero-shot only. Add reference_audio / -voice_profile / prompt_text branches in a follow-up (mirror voxcpm2 -``build_voxcpm2_prompt``). -""" - -from __future__ import annotations - -from typing import Any - -from vllm.inputs import tokens_input - -# Upstream rednote-hilab/dots.tts @ a393d2e -# data/pipelines/tts_pipeline.py:17-18 -_TTS_TEXT_PREFIX = "[文本]" -_TTS_AUDIO_PREFIX = "[文本对应语音]" -_AUDIO_GEN_START_TOKEN = "<|audio_gen_start|>" - - -def build_dots_tts_prompt( - tokenizer: Any, - text: str, -) -> dict[str, Any]: - """Build a dots.tts prefill prompt dict for ``Omni.generate()``. - - Returns ``{"prompt_token_ids": [...], "additional_information": {...}}`` - where ``prompt_token_ids`` is the real Qwen2 token sequence: - - [text_tokens of "[文本][文本对应语音]", audio_gen_start_id] - - talker's preprocess() embeds these into ``hidden_states`` for the AR - base LM. Decode then drives via ``state.curr_embed_for_next`` (DiT - → patch_encoder AR loopback). - """ - if not isinstance(text, str) or not text.strip(): - raise ValueError(f"text must be a non-empty string, got {text!r}") - - prompt_text = f"{_TTS_TEXT_PREFIX}{text}{_TTS_AUDIO_PREFIX}" - text_ids = tokenizer.encode(prompt_text, add_special_tokens=False) - - audio_gen_start_id = tokenizer.convert_tokens_to_ids(_AUDIO_GEN_START_TOKEN) - if not isinstance(audio_gen_start_id, int) or audio_gen_start_id == tokenizer.unk_token_id: - raise ValueError( - f"Tokenizer does not know {_AUDIO_GEN_START_TOKEN!r} " - f"(got id={audio_gen_start_id!r}). Did you load the " - f"dots.tts tokenizer (with added_tokens.json)?" - ) - prompt_token_ids = list(text_ids) + [audio_gen_start_id] - - prompt = tokens_input(prompt_token_ids=prompt_token_ids) - prompt["additional_information"] = { - # Reserved for future voice clone (ref_audio / speaker x-vector). - # V1 zero-shot leaves this empty; preprocess() ignores it. - } - return prompt diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py deleted file mode 100644 index 09961d0ab..000000000 --- a/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py +++ /dev/null @@ -1,732 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright 2026 The vLLM-Omni team. -# Copyright (c) rednote-hilab. All rights reserved. -# Portions Copyright 3D-Speaker -# (https://github.com/alibaba-damo-academy/3D-Speaker), Apache-2.0. -# Adapted from: -# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/speaker -# (whose CAM++ implementation is itself adapted from 3D-Speaker). -# -# Merged from four upstream files (preserved as-is, inference behaviour -# unchanged; only internal cross-module imports were folded inline): -# - modules/speaker/encoder.py -> SpeakerXVectorFeatures -# - modules/speaker/fbank.py -> fbank constants + extract_speaker_fbank -# - modules/speaker/campplus.py -> FCM, CAMPPlus -# - modules/speaker/campplus_layers.py -> get_nonlinear, statistics_pooling, -# StatsPool, TDNNLayer, CAMLayer, -# CAMDenseTDNNLayer, CAMDenseTDNNBlock, -# TransitLayer, DenseLayer, BasicResBlock -# -# Plus two tiny helpers folded from utils/audio.py (used by -# extract_speaker_fbank): -# - high_quality_resample, extract_fbank -# -# Nothing in this file diverges from upstream behaviour. - -from __future__ import annotations - -import math -import random -from collections import OrderedDict - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.utils.checkpoint as cp -import torchaudio -import torchaudio.compliance.kaldi as Kaldi -import torchaudio.functional as AF -from torch.nn.utils.rnn import pad_sequence - -# ============================================================================ -# Constants (from modules/speaker/fbank.py) -# ============================================================================ - -_SPEAKER_FBANK_SAMPLE_RATE = 16000 -_SPEAKER_FBANK_N_MELS = 80 -_SPEAKER_FBANK_MEAN_NORM = True -_SPEAKER_FBANK_DITHER = 0.0 - - -# ============================================================================ -# Audio helpers (folded from utils/audio.py — used by extract_speaker_fbank) -# ============================================================================ - - -def high_quality_resample(x, orig_sr, target_sr): - return AF.resample( - x, - orig_freq=orig_sr, - new_freq=target_sr, - lowpass_filter_width=64, - rolloff=0.95, - resampling_method="sinc_interp_kaiser", - ) - - -def extract_fbank( - waveform: torch.Tensor, - *, - sample_rate: int, - n_mels: int, - dither: float = 0.0, - mean_norm: bool = False, -) -> torch.Tensor: - if waveform.ndim == 1: - feature_input = waveform.unsqueeze(0) - elif waveform.ndim == 2: - feature_input = waveform if waveform.size(0) == 1 else waveform[0:1, :] - else: - raise ValueError(f"FBank expects a 1D or 2D waveform, got shape {tuple(waveform.shape)}.") - features = Kaldi.fbank( - feature_input, - num_mel_bins=n_mels, - sample_frequency=sample_rate, - dither=dither, - ) - if mean_norm: - features = features - features.mean(dim=0, keepdim=True) - return features - - -# ============================================================================ -# Speaker fbank entry (from modules/speaker/fbank.py) -# ============================================================================ - - -def extract_speaker_fbank( - waveform: torch.Tensor, - *, - sample_rate: int, -) -> torch.Tensor: - feature_input = waveform - if sample_rate != _SPEAKER_FBANK_SAMPLE_RATE: - feature_input = high_quality_resample( - waveform, - orig_sr=sample_rate, - target_sr=_SPEAKER_FBANK_SAMPLE_RATE, - ) - return extract_fbank( - feature_input, - sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, - n_mels=_SPEAKER_FBANK_N_MELS, - dither=_SPEAKER_FBANK_DITHER, - mean_norm=_SPEAKER_FBANK_MEAN_NORM, - ) - - -# ============================================================================ -# CAM++ building blocks (from modules/speaker/campplus_layers.py) -# ---------------------------------------------------------------------------- -# 3D-Speaker, Apache-2.0. -# ============================================================================ - - -def get_nonlinear(config_str, channels): - nonlinear = nn.Sequential() - for name in config_str.split("-"): - if name == "relu": - nonlinear.add_module("relu", nn.ReLU(inplace=True)) - elif name == "prelu": - nonlinear.add_module("prelu", nn.PReLU(channels)) - elif name == "batchnorm": - nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels)) - elif name == "batchnorm_": - nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels, affine=False)) - else: - raise ValueError(f"Unexpected module ({name}).") - return nonlinear - - -def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, _eps=1e-2): - mean = x.mean(dim=dim) - std = x.std(dim=dim, unbiased=unbiased) - stats = torch.cat([mean, std], dim=-1) - if keepdim: - stats = stats.unsqueeze(dim=dim) - return stats - - -class StatsPool(nn.Module): - def forward(self, x): - return statistics_pooling(x) - - -class TDNNLayer(nn.Module): - def __init__( - self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - dilation=1, - bias=False, - config_str="batchnorm-relu", - ): - super().__init__() - if padding < 0: - assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" - padding = (kernel_size - 1) // 2 * dilation - self.linear = nn.Conv1d( - in_channels, - out_channels, - kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - bias=bias, - ) - self.nonlinear = get_nonlinear(config_str, out_channels) - - def forward(self, x): - x = self.linear(x) - return self.nonlinear(x) - - -class CAMLayer(nn.Module): - def __init__( - self, - bn_channels, - out_channels, - kernel_size, - stride, - padding, - dilation, - bias, - reduction=2, - ): - super().__init__() - self.linear_local = nn.Conv1d( - bn_channels, - out_channels, - kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - bias=bias, - ) - self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1) - self.relu = nn.ReLU(inplace=True) - self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1) - self.sigmoid = nn.Sigmoid() - - def forward(self, x): - y = self.linear_local(x) - context = x.mean(-1, keepdim=True) + self.seg_pooling(x) - context = self.relu(self.linear1(context)) - m = self.sigmoid(self.linear2(context)) - return y * m - - def seg_pooling(self, x, seg_len=100, stype="avg"): - if stype == "avg": - seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) - elif stype == "max": - seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) - else: - raise ValueError("Wrong segment pooling type.") - shape = seg.shape - seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1) - return seg[..., : x.shape[-1]] - - -class CAMDenseTDNNLayer(nn.Module): - def __init__( - self, - in_channels, - out_channels, - bn_channels, - kernel_size, - stride=1, - dilation=1, - bias=False, - config_str="batchnorm-relu", - memory_efficient=False, - ): - super().__init__() - assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" - padding = (kernel_size - 1) // 2 * dilation - self.memory_efficient = memory_efficient - self.nonlinear1 = get_nonlinear(config_str, in_channels) - self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False) - self.nonlinear2 = get_nonlinear(config_str, bn_channels) - self.cam_layer = CAMLayer( - bn_channels, - out_channels, - kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - bias=bias, - ) - - def bn_function(self, x): - return self.linear1(self.nonlinear1(x)) - - def forward(self, x): - if self.training and self.memory_efficient: - x = cp.checkpoint(self.bn_function, x) - else: - x = self.bn_function(x) - return self.cam_layer(self.nonlinear2(x)) - - -class CAMDenseTDNNBlock(nn.ModuleList): - def __init__( - self, - num_layers, - in_channels, - out_channels, - bn_channels, - kernel_size, - stride=1, - dilation=1, - bias=False, - config_str="batchnorm-relu", - memory_efficient=False, - ): - super().__init__() - for i in range(num_layers): - layer = CAMDenseTDNNLayer( - in_channels=in_channels + i * out_channels, - out_channels=out_channels, - bn_channels=bn_channels, - kernel_size=kernel_size, - stride=stride, - dilation=dilation, - bias=bias, - config_str=config_str, - memory_efficient=memory_efficient, - ) - self.add_module(f"tdnnd{i + 1}", layer) - - def forward(self, x): - for layer in self: - x = torch.cat([x, layer(x)], dim=1) - return x - - -class TransitLayer(nn.Module): - def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"): - super().__init__() - self.nonlinear = get_nonlinear(config_str, in_channels) - self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) - - def forward(self, x): - x = self.nonlinear(x) - return self.linear(x) - - -class DenseLayer(nn.Module): - def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"): - super().__init__() - self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) - self.nonlinear = get_nonlinear(config_str, out_channels) - - def forward(self, x): - if len(x.shape) == 2: - x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1) - else: - x = self.linear(x) - return self.nonlinear(x) - - -class BasicResBlock(nn.Module): - expansion = 1 - - def __init__(self, in_planes, planes, stride=1): - super().__init__() - self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False) - self.bn1 = nn.BatchNorm2d(planes) - self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) - self.bn2 = nn.BatchNorm2d(planes) - - self.shortcut = nn.Sequential() - if stride != 1 or in_planes != self.expansion * planes: - self.shortcut = nn.Sequential( - nn.Conv2d( - in_planes, - self.expansion * planes, - kernel_size=1, - stride=(stride, 1), - bias=False, - ), - nn.BatchNorm2d(self.expansion * planes), - ) - - def forward(self, x): - out = F.relu(self.bn1(self.conv1(x))) - out = self.bn2(self.conv2(out)) - out += self.shortcut(x) - return F.relu(out) - - -# ============================================================================ -# CAM++ model (from modules/speaker/campplus.py) -# ---------------------------------------------------------------------------- -# 3D-Speaker, Apache-2.0. -# ============================================================================ - - -class FCM(nn.Module): - def __init__( - self, - block=BasicResBlock, - num_blocks=(2, 2), - m_channels=32, - feat_dim=_SPEAKER_FBANK_N_MELS, - ): - super().__init__() - self.in_planes = m_channels - self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) - self.bn1 = nn.BatchNorm2d(m_channels) - - self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2) - self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2) - - self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False) - self.bn2 = nn.BatchNorm2d(m_channels) - self.out_channels = m_channels * (feat_dim // 8) - - def _make_layer(self, block, planes, num_blocks, stride): - strides = [stride] + [1] * (num_blocks - 1) - layers = [] - for stride in strides: - layers.append(block(self.in_planes, planes, stride)) - self.in_planes = planes * block.expansion - return nn.Sequential(*layers) - - def forward(self, x): - x = x.unsqueeze(1) - out = F.relu(self.bn1(self.conv1(x))) - out = self.layer1(out) - out = self.layer2(out) - out = F.relu(self.bn2(self.conv2(out))) - - shape = out.shape - return out.reshape(shape[0], shape[1] * shape[2], shape[3]) - - -class CAMPPlus(nn.Module): - _TDNN_KERNEL_SIZE = 5 - _TDNN_STRIDE = 2 - _TDNN_PADDING = 2 - - def __init__( - self, - feat_dim=_SPEAKER_FBANK_N_MELS, - embedding_size=512, - growth_rate=32, - bn_size=4, - init_channels=128, - config_str="batchnorm-relu", - memory_efficient=True, - ): - super().__init__() - - self.head = FCM(feat_dim=feat_dim) - channels = self.head.out_channels - - self.xvector = nn.Sequential( - OrderedDict( - [ - ( - "tdnn", - TDNNLayer( - channels, - init_channels, - self._TDNN_KERNEL_SIZE, - stride=self._TDNN_STRIDE, - dilation=1, - padding=-1, - config_str=config_str, - ), - ), - ] - ) - ) - channels = init_channels - for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2), strict=True)): - block = CAMDenseTDNNBlock( - num_layers=num_layers, - in_channels=channels, - out_channels=growth_rate, - bn_channels=bn_size * growth_rate, - kernel_size=kernel_size, - dilation=dilation, - config_str=config_str, - memory_efficient=memory_efficient, - ) - self.xvector.add_module(f"block{i + 1}", block) - channels = channels + num_layers * growth_rate - self.xvector.add_module( - f"transit{i + 1}", - TransitLayer(channels, channels // 2, bias=False, config_str=config_str), - ) - channels //= 2 - - self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels)) - - self.xvector.add_module("stats", StatsPool()) - self.xvector.add_module("dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_")) - - for m in self.modules(): - if isinstance(m, (nn.Conv1d, nn.Linear)): - nn.init.kaiming_normal_(m.weight.data) - if m.bias is not None: - nn.init.zeros_(m.bias) - - @staticmethod - def _conv_output_lengths(lengths, kernel_size, stride=1, padding=0, dilation=1): - return ( - torch.div( - lengths + 2 * padding - dilation * (kernel_size - 1) - 1, - stride, - rounding_mode="floor", - ) - + 1 - ) - - @staticmethod - def _make_length_mask(lengths, max_len, device): - lengths = lengths.to(device=device, dtype=torch.long).clamp(min=0, max=max_len) - return torch.arange(max_len, device=device).unsqueeze(0) < lengths.unsqueeze(1) - - def _masked_stats_pooling(self, x, lengths, unbiased=True, eps=1e-2): - lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1, max=x.size(-1)) - mask = self._make_length_mask(lengths, x.size(-1), x.device).unsqueeze(1) - mask = mask.to(dtype=x.dtype) - - denom = lengths.to(dtype=x.dtype).view(-1, 1).clamp_min(1.0) - mean = (x * mask).sum(dim=-1) / denom - - centered = (x - mean.unsqueeze(-1)) * mask - var_denom = (lengths - 1).clamp_min(1).to(dtype=x.dtype).view(-1, 1) if unbiased else denom - var = centered.pow(2).sum(dim=-1) / var_denom - std = torch.sqrt(var.clamp_min(eps)) - return torch.cat([mean, std], dim=1) - - def forward(self, x, lengths=None): - x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) - x = self.head(x) - if lengths is not None: - lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1) - - for name, module in self.xvector.named_children(): - if name == "stats": - x = self._masked_stats_pooling(x, lengths) if lengths is not None else module(x) - continue - - x = module(x) - if name == "tdnn" and lengths is not None: - lengths = self._conv_output_lengths( - lengths, - kernel_size=self._TDNN_KERNEL_SIZE, - stride=self._TDNN_STRIDE, - padding=self._TDNN_PADDING, - ) - - return x - - -# ============================================================================ -# SpeakerXVectorFeatures wrapper (from modules/speaker/encoder.py) -# ============================================================================ - - -class SpeakerXVectorFeatures(nn.Module): - """Speaker embedding extractor based on 3D-Speaker CAM++.""" - - def __init__( - self, - sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, - campplus_embedding_size=512, - max_audio_seconds=10.0, - ): - super().__init__() - - self.sample_rate = sample_rate - self.max_audio_seconds = float(max_audio_seconds) - self.model = CAMPPlus( - feat_dim=_SPEAKER_FBANK_N_MELS, - embedding_size=campplus_embedding_size, - ) - self.resample = None - if self.sample_rate != _SPEAKER_FBANK_SAMPLE_RATE: - self.resample = torchaudio.transforms.Resample( - orig_freq=sample_rate, - new_freq=_SPEAKER_FBANK_SAMPLE_RATE, - ) - - for param in self.model.parameters(): - param.requires_grad = False - - @staticmethod - def _normalize_lengths(lengths, batch_size, max_length, device, *, min_length): - if lengths is None: - return torch.full( - (batch_size,), - max_length, - device=device, - dtype=torch.long, - ) - return lengths.to(device=device, dtype=torch.long).clamp( - min=min_length, - max=max_length, - ) - - def _crop_audio(self, audio, audio_lengths=None): - original_lengths = self._normalize_lengths( - audio_lengths, - audio.size(0), - audio.size(-1), - audio.device, - min_length=0, - ) - if self.max_audio_seconds <= 0: - return audio, original_lengths, original_lengths, torch.zeros_like(original_lengths) - - max_input_length = round(self.sample_rate * self.max_audio_seconds) - cropped_audio = [] - cropped_lengths = [] - starts = [] - - for index, total_length_tensor in enumerate(original_lengths): - total_length = int(total_length_tensor.item()) - cropped_length = min(total_length, max_input_length) - start = random.randint(0, total_length - cropped_length) if total_length > cropped_length else 0 - cropped_audio.append(audio[index, start : start + cropped_length]) - cropped_lengths.append(cropped_length) - starts.append(start) - - return ( - pad_sequence( - cropped_audio, - batch_first=True, - padding_value=0.0, - ), - original_lengths, - torch.tensor( - cropped_lengths, - device=audio.device, - dtype=torch.long, - ), - torch.tensor(starts, device=audio.device, dtype=torch.long), - ) - - def _crop_fbank( - self, - fbank, - fbank_lengths, - original_audio_lengths, - cropped_audio_lengths, - starts, - ): - original_fbank_lengths = self._normalize_lengths( - fbank_lengths, - fbank.size(0), - fbank.size(1), - fbank.device, - min_length=1, - ) - cropped_fbank = [] - cropped_fbank_lengths = [] - - for index, total_feat_length_tensor in enumerate(original_fbank_lengths): - total_audio_length = int(original_audio_lengths[index].item()) - total_feat_length = int(total_feat_length_tensor.item()) - start_audio = int(starts[index].item()) - end_audio = start_audio + int(cropped_audio_lengths[index].item()) - - if total_audio_length > 0: - start_feat = math.floor(start_audio * total_feat_length / total_audio_length) - end_feat = math.ceil(end_audio * total_feat_length / total_audio_length) - else: - start_feat = 0 - end_feat = 1 - - start_feat = min(start_feat, total_feat_length - 1) - end_feat = min(max(end_feat, start_feat + 1), total_feat_length) - cropped_fbank.append(fbank[index, start_feat:end_feat]) - cropped_fbank_lengths.append(end_feat - start_feat) - - return pad_sequence( - cropped_fbank, - batch_first=True, - padding_value=0.0, - ), torch.tensor( - cropped_fbank_lengths, - device=fbank.device, - dtype=torch.long, - ) - - def _extract_fbank_batch(self, audio, audio_lengths): - if self.resample is not None: - audio = self.resample(audio) - audio_lengths = torch.ceil(audio_lengths.float() * (_SPEAKER_FBANK_SAMPLE_RATE / self.sample_rate)).long() - - audio_cpu = audio.detach().cpu() - features = [] - - for index, valid_length_tensor in enumerate(audio_lengths): - valid_length = int(valid_length_tensor.item()) - waveform = audio_cpu[index, :valid_length] - if waveform.numel() == 0: - waveform = audio_cpu.new_zeros(1) - features.append( - extract_speaker_fbank( - waveform, - sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, - ) - ) - - fbank_lengths = torch.tensor( - [feature.size(0) for feature in features], - device=audio.device, - dtype=torch.long, - ) - fbank = pad_sequence( - features, - batch_first=True, - padding_value=0.0, - ).to(device=audio.device, dtype=audio.dtype) - return fbank, fbank_lengths - - @torch.no_grad() - @torch.autocast(enabled=False, device_type="cuda") - def forward(self, audio, audio_lengths=None, fbank=None, fbank_lengths=None, **_kwargs): - self.model.eval() - audio = audio.float() - if audio.dim() == 3: - if audio.size(1) != 1: - raise ValueError(f"Speaker encoder expects mono audio, got shape {tuple(audio.shape)}.") - audio = audio[:, 0] - elif audio.dim() != 2: - raise ValueError(f"Speaker encoder expects a 2D or 3D audio tensor, got shape {tuple(audio.shape)}.") - - audio, original_audio_lengths, cropped_audio_lengths, starts = self._crop_audio( - audio, - audio_lengths=audio_lengths, - ) - - if fbank is None: - fbank, fbank_lengths = self._extract_fbank_batch( - audio, - cropped_audio_lengths, - ) - else: - if not isinstance(fbank, torch.Tensor): - raise TypeError("Speaker encoder expects `fbank` to be a torch.Tensor.") - if fbank.dim() != 3 or fbank.size(0) != audio.size(0): - raise ValueError( - f"Speaker encoder expects `fbank` with shape (B, T, F) and matching batch size, got {tuple(fbank.shape)}." # noqa: E501 - ) - fbank, fbank_lengths = self._crop_fbank( - fbank.to(device=audio.device, dtype=torch.float32), - fbank_lengths, - original_audio_lengths, - cropped_audio_lengths, - starts, - ) - - return self.model(fbank, lengths=fbank_lengths) diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py deleted file mode 100644 index 862af8e17..000000000 --- a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py +++ /dev/null @@ -1,1426 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""dots.tts talker — vLLM-native AR base LM + audio side path. - -Mirrors upstream rednote-hilab/dots.tts (pinned @ a393d2e): - - * Qwen2.5-1.5B base LM runs under vLLM PagedAttention (stock - ``Qwen2Model``); weights come from the checkpoint's ``llm.model.*`` - namespace. ``DotsTTSConfig`` hoists the LM fields to top-level - attributes. - * Everything else is side computation, invisible to the vLLM - scheduler and driven once per step by ``_finish_decode``: - DiT flow-matching head (N-step Euler, fp32 integration state, - deterministic per-request noise) → io_helper denormalize → - patch_encoder AR loopback (produces the next step's LLM input - embedding) + streaming AudioVAE decode (per-request - ``BigVGANStreamState`` sliding window, so every latent patch is - decoded with real left conv context) → 48 kHz wav chunks pushed to - ``_audio_queue`` → eos_proj stop signal via ``_results_queue`` - (threshold 0.8, upstream default; the stop step also drains the - vocoder's 2-frame lookahead tail via ``stream_flush``). - * Cross-step state is isolated per request in ``_RequestState`` - (keyed by request_id); eviction is deferred to the end of - ``forward()`` so a finishing request's audio still drains. - * The CAM++ speaker encoder is constructed and its weights load, but - voice cloning is not exposed yet — generation is zero-shot and the - DiT falls back to its null conditioning (``fm_null_g_cond``). - -Debugging: set ``DOTS_TTS_BETA_TRACE=1`` for an env-gated per-patch -rms/max trace across the full side path (zero overhead when unset). - -Reference implementations: ``vllm_omni/model_executor/models/voxcpm2/`` -(architectural template: vLLM-native base LM + side path + per-request -state) and ``ming_flash_omni`` (AudioVAE weight-loading pattern). -""" - -from __future__ import annotations - -import dataclasses -import hashlib -import os -from collections.abc import Iterable -from dataclasses import dataclass -from typing import Any - -import torch -import torch.nn as nn -from vllm.config import VllmConfig -from vllm.logger import init_logger -from vllm.model_executor.models.qwen2 import Qwen2Model -from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix -from vllm.sequence import IntermediateTensors - -from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import DiT -from vllm_omni.model_executor.models.dots_tts.dots_tts_patch_encoder import ( - VAESemanticEncoder, -) -from vllm_omni.model_executor.models.dots_tts.dots_tts_speaker_encoder import ( - SpeakerXVectorFeatures, -) -from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import ( - AudioVAE, - AudioVAEConfig, -) -from vllm_omni.model_executor.models.output_templates import OmniOutput - -logger = init_logger(__name__) - - -def _require_config_block(hf_config: Any, name: str) -> dict[str, Any]: - """Fetch a sub-config dict (``vocoder`` / ``DiT`` / ``PatchEncoder``) - from the checkpoint's config.json, failing loudly when absent — - silently falling back to another checkpoint's values would produce - garbage audio, not an error.""" - block = getattr(hf_config, name, None) - if not isinstance(block, dict) or not block: - raise ValueError( - f"dots.tts checkpoint config.json has no '{name}' block — cannot " - "build the audio side path. Expected the upstream dots.tts config " - "layout (e.g. dots-studio/dots.tts-soar)." - ) - return block - - -def _build_audio_vae_config(hf_config: Any) -> AudioVAEConfig: - # AudioVAEConfig declares exactly the checkpoint's 17 ``vocoder`` keys, - # so an unknown / missing key raises TypeError instead of guessing. - return AudioVAEConfig(**_require_config_block(hf_config, "vocoder")) - - -# Architecture constants used by the per-request FM workspace math. All -# released dots.tts checkpoints (soar / base / mf) share these values; -# _validate_architecture_constants refuses a checkpoint that disagrees -# instead of silently generating garbage. -_FM_HIDDEN = 1024 # DiT.hidden_size (also LLM↔DiT projection target) -_LATENT_DIM = 128 # AudioVAE.latent_dim -_LATENT_PATCH_SIZE = 4 # config.patch_size — DiT samples 4 latent frames / audio patch -_HIDDEN_PATCH_SIZE = 1 # upstream core.py hardcodes self.hidden_patch_size = 1 -_MAX_AUDIO_PATCHES = 1024 # ~164 s @ 4×1920/48k; bounds per-request FM static buffer -_DIT_NUM_STEPS = int( - os.environ.get("DOTS_TTS_DIT_NUM_STEPS", "10") -) # env-gated; upstream default 10 for fixed-step Euler -_DIT_GUIDANCE_SCALE = 1.2 # upstream default guidance_scale for soar -_DIT_NOISE_SEED = 20260601 # base seed for per-request FM noise (voxcpm2 parity) -_PATCH_ENCODER_OUT_DS_RATE = 2 # patch_size / in_ds_rate = 4 / 2 (VAESemanticEncoder hardcodes in_ds_rate=2) - - -class _IOHelper: - """Latent stats wrapper. Holds (mean, var) over the 128-dim VAE latent - space and normalizes / denormalizes between DiT-internal space (which - operates on normalized latents) and AudioVAE / patch_encoder input space - (which expects raw latents). - - Mirrors upstream IOHelper (rednote-hilab/dots.tts @ a393d2e - models/dots_tts/core.py:684). Plain class (not nn.Module) so it stays - out of state_dict() and load_weights() doesn't have to mute its buffers. - Stats are loaded from latent_stats.pt — independent of safetensors. - """ - - def __init__(self, latent_stats_path: str | None = None) -> None: - if latent_stats_path is None: - self.global_mean: torch.Tensor | None = None - self.global_var: torch.Tensor | None = None - return - stats = torch.load(latent_stats_path, weights_only=False) - self.global_mean = torch.as_tensor(stats["mean"]) - self.global_var = torch.as_tensor(stats["var"]) - - def normalize(self, x: torch.Tensor) -> torch.Tensor: - if self.global_mean is None or self.global_var is None: - return x - mean = self.global_mean.to(device=x.device, dtype=x.dtype) - var = self.global_var.to(device=x.device, dtype=x.dtype) - return (x - mean) / torch.sqrt(var) - - def denormalize(self, x: torch.Tensor) -> torch.Tensor: - if self.global_mean is None or self.global_var is None: - return x - mean = self.global_mean.to(device=x.device, dtype=x.dtype) - var = self.global_var.to(device=x.device, dtype=x.dtype) - return x * torch.sqrt(var) + mean - - -# Per-request state container for AR-loop continuity. All cross-step -# side-path state lives here, keyed by request_id, so concurrent requests -# stay fully isolated (the voxcpm2 _RequestState pattern). -@dataclass -class _RequestState: - request_id: str - # AR loopback: patch_encoder's previous decode output, fed to the LM - # as next-step inputs_embeds. Shape: [1, llm_hidden] = [1, 1536]. - # (patch_encoder collapses its internal _PATCH_ENCODER_OUT_DS_RATE - # positions into one LLM token via out_proj; see - # VAESemanticEncoder._project_embeddings.) _finish_decode writes a - # fresh tensor every call. - curr_embed_for_next: torch.Tensor | None = None - # Stop-signal cache: _finish_decode populates precomputed_stop_logits; - # compute_logits drains it. - precomputed_stop_logits: torch.Tensor | None = None - is_stopping: bool = False - prefill_completed: bool = False - # Per-request FM noise counter: draw #n of this request hashes to a - # deterministic Generator seed (see _run_dit_n_step_euler), so outputs - # are reproducible run-to-run and concurrent requests cannot perturb - # each other's noise streams (voxcpm2 _fill_deterministic_cfm_noise). - noise_step: int = 0 - # Per-request FM static workspace (lazy-allocated by - # _initialize_request_fm_state on first _finish_decode call). Sized for - # _MAX_AUDIO_PATCHES × (_HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE) = 1024 × 5 - # = 5120 positions. - fm_sequence: torch.Tensor | None = None # [1, fm_capacity, _FM_HIDDEN] - fm_cfg_sequence: torch.Tensor | None = None # [1, fm_capacity, _FM_HIDDEN] - fm_null_g_cond: torch.Tensor | None = None # [1, _FM_HIDDEN] - fm_seq_len: int = 0 - fm_capacity: int = 0 - # Speaker x-vector after _xvec_proj. Zero-shot generation leaves this - # None → DiT falls back to fm_null_g_cond; voice-clone wiring is a - # follow-up. - g_cond: torch.Tensor | None = None - # AR-loop bookkeeping. - # patch_encoder_state holds conv_tail + per-layer KV caches for the - # patch_encoder's streaming decode (mirrors upstream state.patch_encoder_state). - # Lazy-allocated by _run_patch_encoder_loopback on first audio step. - patch_encoder_state: Any | None = None - # Review M2: streaming vocoder state (BigVGANStreamState — LSTM hidden - # + decoder sliding latent window). Gives each patch real left conv - # context; output lags input by decoder.stream_lookahead (2 frames = - # 3840 samples). Lazy-allocated by _run_vocoder_stream_step. - vocoder_stream_state: Any | None = None - - -# DiT transformer hyperparameters from dots-studio/dots.tts-soar config.json -# (the ``DiT`` block). Mirrors the shape of an upstream pydantic config: DiT's -# constructor calls ``transformer_config.to_dict()`` and reads ``.hidden_size`` -# / ``.num_layers``, so this @dataclass plus the ``to_dict`` helper satisfies -# that minimal interface without dragging in upstream's config system. -# -# Note: ``attn_dropout`` in the JSON spells differently from MultiHeadAttention's -# ``attn_drop`` kwarg — the value gets absorbed by ``**_kwargs`` and the MHA -# default 0.0 applies. Same numeric outcome for the soar checkpoint (both 0.0), -# so this is benign; flagging it because it would mask non-default values if -# upstream ever raised this field. -@dataclass -class _DiTConfig: - num_layers: int = 18 - num_heads: int = 16 - hidden_size: int = 1024 - ffn_hidden_size: int = 4096 - modulation: bool = True - qkv_bias: bool = False - qk_norm: bool = True - attn_dropout: float = 0.0 - dropout: float = 0.0 - norm_layer: str = "RMSNorm" - alibi_bias: bool = False - rotary_bias: bool = True - rotary_theta: float = 10000.0 - - def to_dict(self) -> dict[str, Any]: - return dataclasses.asdict(self) - - -def _build_dit_config(hf_config: Any) -> _DiTConfig: - # _DiTConfig declares exactly the checkpoint's 13 ``DiT`` keys, so an - # unknown key raises TypeError instead of being silently dropped. - return _DiTConfig(**_require_config_block(hf_config, "DiT")) - - -# Patch encoder hyperparameters from dots-studio/dots.tts-soar config.json -# (the top-level ``patch_size`` and the ``PatchEncoder`` sub-block). -# -# Two layers of config: -# * Outer (``_PatchEncoderConfig``): VAESemanticEncoder reads ``.patch_size`` -# and forwards ``.PatchEncoder`` to its inner SuperviseEncoder. -# * Inner (``_PatchEncoderInner``): the constructor uses BOTH attribute -# access (``config.PatchEncoder.hidden_size``) and dict-style -# ``config.get("hidden_size", 1024)``. Upstream's pydantic ConfigBase -# supports both natively; we add a tiny ``.get()`` helper to a dataclass. -# -# Upstream's ``_EncoderConfig`` also defines qk_norm / rotary_bias / -# rotary_theta / input_dim, but ``SuperviseEncoder.__init__`` only reads six -# keys (num_layers / num_heads / hidden_size / ffn_hidden_size / norm_layer / -# causal). The rest are swallowed and MultiHeadAttention falls back to its -# own defaults — matching upstream behaviour means *not* plumbing them through. -@dataclass -class _PatchEncoderInner: - num_layers: int = 24 - num_heads: int = 16 - hidden_size: int = 1024 - ffn_hidden_size: int = 4096 - norm_layer: str = "RMSNorm" - causal: bool = True - - def get(self, key: str, default: Any = None) -> Any: - return getattr(self, key, default) - - -@dataclass -class _PatchEncoderConfig: - patch_size: int = 4 - PatchEncoder: _PatchEncoderInner = dataclasses.field(default_factory=_PatchEncoderInner) - - -def _build_patch_encoder_config(hf_config: Any) -> _PatchEncoderConfig: - block = _require_config_block(hf_config, "PatchEncoder") - # SuperviseEncoder reads only the six declared fields; upstream's - # pydantic config swallows the other checkpoint keys, so filtering - # here matches upstream behaviour exactly. - inner_fields = {f.name for f in dataclasses.fields(_PatchEncoderInner)} - inner = _PatchEncoderInner(**{k: v for k, v in block.items() if k in inner_fields}) - return _PatchEncoderConfig( - patch_size=int(hf_config.patch_size), - PatchEncoder=inner, - ) - - -def _validate_architecture_constants(hf_config: Any, dit_config: _DiTConfig) -> None: - """The FM workspace helpers size buffers off module-level constants; - a checkpoint that disagrees would corrupt the DiT conditioning, so - refuse it up front (making these fully config-driven is deferred - until a released checkpoint actually differs).""" - checks = [ - ("latent_dim", getattr(hf_config, "latent_dim", None), _LATENT_DIM), - ("patch_size", getattr(hf_config, "patch_size", None), _LATENT_PATCH_SIZE), - ("DiT.hidden_size", dit_config.hidden_size, _FM_HIDDEN), - ] - mismatches = [f"{name}={got!r} (expected {want})" for name, got, want in checks if got != want] - if mismatches: - raise ValueError( - "dots.tts checkpoint config disagrees with this integration's " - "architecture constants: " + ", ".join(mismatches) - ) - - -class DotsTTSForConditionalGeneration(nn.Module): - """dots.tts AR talker: Qwen2.5 base LM + DiT / AudioVAE / CAM++ side path.""" - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - - self.vllm_config = vllm_config - self.config = vllm_config.model_config.hf_config - - # vllm-omni capability flags — gpu_model_runner checks these to decide - # whether to call preprocess() / postprocess() / accept multimodal - # outputs. Without these, the runner bypasses our preprocess hook - # and _pending_requests stays empty (forward never dispatches the - # side-path). Voxcpm2 sets the same three. - self.have_multimodal_outputs = True - self.has_preprocess = True - self.has_postprocess = True - - self.model = Qwen2Model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "model"), - ) - self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors - - # Side-path module hyperparameters come off the checkpoint's - # config.json blocks (fail loudly if a block is absent); the FM - # workspace constants are cross-checked against the same config. - dit_config = _build_dit_config(self.config) - _validate_architecture_constants(self.config, dit_config) - llm_hidden = self.config.hidden_size - fm_hidden = dit_config.hidden_size - latent_dim = int(self.config.latent_dim) - xvec_dim = int(getattr(self.config, "campplus_embedding_size", 512)) - - self._audio_vae = AudioVAE(_build_audio_vae_config(self.config)) - # Upstream serializes vocoder.safetensors with weight_norm folded into - # plain `weight` on the decoder (encoder kept weight_norm — it's not in - # the synthesis hot path). Match that layout before load_weights runs, - # so checkpoint keys align with our state_dict 1:1. - self._audio_vae.remove_weight_norm() - # AudioVAE / speaker encoder stay fp32 — upstream only casts - # `model.core` to bf16 (runtime.py:81), leaving vocoder + xvector - # extractor in fp32 because their kernels (conv1d / fbank) need - # higher precision. Without this _audio_vae would inherit talker's - # bf16 dtype via vLLM's auto-cast and dtype-mismatch under - # stream_step (input .float() vs bf16 weights). - - # DiT flow-matching head. Dimensions per upstream core.py:101-104: - # in_dim = config.DiT.hidden_size (DiT internal space) - # out_dim = config.latent_dim (AudioVAE input space) - # mode is "flow_matching" for soar/base; "meanflow" only for the mf - # checkpoint (handled by a separate factory in a later step). - self._head = DiT( - in_dim=fm_hidden, - out_dim=latent_dim, - transformer_config=dit_config, - mode="flow_matching", - ) - - # Patch encoder: closes the AR loop by mapping each DiT-emitted audio - # latent patch back into the LLM hidden space, so Qwen2 "hears" what it - # has already generated. Upstream core.py:75 calls this with - # in_dim=latent_dim and out_dim=llm.config.hidden_size. - self._patch_encoder = VAESemanticEncoder( - in_dim=latent_dim, - out_dim=llm_hidden, - config=_build_patch_encoder_config(self.config), - ) - - # Five thin projectors (upstream core.py:82-113). Soar dimensions: - # llm_hidden=1536, fm_hidden=1024, latent_dim=128, xvec_dim=512. - # These are kept as direct attributes (no vendor file) — they are flat - # nn.Linear / nn.Sequential, and live at the top of the checkpoint - # namespace under their own short prefixes (see load_weights). - self._hidden_proj = nn.Linear(llm_hidden, fm_hidden) - self._latent_proj = nn.Linear(latent_dim, fm_hidden) - self._coordinate_proj = nn.Linear(latent_dim, fm_hidden) - self._xvec_proj = nn.Sequential( - nn.Linear(xvec_dim, fm_hidden), - nn.LayerNorm(fm_hidden), - ) - # eos_proj is the stop predictor — feeds LLM hidden through a 2-layer - # MLP and produces a 2-way logit (continue / stop). - self._eos_proj = nn.Sequential( - nn.Linear(llm_hidden, llm_hidden), - nn.SiLU(), - nn.Linear(llm_hidden, 2), - ) - - # CAM++ speaker encoder (3D-Speaker), produces the x-vector from - # reference audio for voice cloning. sample_rate matches AudioVAE's - # input space (48 kHz); the wrapper resamples to 16k internally for - # CAM++. Weights are frozen (upstream does `requires_grad = False`). - self._speaker_encoder = SpeakerXVectorFeatures( - sample_rate=self._audio_vae.sample_rate, - campplus_embedding_size=xvec_dim, - max_audio_seconds=float(getattr(self.config, "xvec_max_audio_seconds", 10.0)), - ) - - # Pin AudioVAE + speaker encoder to fp32 (mirror upstream - # runtime.py:81 where only `model.core` is cast to bf16). vLLM's - # auto-cast would otherwise lower these to bf16, breaking - # stream_step (input .float() vs bf16 weights) and the fbank - # pipeline. - self._audio_vae.float() - self._speaker_encoder.float() - - # Latent stats — independent of safetensors (torch.load on - # latent_stats.pt sitting next to the checkpoint files). Used by - # _finish_decode to bridge between DiT's normalized latent space and - # the raw latent space expected by patch_encoder + AudioVAE. - latent_stats_path = self._resolve_latent_stats_path(vllm_config.model_config.model) - if latent_stats_path is not None: - self._io_helper = _IOHelper(latent_stats_path) - else: - logger.warning( - "latent_stats.pt not resolvable for %s; running with " - "identity normalize/denormalize. Wav will be ~10x quiet " - "and distorted (no sqrt(var) denormalize).", - vllm_config.model_config.model, - ) - self._io_helper = _IOHelper() - - # (β) diagnostic: env-gated per-patch rms/max trace across the full - # side path (LLM hidden → DiT → denormalize → patch_encoder → wav → - # eos). Zero overhead when DOTS_TTS_BETA_TRACE is unset (single - # bool check per step); kept as the model's debugging instrument - # (cf. voxcpm2's enable_profiling runtime knob). - self._beta_trace = bool(os.environ.get("DOTS_TTS_BETA_TRACE")) - - # Per-request state plumbing. - # _active_states: per-request state keyed by request_id (created in - # preprocess prefill, evicted by _flush_deferred_cleanup). - # _pending_requests: per-step list of (req_id, is_prefill, embeds, - # span_len) populated by preprocess(), consumed by forward(). - # _deferred_cleanup_ids: requests vLLM has marked finished but whose - # side-path audio still needs draining in the current forward step; - # evicted from _active_states at the end of forward(). - self._active_states: dict[str, _RequestState] = {} - self._pending_requests: list[tuple[str, bool, torch.Tensor, int]] = [] - self._deferred_cleanup_ids: set[str] = set() - - # Audio side-channel: - # _audio_queue: _finish_decode pushes (req_id, audio_tensor) tuples - # here; make_omni_output drains them into multimodal_outputs each - # step. - # _sample_rate: cached from AudioVAE config (= 48000 for soar). - self._audio_queue: list[tuple[str, torch.Tensor]] = [] - self._sample_rate = self._audio_vae.sample_rate - - # Stop-signal side-channel: - # _results_queue: _finish_decode pushes (req_id, stop_logits) tuples - # here; compute_logits drains them into the [bsz, vocab_size] - # logits tensor each step. - self._results_queue: list[tuple[str, torch.Tensor | None]] = [] - - logger.info( - "DotsTTS talker built (base_lm=Qwen2[28L,12H,1536d], " - "audio_vae=AudioVAE, dit=DiT[18L,16H,1024d], " - "patch_encoder=VAESemanticEncoder[24L,16H,1024d], " - "5 projectors, speaker_encoder=CAM++[7.2M], " - "io_helper=%s, model=%s, dit_steps=%d); " - "_finish_decode: DiT → denormalize → patch_encoder AR + " - "streaming AudioVAE decode → wav + eos_proj → stop signal.", - "loaded" if self._io_helper.global_mean is not None else "identity", - vllm_config.model_config.model, - _DIT_NUM_STEPS, - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, - ) -> torch.Tensor | IntermediateTensors: - """Drive the base LM, then the per-request audio side path. - - One vLLM step: base LM forward under PagedAttention, then slice - the scaffold hidden states per request (in _pending_requests - order) and run _finish_decode on each — DiT sampling, streaming - AudioVAE decode, AR loopback, and stop signal all happen there. - vLLM only ever sees the returned hidden states; audio leaves - through the _audio_queue side channel. - """ - output = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) - if isinstance(output, IntermediateTensors): - return output - if isinstance(output, tuple): - output = output[0] - - # Per-request side-path: slice scaffold_hidden by token span_len - # in the order preprocess() appended to _pending_requests. - token_offset = 0 - for req_id, is_prefill, _embeds, span_len in self._pending_requests: - req_hidden = output[token_offset : token_offset + span_len] - token_offset += span_len - self._finish_decode(req_id, req_hidden, is_prefill) - - # End-of-step cleanup - self._pending_requests.clear() - self._flush_deferred_cleanup() - - return output - - # ── vllm-omni protocol methods ── - - @staticmethod - def _resolve_latent_stats_path(model_arg: str) -> str | None: - """Resolve latent_stats.pt to a local filesystem path. - - ``vllm_config.model_config.model`` is typically an HF repo ID like - ``"dots-studio/dots.tts-soar"`` — joining it with ``"latent_ - stats.pt"`` gives a non-existent path. Try the local-dir form - first (covers ``--model /local/path``), then fall back to HF - cache lookup with ``local_files_only=True`` so we never trigger - a download from this side (the file was fetched alongside the - safetensors at load time). - """ - local = os.path.join(model_arg, "latent_stats.pt") - if os.path.exists(local): - return local - try: - from vllm_omni.transformers_utils.repo_utils import hf_api - - return hf_api().hf_hub_download( - repo_id=model_arg, - filename="latent_stats.pt", - local_files_only=True, - ) - except Exception: - return None - - def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: - """Map token ids to embeddings via the base LM's embed_tokens. - - Required by vLLM's ``is_text_generation_model`` Protocol check - (interfaces_base._check_vllm_model_embed_input_ids). Without this - the model fails ``--runner generate`` validation at engine startup. - """ - return self.model.embed_tokens(input_ids) - - def preprocess( - self, - input_ids: torch.Tensor, - input_embeds: torch.Tensor | None = None, - **info_dict: Any, - ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - """Per-step embedding + per-request state registration. - - Prefill: embed the wrapped prompt - ``[文本][文本对应语音]`` via the LM's - embed_tokens. Caller (serving layer or test harness) is - responsible for the wrap — talker just embeds whatever it gets. - - Decode (span_len == 1): use ``state.curr_embed_for_next`` from - the previous step's AR loopback (DiT → patch_encoder output); - zero fallback only if the loopback hasn't produced one yet. - - Voice clone is intentionally not wired here — zero-shot only. - - Returns ``(input_ids unchanged, embeds [span_len, hidden_size], - {})``. ``forward()`` consumes ``self._pending_requests`` to - slice scaffold_hidden per request. - """ - # voxcpm2 protocol compatibility: flatten additional_information. - additional = info_dict.get("additional_information") - if isinstance(additional, dict): - merged = {k: v for k, v in info_dict.items() if k != "additional_information"} - for k, v in additional.items(): - merged.setdefault(k, v) - info_dict = merged - - span_len = int(input_ids.shape[0]) - dev = input_ids.device - req_id = info_dict.get("request_id", "default") - state = self._active_states.get(req_id) - # Prefill is detected by state lifecycle (not span_len), so decode - # can use span_len=1 with a multi-position curr_embed_for_next. - is_prefill = state is None or not state.prefill_completed - - if is_prefill: - if state is None: - state = _RequestState(request_id=req_id) - self._active_states[req_id] = state - # Reset per-step state on every prefill — previous run is stale. - state.curr_embed_for_next = None - state.precomputed_stop_logits = None - state.is_stopping = False - state.prefill_completed = False - state.noise_step = 0 - # Reset AR-loop state on every prefill. - state.patch_encoder_state = None - state.vocoder_stream_state = None - embeds = self.model.embed_tokens(input_ids) - else: - curr = state.curr_embed_for_next - if curr is not None: - embeds = curr.to(dev) # [1, llm_hidden] - else: - # First decode step before AR loop closed — zero fallback. - embed_dtype = self.model.embed_tokens.weight.dtype - embeds = torch.zeros( - 1, - self.config.hidden_size, - device=dev, - dtype=embed_dtype, - ) - - self._pending_requests.append((req_id, is_prefill, embeds, span_len)) - return input_ids, embeds, {} - - def postprocess(self, *args: Any, **kwargs: Any) -> dict: - return {} - - def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: - """Mark finished requests for deferred eviction. - - vLLM scheduler calls this BEFORE forward() to notify that certain - requests have completed. We can't drop their _RequestState yet — - the current forward step still needs to drain side-path audio - chunks. Real eviction happens at the end of forward() via - _flush_deferred_cleanup. - """ - for req_id in finished_req_ids: - if req_id in self._active_states: - self._deferred_cleanup_ids.add(req_id) - - def _flush_deferred_cleanup(self) -> None: - """Evict states marked by on_requests_finished. Called at the - tail of forward().""" - for req_id in self._deferred_cleanup_ids: - self._active_states.pop(req_id, None) - self._deferred_cleanup_ids.clear() - - def _finish_decode( - self, - req_id: str, - req_hidden: torch.Tensor, - is_prefill: bool, - ) -> None: - """One request's per-step side path: DiT → AR loopback → wav → stop. - - One vLLM step = one audio patch. patch_encoder.decode_patch - collapses its internal _PATCH_ENCODER_OUT_DS_RATE positions into - a single LLM input embedding (see VAESemanticEncoder._project_ - embeddings), so vLLM's span_len=1 decode matches naturally — no - slot bookkeeping needed. - - Per-call flow (prefill or decode): - 1. _hidden_proj(last_hidden) → fm_sequence (+1 position) - 2. DiT N-step Euler → audio_patch [1, 4, 128] (normalized) - 3. _latent_proj(audio_patch) → fm_sequence (+4 positions, normalized) - 4. io_helper.denormalize(audio_patch) → audio_patch_raw - 5. patch_encoder(audio_patch_raw) → next-step embed [1, 1536] - 6. AudioVAE stream_step(audio_patch_raw) → wav chunk pushed to - _audio_queue (prefill's patch #0 included; output lags 2 - frames behind input) - 7. (decode only) eos_proj(last_hidden) → stop_logits pushed to - _results_queue; threshold > 0.8 flips state.is_stopping and - stream_flush drains the decoder's lookahead tail. - """ - state = self._active_states.get(req_id) - if state is None: - return - - if state.fm_sequence is None: - self._initialize_request_fm_state( - state, - device=req_hidden.device, - dtype=req_hidden.dtype, - ) - - # This step's append would overflow the per-request FM buffer. - # Stop gracefully (same path as a model-decided stop below) instead - # of raising past _append_hidden_chunk/_append_history_chunk — - # those raises are engine-fatal, not request-fatal, and would take - # every other in-flight and future request down with this one. - if state.fm_seq_len + _HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE > state.fm_capacity: - state.is_stopping = True - stop_logits = torch.tensor([[0.0, 1.0]], device=req_hidden.device, dtype=req_hidden.dtype) - state.precomputed_stop_logits = stop_logits - tail = self._audio_vae.stream_flush(state.vocoder_stream_state) - if tail.size(-1) > 0: - self._audio_queue.append((req_id, tail.reshape(-1))) - self._results_queue.append((req_id, stop_logits)) - return - - # 1. Append last LLM hidden to fm_sequence (+1 position). - last_hidden = req_hidden[-_HIDDEN_PATCH_SIZE:].unsqueeze(0) - - self._append_hidden_chunk(state, last_hidden) - - # 2. DiT N-step Euler → audio latent patch [1, 4, 128] (normalized). - audio_patch = self._run_dit_n_step_euler(state) - - # 3. Append latent to fm_sequence (+4 positions). Stays in - # normalized space — DiT's KV cache lives in that space. - self._append_history_chunk(state, audio_patch) - - # 4. Denormalize for downstream consumers (patch_encoder, AudioVAE). - audio_patch_raw = self._io_helper.denormalize(audio_patch) - - # 5. patch_encoder AR loopback → next-step inputs_embeds [1, 1536]. - next_embeds = self._run_patch_encoder_loopback(state, audio_patch_raw) - state.curr_embed_for_next = next_embeds.squeeze(0).detach() - - # 6. Streaming AudioVAE decode → wav chunk. Runs on prefill too: - # upstream zero-shot emits every DiT patch including the first - # (_generate_latents_stream drops one patch only under prompt - # prefill / voice clone), and our prefill patch #0 corresponds - # to upstream's first decode-loop patch. - wav = self._run_vocoder_stream_step(state, audio_patch_raw) - if wav.size(-1) > 0: - self._audio_queue.append((req_id, wav.reshape(-1))) - - if is_prefill: - if self._beta_trace: - self._beta_trace_log( - state, - is_prefill=True, - last_hidden=last_hidden, - audio_patch=audio_patch, - audio_patch_raw=audio_patch_raw, - next_embeds=next_embeds, - wav=wav if wav.size(-1) > 0 else None, - prob_stop=None, - ) - state.prefill_completed = True - # Keep len(_results_queue) == bsz: compute_logits pairs queue - # entries with batch rows positionally, so every request in the - # batch must push exactly one entry per step (voxcpm2_talker.py - # prefill placeholder pattern). None → forced continue via the - # stop_logits-is-None branch in compute_logits. - self._results_queue.append((req_id, None)) - return - - # 7. eos_proj → stop probability. Mirrors upstream - # model.py:1681 — softmax([continue, stop]) on detached LLM - # hidden, threshold the stop slot at 0.8 (upstream default - # eos_threshold). When stop fires for this step, compute_logits - # drains the high-stop logit on the next dispatch and vLLM - # terminates the request. - eos_logits = self._eos_proj(last_hidden.detach()).softmax(dim=-1) - stop_logits = eos_logits.squeeze(0) - state.precomputed_stop_logits = stop_logits - if eos_logits[0, -1, 1].item() > 0.8: - state.is_stopping = True - # Upstream flushes after its decode loop (model.py:1898); our - # request ends when vLLM sees this step's stop token, so drain - # the decoder's 2-frame lookahead tail in the same step. - # Requests cut by max_tokens/abort never reach this path and - # lose the 40 ms tail — no engine hook fires late enough to - # flush after the final step. - tail = self._audio_vae.stream_flush(state.vocoder_stream_state) - if tail.size(-1) > 0: - self._audio_queue.append((req_id, tail.reshape(-1))) - self._results_queue.append((req_id, stop_logits)) - - if self._beta_trace: - self._beta_trace_log( - state, - is_prefill=False, - last_hidden=last_hidden, - audio_patch=audio_patch, - audio_patch_raw=audio_patch_raw, - next_embeds=next_embeds, - wav=wav if wav.size(-1) > 0 else None, - prob_stop=eos_logits[0, -1, 1].item(), - ) - - def _beta_trace_log( - self, - state: _RequestState, - *, - is_prefill: bool, - last_hidden: torch.Tensor, - audio_patch: torch.Tensor, - audio_patch_raw: torch.Tensor, - next_embeds: torch.Tensor, - wav: torch.Tensor | None, - prob_stop: float | None, - ) -> None: - """(β) diagnostic — patch label matches output.wav indexing in - decode rows; prefill row labels itself ``pre``.""" - - def rmsmax(t: torch.Tensor) -> tuple[float, float]: - f = t.detach().float() - return f.pow(2).mean().sqrt().item(), f.abs().max().item() - - patch_label = "pre" if is_prefill else f"{state.fm_seq_len // 5 - 2:2d}" - req_tail = state.request_id[-8:] - hid_rms, hid_max = rmsmax(last_hidden) - dit_rms, dit_max = rmsmax(audio_patch) - raw_rms, raw_max = rmsmax(audio_patch_raw) - nxt_rms, nxt_max = rmsmax(next_embeds) - if wav is not None: - w_rms, w_max = rmsmax(wav) - wav_part = f"wav:rms={w_rms:.3f} max={w_max:.3f}" - else: - wav_part = "wav:rms= - max= - " - stop_part = f"stop={prob_stop:.3f}" if prob_stop is not None else "stop= - " - logger.info( - f"[β req={req_tail} patch={patch_label}] " - f"hid:rms={hid_rms:.3f} max={hid_max:.3f} | " - f"dit:rms={dit_rms:.3f} max={dit_max:.3f} | " - f"raw:rms={raw_rms:.3f} max={raw_max:.3f} | " - f"next:rms={nxt_rms:.3f} max={nxt_max:.3f} | " - f"{wav_part} | {stop_part}" - ) - - # ── FM helpers (per-request workspace + DiT N-step Euler) ── - - def _initialize_request_fm_state( - self, - state: _RequestState, - *, - device: torch.device, - dtype: torch.dtype, - ) -> None: - """Lazy-allocate per-request FM workspace. - - Mirrors upstream _allocate_generate_state (model.py:413) but per- - request (no workspace sharing) so concurrent requests stay - isolated. - """ - fm_capacity = _MAX_AUDIO_PATCHES * (_HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE) - state.fm_capacity = fm_capacity - state.fm_sequence = torch.zeros( - (1, fm_capacity, _FM_HIDDEN), - device=device, - dtype=dtype, - ) - state.fm_cfg_sequence = torch.zeros( - (1, fm_capacity, _FM_HIDDEN), - device=device, - dtype=dtype, - ) - state.fm_null_g_cond = torch.zeros( - (1, _FM_HIDDEN), - device=device, - dtype=dtype, - ) - state.fm_seq_len = 0 - - def _append_hidden_chunk( - self, - state: _RequestState, - hidden_chunk: torch.Tensor, - ) -> None: - """Project last _HIDDEN_PATCH_SIZE LLM hiddens via _hidden_proj into fm_sequence. - - Mirrors upstream _append_hidden_chunk (model.py:1308): cfg buffer - gets the null-projected counterpart (``_hidden_proj`` of zero) so - CFG dropout sees no past hidden, only past audio history. - """ - last_hidden = hidden_chunk[:, -_HIDDEN_PATCH_SIZE:, :] - projected = self._hidden_proj(last_hidden) - null_projected = self._hidden_proj(torch.zeros_like(last_hidden)) - end = state.fm_seq_len + projected.size(1) - if end > state.fm_capacity: - raise RuntimeError(f"FM buffer overflow on hidden append: end={end} cap={state.fm_capacity}") - state.fm_sequence[:, state.fm_seq_len : end].copy_(projected.to(state.fm_sequence.dtype)) - state.fm_cfg_sequence[:, state.fm_seq_len : end].copy_(null_projected.to(state.fm_cfg_sequence.dtype)) - state.fm_seq_len = end - - def _append_history_chunk( - self, - state: _RequestState, - latent_chunk: torch.Tensor, - ) -> None: - """Project audio latent via _latent_proj into fm_sequence (history append). - - Mirrors upstream _append_history_chunk (model.py:1325): cfg - buffer gets the same value (latent history is shared between - conditional and unconditional CFG paths — only the hidden chunk - gets nulled). - """ - history_latent = self._latent_proj(latent_chunk) - end = state.fm_seq_len + history_latent.size(1) - if end > state.fm_capacity: - raise RuntimeError(f"FM buffer overflow on latent append: end={end} cap={state.fm_capacity}") - state.fm_sequence[:, state.fm_seq_len : end].copy_(history_latent.to(state.fm_sequence.dtype)) - state.fm_cfg_sequence[:, state.fm_seq_len : end].copy_(history_latent.to(state.fm_cfg_sequence.dtype)) - state.fm_seq_len = end - - def _build_fm_attn_mask( - self, - state: _RequestState, - total_len: int, - device: torch.device, - ) -> torch.Tensor: - """Build DiT attention mask for one decode step. - - Mirrors upstream _build_fm_attn_mask (model.py:1203). Layout: - * [0, fm_seq_len) — committed history. Causal over the past - up to a non-causal tail block of size _HIDDEN_PATCH_SIZE that - sees itself and the new noise patch (the just-appended hidden). - * [latent_start, total_len) — new noise patch, bidirectional - with the committed history and itself. - """ - attn_mask = torch.zeros((1, total_len, total_len), device=device, dtype=torch.bool) - latent_start = total_len - _LATENT_PATCH_SIZE - block_start = state.fm_seq_len - _HIDDEN_PATCH_SIZE - if block_start > 0: - causal_mask = torch.ones((block_start, block_start), device=device, dtype=torch.bool).triu(1).logical_not() - attn_mask[:, :block_start, :block_start] = causal_mask - attn_mask[:, block_start : state.fm_seq_len, : state.fm_seq_len] = True - attn_mask[:, block_start : state.fm_seq_len, latent_start:] = True - attn_mask[:, latent_start:, : state.fm_seq_len] = True - attn_mask[:, latent_start:, latent_start:] = True - if latent_start > state.fm_seq_len: - padding_indices = torch.arange(state.fm_seq_len, latent_start, device=device) - attn_mask[:, padding_indices, padding_indices] = True - return attn_mask - - def _build_fm_pos_ids( - self, - state: _RequestState, - total_len: int, - device: torch.device, - ) -> torch.Tensor: - """Build DiT position ids for one decode step. - - Mirrors upstream _build_fm_pos_ids (model.py:1236). Position ids - stay fp32 — DiT's RotaryEmbedding does its own dtype handling. - """ - pos_ids = torch.zeros((1, total_len), device=device, dtype=torch.float32) - latent_start = total_len - _LATENT_PATCH_SIZE - if state.fm_seq_len > 0: - pos_ids[:, : state.fm_seq_len] = torch.arange( - state.fm_seq_len, - device=device, - dtype=pos_ids.dtype, - ) - pos_ids[:, latent_start:] = torch.arange( - state.fm_seq_len, - state.fm_seq_len + _LATENT_PATCH_SIZE, - device=device, - dtype=pos_ids.dtype, - ) - return pos_ids - - def _run_dit_n_step_euler( - self, - state: _RequestState, - *, - num_steps: int = _DIT_NUM_STEPS, - guidance_scale: float = _DIT_GUIDANCE_SCALE, - ) -> torch.Tensor: - """DiT N-step Euler integration → audio latent patch [1, 4, 128]. - - Manual Euler loop (upstream uses torchdyn.odeint; we manualize to - avoid the dep and keep control flow explicit). Each step: - 1. z_proj = _coordinate_proj(z), placed into the last - _LATENT_PATCH_SIZE positions of input_sequence / cfg_sequence. - 2. Batch [conditional, unconditional] → DiT velocity field. - 3. CFG blend: v_c + guidance_scale * (v_c - v_u). - 4. Euler: z += v * dt. - - Mirrors upstream fm_solver_step (core.py:295) + _flow_matching_step_fm - (core.py:463), reduced to a fixed-step Euler integrator. - """ - assert state.fm_sequence is not None - assert state.fm_cfg_sequence is not None - assert state.fm_null_g_cond is not None - device = state.fm_sequence.device - dtype = state.fm_sequence.dtype - total_len = state.fm_seq_len + _LATENT_PATCH_SIZE - - # Workspace: committed history + 4 noise slots (overwritten per step). - input_sequence = torch.zeros((1, total_len, _FM_HIDDEN), device=device, dtype=dtype) - input_sequence[:, : state.fm_seq_len] = state.fm_sequence[:, : state.fm_seq_len] - cfg_sequence = torch.zeros((1, total_len, _FM_HIDDEN), device=device, dtype=dtype) - cfg_sequence[:, : state.fm_seq_len] = state.fm_cfg_sequence[:, : state.fm_seq_len] - - attn_mask = self._build_fm_attn_mask(state, total_len, device) - pos_ids = self._build_fm_pos_ids(state, total_len, device) - - g_cond = state.g_cond if state.g_cond is not None else state.fm_null_g_cond - g_cond = g_cond.to(device=device, dtype=dtype) - g_cond_batched = torch.cat([g_cond, torch.zeros_like(g_cond)], dim=0) - - latent_start = total_len - _LATENT_PATCH_SIZE - # fp32 ODE integration (same fix as Ming's CFM sampler, PR #4341): - # bf16's 7-bit mantissa loses the small per-step increments in the - # z += v*dt accumulation (hidden-state cos drift vs upstream past - # step ~7) and quantizes the linspace timesteps. Integration - # state stays fp32; DiT matmuls still run bf16 under the autocast - # block below, and the result is cast back to the sequence dtype - # on return. - # Deterministic per-request noise (voxcpm2 _fill_deterministic_cfm_ - # noise pattern): hash seed:request_key:draw# into a private - # Generator instead of the global CUDA RNG, so outputs reproduce - # run-to-run and concurrent requests cannot perturb each other's - # noise streams. request_key strips the engine's per-run "_" - # uuid suffix so replay across runs keys on the stable batch index. - request_key = state.request_id.split("_", 1)[0] - if not request_key.isdigit(): - request_key = state.request_id - noise_key = f"{_DIT_NOISE_SEED}:{request_key}:{state.noise_step}".encode() - digest = hashlib.blake2b(noise_key, digest_size=8).digest() - gen = torch.Generator(device=device) - gen.manual_seed(int.from_bytes(digest, "little") & 0x7FFF_FFFF_FFFF_FFFF) - z = torch.empty((1, _LATENT_PATCH_SIZE, _LATENT_DIM), device=device, dtype=torch.float32) - z.normal_(generator=gen) - state.noise_step += 1 - dt = 1.0 / num_steps - times = torch.linspace(0.0, 1.0, num_steps + 1, device=device, dtype=torch.float32) - - # Match upstream's autocast wrapper (model.py:280). DiT's - # TimestepEmbedder internally forces fp32 (`.float()` on freqs); - # without autocast the fp32 output hits a bf16 Linear → dtype - # mismatch. autocast bridges this for matmuls inside DiT. - use_amp = device.type == "cuda" and dtype in (torch.float16, torch.bfloat16) - with torch.autocast( - device_type=device.type if device.type == "cuda" else "cuda", - dtype=dtype if use_amp else torch.float32, - enabled=use_amp, - ): - for step in range(num_steps): - t = times[step].reshape(1) - z_proj = self._coordinate_proj(z) - z_c = input_sequence.clone() - z_c[:, latent_start:] = z_proj - z_u = cfg_sequence.clone() - z_u[:, latent_start:] = z_proj - z_batched = torch.cat([z_c, z_u], dim=0) - t_batched = t.repeat(2) - vt = self._head( - x=z_batched, - timesteps=t_batched, - attn_mask=attn_mask, - pos_ids=pos_ids, - g_cond=g_cond_batched, - ) - # Upcast DiT output before CFG blend + Euler update so the - # accumulation arithmetic stays fp32 (autocast only affects - # matmuls; these pointwise ops keep their input dtype). - vt = vt[:, latent_start:].float() - vt_c, vt_u = vt[0:1], vt[1:2] - velocity = vt_c + guidance_scale * (vt_c - vt_u) - z = z + velocity * dt - return z.to(dtype) - - def _run_patch_encoder_loopback( - self, - state: _RequestState, - audio_patch_raw: torch.Tensor, - ) -> torch.Tensor: - """Decode one audio latent patch → 1 LLM input embed. - - Mirrors upstream _consume_audio_patch (model.py:1557) — patch_encoder - side only. Lazy-allocates state.patch_encoder_state on first call. - Returns ``llm_embedding`` of shape ``[1, 1, llm_hidden]`` — the - _PATCH_ENCODER_OUT_DS_RATE encoder positions fold into the feature - axis before ``out_proj``, not into separate output positions. - - ``audio_patch_raw`` must already be denormalized (raw VAE latent - space). Caller (_finish_decode) runs io_helper.denormalize on - DiT's normalized output before passing in. - """ - if state.patch_encoder_state is None: - state.patch_encoder_state = self._patch_encoder.init_decode_state( - max_audio_patch_count=_MAX_AUDIO_PATCHES, - batch_size=1, - device=audio_patch_raw.device, - dtype=audio_patch_raw.dtype, - ) - pe_state = state.patch_encoder_state - positions = ( - torch.arange( - _PATCH_ENCODER_OUT_DS_RATE, - device=audio_patch_raw.device, - dtype=torch.long, - ) - + pe_state.seq_len - ) - llm_embedding, conv_tail = self._patch_encoder.decode_patch( - audio_patch_raw, - pe_state.conv_tail, - pe_state.layer_caches, - positions, - ) - pe_state.conv_tail.copy_(conv_tail) - pe_state.seq_len += _PATCH_ENCODER_OUT_DS_RATE - return llm_embedding - - def _run_vocoder_stream_step( - self, - state: _RequestState, - audio_patch_raw: torch.Tensor, - ) -> torch.Tensor: - """Streaming AudioVAE decode of one latent patch (review M2). - - Mirrors upstream generate_audio_stream (model.py:1881-1900): a - per-request BigVGANStreamState carries the decoder's LSTM hidden - and sliding latent window across steps, so each patch is decoded - with real left context instead of inference_from_latents' - zero-padded isolated boundary. Output lags input by - decoder.stream_lookahead (2 frames): the first call emits 2 frames - (3840 samples), steady state 4 (7680); the stop path in - _finish_decode drains the tail via stream_flush. - - audio_patch_raw: [1, T=4, D=128] denormalized latents. - Returns wav chunk [1, 1, n_samples] (n_samples may be 0). - """ - if state.vocoder_stream_state is None: - state.vocoder_stream_state = self._audio_vae.init_stream_state( - batch_size=1, - chunk_size=_LATENT_PATCH_SIZE, - ) - # stream_step expects [B, latent_dim, T]; the vocoder is pinned - # fp32 (__init__), so cast the (possibly bf16) latents up. - return self._audio_vae.stream_step( - audio_patch_raw.transpose(1, 2).float(), - state.vocoder_stream_state, - ) - - def make_omni_output( - self, - model_outputs: torch.Tensor | OmniOutput, - **kwargs: Any, - ) -> OmniOutput: - """Drain self._audio_queue into the OmniOutput's multimodal_outputs. - - Output contract is **delta** (SKILL I1): each call emits only - audio produced since the previous step; the engine's - _consolidate_multimodal_tensors concatenates across steps for - offline consumers. - - Same-request multi-chunk: chunks for the same req_id within one - step are cat'd together. Multi-request: one entry per req_id - in multimodal_outputs["model_outputs"]. - - Note: ``sr`` is not a declared OmniPayload key but the engine / - serving consumer reads it (voxcpm2 convention). - """ - if isinstance(model_outputs, OmniOutput): - return model_outputs - - mm: dict[str, Any] = {} - if self._audio_queue: - audio_by_req: dict[str, torch.Tensor] = {} - for req_id, audio in self._audio_queue: - if audio is None: - continue - if req_id in audio_by_req: - audio_by_req[req_id] = torch.cat( - [audio_by_req[req_id].reshape(-1), audio.reshape(-1)], - dim=0, - ) - else: - audio_by_req[req_id] = audio - if audio_by_req: - sr = torch.tensor(self._sample_rate, dtype=torch.int32) - ready_req_ids = list(audio_by_req) - chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] - mm["model_outputs"] = chunks - mm["sr"] = [sr for _ in ready_req_ids] - # sparse_audio: ["1"] flags GPUARModelRunner to skip the default - # payload["hidden"] = scaffold_hidden injection (gpu_ar_model_ - # runner.py:1114). Otherwise prefill scaffold (24 hidden × 1536 - # = 29184 numbers observed) leaks into mm["audio"] via the - # output_processor "hidden" → target_key="audio" rename - # (output_processor.py:84), prepending ~0.61 s of noise. - mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} - self._audio_queue.clear() - else: - # Empty-audio step (prefill) still needs the sparse marker so the - # engine doesn't bleed scaffold_hidden into mm["audio"]. - mm["model_outputs"] = [] - mm["sr"] = [] - mm["meta"] = {"req_id": [], "sparse_audio": ["1"]} - - if self._beta_trace: - payload = mm.get("model_outputs") - if payload: - for i, t in enumerate(payload): - logger.info( - f"[β-emit] mm[model_outputs][{i}]: shape={tuple(t.shape)} " - f"dtype={t.dtype} samples={t.numel()} " - f"rms={t.float().pow(2).mean().sqrt().item():.4f} " - f"max={t.abs().max().item():.4f}" - ) - else: - logger.info("[β-emit] mm empty / no model_outputs") - mo_kind = type(model_outputs).__name__ - mo_info = f"shape={tuple(model_outputs.shape)}" if hasattr(model_outputs, "shape") else "no_shape" - logger.info(f"[β-emit] text_hidden_states arg: {mo_kind} {mo_info}") - - return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) - - def compute_logits( - self, - hidden_states: torch.Tensor | OmniOutput, - sampling_metadata: Any = None, - ) -> torch.Tensor | None: - """Per-step stop signal logits (vLLM sampling protocol). - - Encodes continue/stop into ``logits[i, 0]`` / ``logits[i, 1]`` - (voxcpm2 convention — rest of vocab stays at -inf so the sampler - picks between the two slots). Source of stop logits is - ``self._results_queue``, populated by forward()'s _finish_decode. - When the queue is empty (preprocess-only steps), defaults to - all-continue. - """ - if isinstance(hidden_states, OmniOutput): - hidden_states = hidden_states.text_hidden_states - if hidden_states is None: - return None - - bsz = hidden_states.shape[0] - logits = torch.full( - (bsz, self.config.vocab_size), - float("-inf"), - device=hidden_states.device, - dtype=hidden_states.dtype, - ) - - if self._results_queue: - for i, (req_id, stop_logits) in enumerate(self._results_queue): - if i >= bsz: - break - state = self._active_states.get(req_id) - if stop_logits is not None: - if state is not None and state.is_stopping: - logits[i, 0] = 0.0 - logits[i, 1] = 1.0 - else: - # is_stopping=False means prob_stop <= 0.8 — force - # continue. Feeding the raw softmax (continue, stop) - # pair as logits would let the greedy sampler compare - # them directly, silently lowering the effective stop - # threshold to 0.5 (early stop, swallowed endings). - logits[i, 0] = 1.0 - if state is not None: - state.precomputed_stop_logits = None - else: - # No stop signal pushed this step (prefill placeholder). - # Default to continue — stop is driven by eos_proj. - logits[i, 0] = 1.0 # continue - self._results_queue.clear() - else: - logits[:, 0] = 1.0 # all continue - - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load all dots.tts checkpoint weights — step 6b. - - Routing branches in a single pass over the input iterator (it can only - be consumed once): - - * **AudioVAE**: keys from ``vocoder.safetensors`` (no parent prefix). - Matched via set membership against ``self._audio_vae.state_dict()``; - acts as the final catch-all so it must run last. - * **DiT**: keys from ``model.safetensors`` under ``velocity_field_predictor.``. - * **patch_encoder**: keys from ``model.safetensors`` under ``patch_encoder.``. - * **Five projectors** (``hidden_proj`` / ``latent_proj`` / - ``coordinate_proj`` / ``xvec_proj`` / ``eos_proj``): each lives at - its own top-level prefix; dispatched through a shared table. - * **Qwen2 base LM**: keys under ``llm.model.``, handed off to vLLM's - Qwen2Model fuser (q/k/v_proj -> qkv_proj, gate/up_proj -> - gate_up_proj). ``llm.lm_head.*`` is silently dropped (soar uses - ``tie_word_embeddings=True`` and our ``self.model`` is a Qwen2Model - with no lm_head module). - * **CAM++ speaker encoder**: keys from ``speaker_encoder.safetensors`` - (no parent prefix, same convention as AudioVAE). Routed by set - membership against ``self._speaker_encoder.state_dict()`` (keys are - ``model.*`` / ``resample.kernel``, no distinguishing top-level - prefix). **Placed before the AudioVAE fallback** so its keys don't - fall through to the catch-all. - """ - vae_state_keys = set(self._audio_vae.state_dict().keys()) - dit_state_keys = set(self._head.state_dict().keys()) - patch_state_keys = set(self._patch_encoder.state_dict().keys()) - speaker_state_keys = set(self._speaker_encoder.state_dict().keys()) - - # (top-level prefix in checkpoint, local attr name, local nn.Module) - projector_specs: list[tuple[str, str, nn.Module]] = [ - ("hidden_proj.", "_hidden_proj", self._hidden_proj), - ("latent_proj.", "_latent_proj", self._latent_proj), - ("coordinate_proj.", "_coordinate_proj", self._coordinate_proj), - ("xvec_proj.", "_xvec_proj", self._xvec_proj), - ("eos_proj.", "_eos_proj", self._eos_proj), - ] - projector_state_keys = {prefix: set(mod.state_dict().keys()) for prefix, _, mod in projector_specs} - projector_matched: dict[str, list[tuple[str, torch.Tensor]]] = {prefix: [] for prefix, _, _ in projector_specs} - - matched_vae: list[tuple[str, torch.Tensor]] = [] - matched_dit: list[tuple[str, torch.Tensor]] = [] - matched_patch: list[tuple[str, torch.Tensor]] = [] - matched_llm: list[tuple[str, torch.Tensor]] = [] - matched_speaker: list[tuple[str, torch.Tensor]] = [] - skipped_lm_head = 0 - - DIT_PREFIX = "velocity_field_predictor." - PATCH_PREFIX = "patch_encoder." - LLM_MODEL_PREFIX = "llm.model." - LLM_LM_HEAD_PREFIX = "llm.lm_head." - for name, tensor in weights: - if name.startswith(DIT_PREFIX): - candidate = name[len(DIT_PREFIX) :] - if candidate in dit_state_keys: - matched_dit.append((candidate, tensor)) - continue - if name.startswith(PATCH_PREFIX): - candidate = name[len(PATCH_PREFIX) :] - if candidate in patch_state_keys: - matched_patch.append((candidate, tensor)) - continue - if name.startswith(LLM_LM_HEAD_PREFIX): - # Qwen2.5-1.5B-Base has tie_word_embeddings=True, and our - # ``self.model`` is a Qwen2Model (no lm_head submodule). Drop. - skipped_lm_head += 1 - continue - if name.startswith(LLM_MODEL_PREFIX): - candidate = name[len(LLM_MODEL_PREFIX) :] - # Hand off untouched — Qwen2Model.load_weights does its own - # fusion (q/k/v_proj -> qkv_proj, gate/up_proj -> gate_up_proj) - # and tolerates unknown keys, so no local key set to gate. - matched_llm.append((candidate, tensor)) - continue - # Projectors: 5 candidate top-level prefixes, first hit wins. - hit_projector = False - for prefix, _attr, _mod in projector_specs: - if name.startswith(prefix): - candidate = name[len(prefix) :] - if candidate in projector_state_keys[prefix]: - projector_matched[prefix].append((candidate, tensor)) - hit_projector = True - break - if hit_projector: - continue - # Speaker encoder: set-membership match (no distinguishing prefix - # in checkpoint; placed before the AudioVAE fallback so its keys - # don't get silently dropped by the catch-all). - if name in speaker_state_keys: - matched_speaker.append((name, tensor)) - continue - candidate = name[len("vocoder.") :] if name.startswith("vocoder.") else name - if candidate in vae_state_keys: - matched_vae.append((candidate, tensor)) - - loaded: set[str] = set() - - if matched_vae: - vae_loader = AutoWeightsLoader(self._audio_vae) - loaded_vae = vae_loader.load_weights(iter(matched_vae)) - loaded.update(f"_audio_vae.{name}" for name in loaded_vae) - logger.info( - "DotsTTS load_weights: loaded %d/%d AudioVAE tensors.", - len(loaded_vae), - len(vae_state_keys), - ) - - if matched_dit: - dit_loader = AutoWeightsLoader(self._head) - loaded_dit = dit_loader.load_weights(iter(matched_dit)) - loaded.update(f"_head.{name}" for name in loaded_dit) - logger.info( - "DotsTTS load_weights: loaded %d/%d DiT tensors.", - len(loaded_dit), - len(dit_state_keys), - ) - - if matched_patch: - patch_loader = AutoWeightsLoader(self._patch_encoder) - loaded_patch = patch_loader.load_weights(iter(matched_patch)) - loaded.update(f"_patch_encoder.{name}" for name in loaded_patch) - logger.info( - "DotsTTS load_weights: loaded %d/%d patch_encoder tensors.", - len(loaded_patch), - len(patch_state_keys), - ) - - any_projector_matched = False - for prefix, attr_name, mod in projector_specs: - matched = projector_matched[prefix] - if not matched: - continue - any_projector_matched = True - proj_loader = AutoWeightsLoader(mod) - loaded_proj = proj_loader.load_weights(iter(matched)) - loaded.update(f"{attr_name}.{name}" for name in loaded_proj) - logger.info( - "DotsTTS load_weights: loaded %d/%d %s tensors.", - len(loaded_proj), - len(projector_state_keys[prefix]), - prefix[:-1], - ) - - if matched_llm: - # vLLM's Qwen2Model.load_weights does its own fusion + tolerates - # unknown keys. It returns the set of param names it actually - # accepted; re-prefix with ``model.`` for our return value (our - # local attr is ``self.model``). - loaded_llm = self.model.load_weights(iter(matched_llm)) - loaded.update(f"model.{name}" for name in loaded_llm) - logger.info( - "DotsTTS load_weights: loaded %d Qwen2 tensors (fed %d; %d llm.lm_head.* skipped — tied embeddings).", - len(loaded_llm), - len(matched_llm), - skipped_lm_head, - ) - - if matched_speaker: - speaker_loader = AutoWeightsLoader(self._speaker_encoder) - loaded_speaker = speaker_loader.load_weights(iter(matched_speaker)) - loaded.update(f"_speaker_encoder.{name}" for name in loaded_speaker) - logger.info( - "DotsTTS load_weights: loaded %d/%d CAM++ speaker_encoder tensors.", - len(loaded_speaker), - len(speaker_state_keys), - ) - - if ( - not matched_vae - and not matched_dit - and not matched_patch - and not any_projector_matched - and not matched_llm - and not matched_speaker - ): - logger.warning( - "DotsTTS load_weights: no AudioVAE / DiT / " - "patch_encoder / projector / Qwen2 / speaker_encoder keys matched." - ) - - return loaded diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py deleted file mode 100644 index fe17b5c0b..000000000 --- a/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py +++ /dev/null @@ -1,1434 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright 2026 The vLLM-Omni team. -# Copyright (c) rednote-hilab. All rights reserved. -# Adapted from: -# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules -# -# Merged from six upstream files (preserved as-is, inference behaviour -# unchanged; only internal imports were folded inline): -# - modules/vocoder/config.py -> AudioVAEConfig -# - modules/backbone/layers.py -> Conv1d, ConvTranspose1d -# (Mlp/MultiHeadAttention/RotaryEmbedding -# are not vendored here; they belong -# to the DiT side-path module added later) -# - modules/vocoder/alias_free_filter.py -> sinc, kaiser_sinc_filter1d, -# LowPassFilter1d -# - modules/vocoder/alias_free_resample.py-> UpSample1d, DownSample1d -# - modules/vocoder/alias_free_act.py -> Snake, SnakeBeta, Activation1d -# - modules/vocoder/bigvgan.py -> AudioVAE and all its sub-modules -# -# Per-section upstream attributions are preserved at each section header. -# AudioVAEConfig is the only piece that diverges from upstream: pydantic -# ConfigBase replaced with a stdlib @dataclass (extra fields not exercised by -# any shipped checkpoint). - -from __future__ import annotations - -import itertools -import math -from dataclasses import dataclass, field -from fractions import Fraction - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import pow, sin -from torch.nn import Parameter -from torch.nn.utils import remove_weight_norm, weight_norm - -# ============================================================================ -# Configuration (adapted from modules/vocoder/config.py) -# ---------------------------------------------------------------------------- -# Upstream subclasses a pydantic ``ConfigBase`` with ``extra="allow"``; both the -# dots.tts-soar and dots.tts-mf checkpoint ``config.json`` pass only the 17 -# declared fields below, so the escape hatch is never exercised. ``.get()`` -# helper preserves bigvgan.py's ``h.get(...)`` call sites without modification. -# ``num_decoder_lookahead`` / ``num_encoder_lookahead`` are looked up via -# ``h.get(...)`` but not present in any shipped config; they fall back to 2. -# ============================================================================ - - -@dataclass -class AudioVAEConfig: - sample_rate: int = 24000 - upsample_rates: list[int] = field(default_factory=list) - upsample_kernel_sizes: list[int] = field(default_factory=list) - upsample_initial_channel: int = 1536 - resblock: str = "1" - resblock_kernel_sizes: list[int] = field(default_factory=list) - resblock_dilation_sizes: list[list[int]] = field(default_factory=list) - downsample_rates: list[int] = field(default_factory=list) - downsample_channels: list[int] = field(default_factory=list) - activation: str = "snakebeta" - snake_logscale: bool = True - latent_dim: int = 128 - causal: bool = False - mi_num_layers: int = 4 - causal_encoder: bool = False - use_bias_at_final: bool = True - use_tanh_at_final: bool = True - - num_decoder_lookahead: int = 2 - num_encoder_lookahead: int = 2 - - def get(self, key: str, default=None): - return getattr(self, key, default) - - -# ============================================================================ -# Conv layers (vendored from modules/backbone/layers.py) -# ---------------------------------------------------------------------------- -# Only the Conv1d / ConvTranspose1d helpers are needed by the AudioVAE. -# Mlp / MultiHeadAttention / RotaryEmbedding from the same upstream file are -# reserved for the DiT flow-matching head, which gets its own vendored module -# in a later step. -# ============================================================================ - - -class Conv1d(nn.Conv1d): - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int = 1, - stride: int = 1, - dilation: int = 1, - groups: int = 1, - padding_mode: str = "zeros", - bias: bool = True, - padding=None, - causal: bool = False, - **_kwargs, - ): - self.causal = causal - if padding is None: - if causal: - padding = 0 - self.left_padding = dilation * (kernel_size - 1) - else: - padding = int((kernel_size * dilation - dilation) / 2) - - super().__init__( - in_channels, - out_channels, - kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - groups=groups, - padding_mode=padding_mode, - bias=bias, - ) - - self.in_channels = in_channels - - def forward(self, x): - if self.causal: - x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2) - return super().forward(x) - - -class ConvTranspose1d(nn.ConvTranspose1d): - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - stride: int = 1, - output_padding: int = 0, - groups: int = 1, - bias: bool = True, - dilation: int = 1, - padding=None, - padding_mode: str = "zeros", - causal: bool = False, - **_kwargs, - ): - if padding is None: - padding = 0 if causal else (kernel_size - stride) // 2 - if causal: - assert padding == 0, "padding is not allowed in causal ConvTranspose1d." - assert kernel_size == 2 * stride, "kernel_size must be equal to 2*stride in Causal ConvTranspose1d." - - super().__init__( - in_channels, - out_channels, - kernel_size, - stride=stride, - padding=padding, - output_padding=output_padding, - groups=groups, - bias=bias, - dilation=dilation, - padding_mode=padding_mode, - ) - - self.causal = causal - self.stride = stride - - def forward(self, x): - x = super().forward(x) - if self.causal: - x = x[:, :, : -self.stride] - return x - - -# ============================================================================ -# Alias-free low-pass filter (vendored from modules/vocoder/alias_free_filter.py) -# ---------------------------------------------------------------------------- -# Adapted from https://github.com/junjun3518/alias-free-torch under -# the Apache License 2.0. The sinc fallback is adopted from adefossez's -# julius.core.sinc (MIT). -# ============================================================================ - - -if "sinc" in dir(torch): - sinc = torch.sinc -else: - - def sinc(x: torch.Tensor): - """ - Implementation of sinc, i.e. sin(pi * x) / (pi * x) - __Warning__: Different to julius.sinc, the input is multiplied by `pi`! - """ - return torch.where( - x == 0, - torch.tensor(1.0, device=x.device, dtype=x.dtype), - torch.sin(math.pi * x) / math.pi / x, - ) - - -def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size] - even = kernel_size % 2 == 0 - half_size = kernel_size // 2 - - # For kaiser window - delta_f = 4 * half_width - A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 - if A > 50.0: - beta = 0.1102 * (A - 8.7) - elif A >= 21.0: - beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) - else: - beta = 0.0 - window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) - - # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio - if even: - time = torch.arange(-half_size, half_size) + 0.5 - else: - time = torch.arange(kernel_size) - half_size - if cutoff == 0: - filter_ = torch.zeros_like(time) - else: - filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) - # Normalize filter to have sum = 1, otherwise we will have a small leakage - # of the constant component in the input signal. - filter_ /= filter_.sum() - filter = filter_.view(1, 1, kernel_size) - - return filter - - -class LowPassFilter1d(nn.Module): - def __init__( - self, - cutoff=0.5, - half_width=0.6, - stride: int = 1, - padding: bool = True, - padding_mode: str = "replicate", - kernel_size: int = 12, - channels: int = 1, - causal: bool = True, - fixed_filter: bool = False, - ): - # kernel_size should be even number for stylegan3 setup, - # in this implementation, odd number is also possible. - super().__init__() - if cutoff < -0.0: - raise ValueError("Minimum cutoff must be larger than zero.") - if cutoff > 0.5: - raise ValueError("A cutoff above 0.5 does not make sense.") - self.kernel_size = kernel_size - if causal: - self.pad_left = kernel_size - 1 - self.pad_right = 0 - else: - self.even = kernel_size % 2 == 0 - self.pad_left = kernel_size // 2 - int(self.even) - self.pad_right = kernel_size // 2 - self.stride = stride - self.padding = padding - self.padding_mode = padding_mode - self.fixed_filter = fixed_filter - filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) - if fixed_filter: - self.register_buffer("filter", filter) - else: - self.filter = nn.Parameter(filter.expand(channels, -1, -1).clone()) - - # input [B, C, T] - def forward(self, x): - _, C, _ = x.shape - if self.padding: - x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) - if self.fixed_filter: - out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) - else: - out = F.conv1d(x, self.filter, stride=self.stride, groups=C) - return out - - -# ============================================================================ -# Alias-free up/down sampling (vendored from modules/vocoder/alias_free_resample.py) -# ---------------------------------------------------------------------------- -# Adapted from https://github.com/junjun3518/alias-free-torch (Apache 2.0). -# ============================================================================ - - -class UpSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None, channels=None, causal=True, fixed_filter=False): - super().__init__() - self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.stride = ratio - self.channels = channels - self.causal = causal - self.fixed_filter = fixed_filter - if causal: - self.pad = 0 - else: - self.pad = self.kernel_size // ratio - 1 - self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 - self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 - filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) - if self.fixed_filter: - self.register_buffer("filter", filter) - else: - self.filter = nn.Parameter(filter.expand(channels, -1, -1).clone()) - - # x: [B, C, T] - def forward(self, x): - _, C, _ = x.shape - x = F.pad(x, (self.pad, self.pad), mode="replicate") - if self.fixed_filter: - x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) - else: - x = self.ratio * F.conv_transpose1d(x, self.filter, stride=self.stride, groups=C) - if self.causal: - x = x[..., : -(self.kernel_size - self.stride)] - else: - x = x[..., self.pad_left : -self.pad_right] - - return x - - -class DownSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None, channels=None, causal=True, fixed_filter=False): - super().__init__() - self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.lowpass = LowPassFilter1d( - cutoff=0.5 / ratio, - half_width=0.6 / ratio, - stride=ratio, - kernel_size=self.kernel_size, - channels=channels, - causal=causal, - fixed_filter=fixed_filter, - ) - - def forward(self, x): - return self.lowpass(x) - - -# ============================================================================ -# Snake activations + alias-free Activation1d (vendored from -# modules/vocoder/alias_free_act.py) -# ---------------------------------------------------------------------------- -# Adapted from https://github.com/junjun3518/alias-free-torch (Apache 2.0). -# Snake / SnakeBeta paper: https://arxiv.org/abs/2006.08195 -# ============================================================================ - - -class Activation1d(nn.Module): - def __init__( - self, - activation, - up_ratio: int = 2, - down_ratio: int = 2, - up_kernel_size: int = 12, - down_kernel_size: int = 12, - causal=True, - fixed_filter=False, - ): - super().__init__() - # causal=False - self.up_ratio = up_ratio - self.down_ratio = down_ratio - self.act = activation - self.upsample = UpSample1d( - up_ratio, - up_kernel_size, - activation.in_features, - causal=causal, - fixed_filter=fixed_filter, - ) - self.downsample = DownSample1d( - down_ratio, - down_kernel_size, - activation.in_features, - causal=causal, - fixed_filter=fixed_filter, - ) - - # x: [B,C,T] - def forward(self, x): - x = self.upsample(x) - x = self.act(x) - return self.downsample(x) - - -class Snake(nn.Module): - """ - Implementation of a sine-based periodic activation function - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter - """ - - def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): - super().__init__() - self.in_features = in_features - - # initialize alpha - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: # log scale alphas initialized to zeros - self.alpha = Parameter(torch.zeros(in_features) * alpha) - else: # linear scale alphas initialized to ones - self.alpha = Parameter(torch.ones(in_features) * alpha) - - self.alpha.requires_grad = alpha_trainable - - self.no_div_by_zero = 0.000000001 - - def forward(self, x): - # Snake := x + 1/a * sin^2 (xa) - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] - if self.alpha_logscale: - alpha = torch.exp(alpha) - return x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) - - -class SnakeBeta(nn.Module): - """ - A modified Snake function which uses separate parameters for the magnitude - of the periodic components. - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter that controls frequency - - beta - trainable parameter that controls magnitude - """ - - def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): - super().__init__() - self.in_features = in_features - - # initialize alpha - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: # log scale alphas initialized to zeros - self.alpha = Parameter(torch.zeros(in_features) * alpha) - self.beta = Parameter(torch.zeros(in_features) * alpha) - else: # linear scale alphas initialized to ones - self.alpha = Parameter(torch.ones(in_features) * alpha) - self.beta = Parameter(torch.ones(in_features) * alpha) - - self.alpha.requires_grad = alpha_trainable - self.beta.requires_grad = alpha_trainable - - self.no_div_by_zero = 0.000000001 - - def forward(self, x): - # SnakeBeta := x + 1/b * sin^2 (xa) - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] - beta = self.beta.unsqueeze(0).unsqueeze(-1) - if self.alpha_logscale: - alpha = torch.exp(alpha) - beta = torch.exp(beta) - return x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) - - -# ============================================================================ -# BigVGAN AudioVAE (vendored from modules/vocoder/bigvgan.py) -# ---------------------------------------------------------------------------- -# Copyright (c) 2022 NVIDIA CORPORATION. -# Licensed under the MIT license. -# Upstream layout flattened: the only modification is that the three -# original ``from dots_tts.modules.{backbone,vocoder}.* import ...`` lines -# are dropped, since those names are already defined above in this file. -# ============================================================================ - - -@dataclass(slots=True) -class BigVGANStreamState: - lstm_hidden: tuple[torch.Tensor, torch.Tensor] - decoder: DecoderStreamState - - -@dataclass(slots=True) -class DecoderStreamState: - window: torch.Tensor - chunk_size: int - total_frames: int = 0 - emitted_frames: int = 0 - - -def _empty_chunk( - ref: torch.Tensor, - *, - channels: int | None = None, - dtype: torch.dtype | None = None, -) -> torch.Tensor: - return ref.new_zeros( - (ref.size(0), channels or ref.size(1), 0), - dtype=dtype or ref.dtype, - ) - - -def _module_state_device_dtype(module: nn.Module) -> tuple[torch.device, torch.dtype]: - for name in ("weight", "bias", "filter"): - tensor = getattr(module, name, None) - if tensor is not None: - return tensor.device, tensor.dtype - for tensor in itertools.chain(module.parameters(), module.buffers()): - return tensor.device, tensor.dtype - raise RuntimeError(f"Unable to infer state dtype/device for {type(module).__name__}.") - - -def _stream_state_zeros( - batch_size: int, - channels: int, - length: int, - *, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor: - return torch.zeros( - (batch_size, channels, max(0, int(length))), - device=device, - dtype=dtype, - ) - - -def init_weights(m, mean=0.0, std=0.01): - classname = m.__class__.__name__ - if classname.find("Conv") != -1: - m.weight.data.normal_(mean, std) - - -class Conv1d_S(nn.Module): - "Conv1d for spectral normalisation and orthogonal initialisation" - - def __init__( - self, - in_channels, - out_channels, - kernel_size=1, - stride=1, - dilation=1, - groups=1, - causal=False, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.dilation = dilation - self.groups = groups - self.causal = causal - pad = 0 if causal else dilation * (kernel_size - 1) // 2 - self.causal_pad = dilation * (kernel_size - 1) if causal else 0 - - self.layer = weight_norm( - nn.Conv1d( - in_channels, - out_channels, - kernel_size=kernel_size, - stride=stride, - padding=pad, - dilation=dilation, - groups=groups, - ) - ) - - def forward(self, inputs): - if self.causal and self.causal_pad > 0: - inputs = F.pad(inputs, (self.causal_pad, 0)) - return self.layer(inputs) - - -class SLSTM(nn.Module): - """ - LSTM without worrying about the hidden state, nor the layout of the data. - Expects input as convolutional layout. - """ - - def __init__( - self, - dimension: int, - num_layers: int = 2, - skip: bool = True, - bidirectional: bool = False, - ): - super().__init__() - self.skip = skip - self.bidirectional = bidirectional - self.lstm = nn.LSTM( - input_size=dimension, - hidden_size=dimension, - num_layers=num_layers, - bidirectional=bidirectional, - batch_first=True, - ) - self._stream_num_layers = num_layers - self._stream_weight_ih = tuple(getattr(self.lstm, f"weight_ih_l{layer_idx}") for layer_idx in range(num_layers)) - self._stream_weight_hh = tuple(getattr(self.lstm, f"weight_hh_l{layer_idx}") for layer_idx in range(num_layers)) - self._stream_bias_ih = tuple(getattr(self.lstm, f"bias_ih_l{layer_idx}") for layer_idx in range(num_layers)) - self._stream_bias_hh = tuple(getattr(self.lstm, f"bias_hh_l{layer_idx}") for layer_idx in range(num_layers)) - if self.bidirectional: - self.proj_out = nn.Linear(dimension * 2, dimension) - - def forward(self, x): - y, _ = self.lstm(x) - if self.bidirectional: - y = self.proj_out(y) - if self.skip: - y = y + x - return y - - def stream_step( - self, - x: torch.Tensor, - hidden: tuple[torch.Tensor, torch.Tensor] | None, - ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: - if self.bidirectional: - raise RuntimeError("Streaming only supports unidirectional SLSTM.") - - residual = x - if hidden is None: - hidden = self.init_stream_state(x.size(0)) - - hidden_h, hidden_c = hidden - next_hidden_h = [] - next_hidden_c = [] - for layer_idx in range(self._stream_num_layers): - layer_input = x - hx = hidden_h[layer_idx] - cx = hidden_c[layer_idx] - outputs = [] - weight_ih = self._stream_weight_ih[layer_idx] - weight_hh = self._stream_weight_hh[layer_idx] - bias_ih = self._stream_bias_ih[layer_idx] - bias_hh = self._stream_bias_hh[layer_idx] - - for frame_idx in range(x.size(1)): - gates = F.linear(layer_input[:, frame_idx, :], weight_ih, bias_ih) - gates = gates + F.linear(hx, weight_hh, bias_hh) - input_gate, forget_gate, cell_gate, output_gate = gates.chunk(4, dim=-1) - input_gate = torch.sigmoid(input_gate) - forget_gate = torch.sigmoid(forget_gate) - cell_gate = torch.tanh(cell_gate) - output_gate = torch.sigmoid(output_gate) - cx = forget_gate * cx + input_gate * cell_gate - hx = output_gate * torch.tanh(cx) - outputs.append(hx) - - x = torch.stack(outputs, dim=1) - next_hidden_h.append(hx) - next_hidden_c.append(cx) - - y = x - if self.skip: - y = y + residual - return y, (torch.stack(next_hidden_h, dim=0), torch.stack(next_hidden_c, dim=0)) - - def init_stream_state( - self, - batch_size: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - num_directions = 2 if self.bidirectional else 1 - state_shape = ( - self.lstm.num_layers * num_directions, - batch_size, - self.lstm.hidden_size, - ) - weight = self.lstm.weight_hh_l0 - return ( - weight.new_zeros(state_shape), - weight.new_zeros(state_shape), - ) - - -class ResStack(nn.Module): - def __init__(self, channel, kernel_size=3, base=3, nums=4, causal=False): - super().__init__() - - self.layers = nn.ModuleList([]) - for i in range(nums): - dil = base**i - pad1 = dil * (kernel_size - 1) if causal else dil - pad2 = (kernel_size - 1) if causal else 1 - block = [ - nn.LeakyReLU(), - ] - if causal and pad1 > 0: - block.append(nn.ConstantPad1d((pad1, 0), 0.0)) - block.append( - nn.utils.weight_norm( - nn.Conv1d( - channel, - channel, - kernel_size=kernel_size, - dilation=dil, - padding=0 if causal else pad1, - ) - ) - ) - block.append(nn.LeakyReLU()) - if causal and pad2 > 0: - block.append(nn.ConstantPad1d((pad2, 0), 0.0)) - block.append( - nn.utils.weight_norm( - nn.Conv1d( - channel, - channel, - kernel_size=kernel_size, - dilation=1, - padding=0 if causal else pad2, - ) - ) - ) - self.layers.append(nn.Sequential(*block)) - - def forward(self, x): - for layer in self.layers: - x = x + layer(x) - return x - - -class Encoder(nn.Module): - def __init__( - self, - in_channels=1, - out_channels=100, - base_channels=12, - proj_kernel_size=3, - stack_kernel_size=3, - stack_dilation_base=2, - stacks=6, - channels=(12, 24, 48, 96, 192, 384, 768), - down_sample_factors=(2, 2, 2, 2, 4, 4), - causal=False, - lookahead=0, - ): - super().__init__() - - act_slope = 0.2 - layers = [] - # pre proj_layer - layers += [ - Conv1d_S( - in_channels, - base_channels, - kernel_size=proj_kernel_size, - stride=1, - causal=causal, - ), - nn.LeakyReLU(act_slope, True), - ] - - # channels: [512, 256, 128, 64], upsample_factors: [5, 2, 2] - for (in_c, out_c), down_f in zip(itertools.pairwise(channels), down_sample_factors, strict=True): - layers += [ - Conv1d_S(in_c, out_c, kernel_size=down_f * 2, stride=down_f, causal=causal), - ResStack(out_c, stack_kernel_size, stack_dilation_base, stacks, causal=causal), - nn.LeakyReLU(act_slope, True), - ] - - # post layers - if lookahead > 0: - layers += [ - Conv1d_S( - channels[-1], - out_channels, - kernel_size=lookahead * 2 + 1, - stride=1, - causal=False, - ), - ] - else: - layers += [ - Conv1d_S( - channels[-1], - out_channels, - kernel_size=proj_kernel_size, - stride=1, - causal=causal, - ), - ] - self.generator = nn.Sequential(*layers) - - def forward(self, conditions, _z_inputs=None): - return self.generator(conditions) - - -class AMPBlock1(torch.nn.Module): - def __init__( - self, - h, - channels, - kernel_size=3, - dilation=(1, 3, 5), - activation=None, - causal=True, - ): - super().__init__() - self.h = h - - self.convs1 = nn.ModuleList( - [ - weight_norm( - Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[0], - causal=causal, - ) - ), - weight_norm( - Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[1], - causal=causal, - ) - ), - weight_norm( - Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[2], - causal=causal, - ) - ), - ] - ) - self.convs1.apply(init_weights) - - self.convs2 = nn.ModuleList( - [ - weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), - weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), - weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), - ] - ) - self.convs2.apply(init_weights) - - self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers - - if activation == "snake": # periodic nonlinearity with snake function and anti-aliasing - self.activations = nn.ModuleList( - [ - Activation1d( - activation=Snake(channels, alpha_logscale=h.snake_logscale), - causal=causal, - fixed_filter=True, - ) - for _ in range(self.num_layers) - ] - ) - elif activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing - self.activations = nn.ModuleList( - [ - Activation1d( - activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale), - causal=causal, - fixed_filter=True, - ) - for _ in range(self.num_layers) - ] - ) - else: - raise NotImplementedError( - "activation incorrectly specified. check the config file and look for 'activation'." - ) - - def forward(self, x): - acts1, acts2 = self.activations[::2], self.activations[1::2] - for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True): - xt = a1(x) - xt = c1(xt) - xt = a2(xt) - xt = c2(xt) - x = xt + x - - return x - - def remove_weight_norm(self): - for layer in self.convs1: - remove_weight_norm(layer) - for layer in self.convs2: - remove_weight_norm(layer) - - -class AMPBlock2(torch.nn.Module): - def __init__(self, h, channels, kernel_size=3, dilation=(1, 3), activation=None, causal=True): - super().__init__() - self.h = h - - self.convs = nn.ModuleList( - [ - weight_norm( - Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[0], - causal=causal, - ) - ), - weight_norm( - Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[1], - causal=causal, - ) - ), - ] - ) - self.convs.apply(init_weights) - - self.num_layers = len(self.convs) # total number of conv layers - - if activation == "snake": # periodic nonlinearity with snake function and anti-aliasing - self.activations = nn.ModuleList( - [ - Activation1d( - activation=Snake(channels, alpha_logscale=h.snake_logscale), - causal=causal, - fixed_filter=True, - ) - for _ in range(self.num_layers) - ] - ) - elif activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing - self.activations = nn.ModuleList( - [ - Activation1d( - activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale), - causal=causal, - fixed_filter=True, - ) - for _ in range(self.num_layers) - ] - ) - else: - raise NotImplementedError( - "activation incorrectly specified. check the config file and look for 'activation'." - ) - - def forward(self, x): - for c, a in zip(self.convs, self.activations, strict=True): - xt = a(x) - xt = c(xt) - x = xt + x - - return x - - def remove_weight_norm(self): - for layer in self.convs: - remove_weight_norm(layer) - - -class Decoder(nn.Module): - def __init__(self, h): - super().__init__() - self.h = h - causal = h.causal - self._stream_window_sizes: dict[int, int] = {} - - self.num_kernels = len(h.resblock_kernel_sizes) - self.num_upsamples = len(h.upsample_rates) - - num_decoder_lookahead = h.get("num_decoder_lookahead", 2) - # pre conv - self.conv_pre = weight_norm( - Conv1d( - h.latent_dim, - h.upsample_initial_channel, - kernel_size=2 * num_decoder_lookahead + 1, - stride=1, - causal=False, - ) - ) - - # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default - resblock = AMPBlock1 if h.resblock == "1" else AMPBlock2 - - # transposed conv-based upsamplers. does not apply anti-aliasing - self.ups = nn.ModuleList() - for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes, strict=True)): - self.ups.append( - nn.ModuleList( - [ - weight_norm( - ConvTranspose1d( - h.upsample_initial_channel // (2**i), - h.upsample_initial_channel // (2 ** (i + 1)), - k, - u, - causal=causal, - ) - ) - ] - ) - ) - - # residual blocks using anti-aliased multi-periodicity composition modules (AMP) - self.resblocks = nn.ModuleList() - for i in range(len(self.ups)): - ch = h.upsample_initial_channel // (2 ** (i + 1)) - for k, d in zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes, strict=True): - self.resblocks.append(resblock(h, ch, k, d, activation=h.activation, causal=causal)) - - # post conv - if h.activation == "snake": # periodic nonlinearity with snake function and anti-aliasing - activation_post = Snake(ch, alpha_logscale=h.snake_logscale) - self.activation_post = Activation1d(activation=activation_post, causal=causal, fixed_filter=False) - elif h.activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing - activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale) - self.activation_post = Activation1d(activation=activation_post, causal=causal, fixed_filter=False) - else: - raise NotImplementedError( - "activation incorrectly specified. check the config file and look for 'activation'." - ) - - self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal, bias=h.get("use_bias_at_final", True))) - - # weight initialization - for i in range(len(self.ups)): - self.ups[i].apply(init_weights) - self.conv_post.apply(init_weights) - - def forward(self, z): - # pre conv - x = self.conv_pre(z) - - for i in range(self.num_upsamples): - # upsampling - for i_up in range(len(self.ups[i])): - x = self.ups[i][i_up](x) - # AMP blocks - xs = None - for j in range(self.num_kernels): - if xs is None: - xs = self.resblocks[i * self.num_kernels + j](x) - else: - xs += self.resblocks[i * self.num_kernels + j](x) - x = xs / self.num_kernels - - # post conv - x = self.activation_post(x) - x = self.conv_post(x) - if self.h.get("use_tanh_at_final", True): - x = torch.tanh(x) - else: - x = torch.clamp(x, min=-1.0, max=1.0) # Bound the output to [-1, 1] - - return x - - @property - def stream_lookahead(self) -> int: - return (self.conv_pre.kernel_size[0] - 1) // 2 - - @staticmethod - def _conv1d_left_context(layer) -> int: - dilation = layer.dilation[0] if isinstance(layer.dilation, tuple) else layer.dilation - kernel_size = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size - if getattr(layer, "causal", False): - return dilation * (kernel_size - 1) - return layer.padding[0] if isinstance(layer.padding, tuple) else layer.padding - - @staticmethod - def _convtranspose1d_left_context(layer) -> int: - stride = layer.stride[0] if isinstance(layer.stride, tuple) else layer.stride - kernel_size = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size - if not getattr(layer, "causal", False): - raise NotImplementedError("Streaming only supports causal ConvTranspose1d.") - if kernel_size != 2 * stride: - raise ValueError( - "Streaming ConvTranspose1d expects kernel_size == 2 * stride, got " - f"kernel_size={kernel_size} stride={stride}." - ) - return 1 - - @classmethod - def _activation_left_context(cls, activation: Activation1d) -> int: - upsample = activation.upsample - downsample = activation.downsample.lowpass - if not upsample.causal or not downsample.padding or downsample.pad_right != 0: - raise NotImplementedError("Streaming only supports causal alias-free activations.") - ratio = int(upsample.ratio) - if ratio != int(downsample.stride): - raise ValueError( - "Alias-free activation expects matched up/down ratios, got " - f"up_ratio={ratio} down_ratio={downsample.stride}." - ) - total_left = (upsample.kernel_size - 1) + (downsample.kernel_size - 1) - return (total_left + ratio - 1) // ratio - - @classmethod - def _ampblock_left_context(cls, block) -> int: - if isinstance(block, AMPBlock1): - left_context = 0 - acts1 = block.activations[::2] - acts2 = block.activations[1::2] - for conv1, conv2, act1, act2 in zip( - block.convs1, - block.convs2, - acts1, - acts2, - strict=True, - ): - left_context += ( - cls._activation_left_context(act1) - + cls._conv1d_left_context(conv1) - + cls._activation_left_context(act2) - + cls._conv1d_left_context(conv2) - ) - return left_context - if isinstance(block, AMPBlock2): - left_context = 0 - for conv, activation in zip(block.convs, block.activations, strict=True): - left_context += cls._activation_left_context(activation) + cls._conv1d_left_context(conv) - return left_context - raise TypeError(f"Unsupported resblock type: {type(block).__name__}.") - - def _stream_left_context(self) -> int: - left_context = Fraction(self._conv1d_left_context(self.conv_pre), 1) - current_scale = Fraction(1, 1) - for stage_idx, upsample_layers in enumerate(self.ups): - for upsample in upsample_layers: - left_context += current_scale * self._convtranspose1d_left_context(upsample) - stride = upsample.stride[0] if isinstance(upsample.stride, tuple) else upsample.stride - current_scale /= int(stride) - - stage_start = stage_idx * self.num_kernels - stage_end = stage_start + self.num_kernels - stage_context = max(self._ampblock_left_context(block) for block in self.resblocks[stage_start:stage_end]) - left_context += current_scale * stage_context - - left_context += current_scale * self._activation_left_context(self.activation_post) - left_context += current_scale * self._conv1d_left_context(self.conv_post) - return int(left_context.__ceil__()) - - def stream_window_size(self, chunk_size: int) -> int: - if chunk_size < 1: - raise ValueError(f"chunk_size must be >= 1, got {chunk_size}.") - cached = self._stream_window_sizes.get(chunk_size) - if cached is not None: - return cached - - window_size = chunk_size + self.stream_lookahead + self._stream_left_context() - self._stream_window_sizes[chunk_size] = window_size - return window_size - - def remove_weight_norm(self): - for upsample_layers in self.ups: - for upsample_layer in upsample_layers: - remove_weight_norm(upsample_layer) - for resblock in self.resblocks: - resblock.remove_weight_norm() - remove_weight_norm(self.conv_pre) - remove_weight_norm(self.conv_post) - - -class AudioVAE(nn.Module): - def __init__(self, h: AudioVAEConfig): - super().__init__() - self.config = h - - self.h = h - self.hop_size = int(np.prod(h.downsample_rates)) - self.sample_rate = h.sample_rate - self.decoder_lookahead = int(h.get("num_decoder_lookahead", 2)) - - self.audio_encoder = Encoder( - out_channels=h.latent_dim, - down_sample_factors=h.downsample_rates, - channels=h.downsample_channels, - causal=h.causal_encoder, - lookahead=h.get("num_encoder_lookahead", 2), - ) - - intermediate_size = h.latent_dim * 4 - self.enc_mi_layer = nn.Sequential( - nn.Linear(h.latent_dim, intermediate_size), - SLSTM(intermediate_size, num_layers=h.mi_num_layers), - nn.Linear(intermediate_size, h.latent_dim), - ) - self.dec_mi_layer = nn.Sequential( - nn.Linear(h.latent_dim, intermediate_size), - SLSTM(intermediate_size, num_layers=h.mi_num_layers), - nn.Linear(intermediate_size, h.latent_dim), - ) - self.pre_proj = Conv1d( - in_channels=h.latent_dim, - out_channels=h.latent_dim * 2, - kernel_size=1, - stride=1, - ) - self.post_proj = Conv1d(in_channels=h.latent_dim, out_channels=h.latent_dim, kernel_size=1, stride=1) - - self.decoder = Decoder(h) - - def inference(self, data): - latents = self.extract_latents(data["sample"]) - return {"sample": self.inference_from_latents(latents)} - - @torch.autocast(enabled=False, device_type="cuda") - def extract_latents(self, x, do_sample=False): - x = x.float() - x = self.audio_encoder(x) - x = x.permute(0, 2, 1) - x = self.enc_mi_layer(x) - x = x.permute(0, 2, 1) - x = self.pre_proj(x) - if do_sample: - m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1) - x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) - return x - - def inference_from_latents(self, x, do_sample=True, noise_scale=1.0): - if do_sample: - assert x.size(1) == self.h.latent_dim * 2, f"Input must be like [B, D, H], got {x.shape}" - m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1) - x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) * noise_scale - else: - assert x.size(1) == self.h.latent_dim, f"Input must be like [B, D, H], got {x.shape}" - x = self.post_proj(x) - x = x.permute(0, 2, 1) - x = self.dec_mi_layer(x) - x = x.permute(0, 2, 1) - return self.decoder(x) - - def _validate_stream_latents(self, latents: torch.Tensor) -> None: - if latents.ndim != 3: - raise ValueError( - f"Streaming latents must have shape [batch, latent_dim, frames], got {tuple(latents.shape)}." - ) - if latents.size(1) != self.h.latent_dim: - raise ValueError(f"Streaming latent_dim must be {self.h.latent_dim}, got {latents.size(1)}.") - - def init_stream_state( - self, - batch_size: int = 1, - chunk_size: int = 8, - ) -> BigVGANStreamState: - if not self.h.causal: - raise RuntimeError("Strict streaming requires a causal vocoder.") - if batch_size < 1: - raise ValueError(f"batch_size must be >= 1, got {batch_size}.") - window_size = self.decoder.stream_window_size(chunk_size) - device, dtype = _module_state_device_dtype(self.decoder.conv_pre) - return BigVGANStreamState( - lstm_hidden=self.dec_mi_layer[1].init_stream_state(batch_size), - decoder=DecoderStreamState( - window=_stream_state_zeros( - batch_size, - self.h.latent_dim, - window_size, - device=device, - dtype=dtype, - ), - chunk_size=int(chunk_size), - ), - ) - - def _decode_stream_latents( - self, - latents: torch.Tensor, - lstm_hidden: tuple[torch.Tensor, torch.Tensor], - ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: - self._validate_stream_latents(latents) - latents = latents.float() - x = self.post_proj(latents) - x = x.permute(0, 2, 1) - x = self.dec_mi_layer[0](x) - x, lstm_hidden = self.dec_mi_layer[1].stream_step(x, lstm_hidden) - x = self.dec_mi_layer[2](x) - decoder_dtype = self.decoder.conv_pre.weight.dtype - return x.permute(0, 2, 1).to(dtype=decoder_dtype), lstm_hidden - - def _prepare_stream_decoder_input( - self, - latents: torch.Tensor, - state: BigVGANStreamState, - ) -> torch.Tensor: - decoder_input, state.lstm_hidden = self._decode_stream_latents( - latents, - state.lstm_hidden, - ) - return decoder_input - - def _append_stream_decoder_input_tensor( - self, - decoder_input: torch.Tensor, - window: torch.Tensor, - valid_frames: torch.Tensor, - ) -> torch.Tensor: - if window.dtype != decoder_input.dtype: - window = window.to(dtype=decoder_input.dtype) - chunk_size = int(decoder_input.size(-1)) - if chunk_size >= window.size(-1): - raise ValueError(f"decoder window size {window.size(-1)} must be larger than chunk_size {chunk_size}.") - positions = torch.arange( - window.size(-1), - device=window.device, - dtype=valid_frames.dtype, - ) - clipped_valid = valid_frames.clamp(min=0, max=window.size(-1)) - combined = torch.cat( - [window, decoder_input.new_zeros(window.size(0), window.size(1), chunk_size)], - dim=-1, - ) - insert_index = clipped_valid + torch.arange( - chunk_size, - device=window.device, - dtype=valid_frames.dtype, - ) - combined.scatter_( - -1, - insert_index.view(1, 1, -1).expand_as(decoder_input), - decoder_input, - ) - new_valid = (clipped_valid + chunk_size).clamp(max=window.size(-1)) - start = (clipped_valid + chunk_size - window.size(-1)).clamp(min=0) - gather_index = (start + positions).clamp(max=combined.size(-1) - 1) - gathered = combined.gather( - -1, - gather_index.view(1, 1, -1).expand_as(window), - ) - mask = (positions < new_valid).to(dtype=window.dtype).view(1, 1, -1) - return gathered * mask - - def _append_stream_decoder_input( - self, - decoder_input: torch.Tensor, - state: BigVGANStreamState, - ) -> torch.Tensor: - decoder_state = state.decoder - chunk_size = int(decoder_input.size(-1)) - if chunk_size != decoder_state.chunk_size: - raise ValueError(f"Streaming chunk_size must stay fixed at {decoder_state.chunk_size}, got {chunk_size}.") - window = decoder_state.window - valid_frames = min(decoder_state.total_frames, window.size(-1)) - valid_frames_tensor = window.new_tensor(valid_frames, dtype=torch.int64) - new_window = self._append_stream_decoder_input_tensor( - decoder_input, - window, - valid_frames_tensor, - ) - decoder_state.window = new_window - decoder_state.total_frames += chunk_size - return new_window - - def compiled_stream_step( - self, - latents: torch.Tensor, - hidden_h: torch.Tensor, - hidden_c: torch.Tensor, - window: torch.Tensor, - valid_frames: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - decoder_input, (hidden_h, hidden_c) = self._decode_stream_latents( - latents, - (hidden_h, hidden_c), - ) - new_window = self._append_stream_decoder_input_tensor( - decoder_input, - window, - valid_frames, - ) - audio_window = self.decode_stream_window(new_window) - return audio_window, hidden_h, hidden_c, new_window - - def decode_stream_window(self, window: torch.Tensor) -> torch.Tensor: - decoder_dtype = self.decoder.conv_pre.weight.dtype - if window.dtype != decoder_dtype: - window = window.to(dtype=decoder_dtype) - return self.decoder(window) - - def _slice_stream_audio_window( - self, - audio_window: torch.Tensor, - state: BigVGANStreamState, - *, - final: bool, - ) -> torch.Tensor: - decoder_state = state.decoder - stable_end = ( - decoder_state.total_frames if final else max(0, decoder_state.total_frames - self.decoder.stream_lookahead) - ) - if stable_end <= decoder_state.emitted_frames: - return _empty_chunk(audio_window, channels=1) - - window_size = decoder_state.window.size(-1) - valid_frames = min(decoder_state.total_frames, window_size) - window_start = decoder_state.total_frames - valid_frames - if decoder_state.emitted_frames < window_start: - raise RuntimeError("Decoder stream window is too short for fixed-graph decoding.") - - local_start = decoder_state.emitted_frames - window_start - local_end = stable_end - window_start - sample_start = local_start * self.hop_size - sample_end = local_end * self.hop_size - decoder_state.emitted_frames = stable_end - return audio_window[..., sample_start:sample_end] - - def stream_step( - self, - latents: torch.Tensor, - state: BigVGANStreamState, - ) -> torch.Tensor: - decoder_input = self._prepare_stream_decoder_input(latents, state) - window = self._append_stream_decoder_input(decoder_input, state) - audio_window = self.decode_stream_window(window) - return self._slice_stream_audio_window(audio_window, state, final=False) - - def stream_flush(self, state: BigVGANStreamState) -> torch.Tensor: - audio_window = self.decode_stream_window(state.decoder.window) - return self._slice_stream_audio_window(audio_window, state, final=True) - - def remove_weight_norm(self): - self.decoder.remove_weight_norm() - - -__all__ = ["AudioVAEConfig", "AudioVAE"] diff --git a/vllm_omni/model_executor/models/dots_tts/pipeline.py b/vllm_omni/model_executor/models/dots_tts/pipeline.py deleted file mode 100644 index 22b50c629..000000000 --- a/vllm_omni/model_executor/models/dots_tts/pipeline.py +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""dots.tts pipeline topology (frozen). - -Single-stage AR TTS: text → 48 kHz waveform in one pass. -Architecture: Qwen2.5-1.5B base LM (vLLM-native) → DiTAR / MeanFlow head -→ AudioVAE.decode → BigVGAN → 48 kHz mono. - -Speaker conditioning via CAM++ x-vector injected in preprocess(). - -Modeled after voxcpm2's single-stage pipeline (same "vLLM-native base LM -+ side-path computation" pattern), but simplified: no residual LM, no FSQ. -""" - -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) - -DOTS_TTS_PIPELINE = PipelineConfig( - model_type="dots_tts", - model_arch="DotsTTSForConditionalGeneration", - default_deploy_config_name="dots_tts.yaml", - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="latent_generator", - execution_type=StageExecutionType.LLM_AR, - input_sources=(), - final_output=True, - final_output_type="audio", - owns_tokenizer=True, - engine_output_type="audio", - sampling_constraints={ - "detokenize": False, - # Stop signal slot in compute_logits (voxcpm2 convention — - # talker emits logits[i, 1] = 1.0 when prob_stop > 0.8, vLLM - # matches against this list to terminate the request). Not - # Qwen2.5 EOS (151643) — that ID would never come out of the - # 2-slot continue/stop softmax our compute_logits builds. - "stop_token_ids": [1], - }, - ), - ), -) diff --git a/vllm_omni/model_executor/models/dynin_omni/__init__.py b/vllm_omni/model_executor/models/dynin_omni/__init__.py deleted file mode 100644 index fb6738a29..000000000 --- a/vllm_omni/model_executor/models/dynin_omni/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# NOTE: Do not import model classes in this file. Importing any -# submodule in this package triggers __init__.py execution, and -# both the model registry and pipeline registry import submodules -# directly — heavy imports here would be loaded as a side effect -# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni.py deleted file mode 100644 index 9b8baac7d..000000000 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni.py +++ /dev/null @@ -1,747 +0,0 @@ -from __future__ import annotations - -from collections.abc import Iterable, Mapping, Sequence -from functools import cached_property -from importlib import import_module -from typing import Any - -import numpy as np -import torch -import torch.nn as nn -from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions -from vllm.inputs import MultiModalDataDict -from vllm.inputs import MultiModalInput as MultiModalInputs -from vllm.model_executor.models.interfaces import SupportsMultiModal -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import ( - MultiModalFieldConfig, - MultiModalKwargsItems, - PlaceholderRange, -) -from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser -from vllm.multimodal.processing import ( - BaseDummyInputsBuilder, - BaseMultiModalProcessor, - BaseProcessingInfo, - ProcessorInputs, - PromptUpdate, - TimingContext, -) -from vllm.sequence import IntermediateTensors -from vllm.v1.outputs import SamplerOutput -from vllm.v1.sample.metadata import SamplingMetadata -from vllm.v1.sample.sampler import Sampler - -from vllm_omni.model_executor.models.output_templates import OmniOutput - -from .dynin_omni_common import build_zero_input_embeddings - -try: - from PIL import Image as PILImage -except Exception: # pragma: no cover - PILImage = None - - -_MODALITY_ORDER = ("image", "video", "audio") - -_MODALITY_ALIASES = { - "img2img": "image", -} - -_MODALITY_INPUT_KEY_BY_NAME = { - "image": "pixel_values", - "video": "pixel_values_videos", - "audio": "input_audio_features", -} - -_MODALITY_PLACEHOLDER_BY_NAME = { - "image": "<|soi|><|image|><|eoi|>", - "video": "<|sov|><|video|><|eov|>", - "audio": "<|soa|><|audio|><|eoa|>", -} - -_MODALITY_INPUT_ALIASES = { - "image": ("pixel_values", "image_embeds", "img2img"), - "video": ("pixel_values_videos", "video_embeds"), - "audio": ("input_audio_features", "audio_embeds"), -} - - -def _normalize_modality_name(modality: str) -> str: - return _MODALITY_ALIASES.get(modality, modality) - - -def _get_modality_count(mm_counts: Mapping[str, int], modality: str) -> int: - canonical = _normalize_modality_name(modality) - count = mm_counts.get(canonical, 0) - for alias, target in _MODALITY_ALIASES.items(): - if target == canonical: - count += mm_counts.get(alias, 0) - return count - - -def _normalize_mm_data_aliases(mm_data: MultiModalDataDict) -> MultiModalDataDict: - normalized: dict[str, Any] = {} - for modality, value in mm_data.items(): - canonical = _normalize_modality_name(modality) - if canonical in normalized and normalized[canonical] is not None and value is not None: - raise ValueError( - "Dynin received duplicate multimodal inputs for " - f"{canonical!r} via {modality!r}. " - "Provide either the canonical modality or its alias, not both." - ) - if canonical not in normalized or normalized[canonical] is None: - normalized[canonical] = value - return normalized - - -def _get_placeholder_text(modality: str) -> str | None: - modality = _normalize_modality_name(modality) - for base_modality, placeholder in _MODALITY_PLACEHOLDER_BY_NAME.items(): - if modality.startswith(base_modality): - return placeholder - return None - - -class DyninOmniProcessingInfo(BaseProcessingInfo): - def get_data_parser(self) -> MultiModalDataParser: - return DyninOmniMultiModalDataParser( - expected_hidden_size=self._get_expected_hidden_size(), - ) - - def get_supported_mm_limits(self) -> Mapping[str, int | None]: - limits = {modality: 1 for modality in _MODALITY_ORDER} - for alias, target in _MODALITY_ALIASES.items(): - if target in limits: - limits[alias] = limits[target] - return limits - - def get_mm_max_tokens_per_item( - self, - seq_len: int, - mm_counts: Mapping[str, int], - ) -> Mapping[str, int] | None: - del seq_len, mm_counts - limits = {modality: 1 for modality in _MODALITY_ORDER} - for alias, target in _MODALITY_ALIASES.items(): - if target in limits: - limits[alias] = limits[target] - return limits - - -class DyninOmniDummyInputsBuilder(BaseDummyInputsBuilder[DyninOmniProcessingInfo]): - def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: - chunks: list[str] = [] - for modality in _MODALITY_ORDER: - placeholder = _get_placeholder_text(modality) - if placeholder is None: - continue - chunks.extend([placeholder] * _get_modality_count(mm_counts, modality)) - return " ".join(chunks) - - def get_dummy_mm_data( - self, - seq_len: int, - mm_counts: Mapping[str, int], - mm_options: Mapping[str, BaseDummyOptions] | None = None, - ) -> MultiModalDataDict: - del seq_len - - mm_data: dict[str, Any] = {} - - num_images = _get_modality_count(mm_counts, "image") - if num_images > 0: - mm_data["image"] = self._get_dummy_images( - width=224, - height=224, - num_images=num_images, - overrides=mm_options.get("image") if mm_options else None, - ) - - num_videos = _get_modality_count(mm_counts, "video") - if num_videos > 0: - mm_data["video"] = self._get_dummy_videos( - width=224, - height=224, - num_frames=8, - num_videos=num_videos, - overrides=mm_options.get("video") if mm_options else None, - ) - - num_audios = _get_modality_count(mm_counts, "audio") - if num_audios > 0: - mm_data["audio"] = self._get_dummy_audios( - length=16000, - num_audios=num_audios, - overrides=mm_options.get("audio") if mm_options else None, - ) - - return mm_data - - -class DyninOmniMultiModalDataParser(MultiModalDataParser): - def parse_mm_data(self, mm_data: MultiModalDataDict) -> MultiModalDataItems: - normalized = _normalize_mm_data_aliases(mm_data) - mm_items = super().parse_mm_data(normalized) - - for alias, canonical in _MODALITY_ALIASES.items(): - if alias in mm_data and canonical in mm_items and alias not in mm_items: - mm_items[alias] = mm_items[canonical] - - return mm_items - - def _get_audio_with_sr(self, audio: Any) -> tuple[np.ndarray, float | None]: - audio_array, orig_sr = super()._get_audio_with_sr(audio) - if self.audio_resampler.target_sr is None: - return audio_array, None - return audio_array, orig_sr - - -class DyninOmniMultiModalProcessor(BaseMultiModalProcessor[DyninOmniProcessingInfo]): - @staticmethod - def _find_subsequence( - haystack: list[int], - needle: list[int], - start: int, - ) -> int | None: - if not needle: - return None - - max_start = len(haystack) - len(needle) - if max_start < start: - return None - - for idx in range(start, max_start + 1): - if haystack[idx : idx + len(needle)] == needle: - return idx - return None - - @staticmethod - def _make_disabled_embed_mask(length: int) -> torch.Tensor: - return torch.zeros(length, dtype=torch.bool) - - @staticmethod - def _encode_prompt_to_token_ids( - prompt: str | list[int], - tokenizer: Any | None, - ) -> list[int]: - if isinstance(prompt, str): - if tokenizer is None: - raise ValueError("Tokenizer is required to process string prompts for Dynin multimodal inputs.") - return tokenizer.encode(prompt, add_special_tokens=False) - return list(prompt) - - @staticmethod - def _ensure_non_empty_prompt_ids( - prompt_token_ids: list[int], - tokenizer: Any | None, - ) -> list[int]: - if prompt_token_ids: - return prompt_token_ids - - fallback_id = None - if tokenizer is not None: - fallback_id = getattr(tokenizer, "bos_token_id", None) - if fallback_id is None: - fallback_id = getattr(tokenizer, "eos_token_id", None) - if fallback_id is None: - fallback_id = getattr(tokenizer, "pad_token_id", None) - - return [0 if fallback_id is None else int(fallback_id)] - - @classmethod - def _image_to_chw_float_tensor(cls, image: Any) -> torch.Tensor: - if isinstance(image, torch.Tensor): - tensor = image.detach() - elif isinstance(image, np.ndarray): - tensor = torch.from_numpy(image) - elif PILImage is not None and isinstance(image, PILImage.Image): - tensor = torch.from_numpy(np.asarray(image).copy()) - else: - raise TypeError(f"Unsupported image item type: {type(image)!r}") - - if tensor.ndim == 2: - tensor = tensor.unsqueeze(-1) - if tensor.ndim != 3: - raise ValueError(f"Expected 3D image tensor, got shape={tuple(tensor.shape)}") - - if tensor.shape[-1] in (1, 3, 4) and tensor.shape[0] not in (1, 3, 4): - tensor = tensor.permute(2, 0, 1) - - if tensor.shape[0] == 1: - tensor = tensor.repeat(3, 1, 1) - if tensor.shape[0] == 4: - tensor = tensor[:3] - - tensor = tensor.to(dtype=torch.float32) - if tensor.numel() > 0 and torch.max(tensor) > 1.0: - tensor = tensor / 255.0 - return tensor.contiguous() - - @classmethod - def _video_to_tchw_float_tensor(cls, video: Any) -> torch.Tensor: - if isinstance(video, (list, tuple)) and not isinstance(video, torch.Tensor): - frames = [cls._image_to_chw_float_tensor(frame) for frame in video] - if not frames: - return torch.zeros((1, 3, 1, 1), dtype=torch.float32) - return torch.stack(frames, dim=0).contiguous() - - if isinstance(video, torch.Tensor): - tensor = video.detach() - elif isinstance(video, np.ndarray): - tensor = torch.from_numpy(video) - else: - raise TypeError(f"Unsupported video item type: {type(video)!r}") - - if tensor.ndim == 3: - return cls._image_to_chw_float_tensor(tensor).unsqueeze(0).contiguous() - - if tensor.ndim != 4: - raise ValueError(f"Expected 4D video tensor, got shape={tuple(tensor.shape)}") - - if tensor.shape[-1] in (1, 3, 4) and tensor.shape[1] not in (1, 3, 4): - tensor = tensor.permute(0, 3, 1, 2) - - if tensor.shape[1] == 1: - tensor = tensor.repeat(1, 3, 1, 1) - if tensor.shape[1] == 4: - tensor = tensor[:, :3] - - tensor = tensor.to(dtype=torch.float32) - if tensor.numel() > 0 and torch.max(tensor) > 1.0: - tensor = tensor / 255.0 - return tensor.contiguous() - - @staticmethod - def _audio_to_float_tensor(audio: Any) -> torch.Tensor: - if isinstance(audio, tuple) and len(audio) == 2: - audio = audio[0] - - if isinstance(audio, torch.Tensor): - tensor = audio.detach() - elif isinstance(audio, np.ndarray): - tensor = torch.from_numpy(audio) - else: - tensor = torch.as_tensor(audio) - - tensor = tensor.to(dtype=torch.float32).contiguous().view(-1) - if tensor.numel() == 0: - return torch.zeros((16000,), dtype=torch.float32) - - max_abs = torch.max(torch.abs(tensor)) - if max_abs > 1.0: - tensor = tensor / max_abs - - return tensor.contiguous() - - @classmethod - def _convert_modality_item(cls, modality: str, item: Any) -> torch.Tensor: - if modality == "image": - return cls._image_to_chw_float_tensor(item) - if modality == "video": - return cls._video_to_tchw_float_tensor(item) - if modality == "audio": - return cls._audio_to_float_tensor(item) - raise ValueError(f"Unsupported modality for Dynin processor: {modality}") - - def _build_modality_kwargs( - self, - modality: str, - modality_items: Sequence[Any], - ) -> Sequence[Any]: - modality = _normalize_modality_name(modality) - input_key = _MODALITY_INPUT_KEY_BY_NAME[modality] - tensor_items = [self._convert_modality_item(modality, item) for item in modality_items] - mm_kwargs = MultiModalKwargsItems.from_hf_inputs( - {input_key: tensor_items}, - {input_key: MultiModalFieldConfig.batched(modality)}, - ) - return mm_kwargs[modality] - - def _build_placeholder_ranges( - self, - *, - modality: str, - item_count: int, - prompt_token_ids: list[int], - tokenizer: Any | None, - search_start: int, - ) -> tuple[list[PlaceholderRange], int]: - ranges: list[PlaceholderRange] = [] - - for _ in range(item_count): - placeholder_text = _get_placeholder_text(modality) - placeholder_token_ids: list[int] = [] - - if placeholder_text and tokenizer is not None: - placeholder_token_ids = tokenizer.encode( - placeholder_text, - add_special_tokens=False, - ) - - found_offset = None - if placeholder_token_ids: - found_offset = self._find_subsequence( - prompt_token_ids, - placeholder_token_ids, - search_start, - ) - - if found_offset is None: - found_offset = min(search_start, len(prompt_token_ids) - 1) - placeholder_len = 1 - else: - placeholder_len = len(placeholder_token_ids) - - ranges.append( - PlaceholderRange( - offset=found_offset, - length=placeholder_len, - is_embed=self._make_disabled_embed_mask(placeholder_len), - ) - ) - search_start = found_offset + placeholder_len - - return ranges, search_start - - def _get_mm_fields_config( - self, - hf_inputs: Any, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> Mapping[str, MultiModalFieldConfig]: - del hf_inputs, hf_processor_mm_kwargs - return {} - - def _get_prompt_updates( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - out_mm_kwargs: MultiModalKwargsItems, - ) -> Sequence[PromptUpdate]: - del mm_items, hf_processor_mm_kwargs, out_mm_kwargs - return [] - - def apply( - self, - inputs: ProcessorInputs, - timing_ctx: TimingContext, - ) -> MultiModalInputs: - prompt = inputs.prompt - mm_items = inputs.mm_data_items - - with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes( - self.info.model_id, - self.info.ctx.get_mm_config().mm_hasher_algorithm, - ) - - tokenizer = self.info.ctx.tokenizer - prompt_token_ids = self._encode_prompt_to_token_ids(prompt, tokenizer) - prompt_token_ids = self._ensure_non_empty_prompt_ids(prompt_token_ids, tokenizer) - - mm_kwargs_by_modality: dict[str, Sequence[Any]] = {} - mm_placeholders: dict[str, list[PlaceholderRange]] = {} - search_start = 0 - mm_counts = mm_items.get_all_counts() - - for modality in _MODALITY_ORDER: - item_count = mm_counts.get(modality, 0) - if item_count <= 0: - continue - - modality_items = mm_items[modality].get_all() - if len(modality_items) != item_count: - raise RuntimeError( - f"Parsed {len(modality_items)} items but expected {item_count} for modality={modality!r}" - ) - - mm_kwargs_by_modality[modality] = self._build_modality_kwargs( - modality, - modality_items, - ) - - placeholder_ranges, search_start = self._build_placeholder_ranges( - modality=modality, - item_count=item_count, - prompt_token_ids=prompt_token_ids, - tokenizer=tokenizer, - search_start=search_start, - ) - mm_placeholders[modality] = placeholder_ranges - - return MultiModalInputs( - type="multimodal", - prompt_token_ids=prompt_token_ids, - mm_kwargs=MultiModalKwargsItems(mm_kwargs_by_modality), - mm_hashes=mm_hashes, - mm_placeholders=mm_placeholders, - ) - - -class DyninOmniStageBase(nn.Module): - stage_name = "Dynin stage" - - def make_empty_intermediate_tensors( - self, - batch_size: int, - dtype: torch.dtype, - device: torch.device, - ) -> IntermediateTensors: - del batch_size, dtype, device - return IntermediateTensors({}) - - def embed_input_ids( - self, - input_ids: torch.Tensor, - multimodal_embeddings: Any = None, - is_multimodal: torch.Tensor | None = None, - **kwargs: Any, - ) -> torch.Tensor: - del multimodal_embeddings, is_multimodal, kwargs - return build_zero_input_embeddings( - input_ids=input_ids, - hidden_size=self.hidden_size, - stage_name=self.stage_name, - ) - - def load_weights( - self, - weights: Iterable[tuple[str, torch.Tensor]], - ) -> set[str]: - return {name for name, _ in weights} - - def compute_logits( - self, - hidden_states: torch.Tensor | OmniOutput, - sampling_metadata: Any = None, - ) -> torch.Tensor | None: - del hidden_states, sampling_metadata - return None - - -@MULTIMODAL_REGISTRY.register_processor( - DyninOmniMultiModalProcessor, - info=DyninOmniProcessingInfo, - dummy_inputs=DyninOmniDummyInputsBuilder, -) -class DyninOmniForConditionalGeneration(nn.Module, SupportsMultiModal): - supports_multimodal_raw_input_only = True - STAGE_ALIAS = { - "tokenizer": "token2text", - "token2token": "token2text", - "detok_text": "token2text", - "token2img": "token2image", - "token2wav": "token2audio", - "token2speech": "token2audio", - } - - STAGE_IMPL = { - "token2text": (".dynin_omni_token2text", "DyninOmniToken2Text"), - "token2image": (".dynin_omni_token2image", "DyninOmniToken2Image"), - "token2audio": (".dynin_omni_token2audio", "DyninOmniToken2Audio"), - } - - _STAGE_IMPL_CACHE: dict[str, type[nn.Module]] = {} - - @classmethod - def get_placeholder_str(cls, modality: str, i: int) -> str | None: - del i - return _get_placeholder_text(modality) - - @classmethod - def _resolve_stage_impl_class(cls, model_stage: str) -> type[nn.Module]: - impl = cls._STAGE_IMPL_CACHE.get(model_stage) - if impl is not None: - return impl - - module_name, class_name = cls.STAGE_IMPL[model_stage] - module = import_module(module_name, package=__package__) - impl = getattr(module, class_name) - cls._STAGE_IMPL_CACHE[model_stage] = impl - return impl - - @classmethod - def _normalize_stage_name(cls, raw_stage: str) -> str: - normalized = cls.STAGE_ALIAS.get(raw_stage, raw_stage) - if normalized not in cls.STAGE_IMPL: - raise ValueError( - "Unsupported DYNIN omni model_stage: " - f"{raw_stage} (normalized={normalized}). " - f"Supported: {sorted(cls.STAGE_IMPL.keys())}" - ) - return normalized - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - raw_stage = str(getattr(vllm_config.model_config, "model_stage", "token2text")).lower() - self.model_stage = self._normalize_stage_name(raw_stage) - - impl_cls = self._resolve_stage_impl_class(self.model_stage) - self.impl = impl_cls(vllm_config=vllm_config, prefix=prefix) - self.model = self.impl - - self.has_preprocess = False - self.has_postprocess = False - self.have_multimodal_outputs = getattr(self.impl, "have_multimodal_outputs", True) - self.requires_raw_input_tokens = getattr(self.impl, "requires_raw_input_tokens", True) - self.language_model = self._resolve_language_model() - - def _resolve_language_model(self) -> Any | None: - if hasattr(self.impl, "get_language_model"): - language_model = self.impl.get_language_model() - if language_model is not None: - return language_model - - if hasattr(self.impl, "language_model"): - language_model = getattr(self.impl, "language_model") - if language_model is not None: - return language_model - - if self.model_stage == "token2text": - return getattr(self.impl, "model", None) - - return None - - def get_language_model(self) -> Any | None: - return self.language_model - - @cached_property - def sampler(self): - if hasattr(self.model, "sampler"): - return self.model.sampler - if self.language_model is not None and hasattr(self.language_model, "sampler"): - return self.language_model.sampler - return Sampler() - - def init_multi_modal(self, thinker_config: Any = None) -> None: - if hasattr(self.model, "init_multi_modal"): - self.model.init_multi_modal(thinker_config) - - def _collect_multimodal_inputs(self, **kwargs: Any) -> dict[str, Any]: - mm_inputs: dict[str, Any] = {} - for modality, aliases in _MODALITY_INPUT_ALIASES.items(): - for alias in aliases: - if alias in kwargs and kwargs[alias] is not None: - mm_inputs[modality] = kwargs[alias] - break - return mm_inputs - - def _normalize_loaded_weight_names( - self, - loaded: set[str], - expected_param_names: set[str], - ) -> set[str]: - if self.model_stage != "token2text": - return loaded - - normalized_loaded: set[str] = set() - prefixes = ("", "impl.", "impl.model.") - - for name in loaded: - for prefix in prefixes: - candidate = f"{prefix}{name}" if prefix else name - if candidate in expected_param_names: - normalized_loaded.add(candidate) - break - - if len(normalized_loaded) < len(expected_param_names): - normalized_loaded.update(expected_param_names) - - return normalized_loaded - - def forward( - self, - input_ids: torch.Tensor | None = None, - positions: torch.Tensor | None = None, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, - ) -> OmniOutput: - return self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **kwargs, - ) - - def make_empty_intermediate_tensors( - self, - batch_size: int, - dtype: torch.dtype, - device: torch.device, - ) -> IntermediateTensors: - return self.model.make_empty_intermediate_tensors(batch_size, dtype, device) - - def embed_input_ids( - self, - input_ids: torch.Tensor, - multimodal_embeddings: Any = None, - is_multimodal: torch.Tensor | None = None, - **kwargs: Any, - ) -> torch.Tensor: - squeezed_batch = False - staged_input_ids = input_ids - - if input_ids.ndim == 0: - staged_input_ids = input_ids.view(1, 1) - squeezed_batch = True - elif input_ids.ndim == 1: - staged_input_ids = input_ids.unsqueeze(0) - squeezed_batch = True - - embeddings = self.model.embed_input_ids( - staged_input_ids, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - **kwargs, - ) - - if squeezed_batch and isinstance(embeddings, torch.Tensor): - if embeddings.ndim == 3 and embeddings.shape[0] == 1: - return embeddings.squeeze(0) - if embeddings.ndim == 2 and input_ids.ndim == 0 and embeddings.shape[0] == 1: - return embeddings - - return embeddings - - def embed_multimodal(self, **kwargs: Any) -> Any: - if hasattr(self.model, "embed_multimodal"): - return self.model.embed_multimodal(**kwargs) - - self._collect_multimodal_inputs(**kwargs) - return None - - def load_weights( - self, - weights: Iterable[tuple[str, torch.Tensor]], - ) -> set[str]: - loaded = self.model.load_weights(weights) - if loaded is None: - loaded = set() - - expected_param_names = {name for name, _ in self.named_parameters()} - if not expected_param_names: - return loaded - - return self._normalize_loaded_weight_names(loaded, expected_param_names) - - def compute_logits( - self, - hidden_states: torch.Tensor | OmniOutput, - sampling_metadata: Any = None, - ) -> torch.Tensor | None: - return self.model.compute_logits(hidden_states, sampling_metadata=sampling_metadata) - - def sample( - self, - logits: torch.Tensor, - sampling_metadata: SamplingMetadata, - ) -> SamplerOutput | None: - if hasattr(self.model, "sample"): - return self.model.sample(logits, sampling_metadata) - if self.language_model is not None and hasattr(self.language_model, "sample"): - return self.language_model.sample(logits, sampling_metadata) - return None diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py deleted file mode 100644 index 125822c1e..000000000 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py +++ /dev/null @@ -1,1242 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import hashlib -import importlib.util -import os -import sys -import threading -import types -from collections.abc import Iterable -from dataclasses import dataclass -from enum import IntEnum -from functools import lru_cache -from pathlib import Path -from typing import Any - -import torch -from vllm.config import VllmConfig -from vllm.logger import init_logger - -logger = init_logger(__name__) - -try: - from vllm_omni.transformers_utils.repo_utils import hf_api -except Exception: # pragma: no cover - hf_api = None - - -class DetokTarget(IntEnum): - TEXT = 0 - AUDIO = 1 - IMAGE = 2 - - -TASK_TO_DETOK = { - "mmu": DetokTarget.TEXT, - "s2t": DetokTarget.TEXT, - "mmu_fast": DetokTarget.TEXT, - "mmu_fastdllm_v1": DetokTarget.TEXT, - "v2t": DetokTarget.TEXT, - "t2s": DetokTarget.AUDIO, - "t2s_mmu_like": DetokTarget.AUDIO, - "t2s_fixed": DetokTarget.AUDIO, - "s2s": DetokTarget.AUDIO, - "v2s": DetokTarget.AUDIO, - "t2i": DetokTarget.IMAGE, - "i2i": DetokTarget.IMAGE, - "ti2ti": DetokTarget.IMAGE, -} - -DEFAULT_VQ_IMAGE_SOURCE = "snu-aidas/magvitv2" -DEFAULT_VQ_AUDIO_SOURCE = "snu-aidas/emova_speech_tokenizer_vllm" -DEFAULT_MAGVIT_REMOTE_CODE_REPO = "snu-aidas/magvitv2" -DEFAULT_DYNIN_REMOTE_CODE_REPO = "snu-aidas/Dynin-Omni" -DYNIN_PROMPT_SOURCE_KEY = "dynin_prompt_source" -DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT = "offline_prebuilt" - -DYNIN_TASK_DEFAULT_RUNTIME = { - "t2t": ("mmu", "mmu", 0, "text"), - "t2i": ("t2i", "t2i_gen", 2, "image"), - "t2s": ("t2s_mmu_like", "t2s_gen", 1, "audio"), - "i2i": ("i2i", "i2i", 2, "image"), -} - -DYNIN_TASK_RUNTIME_FALLBACKS: dict[str, dict[str, Any]] = { - "t2t": { - "prompt_max_text_len": 1024, - "max_new_tokens": 1024, - "steps": 1024, - "block_length": 16, - "temperature": 0.0, - "cfg_scale": 0.0, - }, - "t2i": { - "prompt_max_text_len": 128, - "image_token_count": 1024, - "mask_token_id": 126336, - "codebook_size": 8192, - "timesteps": 20, - "guidance_scale": 3.5, - "temperature": 1.0, - }, - "i2i": { - "prompt_max_text_len": 128, - "mask_token_id": 126336, - "codebook_size": 8192, - "timesteps": 64, - "guidance_scale": 3.5, - "temperature": 1.0, - "image_resolution": 336, - "use_train_i2i_prompt": True, - }, - "t2s": { - "runtime_task": "t2s_mmu_like", - "prompting_task": "t2s_gen", - "prompt_max_text_len": 1024, - "t2s_token_length": 512, - "mask_token_id": 126336, - "codebook_size": 8192, - "audio_codebook_size": 4096, - "steps": 512, - "block_length": 128, - "temperature": 1.0, - "cfg_scale": 2.5, - "t2s_condition": "gender-female_emotion-neutral_speed-normal_pitch-normal", - }, -} - -DEFAULT_DYNIN_T2S_INSTRUCTION = "Please read the following text naturally." - -DYNIN_SPECIAL_TOKENS = ( - "<|soi|>", - "<|eoi|>", - "<|sov|>", - "<|eov|>", - "<|t2i|>", - "<|mmu|>", - "<|t2v|>", - "<|v2v|>", - "<|lvg|>", - "<|i2i|>", - "<|ti2ti|>", - "<|v2t|>", - "<|v2s|>", - "<|s2t|>", - "<|t2s|>", - "<|s2s|>", - "<|soa|>", - "<|eoa|>", -) - -_DYNIN_ONLINE_PROMPT_TOKEN_BY_TASK = { - "t2i": "<|t2i|>", - "i2i": "<|i2i|>", - "t2s": "<|t2s|>", -} - -_DYNIN_MODALITY_PLACEHOLDERS = ( - "<|soi|><|image|><|eoi|>", - "<|sov|><|video|><|eov|>", - "<|soa|><|audio|><|eoa|>", -) - -_DYNIN_CONFIG_CANDIDATE_RELPATHS = ( - "configs/dynin_omni.yaml", - "models/configs/dynin_omni.yaml", - "vllm_omni/model_executor/models/dynin_omni/configs/dynin_omni.yaml", - "dynin_omni.yaml", -) - -_DYNIN_REMOTE_ALLOW_PATTERNS = ("*.py", "*.json", "*.yaml", "*.yml") - -_DYNIN_REMOTE_CACHE_LOCK = threading.Lock() -_DYNIN_REMOTE_PACKAGE_BY_SNAPSHOT: dict[str, str] = {} -_DYNIN_REMOTE_ATTR_CACHE: dict[tuple[str, str, str, str | None, bool], Any] = {} - - -@dataclass(frozen=True) -class DyninInferSources: - model_source: str - tokenizer_source: str - vq_image_source: str - vq_audio_source: str - model_local_files_only: bool - vq_image_local_files_only: bool - vq_audio_local_files_only: bool - config_path: str | None = None - - @property - def local_files_only(self) -> bool: - return self.model_local_files_only - - -@dataclass(frozen=True) -class RemoteCodeSettings: - default_repo: str - repo_env: str - revision_env: str - local_only_env: str - - -DYNIN_REMOTE_SETTINGS = RemoteCodeSettings( - default_repo=DEFAULT_DYNIN_REMOTE_CODE_REPO, - repo_env="DYNIN_REMOTE_CODE_REPO_ID", - revision_env="DYNIN_REMOTE_CODE_REVISION", - local_only_env="DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY", -) - -MAGVIT_REMOTE_SETTINGS = RemoteCodeSettings( - default_repo=DEFAULT_MAGVIT_REMOTE_CODE_REPO, - repo_env="DYNIN_MAGVIT_REMOTE_CODE_REPO_ID", - revision_env="DYNIN_MAGVIT_REMOTE_CODE_REVISION", - local_only_env="DYNIN_MAGVIT_REMOTE_CODE_LOCAL_FILES_ONLY", -) - - -def unwrap_first_value(value: Any, default: Any = None) -> Any: - if value is None: - return default - if isinstance(value, list): - return default if not value else value[0] - if isinstance(value, torch.Tensor): - if value.numel() == 0: - return default - if value.numel() == 1: - return value.item() - return value - return value - - -def normalize_runtime_info(runtime_additional_information: Any) -> dict[str, Any]: - if isinstance(runtime_additional_information, list): - if not runtime_additional_information: - return {} - first = runtime_additional_information[0] - return first if isinstance(first, dict) else {} - if isinstance(runtime_additional_information, dict): - return runtime_additional_information - return {} - - -def logical_dynin_task(task: Any) -> str: - task_text = str(unwrap_first_value(task, "") or "").strip().lower() - if task_text in ("t2s", "t2s_mmu_like", "t2s_fixed"): - return "t2s" - if task_text in ("t2i", "i2i"): - return task_text - return "t2t" - - -def dynin_runtime_fallback(task: str, key: str, value: Any = None) -> Any: - if isinstance(value, str): - if value.strip() != "": - return value - elif value is not None: - return value - return DYNIN_TASK_RUNTIME_FALLBACKS.get(task, {}).get(key) - - -def coerce_token_ids_1d( - value: Any, - ref_device: torch.device | None = None, -) -> torch.Tensor: - if isinstance(value, tuple): - value = value[0] - - if isinstance(value, list): - if not value: - device = ref_device or torch.device("cpu") - return torch.empty(0, dtype=torch.long, device=device) - if isinstance(value[0], torch.Tensor): - value = value[0] - else: - value = torch.tensor( - value[0] if isinstance(value[0], list) else value, - dtype=torch.long, - ) - - if not isinstance(value, torch.Tensor): - value = torch.tensor(value, dtype=torch.long) - - if value.ndim == 0: - value = value.unsqueeze(0) - if value.ndim > 1: - value = value[0] - - if ref_device is not None and value.device != ref_device: - value = value.to(ref_device) - - return value.to(dtype=torch.long).contiguous() - - -def _first_positive_int(value: Any) -> int | None: - if value is None: - return None - if isinstance(value, torch.Tensor): - if value.numel() != 1: - return None - value = value.item() - try: - value = int(value) - except (TypeError, ValueError): - return None - return value if value > 0 else None - - -def resolve_hidden_size( - *, - vllm_config: VllmConfig, - model: Any | None = None, - default: int = 1024, -) -> int: - if model is not None: - try: - embeddings = model.get_input_embeddings() - weight = getattr(embeddings, "weight", None) - if isinstance(weight, torch.Tensor) and weight.ndim >= 2: - hidden_size = _first_positive_int(weight.shape[-1]) - if hidden_size is not None: - return hidden_size - except Exception: - pass - - model_cfg = getattr(model, "config", None) - for key in ("hidden_size", "d_model", "n_embd", "dim", "model_dim", "embed_dim"): - hidden_size = _first_positive_int(getattr(model_cfg, key, None)) - if hidden_size is not None: - return hidden_size - - for config_obj in ( - getattr(vllm_config.model_config, "hf_config", None), - getattr(vllm_config.model_config, "hf_text_config", None), - ): - if config_obj is None: - continue - for key in ("hidden_size", "d_model", "n_embd", "dim", "model_dim", "embed_dim"): - value = config_obj.get(key) if isinstance(config_obj, dict) else getattr(config_obj, key, None) - hidden_size = _first_positive_int(value) - if hidden_size is not None: - return hidden_size - - return default - - -def build_zero_input_embeddings( - *, - input_ids: torch.Tensor, - hidden_size: int, - stage_name: str, - dtype: torch.dtype = torch.bfloat16, -) -> torch.Tensor: - if input_ids.ndim == 0: - shape = (1, hidden_size) - elif input_ids.ndim == 1: - shape = (input_ids.shape[0], hidden_size) - elif input_ids.ndim == 2: - shape = (input_ids.shape[0], input_ids.shape[1], hidden_size) - else: - raise ValueError(f"Unsupported input_ids rank for {stage_name}: {input_ids.ndim}") - return torch.zeros(shape, dtype=dtype, device=input_ids.device) - - -def _to_bool(value: Any, default: bool = False) -> bool: - if value is None: - return default - if isinstance(value, bool): - return value - if isinstance(value, (int, float)): - return bool(value) - - text = str(value).strip().lower() - if text in ("1", "true", "yes", "y", "on"): - return True - if text in ("0", "false", "no", "n", "off", "", "none", "null"): - return False - return default - - -def _runtime_value(runtime_info: dict[str, Any], key: str) -> Any: - return unwrap_first_value(runtime_info.get(key), None) - - -def _runtime_first_value(runtime_info: dict[str, Any], keys: tuple[str, ...]) -> Any: - for key in keys: - value = _runtime_value(runtime_info, key) - if value is not None: - return value - return None - - -def _node_value(node: Any, key: str, default: Any = None) -> Any: - if node is None: - return default - if isinstance(node, dict): - return node.get(key, default) - try: - return node.get(key, default) - except Exception: - return getattr(node, key, default) - - -def _looks_like_hf_repo_id(value: str | None) -> bool: - if not isinstance(value, str): - return False - if value.count("/") != 1: - return False - org, name = value.split("/", 1) - return bool(org and name) - - -def _find_dynin_config_under_root(root: Path) -> Path | None: - for rel_path in _DYNIN_CONFIG_CANDIDATE_RELPATHS: - candidate = root.expanduser() / rel_path - if candidate.exists(): - return candidate.resolve() - return None - - -@lru_cache(maxsize=16) -def _resolve_dynin_config_from_hf_repo(repo_id: str) -> str | None: - if not _looks_like_hf_repo_id(repo_id) or hf_api is None: - return None - - try: - snapshot_dir = ( - Path( - hf_api().snapshot_download( - repo_id=repo_id, - repo_type="model", - allow_patterns=list(_DYNIN_CONFIG_CANDIDATE_RELPATHS), - local_files_only=True, - ) - ) - .expanduser() - .resolve() - ) - except Exception: - return None - - found = _find_dynin_config_under_root(snapshot_dir) - return str(found) if found is not None else None - - -def _resolve_existing_path(path_like: Any, source_name: str) -> str | None: - if path_like is None: - return None - text = str(path_like).strip() - if not text: - return None - - path = Path(text).expanduser() - if path.is_file(): - return str(path.resolve()) - - logger.warning( - "DYNIN config path from %s does not exist: %s. Falling back to auto-discovery.", - source_name, - path, - ) - return None - - -def _resolve_config_path(vllm_config: VllmConfig, runtime_info: dict[str, Any]) -> str | None: - for value, name in ( - (_runtime_value(runtime_info, "dynin_config_path"), "runtime_info.dynin_config_path"), - (os.getenv("DYNIN_CONFIG_PATH"), "DYNIN_CONFIG_PATH"), - (getattr(vllm_config.model_config, "dynin_config_path", None), "vllm_config.model_config.dynin_config_path"), - ): - resolved = _resolve_existing_path(value, name) - if resolved: - return resolved - - model_source = str(getattr(vllm_config.model_config, "model", "") or "") - tokenizer_source = str(getattr(vllm_config.model_config, "tokenizer", "") or "") - hf_config = getattr(vllm_config.model_config, "hf_config", None) - hf_name_or_path = ( - hf_config.get("_name_or_path") if isinstance(hf_config, dict) else getattr(hf_config, "_name_or_path", None) - ) - - hf_repo_candidates: list[str] = [] - for source in (model_source, tokenizer_source, hf_name_or_path): - if not _looks_like_hf_repo_id(source): - continue - source = str(source) - if source not in hf_repo_candidates: - hf_repo_candidates.append(source) - - for source in hf_repo_candidates: - resolved = _resolve_dynin_config_from_hf_repo(source) - if resolved is not None: - logger.info("Resolved dynin config from Hugging Face cache for %s: %s", source, resolved) - return resolved - - for source in (model_source, tokenizer_source): - source_path = Path(source).expanduser() - if source_path.is_dir(): - found = _find_dynin_config_under_root(source_path) - if found is not None: - return str(found) - - module_root = Path(__file__).resolve().parent - for bundled in ( - module_root / "configs" / "dynin_omni.yaml", - module_root / "models" / "configs" / "dynin_omni.yaml", - ): - if bundled.exists(): - return str(bundled) - - return None - - -@lru_cache(maxsize=16) -def _load_omega_config(config_path: str) -> Any: - try: - from omegaconf import OmegaConf - except ImportError as e: - raise ImportError( - f"omegaconf is required to load Dynin config files. Install it to read config: {config_path}" - ) from e - return OmegaConf.load(config_path) - - -def resolve_dynin_infer_sources( - *, - vllm_config: VllmConfig, - runtime_info: dict[str, Any] | None = None, -) -> DyninInferSources: - runtime_info = runtime_info or {} - - base_model_source = str(getattr(vllm_config.model_config, "model", "")) - base_model_path = Path(base_model_source).expanduser() - local_vllm_model_source = str(base_model_path) if base_model_path.is_dir() else None - - model_source = base_model_source - tokenizer_source = model_source - vq_image_source = DEFAULT_VQ_IMAGE_SOURCE - vq_audio_source = DEFAULT_VQ_AUDIO_SOURCE - model_local_files_only = False - vq_image_local_files_only = False - vq_audio_local_files_only = False - - resolver_source: str | None = base_model_source if base_model_source else None - resolver_local_files_only: bool | None = True if base_model_path.is_dir() else None - resolve_model_pretrained_source_fn = get_dynin_config_resolver_attr( - "resolve_model_pretrained_source", - source=resolver_source, - local_files_only=resolver_local_files_only, - ) - resolve_tokenizer_source_fn = get_dynin_config_resolver_attr( - "resolve_tokenizer_source", - source=resolver_source, - local_files_only=resolver_local_files_only, - ) - resolve_model_local_files_only_fn = get_dynin_config_resolver_attr( - "resolve_model_local_files_only", - source=resolver_source, - local_files_only=resolver_local_files_only, - ) - resolve_vq_cfg_block_fn = get_dynin_config_resolver_attr( - "resolve_vq_cfg_block", - source=resolver_source, - local_files_only=resolver_local_files_only, - ) - resolve_vq_repo_source_fn = get_dynin_config_resolver_attr( - "resolve_vq_repo_source", - source=resolver_source, - local_files_only=resolver_local_files_only, - ) - - config_path = _resolve_config_path(vllm_config, runtime_info) - if config_path: - config_file = Path(config_path).expanduser() - if config_file.exists(): - try: - dynin_cfg = _load_omega_config(str(config_file)) - model_source = resolve_model_pretrained_source_fn( - dynin_cfg, - default=model_source, - ) - tokenizer_source = resolve_tokenizer_source_fn( - dynin_cfg, - default=tokenizer_source, - ) - model_local_files_only = resolve_model_local_files_only_fn( - dynin_cfg, - default=model_local_files_only, - ) - vq_image_cfg = resolve_vq_cfg_block_fn(dynin_cfg, modality="image") - vq_audio_cfg = resolve_vq_cfg_block_fn(dynin_cfg, modality="audio") - vq_image_source = resolve_vq_repo_source_fn( - vq_image_cfg, - default=vq_image_source, - ) - vq_audio_source = resolve_vq_repo_source_fn( - vq_audio_cfg, - default=vq_audio_source, - ) - vq_image_local_files_only = _to_bool( - _node_value(vq_image_cfg, "local_files_only", None), - default=model_local_files_only, - ) - vq_audio_local_files_only = _to_bool( - _node_value(vq_audio_cfg, "local_files_only", None), - default=model_local_files_only, - ) - except Exception as e: - logger.warning( - "Failed to resolve DYNIN inference config from %s: %s", - config_file, - e, - ) - else: - logger.warning("DYNIN config path does not exist: %s", config_file) - - runtime_model_source = _runtime_value(runtime_info, "dynin_model_path") - if runtime_model_source: - model_source = str(runtime_model_source) - - runtime_tokenizer_source = _runtime_value(runtime_info, "tokenizer_path") - if runtime_tokenizer_source: - tokenizer_source = str(runtime_tokenizer_source) - - runtime_vq_image_source = _runtime_value(runtime_info, "vq_model_image_path") - if runtime_vq_image_source is None: - runtime_vq_image_source = _runtime_value(runtime_info, "vq_model_path_image") - if runtime_vq_image_source: - vq_image_source = str(runtime_vq_image_source) - - runtime_vq_audio_source = _runtime_value(runtime_info, "vq_model_audio_path") - if runtime_vq_audio_source is None: - runtime_vq_audio_source = _runtime_value(runtime_info, "vq_model_path_audio") - if runtime_vq_audio_source: - vq_audio_source = str(runtime_vq_audio_source) - - runtime_local_global = _runtime_value(runtime_info, "local_files_only") - runtime_local_model = _runtime_first_value( - runtime_info, - ("model_local_files_only", "local_files_only_model"), - ) - runtime_local_vq_image = _runtime_first_value( - runtime_info, - ("vq_model_image_local_files_only", "local_files_only_vq_image"), - ) - runtime_local_vq_audio = _runtime_first_value( - runtime_info, - ("vq_model_audio_local_files_only", "local_files_only_vq_audio"), - ) - - if runtime_local_global is not None: - global_local = _to_bool(runtime_local_global, default=False) - if runtime_local_model is None: - model_local_files_only = global_local - if runtime_local_vq_image is None: - vq_image_local_files_only = global_local - if runtime_local_vq_audio is None: - vq_audio_local_files_only = global_local - - if runtime_local_model is not None: - model_local_files_only = _to_bool( - runtime_local_model, - default=model_local_files_only, - ) - if runtime_local_vq_image is not None: - vq_image_local_files_only = _to_bool( - runtime_local_vq_image, - default=vq_image_local_files_only, - ) - if runtime_local_vq_audio is not None: - vq_audio_local_files_only = _to_bool( - runtime_local_vq_audio, - default=vq_audio_local_files_only, - ) - - if runtime_local_global is None and runtime_local_model is None and local_vllm_model_source is not None: - model_local_files_only = True - - if local_vllm_model_source is not None: - if not runtime_model_source: - if model_source != local_vllm_model_source: - logger.info( - "DYNIN infer model source overridden to local vLLM model path: %s (from %s)", - local_vllm_model_source, - model_source, - ) - model_source = local_vllm_model_source - if not runtime_tokenizer_source: - tokenizer_source = local_vllm_model_source - - return DyninInferSources( - model_source=model_source, - tokenizer_source=tokenizer_source, - vq_image_source=vq_image_source, - vq_audio_source=vq_audio_source, - model_local_files_only=model_local_files_only, - vq_image_local_files_only=vq_image_local_files_only, - vq_audio_local_files_only=vq_audio_local_files_only, - config_path=config_path, - ) - - -def _resolve_remote_source(source: str | None, settings: RemoteCodeSettings) -> str: - if isinstance(source, str): - stripped = source.strip() - if stripped: - source_path = Path(stripped).expanduser() - if source_path.is_dir(): - return str(source_path.resolve()) - if _looks_like_hf_repo_id(stripped): - return stripped - - env_repo = os.getenv(settings.repo_env) - if _looks_like_hf_repo_id(env_repo): - return str(env_repo).strip() - - return settings.default_repo - - -def _resolve_remote_revision(revision: str | None, settings: RemoteCodeSettings) -> str | None: - if isinstance(revision, str) and revision.strip(): - return revision.strip() - env_revision = os.getenv(settings.revision_env) - if isinstance(env_revision, str) and env_revision.strip(): - return env_revision.strip() - return None - - -def _resolve_remote_local_only(local_files_only: bool | None, settings: RemoteCodeSettings) -> bool: - if local_files_only is not None: - return bool(local_files_only) - return _to_bool(os.getenv(settings.local_only_env), default=False) - - -def _resolve_remote_snapshot_dir( - *, - source: str, - revision: str | None, - local_files_only: bool, -) -> str: - source_path = Path(source).expanduser() - if source_path.is_dir(): - return str(source_path.resolve()) - - if hf_api is None: - raise RuntimeError("huggingface_hub is required to load remote code.") - - kwargs: dict[str, Any] = { - "repo_id": source, - "repo_type": "model", - "allow_patterns": list(_DYNIN_REMOTE_ALLOW_PATTERNS), - "local_files_only": bool(local_files_only), - } - if revision is not None: - kwargs["revision"] = revision - - try: - return str(hf_api().snapshot_download(**kwargs)) - except TypeError: - kwargs.pop("local_files_only", None) - return str(hf_api().snapshot_download(**kwargs)) - - -def _ensure_remote_package(snapshot_dir: str) -> str: - with _DYNIN_REMOTE_CACHE_LOCK: - existing = _DYNIN_REMOTE_PACKAGE_BY_SNAPSHOT.get(snapshot_dir) - if existing is not None: - return existing - - digest = hashlib.sha1(snapshot_dir.encode("utf-8")).hexdigest()[:12] - package_name = f"_dynin_hf_remote_{digest}" - - package = types.ModuleType(package_name) - package.__path__ = [snapshot_dir] # type: ignore[attr-defined] - package.__file__ = str(Path(snapshot_dir) / "__init__.py") - - sys.modules.setdefault(package_name, package) - _DYNIN_REMOTE_PACKAGE_BY_SNAPSHOT[snapshot_dir] = package_name - return package_name - - -def _load_remote_module( - *, - module_name: str, - source: str, - revision: str | None, - local_files_only: bool, -): - snapshot_dir = _resolve_remote_snapshot_dir( - source=source, - revision=revision, - local_files_only=local_files_only, - ) - - module_path = Path(snapshot_dir) / f"{module_name}.py" - if not module_path.is_file(): - raise ImportError(f"Remote code module '{module_name}.py' not found under '{snapshot_dir}'. source={source!r}") - - package_name = _ensure_remote_package(snapshot_dir) - full_name = f"{package_name}.{module_name}" - - existing = sys.modules.get(full_name) - if existing is not None: - return existing - - spec = importlib.util.spec_from_file_location(full_name, module_path) - if spec is None or spec.loader is None: - raise ImportError(f"Failed to create import spec for '{module_path}'.") - - module = importlib.util.module_from_spec(spec) - module.__package__ = package_name - sys.modules[full_name] = module - try: - spec.loader.exec_module(module) - except Exception: - sys.modules.pop(full_name, None) - raise - return module - - -def resolve_remote_attr( - attr_name: str, - *, - module_name: str, - settings: RemoteCodeSettings, - source: str | None = None, - revision: str | None = None, - local_files_only: bool | None = None, - fallback_module_names: Iterable[str] = (), - optional: bool = False, -) -> Any | None: - resolved_source = _resolve_remote_source(source, settings) - resolved_revision = _resolve_remote_revision(revision, settings) - resolved_local_only = _resolve_remote_local_only(local_files_only, settings) - - module_candidates = [module_name, *[m for m in fallback_module_names if m and m != module_name]] - last_error: Exception | None = None - - for candidate in module_candidates: - cache_key = (attr_name, candidate, resolved_source, resolved_revision, resolved_local_only) - cached = _DYNIN_REMOTE_ATTR_CACHE.get(cache_key) - if cached is not None: - return cached - - try: - module = _load_remote_module( - module_name=candidate, - source=resolved_source, - revision=resolved_revision, - local_files_only=resolved_local_only, - ) - if hasattr(module, attr_name): - value = getattr(module, attr_name) - _DYNIN_REMOTE_ATTR_CACHE[cache_key] = value - return value - except Exception as e: - last_error = e - - if optional: - if last_error is not None: - logger.debug( - "Optional remote attr not found: attr=%s source=%s revision=%s err=%s", - attr_name, - resolved_source, - resolved_revision, - last_error, - ) - return None - - raise ImportError( - f"Failed to resolve '{attr_name}' from remote code " - f"(source={resolved_source!r}, revision={resolved_revision!r}, modules={module_candidates})." - ) from last_error - - -_DYNIN_MODELING_REMOTE_EXPORTS = { - "DyninOmniConfig": "DyninOmniConfig", - "DyninOmniModelLM": "DyninOmniModelLM", - "VideoTokenMerger": "VideoTokenMerger", -} - -_DYNIN_SAMPLING_REMOTE_EXPORTS = { - "log": "log", - "gumbel_noise": "gumbel_noise", - "gumbel_sample": "gumbel_sample", - "top_k": "top_k", - "mask_by_random_topk": "mask_by_random_topk", - "cosine_schedule": "cosine_schedule", - "linear_schedule": "linear_schedule", - "pow": "pow", - "sigmoid_schedule": "sigmoid_schedule", - "get_mask_schedule": "get_mask_schedule", - "top_k_top_p_filtering": "top_k_top_p_filtering", -} - -_DYNIN_CONFIG_RESOLVER_REMOTE_EXPORTS = { - "resolve_model_pretrained_source": "resolve_model_pretrained_source", - "resolve_tokenizer_source": "resolve_tokenizer_source", - "resolve_model_local_files_only": "resolve_model_local_files_only", - "resolve_vq_cfg_block": "resolve_vq_cfg_block", - "resolve_vq_repo_source": "resolve_vq_repo_source", -} - -_DYNIN_MAGVIT_REMOTE_EXPORTS = { - "VQGANEncoder": "VQGANEncoder", - "VQGANDecoder": "VQGANDecoder", - "LFQuantizer": "LFQuantizer", - "MAGVITv2": "MAGVITv2", -} - - -def _get_export_attr( - name: str, - export_map: dict[str, str], - *, - module_name: str, - settings: RemoteCodeSettings, - source: str | None = None, - revision: str | None = None, - local_files_only: bool | None = None, - optional: bool = False, -) -> Any | None: - attr_name = export_map.get(name) - if attr_name is None: - raise AttributeError(f"Unsupported export: {name!r}") - - return resolve_remote_attr( - attr_name, - module_name=module_name, - settings=settings, - source=source, - revision=revision, - local_files_only=local_files_only, - optional=optional, - ) - - -def get_dynin_modeling_attr(name: str) -> Any: - return _get_export_attr( - name, - _DYNIN_MODELING_REMOTE_EXPORTS, - module_name="modeling_dynin_omni", - settings=DYNIN_REMOTE_SETTINGS, - ) - - -def get_dynin_sampling_attr(name: str) -> Any: - return _get_export_attr( - name, - _DYNIN_SAMPLING_REMOTE_EXPORTS, - module_name="sampling", - settings=DYNIN_REMOTE_SETTINGS, - ) - - -def get_dynin_config_resolver_attr( - name: str, - *, - source: str | None = None, - revision: str | None = None, - local_files_only: bool | None = None, -) -> Any: - attr_name = _DYNIN_CONFIG_RESOLVER_REMOTE_EXPORTS.get(name) - if attr_name is None: - raise AttributeError(f"Unsupported Dynin config_resolver export: {name!r}") - - if source is not None: - value = resolve_remote_attr( - attr_name, - module_name="config_resolver", - settings=DYNIN_REMOTE_SETTINGS, - source=source, - revision=revision, - local_files_only=local_files_only, - optional=True, - ) - if value is not None: - return value - - return resolve_remote_attr( - attr_name, - module_name="config_resolver", - settings=DYNIN_REMOTE_SETTINGS, - source=DEFAULT_DYNIN_REMOTE_CODE_REPO, - revision=revision, - local_files_only=local_files_only, - optional=False, - ) - - -def get_dynin_magvit_attr( - name: str, - *, - source: str | None = None, - revision: str | None = None, - local_files_only: bool | None = None, -) -> Any: - attr_name = _DYNIN_MAGVIT_REMOTE_EXPORTS.get(name) - if attr_name is None: - raise AttributeError(f"Unsupported Dynin MAGVIT export: {name!r}") - - value = resolve_remote_attr( - attr_name, - module_name="modeling_magvitv2", - settings=MAGVIT_REMOTE_SETTINGS, - source=source, - revision=revision, - local_files_only=local_files_only, - optional=True, - ) - if value is not None: - return value - - resolved_source = _resolve_remote_source(source, MAGVIT_REMOTE_SETTINGS) - resolved_revision = _resolve_remote_revision(revision, MAGVIT_REMOTE_SETTINGS) - resolved_local_only = _resolve_remote_local_only(local_files_only, MAGVIT_REMOTE_SETTINGS) - - if resolved_source != DEFAULT_MAGVIT_REMOTE_CODE_REPO: - return resolve_remote_attr( - attr_name, - module_name="modeling_magvitv2", - settings=MAGVIT_REMOTE_SETTINGS, - source=DEFAULT_MAGVIT_REMOTE_CODE_REPO, - revision=resolved_revision, - local_files_only=resolved_local_only, - optional=False, - ) - - raise ImportError( - f"Failed to resolve MAGVIT attr '{attr_name}' from source={resolved_source!r} (revision={resolved_revision!r})." - ) - - -def build_dynin_chat_prompt(content: str) -> str: - return ( - f"<|start_header_id|>user<|end_header_id|>\n{content}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" - ) - - -def extract_dynin_user_prompt_text(decoded_prompt: str) -> str: - text = str(decoded_prompt or "") - assistant_marker = "<|start_header_id|>assistant<|end_header_id|>" - user_marker = "<|start_header_id|>user<|end_header_id|>" - end_header_marker = "<|end_header_id|>" - eot_marker = "<|eot_id|>" - - if assistant_marker in text: - text = text.rsplit(assistant_marker, 1)[0] - if eot_marker in text: - text = text.rsplit(eot_marker, 1)[0] - if user_marker in text: - text = text.rsplit(user_marker, 1)[-1] - if end_header_marker in text: - text = text.split(end_header_marker, 1)[-1] - return text.strip() - - -def normalize_dynin_online_prompt_text(task: str, decoded_prompt: str) -> str: - text = extract_dynin_user_prompt_text(decoded_prompt) - if not text: - text = str(decoded_prompt or "") - - for placeholder in _DYNIN_MODALITY_PLACEHOLDERS: - text = text.replace(placeholder, " ") - - task_token = _DYNIN_ONLINE_PROMPT_TOKEN_BY_TASK.get(task) - if task_token: - text = text.replace(task_token, " ", 1) - - text = " ".join(text.split()).strip() - - if task == "t2s": - if not text: - text = "Hello. This is a default text-to-speech sample." - text = build_dynin_chat_prompt(f"{DEFAULT_DYNIN_T2S_INSTRUCTION}\n{text}") - elif task in {"t2i", "i2i"} and not text: - text = "A high quality detailed image." - - return text - - -def infer_dynin_online_task( - *, - decoded_prompt: str, - has_image: bool = False, - has_audio: bool = False, - has_video: bool = False, -) -> str: - prompt = str(decoded_prompt or "") - if "<|i2i|>" in prompt: - return "i2i" - if "<|t2i|>" in prompt and not has_audio and not has_video: - return "t2i" - if "<|t2s|>" in prompt and not has_audio and not has_video: - return "t2s" - return "t2t" - - -def build_dynin_prompt_payload( - *, - task: str, - text: str, - image_tokens: torch.Tensor | None, - image_placeholder_tokens: int, - audio_placeholder_tokens: int, - image_token_offset: int, - mask_token_id: int, - use_train_i2i_prompt: bool, -) -> tuple[Any, str]: - _, prompting_task, _, _ = DYNIN_TASK_DEFAULT_RUNTIME[task] - - if task == "t2t": - payload = ([[]], [build_dynin_chat_prompt(text)]) - return payload, prompting_task - - if task == "t2i": - image_placeholder = torch.full( - (1, int(image_placeholder_tokens)), - fill_value=int(mask_token_id), - dtype=torch.long, - ) - payload = ([text], image_placeholder) - return payload, prompting_task - - if task == "i2i": - if image_tokens is None: - raise ValueError("i2i requires image tokens") - src = image_tokens.view(1, -1).long() + int(image_token_offset) - target_len = int(image_placeholder_tokens) if image_placeholder_tokens > 0 else int(src.shape[1]) - image_placeholder = torch.full( - (1, target_len), - fill_value=int(mask_token_id), - dtype=torch.long, - ) - if use_train_i2i_prompt: - labels_placeholder = torch.full( - (1, target_len), - fill_value=-100, - dtype=torch.long, - ) - payload = ([text], src, image_placeholder, labels_placeholder) - return payload, "i2i" - payload = ([text], src, image_placeholder) - return payload, "i2i_gen" - - if task == "t2s": - audio_placeholder = torch.full( - (1, int(audio_placeholder_tokens)), - fill_value=int(mask_token_id), - dtype=torch.long, - ) - payload = ([text], audio_placeholder) - return payload, prompting_task - - raise ValueError(f"Unsupported Dynin online bootstrap task: {task}") - - -def _wrap_runtime_field(value: Any) -> list[Any]: - return [value] - - -def build_dynin_online_runtime_info( - *, - task: str, - text_vocab_size: int, - infer_sources: DyninInferSources, - dynin_config_path: str | None = None, - prompting_input: Any | None = None, - attention_mask: list[int] | None = None, - prompt_length: int | None = None, - uncond_prompting_input: Any | None = None, - image_token_count: int = 0, - t2s_token_length: int | None = None, - use_train_i2i_prompt: bool | None = None, -) -> dict[str, Any]: - runtime_task, prompting_task, detok_id, _ = DYNIN_TASK_DEFAULT_RUNTIME[task] - - prompt_max_text_len = int(dynin_runtime_fallback(task, "prompt_max_text_len", None) or 1024) - max_new_tokens = int(dynin_runtime_fallback(task, "max_new_tokens", None) or 256) - steps = int(dynin_runtime_fallback(task, "steps", None) or 256) - block_length = int(dynin_runtime_fallback(task, "block_length", None) or 2) - temperature = float(dynin_runtime_fallback(task, "temperature", None) or 0.0) - cfg_scale = float(dynin_runtime_fallback(task, "cfg_scale", None) or 0.0) - remasking = str(dynin_runtime_fallback(task, "remasking", None) or "low_confidence") - timesteps = int(dynin_runtime_fallback(task, "timesteps", None) or 20) - guidance_scale = float(dynin_runtime_fallback(task, "guidance_scale", None) or 0.0) - mask_token_id = int(dynin_runtime_fallback(task, "mask_token_id", None) or 126336) - codebook_size = int(dynin_runtime_fallback(task, "codebook_size", None) or 8192) - audio_codebook_size = int(dynin_runtime_fallback(task, "audio_codebook_size", None) or 4096) - image_resolution = int(dynin_runtime_fallback(task, "image_resolution", None) or 336) - if image_token_count <= 0 and task in {"t2i", "i2i"}: - fallback_count = dynin_runtime_fallback(task, "image_token_count", None) - if fallback_count is not None: - image_token_count = int(fallback_count) - else: - image_token_count = max(1, (image_resolution // 16) ** 2) - - if t2s_token_length is None: - t2s_token_length = int(dynin_runtime_fallback(task, "t2s_token_length", None) or 383) - t2s_condition = str( - dynin_runtime_fallback( - task, - "t2s_condition", - None, - ) - or "gender-female_emotion-neutral_speed-normal_pitch-normal" - ) - if use_train_i2i_prompt is None: - use_train_i2i_prompt = bool(dynin_runtime_fallback(task, "use_train_i2i_prompt", task == "i2i")) - - runtime_info: dict[str, Any] = { - "task": _wrap_runtime_field(runtime_task), - "prompting_task": _wrap_runtime_field(prompting_task), - "detok_id": _wrap_runtime_field(int(detok_id)), - "prompt_max_text_len": _wrap_runtime_field(prompt_max_text_len), - "prompting_max_text_len": _wrap_runtime_field(prompt_max_text_len), - "cond_dropout_prob": _wrap_runtime_field(0.0), - "prompting_cond_dropout_prob": _wrap_runtime_field(0.0), - "tokenizer_path": _wrap_runtime_field(str(infer_sources.tokenizer_source)), - "text_vocab_size": _wrap_runtime_field(int(text_vocab_size)), - "model_local_files_only": _wrap_runtime_field(bool(infer_sources.model_local_files_only)), - "max_new_tokens": _wrap_runtime_field(int(t2s_token_length if task == "t2s" else max_new_tokens)), - "steps": _wrap_runtime_field(steps), - "block_length": _wrap_runtime_field(block_length), - "temperature": _wrap_runtime_field(temperature), - "cfg_scale": _wrap_runtime_field(cfg_scale), - "remasking": _wrap_runtime_field(remasking), - "mask_id": _wrap_runtime_field(mask_token_id), - "mask_token_id": _wrap_runtime_field(mask_token_id), - "codebook_size": _wrap_runtime_field(codebook_size), - "audio_codebook_size": _wrap_runtime_field(audio_codebook_size), - "timesteps": _wrap_runtime_field(timesteps), - "guidance_scale": _wrap_runtime_field(guidance_scale), - "noise_type": _wrap_runtime_field("mask"), - "noise_schedule_name": _wrap_runtime_field("cosine"), - "noise_schedule_params": _wrap_runtime_field({}), - "seq_len": _wrap_runtime_field(int(image_token_count)), - "condition": _wrap_runtime_field(t2s_condition), - "t2s_condition": _wrap_runtime_field(t2s_condition), - "vq_model_image_path": _wrap_runtime_field(str(infer_sources.vq_image_source)), - "vq_model_image_local_files_only": _wrap_runtime_field(bool(infer_sources.vq_image_local_files_only)), - "vq_model_audio_path": _wrap_runtime_field(str(infer_sources.vq_audio_source)), - "vq_model_audio_local_files_only": _wrap_runtime_field(bool(infer_sources.vq_audio_local_files_only)), - "image_resolution": _wrap_runtime_field(image_resolution), - "t2s_token_length": _wrap_runtime_field(int(t2s_token_length)), - "use_train_i2i_prompt": _wrap_runtime_field(bool(use_train_i2i_prompt)), - } - - if dynin_config_path: - runtime_info["dynin_config_path"] = _wrap_runtime_field(str(dynin_config_path)) - if prompting_input is not None: - runtime_info["prompting_input"] = _wrap_runtime_field(prompting_input) - if uncond_prompting_input is not None: - runtime_info["uncond_prompting_input"] = _wrap_runtime_field(uncond_prompting_input) - if attention_mask: - runtime_info["attention_mask"] = _wrap_runtime_field(list(attention_mask)) - if prompt_length is None and attention_mask: - prompt_length = len(attention_mask) - if prompt_length is not None: - runtime_info["prompt_length"] = _wrap_runtime_field(int(prompt_length)) - - return runtime_info diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py deleted file mode 100644 index dbbb63524..000000000 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py +++ /dev/null @@ -1,313 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import os -import tempfile -from pathlib import Path -from typing import Any - -import torch -from vllm.config import VllmConfig -from vllm.logger import init_logger -from vllm.sequence import IntermediateTensors - -from vllm_omni.model_executor.models.output_templates import OmniOutput - -from .dynin_omni import DyninOmniStageBase -from .dynin_omni_common import ( - DetokTarget, - _looks_like_hf_repo_id, - coerce_token_ids_1d, - normalize_runtime_info, - resolve_dynin_infer_sources, - resolve_hidden_size, - unwrap_first_value, -) - -logger = init_logger(__name__) - - -def _get_hf_token() -> str | None: - return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") - - -def _ensure_transformers_tied_weights_compat() -> None: - """Provide the transformers>=5 tied-weight attr expected during loading. - - The DYNIN audio detokenizer uses Hugging Face remote code that still only - defines the older ``_tied_weights_keys`` attribute. transformers 5.x now - reads ``all_tied_weights_keys`` inside ``from_pretrained`` before the model - instance is returned, so patching the loaded object is too late. - """ - try: - from transformers import PreTrainedModel - except Exception: - return - - if hasattr(PreTrainedModel, "all_tied_weights_keys"): - return - - def all_tied_weights_keys(self: Any) -> dict[str, Any]: - keys: dict[str, Any] = {} - for attr_name in ("_tied_weights_keys", "_dynamic_tied_weights_keys"): - raw_keys = getattr(self, attr_name, None) - if raw_keys is None: - continue - if isinstance(raw_keys, dict): - keys.update({str(k): v for k, v in raw_keys.items()}) - continue - if isinstance(raw_keys, str): - keys[raw_keys] = None - continue - for key in raw_keys: - keys[str(key)] = None - return keys - - PreTrainedModel.all_tied_weights_keys = property(all_tied_weights_keys) # type: ignore[attr-defined] - - -def _ensure_remote_s2u_vendor_root( - *, - repo_id: str, - local_files_only: bool, -) -> str | None: - if local_files_only or not _looks_like_hf_repo_id(repo_id): - return None - - existing = os.environ.get("DYNIN_S2U_VENDOR_ROOT") - if existing: - existing_path = Path(existing).expanduser().resolve() - if existing_path.is_dir(): - return str(existing_path) - - try: - from vllm_omni.transformers_utils.repo_utils import hf_api - except Exception as e: - logger.warning("huggingface_hub unavailable; cannot fetch s2u_vendor from %s: %s", repo_id, e) - return None - - token = _get_hf_token() - last_error: Exception | None = None - revisions: list[str | None] = [None] - - for revision in revisions: - try: - snapshot_dir = hf_api().snapshot_download( - repo_id=repo_id, - revision=revision, - allow_patterns=["s2u_vendor/**"], - token=token, - ) - except TypeError: - try: - snapshot_dir = hf_api().snapshot_download( - repo_id=repo_id, - revision=revision, - allow_patterns=["s2u_vendor/**"], - ) - except Exception as e: - last_error = e - continue - except Exception as e: - last_error = e - continue - - vendor_root = (Path(snapshot_dir) / "s2u_vendor").resolve() - if vendor_root.is_dir(): - os.environ["DYNIN_S2U_VENDOR_ROOT"] = str(vendor_root) - logger.info("Using remote S2U vendor root: %s", vendor_root) - return str(vendor_root) - - if last_error is not None: - logger.warning("Failed to download remote s2u_vendor from %s: %s", repo_id, last_error) - return None - - -class DyninOmniToken2Audio(DyninOmniStageBase): - """Stage-3: token detokenization to speech (or pass-through).""" - - stage_name = "Dynin token2audio" - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - del prefix - super().__init__() - self.vllm_config = vllm_config - self.have_multimodal_outputs = True - self.requires_raw_input_tokens = True - self.hidden_size = resolve_hidden_size(vllm_config=vllm_config) - self._vq_audio = None - self._vq_audio_path: str | None = None - self._vq_audio_local_files_only: bool | None = None - - def forward( - self, - input_ids: torch.Tensor | None = None, - positions: torch.Tensor | None = None, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, - ) -> OmniOutput: - del positions, intermediate_tensors, inputs_embeds - if input_ids is None: - raise ValueError("token2audio stage requires input_ids") - - runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information")) - detok_id = int(unwrap_first_value(runtime_info.get("detok_id"), 0)) - tokens = coerce_token_ids_1d(input_ids) - - if detok_id != DetokTarget.AUDIO: - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={ - "token_ids": tokens, - "detok_id": torch.tensor([detok_id], dtype=torch.long, device=tokens.device), - }, - ) - - audio, sample_rate = self._decode_audio_tokens(tokens, runtime_info=runtime_info) - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={ - "speech": audio, - "audio": audio, - "sr": torch.tensor([sample_rate], dtype=torch.int, device=audio.device), - "detok_id": torch.tensor([detok_id], dtype=torch.long, device=audio.device), - }, - ) - - def _decode_audio_tokens(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> tuple[torch.Tensor, int]: - # Follow DYNIN validation path: - # token list -> "<|speech_x|>" string -> vq_model_audio.decode(...). - vq_audio = self._ensure_vq_audio(runtime_info=runtime_info, ref_device=tokens.device) - - audio_codebook_size = int(unwrap_first_value(runtime_info.get("audio_codebook_size"), 4096)) - audio_vocab_offset = unwrap_first_value( - runtime_info.get("audio_vocab_offset"), - unwrap_first_value(runtime_info.get("t2s_vocab_start"), None), - ) - - token_ids = tokens.to(torch.long) - if audio_vocab_offset is not None: - off = int(audio_vocab_offset) - token_ids = torch.where(token_ids >= off, token_ids - off, token_ids) - token_ids = token_ids[(token_ids >= 0) & (token_ids < audio_codebook_size)] - if token_ids.numel() == 0: - raise RuntimeError("Audio detokenizer got no valid audio token ids.") - - speech_unit_str = " ".join(map(str, token_ids.detach().cpu().tolist())) - speech_unit_for_decode = "".join(f"<|speech_{unit}|>" for unit in speech_unit_str.split(" ") if unit != "") - - condition = unwrap_first_value( - runtime_info.get("condition"), - unwrap_first_value(runtime_info.get("t2s_condition"), None), - ) - output_wav_file = unwrap_first_value(runtime_info.get("output_wav_file"), None) - created_tmp = False - if output_wav_file is None: - fd, tmp_wav = tempfile.mkstemp(prefix="dynin_t2s_", suffix=".wav") - os.close(fd) - output_wav_file = tmp_wav - created_tmp = True - - audio_array = vq_audio.decode(speech_unit_for_decode, condition=condition, output_wav_file=output_wav_file) - if created_tmp: - try: - os.remove(output_wav_file) - except Exception: - pass - if not isinstance(audio_array, torch.Tensor): - audio_array = torch.as_tensor(audio_array, dtype=torch.float32, device=tokens.device) - else: - audio_array = audio_array.to(device=tokens.device, dtype=torch.float32) - - if audio_array.ndim > 1: - audio_array = audio_array.reshape(-1) - audio_array = audio_array.contiguous() - - sample_rate = int( - unwrap_first_value( - runtime_info.get("sr"), - unwrap_first_value(runtime_info.get("sample_rate"), 24000), - ) - ) - try: - cfg = getattr(vq_audio, "u2s_config", None) - cfg_sr = getattr(cfg, "sampling_rate", None) - if cfg_sr is None: - cfg_sr = getattr(getattr(cfg, "data", None), "sampling_rate", None) - if cfg_sr is not None: - sample_rate = int(cfg_sr) - except Exception: - pass - return audio_array, sample_rate - - def _ensure_vq_audio(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any: - sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info) - model_path = str(sources.vq_audio_source) - local_files_only = bool(sources.vq_audio_local_files_only) - - _ensure_remote_s2u_vendor_root( - repo_id=model_path, - local_files_only=local_files_only, - ) - - if ( - self._vq_audio is None - or self._vq_audio_path != model_path - or self._vq_audio_local_files_only != local_files_only - ): - logger.info( - "Loading DYNIN audio detokenizer from %s (local_files_only=%s)", - model_path, - local_files_only, - ) - try: - from transformers import AutoModel - except Exception as e: - raise RuntimeError( - "transformers is required to load EMOVASpeechTokenizer remote code from Hugging Face." - ) from e - - _ensure_transformers_tied_weights_compat() - try: - self._vq_audio = AutoModel.from_pretrained( - model_path, - trust_remote_code=True, - local_files_only=local_files_only, - low_cpu_mem_usage=False, - ) - except TypeError: - try: - self._vq_audio = AutoModel.from_pretrained( - model_path, - trust_remote_code=True, - local_files_only=local_files_only, - ) - except TypeError: - self._vq_audio = AutoModel.from_pretrained( - model_path, - trust_remote_code=True, - ) - except Exception as e: - raise RuntimeError( - f"Failed to load EMOVASpeechTokenizer from Hugging Face remote code for model path '{model_path}'." - ) from e - - if not hasattr(self._vq_audio, "decode"): - raise RuntimeError( - "Loaded audio tokenizer does not expose decode(). " - "Check HF config.json auto_map/model_type and ensure trust_remote_code=True." - ) - self._vq_audio.eval() - self._vq_audio.requires_grad_(False) - self._vq_audio_path = model_path - self._vq_audio_local_files_only = local_files_only - if hasattr(self._vq_audio, "to"): - self._vq_audio = self._vq_audio.to(ref_device) - return self._vq_audio - - def embed_multimodal(self, **kwargs: Any) -> Any: - del kwargs - return None diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2image.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2image.py deleted file mode 100644 index 6b5110a77..000000000 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2image.py +++ /dev/null @@ -1,150 +0,0 @@ -from __future__ import annotations - -import os -from typing import Any - -import torch -from vllm.config import VllmConfig -from vllm.logger import init_logger -from vllm.sequence import IntermediateTensors - -from vllm_omni.model_executor.models.output_templates import OmniOutput - -from .dynin_omni import DyninOmniStageBase -from .dynin_omni_common import ( - DetokTarget, - _to_bool, - coerce_token_ids_1d, - get_dynin_magvit_attr, - normalize_runtime_info, - resolve_dynin_infer_sources, - resolve_hidden_size, - unwrap_first_value, -) - -logger = init_logger(__name__) - - -class DyninOmniToken2Image(DyninOmniStageBase): - """Stage-2: token detokenization to image (or pass-through).""" - - stage_name = "Dynin token2image" - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - del prefix - super().__init__() - - self.vllm_config = vllm_config - self.have_multimodal_outputs = True - self.requires_raw_input_tokens = True - self.hidden_size = resolve_hidden_size(vllm_config=vllm_config) - self._vq_model = None - self._vq_model_path: str | None = None - self._vq_local_files_only: bool | None = None - - def forward( - self, - input_ids: torch.Tensor | None = None, - positions: torch.Tensor | None = None, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, - ) -> OmniOutput: - del positions, intermediate_tensors, inputs_embeds - if input_ids is None: - raise ValueError("token2image stage requires input_ids") - runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information")) - detok_id = int(unwrap_first_value(runtime_info.get("detok_id"), 0)) - tokens = coerce_token_ids_1d(input_ids) - - if detok_id != DetokTarget.IMAGE: - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={ - "token_ids": tokens, - "detok_id": torch.tensor([detok_id], dtype=torch.long, device=tokens.device), - }, - ) - - image = self._decode_image_tokens(tokens, runtime_info=runtime_info) - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={ - "image": image, - "detok_id": torch.tensor([detok_id], dtype=torch.long, device=image.device), - }, - ) - - def _decode_image_tokens(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> torch.Tensor: - # Follow DYNIN validation path: - # tokens -> clamp -> vq_model.decode_code -> (x+1)/2 -> [0,1]. - vq_model = self._ensure_vq_model(runtime_info=runtime_info, ref_device=tokens.device) - codebook_size = int(unwrap_first_value(runtime_info.get("codebook_size"), 8192)) - image_vocab_offset = unwrap_first_value(runtime_info.get("image_vocab_offset"), None) - if image_vocab_offset is None: - text_vocab_size = unwrap_first_value(runtime_info.get("text_vocab_size"), None) - num_new_special_tokens = int(unwrap_first_value(runtime_info.get("num_new_special_tokens"), 0)) - if text_vocab_size is not None: - image_vocab_offset = int(text_vocab_size) + num_new_special_tokens - - token_ids = tokens.to(torch.long) - if image_vocab_offset is not None: - off = int(image_vocab_offset) - token_ids = torch.where(token_ids >= off, token_ids - off, token_ids) - token_ids = torch.clamp(token_ids, min=0, max=max(0, codebook_size - 1)) - token_ids = token_ids.unsqueeze(0) - - decoded = vq_model.decode_code(token_ids) - decoded = torch.clamp((decoded + 1.0) / 2.0, min=0.0, max=1.0) - if decoded.ndim != 4 or decoded.shape[0] == 0: - raise RuntimeError(f"Unexpected MAGVIT decode output shape: {tuple(decoded.shape)}") - return decoded[0].contiguous() - - def _ensure_vq_model(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any: - sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info) - model_path = str(sources.vq_image_source) - local_files_only = bool(sources.vq_image_local_files_only) - if self._vq_model is None or self._vq_model_path != model_path or self._vq_local_files_only != local_files_only: - disable_xet = unwrap_first_value( - runtime_info.get("hf_hub_disable_xet"), - unwrap_first_value(runtime_info.get("disable_hf_xet"), True), - ) - if _to_bool(disable_xet, default=True): - os.environ.setdefault("HF_HUB_DISABLE_XET", "1") - logger.info( - "Loading DYNIN image detokenizer from %s (local_files_only=%s)", - model_path, - local_files_only, - ) - try: - MAGVITv2 = get_dynin_magvit_attr( - "MAGVITv2", - source=model_path, - local_files_only=local_files_only, - ) - try: - self._vq_model = MAGVITv2.from_pretrained( - model_path, - local_files_only=local_files_only, - ) - except TypeError: - self._vq_model = MAGVITv2.from_pretrained(model_path) - except Exception as e: - raise RuntimeError( - "Failed to load MAGVITv2 from local DYNIN submodel implementation " - f"for model path '{model_path}'. " - "If your environment cannot access huggingface.co, set " - "additional_information.vq_model_image_path to a local MAGVITv2 directory " - "and set additional_information.vq_model_image_local_files_only=true." - ) from e - self._vq_model.eval() - self._vq_model.requires_grad_(False) - self._vq_model_path = model_path - self._vq_local_files_only = local_files_only - if hasattr(self._vq_model, "to"): - self._vq_model = self._vq_model.to(ref_device) - return self._vq_model - - def embed_multimodal(self, **kwargs: Any) -> Any: - del kwargs - return None diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2text.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2text.py deleted file mode 100644 index 38caa84f7..000000000 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2text.py +++ /dev/null @@ -1,1640 +0,0 @@ -from __future__ import annotations - -import inspect -import json -from contextlib import contextmanager -from typing import Any - -import torch -import torch.nn.functional as F -from transformers import AutoTokenizer -from vllm.config import VllmConfig -from vllm.logger import init_logger -from vllm.sequence import IntermediateTensors - -from vllm_omni.model_executor.models.output_templates import OmniOutput - -from .dynin_omni import DyninOmniStageBase -from .dynin_omni_common import ( - DYNIN_PROMPT_SOURCE_KEY, - DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT, - DYNIN_REMOTE_SETTINGS, - DYNIN_SPECIAL_TOKENS, - TASK_TO_DETOK, - DetokTarget, - _to_bool, - build_dynin_online_runtime_info, - build_dynin_prompt_payload, - coerce_token_ids_1d, - dynin_runtime_fallback, - get_dynin_magvit_attr, - get_dynin_modeling_attr, - get_dynin_sampling_attr, - infer_dynin_online_task, - logical_dynin_task, - normalize_dynin_online_prompt_text, - normalize_runtime_info, - resolve_dynin_infer_sources, - resolve_hidden_size, - resolve_remote_attr, - unwrap_first_value, -) - -logger = init_logger(__name__) - -TASK_TO_PROMPTING_TASK = { - "t2i": "t2i_gen", - "i2i": "i2i_gen", - "ti2ti": "ti2ti_gen", - "t2s": "t2s_gen", - "t2s_mmu_like": "t2s_gen", - "t2s_fixed": "t2s_fixed_gen", - "s2s": "s2s_gen", - "v2s": "v2s_gen", - "mmu": "mmu", - "mmu_fast": "mmu", - "mmu_fastdllm_v1": "mmu", - "s2t": "s2t", - "v2t": "v2t", -} - -TASK_TO_GENERATE_FN = { - "t2i": "t2i_generate", - "i2i": "i2i_generate", - "ti2ti": "ti2ti_generate", - "t2s": "t2s_generate", - "t2s_mmu_like": "t2s_generate_mmu_like", - "t2s_fixed": "t2s_fixed_generate", - "s2s": "t2s_generate_mmu_like", - "v2s": "t2s_generate_mmu_like", - "s2t": "s2t_generate", - "mmu": "mmu_generate", - "t2t": "generate", - "mmu_fast": "mmu_generate_fast", - "mmu_fastdllm_v1": "mmu_generate_fastdllm_v1", - "v2t": "mmu_generate", -} - -TASKS_USING_UNI_PROMPTING = set(TASK_TO_PROMPTING_TASK.keys()) -PROMPT_PAYLOAD_REQUIRED_TASKS = { - "t2i", - "i2i", - "ti2ti", - "t2s", - "t2s_mmu_like", - "t2s_fixed", - "s2s", - "v2s", -} - -GENERATE_RUNTIME_KWARG_KEYS = ( - "uncond_input_ids", - "uncond_attention_mask", - "noise_schedule", - "generator", - "config", - "uni_prompting", - "resolution", - "max_new_tokens", - "steps", - "block_length", - "temperature", - "top_k", - "eot_token", - "cfg_scale", - "remasking", - "mask_id", - "attention_mask", - "timesteps", - "guidance_scale", - "noise_type", - "seq_len", - "mask_token_id", - "codebook_size", - "audio_codebook_size", - "use_cache", - "threshold", - "factor", -) - -PASSTHROUGH_GENERATE_KWARG_KEYS = ( - "attention_mask", - "uncond_input_ids", - "uncond_attention_mask", - "noise_schedule", - "uni_prompting", - "generator", - "noise_type", -) - -PROMPTING_PAYLOAD_KEYS = ( - "prompting_input", - "prompting_inputs", - "dynin_inputs", - "model_inputs", - "raw_inputs", -) - -UNCOND_PROMPTING_PAYLOAD_KEYS = ( - "uncond_prompting_input", - "uncond_prompting_inputs", -) - -PROMPTING_META_KEYS = ( - "uncond_prompting_input", - "uncond_prompting_inputs", - "uni_prompting", - "prompting_task", - "prompting_config", -) - -MM_INPUT_ALIASES = { - "image": ("pixel_values", "image_embeds", "img2img"), - "video": ("pixel_values_videos", "video_embeds"), - "audio": ("input_audio_features", "audio_embeds"), -} - - -class DyninOmniToken2Text(DyninOmniStageBase): - """Stage-1: DYNIN generation + text detokenization or pass-through.""" - - stage_name = "Dynin token2text" - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - del prefix - super().__init__() - - self.vllm_config = vllm_config - self.have_multimodal_outputs = True - self.requires_raw_input_tokens = True - - self._infer_sources = resolve_dynin_infer_sources(vllm_config=vllm_config) - if self._infer_sources.config_path: - logger.info( - "DYNIN token2text using inference config: %s", - self._infer_sources.config_path, - ) - - self.model = self._load_text_model( - self._infer_sources.model_source, - local_files_only=self._infer_sources.model_local_files_only, - ) - self.model.eval() - self.model.requires_grad_(False) - - self.hidden_size = resolve_hidden_size( - vllm_config=vllm_config, - model=self.model, - ) - - self.tokenizer: Any | None = None - self._tokenizer_path: str | None = None - self._uni_prompting: Any | None = None - self._uni_prompting_init_spec: tuple[Any, ...] | None = None - self._prompt_vq_model: Any | None = None - self._prompt_vq_model_path: str | None = None - self._prompt_vq_local_files_only: bool | None = None - self._cached_mm_inputs: dict[str, Any] = {} - - try: - self._set_tokenizer( - self._infer_sources.tokenizer_source, - local_files_only=self._infer_sources.model_local_files_only, - ) - except Exception: - self.tokenizer = None - self._tokenizer_path = None - - @staticmethod - def _load_text_model(model_path: str, *, local_files_only: bool = False) -> Any: - try: - dynin_model_cls = get_dynin_modeling_attr("DyninOmniModelLM") - # Remote Dynin `DyninOmniConfig` may omit `use_cache`, but bundled - # `modeling_llada.forward` reads `self.config.use_cache` like a - # standard HF causal LM config (see kernel warmup / mmu_generate). - dynin_config_cls = get_dynin_modeling_attr("DyninOmniConfig") - if not getattr(dynin_config_cls, "_dynin_use_cache_shim_applied", False): - _orig_cfg_init = dynin_config_cls.__init__ - - def _cfg_init_shim(self, *args, **kwargs): - _orig_cfg_init(self, *args, **kwargs) - if not hasattr(self, "use_cache"): - setattr(self, "use_cache", True) - - dynin_config_cls.__init__ = _cfg_init_shim # type: ignore[method-assign] - dynin_config_cls._dynin_use_cache_shim_applied = True - # transformers>=5.0 renamed the weight-tying manifest accessor from - # `_tied_weights_keys` (a class attribute) to `all_tied_weights_keys` - # (a property returning a {module_path: [tied_keys]} mapping) and - # invokes it during `from_pretrained -> caching_allocator_warmup - # -> get_total_byte_count`. Remote Dynin code pre-dates the rename, - # so patch a minimal shim onto the class if it's missing. Keeping - # the mapping empty is safe here — the root module has no weight - # tying; the real tied-weights (if any) live on nested submodules - # that transformers discovers via their own `_tied_weights_keys`. - if not hasattr(dynin_model_cls, "all_tied_weights_keys"): - dynin_model_cls.all_tied_weights_keys = property(lambda self: {}) - # transformers>=5.0 also calls `get_total_byte_count` which does - # `len(model._tp_plan)`; PreTrainedModel's class default for - # `_tp_plan` is `None`, and unless `self.config.base_model_tp_plan` - # was populated during __init__ (via `post_init`) the instance - # attribute remains None, crashing with - # `TypeError: object of type 'NoneType' has no len()`. The remote - # Dynin model class bypasses that path, so shim an empty dict at - # the class level as a safe fallback (PreTrainedModel sets the - # instance attribute when the config actually carries a plan). - if getattr(dynin_model_cls, "_tp_plan", None) is None: - dynin_model_cls._tp_plan = {} - # transformers>=5.0 `_finalize_model_loading` calls - # `model.tie_weights(missing_keys=..., recompute_mapping=...)` while - # the remote Dynin code still defines `tie_weights(self)` with no - # keyword arguments, raising - # `TypeError: LLaDAModelLM.tie_weights() got an unexpected keyword - # argument 'missing_keys'`. Wrap the method once to silently drop - # the new kwargs; the remote implementation is a no-op w.r.t. the - # missing-keys bookkeeping transformers would otherwise perform. - _orig_tie_weights = dynin_model_cls.tie_weights - if not getattr(_orig_tie_weights, "_dynin_accepts_v5_kwargs", False): - _tie_sig = inspect.signature(_orig_tie_weights) - _tie_accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in _tie_sig.parameters.values()) - _tie_accepts_new_kwargs = _tie_accepts_var_kw or ( - "missing_keys" in _tie_sig.parameters and "recompute_mapping" in _tie_sig.parameters - ) - if not _tie_accepts_new_kwargs: - - def _tie_weights_shim(self, *args, **kwargs): - kwargs.pop("missing_keys", None) - kwargs.pop("recompute_mapping", None) - return _orig_tie_weights(self, *args, **kwargs) - - _tie_weights_shim._dynin_accepts_v5_kwargs = True # type: ignore[attr-defined] - dynin_model_cls.tie_weights = _tie_weights_shim - try: - return dynin_model_cls.from_pretrained( - model_path, - torch_dtype=torch.bfloat16, - local_files_only=local_files_only, - ) - except TypeError: - return dynin_model_cls.from_pretrained( - model_path, - torch_dtype=torch.bfloat16, - ) - except Exception as e: - raise RuntimeError( - f"Failed to load DyninOmniModelLM via remote Dynin code for model path '{model_path}'." - ) from e - - @staticmethod - def _load_tokenizer_from_source( - source: str, - *, - local_files_only: bool = False, - trust_remote_code: bool = False, - ) -> Any: - load_kwargs = { - "trust_remote_code": trust_remote_code, - "local_files_only": _to_bool(local_files_only, default=False), - } - try: - return AutoTokenizer.from_pretrained(source, **load_kwargs) - except TypeError: - load_kwargs.pop("local_files_only", None) - return AutoTokenizer.from_pretrained(source, **load_kwargs) - - def _set_tokenizer(self, source: str, *, local_files_only: bool) -> None: - try: - tokenizer = self._load_tokenizer_from_source( - source, - local_files_only=local_files_only, - trust_remote_code=False, - ) - except Exception as e: - logger.info( - "Falling back to trust_remote_code=True tokenizer loading for %s: %s", - source, - e, - ) - tokenizer = self._load_tokenizer_from_source( - source, - local_files_only=local_files_only, - trust_remote_code=True, - ) - - self.tokenizer = tokenizer - self._tokenizer_path = source - self._reset_uni_prompting_cache() - - def _reset_uni_prompting_cache(self) -> None: - self._uni_prompting = None - self._uni_prompting_init_spec = None - - def get_language_model(self) -> Any: - return self.model - - @staticmethod - def _merge_runtime_info_missing_values( - runtime_info: dict[str, Any], - fallback_info: dict[str, Any], - ) -> dict[str, Any]: - merged = dict(runtime_info) - for key, value in fallback_info.items(): - if unwrap_first_value(merged.get(key), None) is None: - merged[key] = value - return merged - - def _runtime_info_needs_bootstrap( - self, - runtime_info: dict[str, Any], - logical_task_name: str, - ) -> bool: - task = str(unwrap_first_value(runtime_info.get("task"), "") or "").lower() - detok_id = unwrap_first_value(runtime_info.get("detok_id"), None) - prompt_length = unwrap_first_value(runtime_info.get("prompt_length"), None) - - if not task or detok_id is None: - return True - if prompt_length is None: - return True - if ( - task in PROMPT_PAYLOAD_REQUIRED_TASKS - and self._find_first_payload( - runtime_info=runtime_info, - kwargs={}, - keys=PROMPTING_PAYLOAD_KEYS, - ) - is None - ): - return True - if logical_task_name in {"t2i", "i2i"}: - for key in ("codebook_size", "text_vocab_size", "vq_model_image_path"): - if unwrap_first_value(runtime_info.get(key), None) is None: - return True - if logical_task_name == "t2s": - for key in ("audio_codebook_size", "condition", "vq_model_audio_path"): - if unwrap_first_value(runtime_info.get(key), None) is None: - return True - return False - - def _decode_prompt_for_bootstrap( - self, - input_ids: torch.Tensor, - runtime_info: dict[str, Any], - ) -> str: - self._maybe_load_runtime_tokenizer(runtime_info) - if self.tokenizer is None: - return "" - token_ids = coerce_token_ids_1d(input_ids).detach().cpu().tolist() - try: - return str(self.tokenizer.decode(token_ids, skip_special_tokens=False)) - except Exception: - return "" - - def _bootstrap_runtime_info_if_needed( - self, - *, - input_ids: torch.Tensor, - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - ) -> dict[str, Any]: - if unwrap_first_value(runtime_info.get(DYNIN_PROMPT_SOURCE_KEY), None) == DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT: - return runtime_info - - mm_inputs = self._collect_mm_inputs(**kwargs) - decoded_prompt = "" - - task_value = unwrap_first_value(runtime_info.get("task"), None) - if task_value is None: - decoded_prompt = self._decode_prompt_for_bootstrap(input_ids, runtime_info) - logical_task_name = infer_dynin_online_task( - decoded_prompt=decoded_prompt, - has_image="image" in mm_inputs, - has_audio="audio" in mm_inputs, - has_video="video" in mm_inputs, - ) - else: - logical_task_name = logical_dynin_task(task_value) - - if not self._runtime_info_needs_bootstrap(runtime_info, logical_task_name): - return runtime_info - - self._maybe_load_runtime_tokenizer(runtime_info) - if self.tokenizer is None: - logger.warning("Unable to bootstrap Dynin runtime info because tokenizer is unavailable.") - return runtime_info - - if not decoded_prompt: - decoded_prompt = self._decode_prompt_for_bootstrap(input_ids, runtime_info) - - text_vocab_size = int(len(self.tokenizer)) - prompt_len = int(coerce_token_ids_1d(input_ids).numel()) - dynin_config_path = self._infer_sources.config_path - - base_runtime_info = build_dynin_online_runtime_info( - task=logical_task_name, - text_vocab_size=text_vocab_size, - infer_sources=self._infer_sources, - dynin_config_path=dynin_config_path, - attention_mask=([1] * prompt_len) if logical_task_name == "t2t" else None, - prompt_length=prompt_len if logical_task_name == "t2t" else None, - ) - merged_runtime_info = self._merge_runtime_info_missing_values(runtime_info, base_runtime_info) - - payload_required = logical_task_name in {"t2i", "i2i", "t2s"} - existing_prompt_payload = self._find_first_payload( - runtime_info=merged_runtime_info, - kwargs=kwargs, - keys=PROMPTING_PAYLOAD_KEYS, - ) - has_prompt_payload = existing_prompt_payload is not None - needs_prompt_length = unwrap_first_value(merged_runtime_info.get("prompt_length"), None) is None - if not payload_required: - return merged_runtime_info - - use_train_i2i_prompt = _to_bool( - unwrap_first_value( - merged_runtime_info.get("use_train_i2i_prompt"), - dynin_runtime_fallback(logical_task_name, "use_train_i2i_prompt", logical_task_name == "i2i"), - ), - default=logical_task_name == "i2i", - ) - t2s_token_length = int( - dynin_runtime_fallback( - logical_task_name, - "t2s_token_length", - unwrap_first_value(merged_runtime_info.get("t2s_token_length"), None), - ) - or 383 - ) - image_resolution = int( - dynin_runtime_fallback( - logical_task_name, - "image_resolution", - unwrap_first_value(merged_runtime_info.get("image_resolution"), None), - ) - or 336 - ) - - image_token_count = int( - dynin_runtime_fallback( - logical_task_name, - "image_token_count", - unwrap_first_value(merged_runtime_info.get("seq_len"), None), - ) - or 0 - ) - image_tokens: torch.Tensor | None = None - if logical_task_name == "i2i" and (not has_prompt_payload or image_token_count <= 0): - image_tokens = self._encode_prompt_image_tokens( - runtime_info=merged_runtime_info, - mm_inputs=mm_inputs, - resolution=image_resolution, - ) - image_token_count = int(image_tokens.numel()) - - mask_token_id = int(unwrap_first_value(merged_runtime_info.get("mask_token_id"), 126336)) - prompting_input = self._unwrap_singleton(existing_prompt_payload) - prompting_task = str( - unwrap_first_value( - merged_runtime_info.get("prompting_task"), - TASK_TO_PROMPTING_TASK.get( - str(unwrap_first_value(merged_runtime_info.get("task"), "mmu")).lower(), - "mmu", - ), - ) - ) - if not has_prompt_payload: - prompt_text = normalize_dynin_online_prompt_text(logical_task_name, decoded_prompt) - prompting_input, prompting_task = build_dynin_prompt_payload( - task=logical_task_name, - text=prompt_text, - image_tokens=image_tokens, - image_placeholder_tokens=image_token_count, - audio_placeholder_tokens=t2s_token_length, - image_token_offset=text_vocab_size, - mask_token_id=mask_token_id, - use_train_i2i_prompt=use_train_i2i_prompt, - ) - - prompt_runtime_info = build_dynin_online_runtime_info( - task=logical_task_name, - text_vocab_size=text_vocab_size, - infer_sources=self._infer_sources, - dynin_config_path=dynin_config_path, - image_token_count=image_token_count, - t2s_token_length=t2s_token_length, - use_train_i2i_prompt=use_train_i2i_prompt, - ) - prompt_runtime_info["prompting_task"] = [str(prompting_task)] - prompt_runtime_info["prompting_input"] = [prompting_input] - merged_runtime_info = self._merge_runtime_info_missing_values(merged_runtime_info, prompt_runtime_info) - - if not needs_prompt_length and has_prompt_payload: - return merged_runtime_info - - uni_prompting = self._get_or_create_uni_prompting( - runtime_info=merged_runtime_info, - kwargs=kwargs, - ) - if uni_prompting is not None: - prepared_input_ids, prepared_attention_mask = self._prepare_prompting_input( - payload=prompting_input, - task=str(unwrap_first_value(merged_runtime_info.get("task"), "mmu")), - runtime_info=merged_runtime_info, - kwargs=kwargs, - uni_prompting=uni_prompting, - ref_device=input_ids.device, - ) - if prepared_input_ids is not None: - prepared_prompt_len = int(prepared_input_ids.shape[-1]) - prepared_attention_list: list[int] | None = None - if prepared_attention_mask is not None: - prepared_attention_list = prepared_attention_mask.view(-1).detach().cpu().tolist() - final_runtime_info = build_dynin_online_runtime_info( - task=logical_task_name, - text_vocab_size=text_vocab_size, - infer_sources=self._infer_sources, - dynin_config_path=dynin_config_path, - prompting_input=prompting_input, - attention_mask=prepared_attention_list, - prompt_length=prepared_prompt_len, - image_token_count=image_token_count, - t2s_token_length=t2s_token_length, - use_train_i2i_prompt=use_train_i2i_prompt, - ) - final_runtime_info["prompting_task"] = [str(prompting_task)] - - guidance_scale = float(unwrap_first_value(merged_runtime_info.get("guidance_scale"), 0.0)) - if logical_task_name in {"t2i", "i2i"} and guidance_scale > 0: - uncond_prompting_input, _ = build_dynin_prompt_payload( - task=logical_task_name, - text="", - image_tokens=image_tokens, - image_placeholder_tokens=image_token_count, - audio_placeholder_tokens=t2s_token_length, - image_token_offset=text_vocab_size, - mask_token_id=mask_token_id, - use_train_i2i_prompt=use_train_i2i_prompt, - ) - final_runtime_info["uncond_prompting_input"] = [uncond_prompting_input] - - merged_runtime_info = self._merge_runtime_info_missing_values( - merged_runtime_info, - final_runtime_info, - ) - - return merged_runtime_info - - @staticmethod - def _build_downstream_runtime_info(runtime_info: dict[str, Any]) -> dict[str, Any]: - bridge_keys = ( - "task", - "detok_id", - "dynin_config_path", - "codebook_size", - "audio_codebook_size", - "text_vocab_size", - "num_new_special_tokens", - "image_vocab_offset", - "audio_vocab_offset", - "t2s_vocab_start", - "condition", - "t2s_condition", - "vq_model_image_path", - "vq_model_image_local_files_only", - "vq_model_audio_path", - "vq_model_audio_local_files_only", - "model_local_files_only", - "local_files_only", - "hf_hub_disable_xet", - "disable_hf_xet", - ) - return {key: runtime_info[key] for key in bridge_keys if key in runtime_info} - - @staticmethod - def _jsonify_runtime_value(value: Any) -> Any: - if isinstance(value, torch.Tensor): - return value.detach().cpu().tolist() - if isinstance(value, (list, tuple)): - return [DyninOmniToken2Text._jsonify_runtime_value(item) for item in value] - if isinstance(value, dict): - return {str(key): DyninOmniToken2Text._jsonify_runtime_value(val) for key, val in value.items()} - if isinstance(value, (str, int, float, bool)) or value is None: - return value - return str(value) - - def _encode_runtime_info_tensor( - self, - runtime_info: dict[str, Any], - *, - device: torch.device, - ) -> torch.Tensor | None: - if not runtime_info: - return None - payload = {key: self._jsonify_runtime_value(value) for key, value in runtime_info.items()} - encoded = json.dumps( - payload, - ensure_ascii=False, - separators=(",", ":"), - sort_keys=True, - ).encode("utf-8") - if not encoded: - return None - return torch.tensor(list(encoded), dtype=torch.uint8, device=device) - - def forward( - self, - input_ids: torch.Tensor | None = None, - positions: torch.Tensor | None = None, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, - ) -> OmniOutput: - del positions, intermediate_tensors, inputs_embeds - - if input_ids is None: - raise ValueError("token2text stage requires input_ids") - try: - runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information")) - runtime_info = self._bootstrap_runtime_info_if_needed( - input_ids=input_ids, - runtime_info=runtime_info, - kwargs=kwargs, - ) - task = str(unwrap_first_value(runtime_info.get("task"), "mmu")).lower() - - detok_id = int( - unwrap_first_value( - runtime_info.get("detok_id"), - TASK_TO_DETOK.get(task, DetokTarget.TEXT), - ) - ) - - token_ids = self._generate_token_ids( - task=task, - input_ids=input_ids, - runtime_info=runtime_info, - kwargs=kwargs, - ) - bridge_runtime_info = self._build_downstream_runtime_info(runtime_info) - runtime_info_tensor = self._encode_runtime_info_tensor( - bridge_runtime_info, - device=token_ids.device, - ) - - if detok_id != int(DetokTarget.TEXT): - multimodal_outputs = { - "token_ids": token_ids, - "detok_id": torch.tensor( - [detok_id], - dtype=torch.long, - device=token_ids.device, - ), - } - if runtime_info_tensor is not None: - multimodal_outputs["runtime_info_json"] = runtime_info_tensor - return OmniOutput( - text_hidden_states=None, - multimodal_outputs=multimodal_outputs, - ) - - decode_tokens = self._extract_decode_tokens(token_ids, runtime_info=runtime_info) - multimodal_outputs = { - "token_ids": token_ids, - "text_tokens": decode_tokens, - "detok_id": torch.tensor( - [detok_id], - dtype=torch.long, - device=token_ids.device, - ), - } - if runtime_info_tensor is not None: - multimodal_outputs["runtime_info_json"] = runtime_info_tensor - - return OmniOutput( - text_hidden_states=None, - multimodal_outputs=multimodal_outputs, - ) - finally: - self._cached_mm_inputs = {} - - def _generate_token_ids( - self, - task: str, - input_ids: torch.Tensor, - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - ) -> torch.Tensor: - precomputed = self._get_precomputed_token_ids(runtime_info) - if precomputed is not None: - return coerce_token_ids_1d(precomputed, ref_device=input_ids.device) - - gen_fn_name = TASK_TO_GENERATE_FN.get(task, "mmu_generate") - gen_fn = self._resolve_generate_fn(gen_fn_name) - - gen_kwargs = self._collect_generate_kwargs(runtime_info=runtime_info, kwargs=kwargs) - - if "noise_schedule" not in gen_kwargs: - noise_schedule = self._resolve_noise_schedule( - runtime_info=runtime_info, - kwargs=kwargs, - ) - if noise_schedule is not None: - gen_kwargs["noise_schedule"] = noise_schedule - - if task in TASKS_USING_UNI_PROMPTING and "uni_prompting" not in gen_kwargs: - uni_prompting = self._get_or_create_uni_prompting( - runtime_info=runtime_info, - kwargs=kwargs, - ) - if uni_prompting is not None: - gen_kwargs["uni_prompting"] = uni_prompting - - should_prepare_prompting_inputs = task in TASKS_USING_UNI_PROMPTING or self._contains_prompting_payload( - runtime_info=runtime_info, kwargs=kwargs - ) - if should_prepare_prompting_inputs: - input_ids, gen_kwargs = self._prepare_prompting_inputs_if_needed( - task=task, - input_ids=input_ids, - runtime_info=runtime_info, - kwargs=kwargs, - gen_kwargs=gen_kwargs, - ) - - input_ids, gen_kwargs = self._normalize_generate_inputs( - input_ids=input_ids, - gen_kwargs=gen_kwargs, - ref_device=input_ids.device, - ) - gen_kwargs = self._filter_supported_generate_kwargs( - gen_fn=gen_fn, - gen_kwargs=gen_kwargs, - fn_name=gen_fn_name, - ) - - generated = self._call_generate_fn( - gen_fn=gen_fn, - input_ids=input_ids, - gen_kwargs=gen_kwargs, - ) - return coerce_token_ids_1d(generated, ref_device=input_ids.device) - - @staticmethod - def _get_precomputed_token_ids(runtime_info: dict[str, Any]) -> Any | None: - precomputed = runtime_info.get("generated_token_ids") - if precomputed is None: - precomputed = runtime_info.get("token_ids") - return precomputed - - def _resolve_generate_fn(self, fn_name: str) -> Any: - if not hasattr(self.model, fn_name): - raise RuntimeError( - f"DYNIN model does not expose '{fn_name}'. " - "Pass additional_information.generated_token_ids or adjust task mapping." - ) - return getattr(self.model, fn_name) - - @staticmethod - def _collect_generate_kwargs( - *, - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - ) -> dict[str, Any]: - gen_kwargs: dict[str, Any] = {} - - for key in GENERATE_RUNTIME_KWARG_KEYS: - if key in runtime_info: - gen_kwargs[key] = unwrap_first_value(runtime_info[key]) - - for key in PASSTHROUGH_GENERATE_KWARG_KEYS: - if key not in gen_kwargs and key in kwargs: - gen_kwargs[key] = kwargs[key] - - return gen_kwargs - - @staticmethod - def _contains_prompting_payload( - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - ) -> bool: - keys = PROMPTING_PAYLOAD_KEYS + PROMPTING_META_KEYS - return any(key in runtime_info for key in keys) or any(key in kwargs for key in keys) - - @staticmethod - def _filter_supported_generate_kwargs( - *, - gen_fn: Any, - gen_kwargs: dict[str, Any], - fn_name: str, - ) -> dict[str, Any]: - if not gen_kwargs: - return gen_kwargs - - try: - signature = inspect.signature(gen_fn) - except (TypeError, ValueError): - return gen_kwargs - - params = signature.parameters - accepts_var_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) - if accepts_var_kwargs: - return gen_kwargs - - allowed_keys = { - name - for name, param in params.items() - if param.kind - in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - ) - } - filtered = {k: v for k, v in gen_kwargs.items() if k in allowed_keys} - - removed_keys = sorted(set(gen_kwargs.keys()) - set(filtered.keys())) - if removed_keys: - logger.debug("Filtered unsupported kwargs for %s: %s", fn_name, removed_keys) - - return filtered - - @staticmethod - def _call_generate_fn( - *, - gen_fn: Any, - input_ids: torch.Tensor, - gen_kwargs: dict[str, Any], - ) -> Any: - try: - signature = inspect.signature(gen_fn) - params = signature.parameters - except (TypeError, ValueError): - params = {} - - if "idx" in params: - return gen_fn(idx=input_ids, **gen_kwargs) - if "input_ids" in params: - return gen_fn(input_ids=input_ids, **gen_kwargs) - - try: - return gen_fn(input_ids, **gen_kwargs) - except TypeError: - try: - return gen_fn(idx=input_ids, **gen_kwargs) - except TypeError: - return gen_fn(input_ids=input_ids, **gen_kwargs) - - def _normalize_generate_inputs( - self, - *, - input_ids: torch.Tensor, - gen_kwargs: dict[str, Any], - ref_device: torch.device, - ) -> tuple[torch.Tensor, dict[str, Any]]: - normalized_input_ids = self._coerce_long_tensor_2d(input_ids, ref_device) - if normalized_input_ids is None: - normalized_input_ids = input_ids - - normalized_kwargs = dict(gen_kwargs) - for key in ("attention_mask", "uncond_input_ids", "uncond_attention_mask"): - if key not in normalized_kwargs: - continue - normalized_value = self._coerce_long_tensor_2d( - normalized_kwargs[key], - ref_device, - ) - if normalized_value is not None: - normalized_kwargs[key] = normalized_value - - return normalized_input_ids, normalized_kwargs - - def _get_or_create_uni_prompting( - self, - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - ) -> Any | None: - runtime_uni_prompting = runtime_info.get("uni_prompting") - if runtime_uni_prompting is not None: - runtime_uni_prompting = self._unwrap_singleton(runtime_uni_prompting) - if runtime_uni_prompting is not None: - return runtime_uni_prompting - - kwargs_uni_prompting = self._unwrap_singleton(kwargs.get("uni_prompting")) - if kwargs_uni_prompting is not None: - return kwargs_uni_prompting - - self._maybe_load_runtime_tokenizer(runtime_info) - if self.tokenizer is None: - return None - - use_reserved_token = _to_bool( - unwrap_first_value( - runtime_info.get("use_reserved_token"), - unwrap_first_value(runtime_info.get("prompting_use_reserved_token"), True), - ), - default=True, - ) - - max_text_len_value = unwrap_first_value( - runtime_info.get("prompt_max_text_len"), - unwrap_first_value( - runtime_info.get("prompting_max_text_len"), - unwrap_first_value(runtime_info.get("max_text_len"), None), - ), - ) - cond_dropout_value = unwrap_first_value( - runtime_info.get("cond_dropout_prob"), - unwrap_first_value(runtime_info.get("prompting_cond_dropout_prob"), None), - ) - max_audio_len_value = unwrap_first_value( - runtime_info.get("max_audio_len"), - unwrap_first_value(runtime_info.get("t2s_token_length"), None), - ) - max_audio_len_short_value = unwrap_first_value( - runtime_info.get("max_audio_len_short"), - None, - ) - - max_text_len: int | None = None - if max_text_len_value is not None: - try: - parsed = int(max_text_len_value) - if parsed > 0: - max_text_len = parsed - except Exception: - pass - - cond_dropout_prob: float | None = None - if cond_dropout_value is not None: - try: - cond_dropout_prob = float(cond_dropout_value) - except Exception: - pass - - max_audio_len: int | None = None - if max_audio_len_value is not None: - try: - parsed = int(max_audio_len_value) - if parsed > 0: - max_audio_len = max(parsed, 512) - except Exception: - pass - - max_audio_len_short: int | None = None - if max_audio_len_short_value is not None: - try: - parsed = int(max_audio_len_short_value) - if parsed > 0: - max_audio_len_short = parsed - except Exception: - pass - elif max_audio_len is not None: - max_audio_len_short = max(256, max_audio_len // 2) - - if self._uni_prompting is not None: - if max_text_len is None and hasattr(self._uni_prompting, "max_text_len"): - try: - existing_max_text_len = int(getattr(self._uni_prompting, "max_text_len")) - if existing_max_text_len > 0: - max_text_len = existing_max_text_len - 1 - except Exception: - pass - if cond_dropout_prob is None and hasattr(self._uni_prompting, "cond_dropout_prob"): - try: - cond_dropout_prob = float(getattr(self._uni_prompting, "cond_dropout_prob")) - except Exception: - pass - - desired_spec = ( - id(self.tokenizer), - use_reserved_token, - max_text_len, - cond_dropout_prob, - max_audio_len, - max_audio_len_short, - ) - - if self._uni_prompting is not None and self._uni_prompting_init_spec != desired_spec: - self._reset_uni_prompting_cache() - - if self._uni_prompting is None: - try: - universal_prompting_cls = resolve_remote_attr( - "UniversalPrompting", - module_name="prompting_utils", - settings=DYNIN_REMOTE_SETTINGS, - source=self._infer_sources.model_source, - local_files_only=self._infer_sources.model_local_files_only, - fallback_module_names=("modeling_dynin_omni",), - optional=True, - ) - except Exception: - universal_prompting_cls = None - - try: - if universal_prompting_cls is None: - raise ImportError("UniversalPrompting is not available in the configured remote Dynin code.") - - init_kwargs: dict[str, Any] = { - "use_reserved_token": use_reserved_token, - "special_tokens": DYNIN_SPECIAL_TOKENS, - "ignore_id": -100, - } - if max_text_len is not None: - init_kwargs["max_text_len"] = max_text_len - if cond_dropout_prob is not None: - init_kwargs["cond_dropout_prob"] = cond_dropout_prob - if max_audio_len is not None: - init_kwargs["max_audio_len"] = max_audio_len - if max_audio_len_short is not None: - init_kwargs["max_audio_len_short"] = max_audio_len_short - - try: - self._uni_prompting = universal_prompting_cls(self.tokenizer, **init_kwargs) - except TypeError: - trimmed_audio_kwargs = dict(init_kwargs) - trimmed_audio_kwargs.pop("max_audio_len", None) - trimmed_audio_kwargs.pop("max_audio_len_short", None) - try: - self._uni_prompting = universal_prompting_cls(self.tokenizer, **trimmed_audio_kwargs) - except TypeError: - minimal_kwargs = dict(trimmed_audio_kwargs) - minimal_kwargs.pop("special_tokens", None) - minimal_kwargs.pop("ignore_id", None) - self._uni_prompting = universal_prompting_cls(self.tokenizer, **minimal_kwargs) - self._uni_prompting_init_spec = desired_spec - except Exception as e: - logger.warning("Failed to initialize UniversalPrompting: %s", e) - self._reset_uni_prompting_cache() - - return self._uni_prompting - - @staticmethod - def _unwrap_singleton(value: Any) -> Any: - if isinstance(value, list) and len(value) == 1: - return value[0] - return value - - @classmethod - def _coerce_schedule_params(cls, value: Any) -> dict[str, Any]: - value = cls._unwrap_singleton(value) - if value is None: - return {} - if isinstance(value, dict): - return {str(k): v for k, v in value.items()} - if hasattr(value, "items"): - try: - return {str(k): v for k, v in dict(value).items()} - except Exception: - return {} - if isinstance(value, str): - text = value.strip() - if not text: - return {} - try: - parsed = json.loads(text) - except Exception: - return {} - if isinstance(parsed, dict): - return {str(k): v for k, v in parsed.items()} - return {} - - def _resolve_noise_schedule( - self, - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - ) -> Any | None: - runtime_noise_schedule = unwrap_first_value( - runtime_info.get("noise_schedule"), - kwargs.get("noise_schedule"), - ) - runtime_noise_schedule = self._unwrap_singleton(runtime_noise_schedule) - if callable(runtime_noise_schedule): - return runtime_noise_schedule - - schedule_name: str | None = None - if isinstance(runtime_noise_schedule, str) and runtime_noise_schedule.strip(): - schedule_name = runtime_noise_schedule.strip() - - if schedule_name is None: - for key in ("noise_schedule_name", "mask_schedule", "schedule"): - value = unwrap_first_value(runtime_info.get(key), None) - if value is None and key in kwargs: - value = self._unwrap_singleton(kwargs.get(key)) - if isinstance(value, str) and value.strip(): - schedule_name = value.strip() - break - - if schedule_name is None: - return None - - schedule_params = self._coerce_schedule_params( - unwrap_first_value( - runtime_info.get("noise_schedule_params"), - kwargs.get("noise_schedule_params"), - ) - ) - - try: - get_mask_schedule = get_dynin_sampling_attr("get_mask_schedule") - return get_mask_schedule(schedule_name, **schedule_params) - except Exception as e: - logger.warning( - "Failed to resolve mask schedule '%s' with params=%s: %s", - schedule_name, - schedule_params, - e, - ) - return None - - @staticmethod - def _coerce_long_tensor_2d( - value: Any, - device: torch.device, - ) -> torch.Tensor | None: - if value is None: - return None - out = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) - if out.ndim == 1: - out = out.unsqueeze(0) - if out.ndim > 2: - out = out.view(out.shape[0], -1) - return out.to(device=device, dtype=torch.long).contiguous() - - @staticmethod - def _config_get(config_obj: Any, key: str) -> Any: - if config_obj is None: - return None - if isinstance(config_obj, dict): - return config_obj.get(key) - if hasattr(config_obj, "get"): - try: - return config_obj.get(key) - except Exception: - return None - return None - - @classmethod - def _is_numeric_token_structure(cls, value: Any) -> bool: - if isinstance(value, torch.Tensor): - return True - if isinstance(value, bool): - return True - if isinstance(value, int): - return True - if isinstance(value, float): - return float(value).is_integer() - if isinstance(value, (list, tuple)): - if not value: - return False - return all(cls._is_numeric_token_structure(v) for v in value) - return False - - @classmethod - def _materialize_prompting_payload(cls, value: Any, ref_device: torch.device) -> Any: - if isinstance(value, torch.Tensor): - return value.to(device=ref_device, dtype=torch.long).contiguous() - if isinstance(value, dict): - return {k: cls._materialize_prompting_payload(v, ref_device) for k, v in value.items()} - if isinstance(value, (list, tuple)): - if cls._is_numeric_token_structure(value): - try: - return torch.as_tensor(value, dtype=torch.long, device=ref_device) - except Exception: - pass - converted = [cls._materialize_prompting_payload(v, ref_device) for v in value] - return tuple(converted) if isinstance(value, tuple) else converted - return value - - @contextmanager - def _temporary_prompting_overrides(self, uni_prompting: Any, prompting_cfg: Any): - restore_values: dict[str, Any] = {} - try: - max_text_len_override = self._config_get(prompting_cfg, "max_text_len_override") - if max_text_len_override is not None and hasattr(uni_prompting, "max_text_len"): - try: - override_int = int(max_text_len_override) - if override_int > 0: - restore_values["max_text_len"] = getattr(uni_prompting, "max_text_len") - setattr(uni_prompting, "max_text_len", override_int + 1) - except Exception: - pass - yield - finally: - for attr_name, original_value in restore_values.items(): - try: - setattr(uni_prompting, attr_name, original_value) - except Exception: - pass - - def _prepare_prompting_input( - self, - *, - payload: Any, - task: str, - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - uni_prompting: Any, - ref_device: torch.device, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - if payload is None: - return None, None - - payload = self._unwrap_singleton(payload) - prompting_task = str( - self._unwrap_singleton( - unwrap_first_value( - runtime_info.get("prompting_task"), - TASK_TO_PROMPTING_TASK.get(task, task), - ) - ) - ) - prompting_cfg = self._unwrap_singleton( - unwrap_first_value( - runtime_info.get("prompting_config"), - kwargs.get("prompting_config"), - ) - ) - - if isinstance(payload, dict): - if payload.get("task") is not None: - prompting_task = str(payload["task"]) - if payload.get("config") is not None: - prompting_cfg = payload["config"] - payload = payload.get("input", payload.get("inputs", payload.get("data", payload))) - - payload = self._materialize_prompting_payload(payload, ref_device) - - try: - with self._temporary_prompting_overrides(uni_prompting, prompting_cfg): - prepared = uni_prompting(payload, prompting_task, config=prompting_cfg) - except Exception as e: - logger.warning( - "UniversalPrompting failed for task=%s prompting_task=%s: %s", - task, - prompting_task, - e, - ) - return None, None - - if isinstance(prepared, tuple): - prepared_input_ids = prepared[0] if len(prepared) > 0 else None - prepared_attention_mask = prepared[1] if len(prepared) > 1 else None - else: - prepared_input_ids = prepared - prepared_attention_mask = None - - return ( - self._coerce_long_tensor_2d(prepared_input_ids, ref_device), - self._coerce_long_tensor_2d(prepared_attention_mask, ref_device), - ) - - def _prepare_prompting_inputs_if_needed( - self, - *, - task: str, - input_ids: torch.Tensor, - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - gen_kwargs: dict[str, Any], - ) -> tuple[torch.Tensor, dict[str, Any]]: - uni_prompting = gen_kwargs.get("uni_prompting") - if uni_prompting is None: - uni_prompting = self._get_or_create_uni_prompting( - runtime_info=runtime_info, - kwargs=kwargs, - ) - if uni_prompting is not None: - gen_kwargs["uni_prompting"] = uni_prompting - - if uni_prompting is None: - return input_ids, gen_kwargs - - payload = self._find_first_payload( - runtime_info=runtime_info, - kwargs=kwargs, - keys=PROMPTING_PAYLOAD_KEYS, - ) - - if payload is not None: - prepared_input_ids, prepared_attention_mask = self._prepare_prompting_input( - payload=payload, - task=task, - runtime_info=runtime_info, - kwargs=kwargs, - uni_prompting=uni_prompting, - ref_device=input_ids.device, - ) - if prepared_input_ids is not None: - input_ids = prepared_input_ids - if prepared_attention_mask is not None and "attention_mask" not in gen_kwargs: - gen_kwargs["attention_mask"] = prepared_attention_mask - - uncond_payload = self._find_first_payload( - runtime_info=runtime_info, - kwargs=kwargs, - keys=UNCOND_PROMPTING_PAYLOAD_KEYS, - ) - if uncond_payload is not None and "uncond_input_ids" not in gen_kwargs: - uncond_input_ids, uncond_attention_mask = self._prepare_prompting_input( - payload=uncond_payload, - task=task, - runtime_info=runtime_info, - kwargs=kwargs, - uni_prompting=uni_prompting, - ref_device=input_ids.device, - ) - if uncond_input_ids is not None: - gen_kwargs["uncond_input_ids"] = uncond_input_ids - if uncond_attention_mask is not None and "uncond_attention_mask" not in gen_kwargs: - gen_kwargs["uncond_attention_mask"] = uncond_attention_mask - - return input_ids, gen_kwargs - - @staticmethod - def _find_first_payload( - *, - runtime_info: dict[str, Any], - kwargs: dict[str, Any], - keys: tuple[str, ...], - ) -> Any | None: - for key in keys: - if key in runtime_info: - return runtime_info[key] - if key in kwargs: - return kwargs[key] - return None - - def _extract_decode_tokens( - self, - tokens: torch.Tensor, - runtime_info: dict[str, Any], - ) -> torch.Tensor: - prompt_len = int( - unwrap_first_value( - runtime_info.get("prompt_length"), - unwrap_first_value( - runtime_info.get("prompt_len"), - unwrap_first_value(runtime_info.get("prompt_token_len"), 0), - ), - ) - ) - - decode_tokens = tokens - if 0 < prompt_len < tokens.numel(): - decode_tokens = tokens[prompt_len:] - - text_vocab_size = unwrap_first_value(runtime_info.get("text_vocab_size"), None) - if text_vocab_size is None and self.tokenizer is not None: - text_vocab_size = len(self.tokenizer) - - if text_vocab_size is not None: - vocab_size = int(text_vocab_size) - valid = decode_tokens[(decode_tokens >= 0) & (decode_tokens < vocab_size)] - if valid.numel() > 0: - decode_tokens = valid - - return decode_tokens.contiguous() - - def _decode_text(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> str: - self._maybe_load_runtime_tokenizer(runtime_info) - if self.tokenizer is None: - return "" - try: - return self.tokenizer.decode( - tokens.detach().cpu().tolist(), - skip_special_tokens=True, - ) - except Exception: - return "" - - def _maybe_load_runtime_tokenizer(self, runtime_info: dict[str, Any]) -> None: - tokenizer_path = unwrap_first_value(runtime_info.get("tokenizer_path"), None) - if tokenizer_path is not None: - tokenizer_path = str(tokenizer_path) - - runtime_local_files_only = unwrap_first_value( - runtime_info.get("local_files_only_model"), - unwrap_first_value( - runtime_info.get("model_local_files_only"), - unwrap_first_value( - runtime_info.get("local_files_only"), - self._infer_sources.model_local_files_only, - ), - ), - ) - local_only = _to_bool( - runtime_local_files_only, - default=self._infer_sources.model_local_files_only, - ) - - if tokenizer_path and tokenizer_path != self._tokenizer_path: - try: - logger.info("Loading DYNIN text tokenizer from %s", tokenizer_path) - self._set_tokenizer(tokenizer_path, local_files_only=local_only) - except Exception as e: - logger.warning("Failed to load tokenizer from %s: %s", tokenizer_path, e) - - def _ensure_prompt_vq_model(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any: - sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info) - model_path = str(sources.vq_image_source) - local_files_only = bool(sources.vq_image_local_files_only) - if ( - self._prompt_vq_model is None - or self._prompt_vq_model_path != model_path - or self._prompt_vq_local_files_only != local_files_only - ): - logger.info( - "Loading DYNIN prompt VQ encoder from %s (local_files_only=%s)", - model_path, - local_files_only, - ) - magvit_cls = get_dynin_magvit_attr( - "MAGVITv2", - source=model_path, - local_files_only=local_files_only, - ) - try: - self._prompt_vq_model = magvit_cls.from_pretrained( - model_path, - local_files_only=local_files_only, - ) - except TypeError: - self._prompt_vq_model = magvit_cls.from_pretrained(model_path) - self._prompt_vq_model.eval() - self._prompt_vq_model.requires_grad_(False) - self._prompt_vq_model_path = model_path - self._prompt_vq_local_files_only = local_files_only - if hasattr(self._prompt_vq_model, "to"): - self._prompt_vq_model = self._prompt_vq_model.to(ref_device) - return self._prompt_vq_model - - @staticmethod - def _prepare_prompt_image_tensor( - image: Any, - *, - resolution: int, - device: torch.device, - ) -> torch.Tensor: - tensor = image if isinstance(image, torch.Tensor) else torch.as_tensor(image) - if tensor.ndim == 4: - tensor = tensor[0] - if tensor.ndim != 3: - raise ValueError(f"Unsupported image tensor shape for Dynin bootstrap: {tuple(tensor.shape)}") - - if tensor.shape[0] not in (1, 3, 4) and tensor.shape[-1] in (1, 3, 4): - tensor = tensor.permute(2, 0, 1) - if tensor.shape[0] == 1: - tensor = tensor.repeat(3, 1, 1) - if tensor.shape[0] == 4: - tensor = tensor[:3] - - tensor = tensor.to(device=device, dtype=torch.float32) - if tensor.numel() > 0 and tensor.max() > 1.0: - tensor = tensor / 255.0 - - tensor = tensor.unsqueeze(0) - _, _, height, width = tensor.shape - short_side = max(1, min(int(height), int(width))) - scale = float(resolution) / float(short_side) - new_height = max(1, int(round(height * scale))) - new_width = max(1, int(round(width * scale))) - tensor = F.interpolate( - tensor, - size=(new_height, new_width), - mode="bicubic", - align_corners=False, - ) - top = max(0, (new_height - resolution) // 2) - left = max(0, (new_width - resolution) // 2) - tensor = tensor[:, :, top : top + resolution, left : left + resolution] - if tensor.shape[-2:] != (resolution, resolution): - tensor = F.interpolate( - tensor, - size=(resolution, resolution), - mode="bicubic", - align_corners=False, - ) - tensor = torch.clamp(tensor, min=0.0, max=1.0) - return ((tensor - 0.5) / 0.5).contiguous() - - def _encode_prompt_image_tokens( - self, - *, - runtime_info: dict[str, Any], - mm_inputs: dict[str, Any], - resolution: int, - ) -> torch.Tensor: - image_value = mm_inputs.get("image") - image_items = self._split_mm_items(image_value) - if not image_items: - raise ValueError("Dynin online i2i bootstrap requires an image input.") - - device = self._default_mm_device() - image_tensor = self._prepare_prompt_image_tensor( - image_items[0], - resolution=resolution, - device=device, - ) - vq_model = self._ensure_prompt_vq_model(runtime_info=runtime_info, ref_device=device) - with torch.no_grad(): - token_ids = vq_model.get_code(image_tensor) - token_ids = torch.as_tensor(token_ids, dtype=torch.long).detach().cpu() - if token_ids.ndim == 2 and token_ids.shape[0] == 1: - token_ids = token_ids[0] - return token_ids.contiguous() - - @staticmethod - def _split_mm_items(value: Any) -> list[Any]: - if value is None: - return [] - if isinstance(value, torch.Tensor): - if value.ndim == 0: - return [value] - return [value[i] for i in range(value.shape[0])] - if isinstance(value, list): - return value - if isinstance(value, tuple): - if len(value) == 2 and isinstance(value[1], (int, float)): - return [value] - return list(value) - return [value] - - def _default_mm_device(self) -> torch.device: - try: - return next(self.model.parameters()).device - except StopIteration: - return torch.device("cpu") - - @staticmethod - def _coerce_mm_item_to_float_tensor( - item: Any, - *, - device: torch.device, - ) -> torch.Tensor: - if isinstance(item, tuple) and len(item) == 2 and isinstance(item[1], (int, float)): - item = item[0] - - if isinstance(item, torch.Tensor): - tensor = item.detach().to(device=device, dtype=torch.float32) - else: - tensor = torch.as_tensor(item, dtype=torch.float32, device=device) - - return tensor.contiguous() - - def _build_deterministic_mm_embedding( - self, - item: Any, - *, - device: torch.device, - ) -> torch.Tensor: - tensor = self._coerce_mm_item_to_float_tensor(item, device=device) - if tensor.numel() == 0: - return torch.zeros((1, self.hidden_size), dtype=torch.bfloat16, device=device) - - flattened = tensor.view(-1) - first = flattened[0] - last = flattened[-1] - mean = flattened.mean() - std = flattened.std(unbiased=False) - abs_mean = flattened.abs().mean() - max_abs = flattened.abs().max() - l2 = torch.linalg.vector_norm(flattened) / max(float(flattened.numel()), 1.0) - - base = torch.stack([first, last, mean, std, abs_mean, max_abs, l2], dim=0) - denom = torch.clamp(base.abs().max(), min=1.0) - base = base / denom - - repeats = (self.hidden_size + base.numel() - 1) // base.numel() - embedding = base.repeat(repeats)[: self.hidden_size].to(dtype=torch.bfloat16) - return embedding.unsqueeze(0).contiguous() - - def _collect_mm_inputs(self, **kwargs: Any) -> dict[str, Any]: - mm_inputs: dict[str, Any] = {} - for modality, aliases in MM_INPUT_ALIASES.items(): - for alias in aliases: - if alias in kwargs and kwargs[alias] is not None: - mm_inputs[modality] = kwargs[alias] - break - for modality, value in self._cached_mm_inputs.items(): - if modality not in mm_inputs and value is not None: - mm_inputs[modality] = value - return mm_inputs - - def embed_multimodal(self, **kwargs: Any) -> Any: - mm_inputs = self._collect_mm_inputs(**kwargs) - self._cached_mm_inputs = dict(mm_inputs) - if not mm_inputs: - return None - - device = self._default_mm_device() - mm_embeddings: list[torch.Tensor] = [] - - for modality in ("image", "video", "audio"): - value = mm_inputs.get(modality) - if value is None: - continue - for item in self._split_mm_items(value): - mm_embeddings.append(self._build_deterministic_mm_embedding(item, device=device)) - - return tuple(mm_embeddings) if mm_embeddings else None diff --git a/vllm_omni/model_executor/models/dynin_omni/pipeline.py b/vllm_omni/model_executor/models/dynin_omni/pipeline.py deleted file mode 100644 index 6f59ebb36..000000000 --- a/vllm_omni/model_executor/models/dynin_omni/pipeline.py +++ /dev/null @@ -1,67 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Dynin-Omni pipeline topology (frozen). - -Stage 0: token2text — multimodal understanding / text generation (comprehension) -Stage 1: token2image — stage-0 tokens → image latents -Stage 2: token2audio — stage-1 tokens → audio latents - -All three stages run on the generation worker (``LLM_GENERATION``); the -inter-stage hand-off uses the worker-connector full-payload data plane -(``*_full_payload`` producers + ``token2text_to_token2image`` / -``token2image_to_token2audio`` consumers). Deploy knobs (devices, GPU memory, -batched tokens, connectors) live in ``vllm_omni/deploy/dynin_omni*.yaml``. -""" - -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) - -_PROC = "vllm_omni.model_executor.stage_input_processors.dynin_omni" - -DYNIN_OMNI_PIPELINE = PipelineConfig( - model_type="dynin_omni", - default_deploy_config_name="dynin_omni.yaml", - model_arch="DyninOmniForConditionalGeneration", - # Arch-fallback safety net: route by hf_config.architectures when the - # auto-detected model_type does not match the registry key exactly. - hf_architectures=("DyninOmniForConditionalGeneration",), - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="token2text", - execution_type=StageExecutionType.LLM_GENERATION, - input_sources=(), - final_output=True, - final_output_type="text", - owns_tokenizer=True, - engine_output_type="latent", - custom_process_next_stage_input_func=f"{_PROC}.token2text_to_token2image_full_payload", - ), - StagePipelineConfig( - stage_id=1, - model_stage="token2image", - execution_type=StageExecutionType.LLM_GENERATION, - input_sources=(0,), - final_output=True, - final_output_type="image", - engine_output_type="latent", - custom_process_input_func=f"{_PROC}.token2text_to_token2image", - custom_process_next_stage_input_func=f"{_PROC}.token2image_to_token2audio_full_payload", - requires_full_payload_input=True, - ), - StagePipelineConfig( - stage_id=2, - model_stage="token2audio", - execution_type=StageExecutionType.LLM_GENERATION, - input_sources=(1,), - final_output=True, - final_output_type="audio", - engine_output_type="latent", - custom_process_input_func=f"{_PROC}.token2image_to_token2audio", - requires_full_payload_input=True, - ), - ), -) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 28accbc51..30fc504b3 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -171,10 +171,9 @@ class Qwen3OmniMoeForConditionalGeneration( # such as the vLLM-text perf benchmark. # # Default to the thinker: it is the text-generation stage, which is what - # a non-staged run of this model is asking for. Same shape as - # dynin_omni ("token2text") and glm_tts ("glm_tts"). An explicitly wrong - # value still reaches the ValueError below rather than being silently - # accepted. + # a non-staged run of this model is asking for. Same shape as glm_tts + # ("glm_tts"). An explicitly wrong value still reaches the ValueError + # below rather than being silently accepted. self.model_stage = getattr(vllm_config.model_config, "model_stage", None) or "thinker" # Staged startup always injects model_stage; its absence means a plain # vLLM run with no talker stage downstream, so no one consumes captured diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index d1638ae9d..468ed5302 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -321,12 +321,6 @@ _OMNI_MODELS = { "voxcpm2_talker", "VoxCPM2TalkerForConditionalGeneration", ), - ## dots.tts - "DotsTTSForConditionalGeneration": ( - "dots_tts", - "dots_tts_talker", - "DotsTTSForConditionalGeneration", - ), ## Voxtral TTS "VoxtralTTSForConditionalGeneration": ( "voxtral_tts", @@ -396,11 +390,6 @@ _OMNI_MODELS = { "modeling_moss_tts_codec", "MossTTSCodecDecoder", ), - "DyninOmniForConditionalGeneration": ( - "dynin_omni", - "dynin_omni", - "DyninOmniForConditionalGeneration", - ), ## IndexTTS2 "IndexTTS2TalkerForConditionalGeneration": ( "indextts2", diff --git a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py b/vllm_omni/model_executor/stage_input_processors/dynin_omni.py deleted file mode 100644 index 5a5804de3..000000000 --- a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py +++ /dev/null @@ -1,276 +0,0 @@ -from __future__ import annotations - -import json -from typing import Any - -import torch -from vllm.inputs import TextPrompt -from vllm.logger import init_logger - -from vllm_omni.inputs.data import OmniTokensPrompt - -logger = init_logger(__name__) - - -def _to_prompt_dict(prompt_item: OmniTokensPrompt | TextPrompt | str | None) -> dict[str, Any]: - if isinstance(prompt_item, dict): - return prompt_item - return {} - - -def _to_token_id_list(value: Any) -> list[int]: - if isinstance(value, torch.Tensor): - value = value.detach().to("cpu") - if value.ndim == 0: - return [int(value.item())] - if value.ndim > 1: - value = value[0] - return [int(x) for x in value.tolist()] - if isinstance(value, list): - if not value: - return [] - if isinstance(value[0], list): - return [int(x) for x in value[0]] - return [int(x) for x in value] - if value is None: - return [] - return [int(value)] - - -def _to_int(value: Any, default: int = 0) -> int: - if isinstance(value, torch.Tensor): - if value.numel() == 0: - return default - return int(value.view(-1)[0].item()) - if isinstance(value, list): - if not value: - return default - return int(value[0]) - if value is None: - return default - return int(value) - - -def _normalize_additional_info(value: Any) -> dict[str, Any]: - if not isinstance(value, dict): - return {} - normalized: dict[str, Any] = {} - for key, val in value.items(): - if isinstance(val, list): - normalized[key] = val - else: - normalized[key] = [val] - return normalized - - -def _decode_runtime_bridge_info(value: Any) -> dict[str, Any]: - if isinstance(value, torch.Tensor): - tensor = value.detach().to("cpu").reshape(-1).to(torch.uint8) - raw = bytes(tensor.tolist()) - elif isinstance(value, (bytes, bytearray)): - raw = bytes(value) - elif isinstance(value, list): - try: - raw = bytes(int(item) for item in value) - except Exception: - return {} - elif value is None: - return {} - else: - return value if isinstance(value, dict) else {} - - if not raw: - return {} - - try: - decoded = json.loads(raw.decode("utf-8")) - except Exception: - return {} - return decoded if isinstance(decoded, dict) else {} - - -def _bridge_tokens( - source_outputs, - prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = False, -): - next_inputs = [] - if not isinstance(prompt, list): - prompt = [prompt] - - prompt_meta_by_reqid = {src_out.request_id: _to_prompt_dict(p) for src_out, p in zip(source_outputs, prompt)} - - for source_output in source_outputs: - output = source_output.outputs[0] - mm_out = getattr(output, "multimodal_output", None) or {} - - token_ids = _to_token_id_list(mm_out.get("token_ids")) - if not token_ids: - token_ids = _to_token_id_list(mm_out.get("text_tokens")) - if not token_ids: - token_ids = list(output.cumulative_token_ids or []) - if not token_ids: - raise RuntimeError(f"Stage output for request {source_output.request_id} has no token_ids") - - detok_id = _to_int(mm_out.get("detok_id"), default=0) - src_prompt = prompt_meta_by_reqid.get(source_output.request_id, {}) - src_additional_info = src_prompt.get("additional_information", {}) or {} - runtime_bridge_info = _decode_runtime_bridge_info(mm_out.get("runtime_info_json")) - if not runtime_bridge_info: - runtime_bridge_info = mm_out.get("runtime_info", {}) or {} - - additional_information: dict[str, Any] = _normalize_additional_info(src_additional_info) - additional_information.update(_normalize_additional_info(runtime_bridge_info)) - additional_information["detok_id"] = [detok_id] - - next_inputs.append( - OmniTokensPrompt( - prompt_token_ids=token_ids, - additional_information=additional_information, - multi_modal_data=(src_prompt.get("multi_modal_data") if requires_multimodal_data else None), - mm_processor_kwargs=None, - ) - ) - - return next_inputs - - -def token2text_to_token2image( - source_outputs, - prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = False, -): - return _bridge_tokens(source_outputs, prompt, requires_multimodal_data) - - -def token2image_to_token2audio( - source_outputs, - prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = False, -): - return _bridge_tokens(source_outputs, prompt, requires_multimodal_data) - - -# ============================================================================ -# Worker-connector data plane (non-async-chunk path). -# ============================================================================ - -# Per-model REPLACE-keys for the full-payload accumulator. dynin_omni's -# producer model emits new chunks per step (token_ids / runtime_info_json), -# all of which use the default CONCAT/replace semantics — no model_outputs -# entry needs explicit REPLACE. -_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset() - - -def _build_full_payload(pooling_output: dict[str, Any] | None, request: Any) -> dict[str, Any] | None: - """Producer-side payload builder: assemble dynin_omni connector payload. - - Reads token_ids from ``pooling_output["token_ids"]`` (preferred) or - ``request.output_token_ids`` (fallback). Reads structured non-tensor - metadata from ``pooling_output["runtime_info_json"]`` (JSON-in-uint8) - if present, falling back to ``pooling_output["runtime_info"]`` dict. - Carries forward ``request.additional_information`` so prompt-side - metadata (speaker / language / detok_id) survives the IPC boundary. - """ - if not isinstance(pooling_output, dict): - pooling_output = {} - - token_ids = _to_token_id_list(pooling_output.get("token_ids")) - if not token_ids: - token_ids = _to_token_id_list(pooling_output.get("text_tokens")) - if not token_ids and request is not None: - token_ids = _to_token_id_list(getattr(request, "output_token_ids", None)) - if not token_ids: - logger.warning( - "dynin_omni._build_full_payload: no token_ids found in pooling_output " - "(keys=%s) or request.output_token_ids for req=%s; consumer wait gate may hang.", - list(pooling_output.keys()), - getattr(request, "request_id", "?"), - ) - return None - - src_additional_info = getattr(request, "additional_information", {}) if request is not None else {} - if not isinstance(src_additional_info, dict): - src_additional_info = {} - - runtime_bridge_info = _decode_runtime_bridge_info(pooling_output.get("runtime_info_json")) - if not runtime_bridge_info: - runtime_bridge_info = pooling_output.get("runtime_info", {}) or {} - - payload = _normalize_additional_info(src_additional_info) - payload.update(_normalize_additional_info(runtime_bridge_info)) - payload["detok_id"] = [_to_int(pooling_output.get("detok_id"), default=_to_int(payload.get("detok_id"), default=0))] - # Use nested OmniPayload shape so the scheduling-metadata extractor in - # OmniConnectorModelRunnerMixin reads codes.audio and meta.finished - # (flat keys at the top level are silently dropped with a warning). - payload["codes"] = {"audio": token_ids} - payload["meta"] = {"finished": torch.tensor(True, dtype=torch.bool)} - return payload - - -def token2text_to_token2image_full_payload( - transfer_manager: Any, - pooling_output: dict[str, Any], - request: Any, -) -> dict[str, Any] | None: - """Producer-side payload builder for the Stage-0 → Stage-1 (text → image) transition.""" - del transfer_manager - return _build_full_payload(pooling_output, request) - - -def token2image_to_token2audio_full_payload( - transfer_manager: Any, - pooling_output: dict[str, Any], - request: Any, -) -> dict[str, Any] | None: - """Producer-side payload builder for the Stage-1 → Stage-2 (image → audio) transition.""" - del transfer_manager - return _build_full_payload(pooling_output, request) - - -def _token_only_from_source(source_outputs: list[Any]) -> list[OmniTokensPrompt]: - """Length-only placeholder list mirroring ``_bridge_tokens`` token counts.""" - inputs: list[OmniTokensPrompt] = [] - for source_output in source_outputs: - output = source_output.outputs[0] - mm_out = getattr(output, "multimodal_output", None) or {} - token_ids = _to_token_id_list(mm_out.get("token_ids")) - if not token_ids: - token_ids = _to_token_id_list(mm_out.get("text_tokens")) - if not token_ids: - token_ids = list(getattr(output, "token_ids", []) or []) - if not token_ids: - token_ids = [0] - inputs.append( - OmniTokensPrompt( - prompt_token_ids=[0] * len(token_ids), - additional_information=None, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - return inputs - - -def token2text_to_token2image_token_only( - stage_list, - engine_input_source, - prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - """Sync-side placeholder for Stage-1 input (token2image).""" - source_stage_id = engine_input_source[0] if engine_input_source else 0 - source_outputs = stage_list[source_stage_id].engine_outputs - return _token_only_from_source(source_outputs) - - -def token2image_to_token2audio_token_only( - stage_list, - engine_input_source, - prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - """Sync-side placeholder for Stage-2 input (token2audio).""" - source_stage_id = engine_input_source[0] if engine_input_source else 0 - source_outputs = stage_list[source_stage_id].engine_outputs - return _token_only_from_source(source_outputs) diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index fdbfb4b60..63e38b01e 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -23,7 +23,6 @@ _CLASS_TO_MODULE: dict[str, str] = { "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", "GLMTTSConfig": "vllm_omni.transformers_utils.configs.glm_tts", "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2", - "DotsTTSConfig": "vllm_omni.transformers_utils.configs.dots_tts", "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts", "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3", "MiniMaxMusic3Config": "vllm_omni.transformers_utils.configs.minimax_music3", @@ -53,7 +52,6 @@ __all__ = [ "FishSpeechFastARConfig", "GLMTTSConfig", "VoxCPM2Config", - "DotsTTSConfig", "VoxtralTTSConfig", "CosyVoice3Config", "MiniMaxMusic3Config", @@ -86,7 +84,6 @@ def __dir__(): from vllm_omni.transformers_utils.configs import audio8_tts as _audio8_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import auk as _auk # noqa: F401, E402 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 -from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402 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 diff --git a/vllm_omni/transformers_utils/configs/dots_tts.py b/vllm_omni/transformers_utils/configs/dots_tts.py deleted file mode 100644 index 185d23ad7..000000000 --- a/vllm_omni/transformers_utils/configs/dots_tts.py +++ /dev/null @@ -1,102 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""HF config for dots.tts. - -dots.tts splits its config across two files in the HF checkpoint: - * ``config.json`` — top-level dots.tts fields (latent_dim / patch_size - / PatchEncoder / DiT / vocoder / etc.). - * ``llm_config.json`` — full Qwen2 LM config (vocab_size / hidden_size / - num_hidden_layers / rope_theta / etc.). - -``AutoConfig.from_pretrained`` only consumes ``config.json``: its dots.tts -sub-blocks (``vocoder`` / ``DiT`` / ``PatchEncoder``) arrive through -``**kwargs`` and land as attributes, which the talker's config factories -consume directly. The Qwen2 LM fields live in ``llm_config.json``, which -AutoConfig never reads — this class hoists them to top-level attributes -with defaults matching the released checkpoints (soar / base / mf all ship -the same Qwen2.5-1.5B backbone). vLLM's ``Qwen2Model.__init__`` then reads -``config.hidden_size`` / ``config.vocab_size`` / etc. through -``get_text_config()``, the same pattern voxcpm2 follows. - -Future work: read ``llm_config.json`` during ``from_pretrained`` so a -checkpoint with a different LM shape works without editing this file. - -Upstream reference: - https://huggingface.co/dots-studio/dots.tts-soar/blob/main/llm_config.json -""" - -from __future__ import annotations - -from transformers import AutoConfig -from transformers.configuration_utils import PretrainedConfig - - -class DotsTTSConfig(PretrainedConfig): - """Configuration for dots.tts integration (Qwen2.5-1.5B-Base AR backbone).""" - - model_type = "dots_tts" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - # -- top-level dots.tts params -- - lm_config: dict | None = None, - # -- LM defaults (overridden by lm_config dict if present) -- - # All values match dots.tts-soar's llm_config.json. - vocab_size: int = 151672, - hidden_size: int = 1536, - intermediate_size: int = 8960, - num_hidden_layers: int = 28, - num_attention_heads: int = 12, - num_key_value_heads: int = 2, - max_position_embeddings: int = 131072, - rms_norm_eps: float = 1e-6, - rope_theta: float = 1_000_000.0, - tie_word_embeddings: bool = True, - hidden_act: str = "silu", - bos_token_id: int = 151643, - eos_token_id: int = 151643, - **kwargs, - ): - super().__init__( - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) - # The dots.tts sub-blocks (vocoder / DiT / PatchEncoder / latent_dim - # / patch_size / ...) arrive via **kwargs and are stored as - # attributes by PretrainedConfig; the talker's config factories - # read them from there. - self.lm_config = lm_config or {} - - # -- Hoist LM parameters to top-level for vLLM's Qwen2Model -- - lm = self.lm_config - self.vocab_size = lm.get("vocab_size", vocab_size) - self.hidden_size = lm.get("hidden_size", hidden_size) - self.intermediate_size = lm.get("intermediate_size", intermediate_size) - self.num_hidden_layers = lm.get("num_hidden_layers", num_hidden_layers) - self.num_attention_heads = lm.get("num_attention_heads", num_attention_heads) - self.num_key_value_heads = lm.get("num_key_value_heads", num_key_value_heads) - self.max_position_embeddings = lm.get("max_position_embeddings", max_position_embeddings) - self.rms_norm_eps = lm.get("rms_norm_eps", rms_norm_eps) - self.rope_theta = lm.get("rope_theta", rope_theta) - self.hidden_act = lm.get("hidden_act", hidden_act) - self.head_dim = self.hidden_size // self.num_attention_heads - - # vLLM's Qwen2Model reads ``max_window_layers`` and checks against - # ``num_hidden_layers`` only when ``use_sliding_window`` is true. - # soar disables sliding window — set both to the same value so the - # assertion (if executed) stays vacuous. - self.use_sliding_window = lm.get("use_sliding_window", False) - self.sliding_window = lm.get("sliding_window", None) - self.max_window_layers = lm.get("max_window_layers", self.num_hidden_layers) - - def get_text_config(self, **kwargs): - """Return self — LM fields are hoisted to top-level attributes.""" - return self - - -AutoConfig.register("dots_tts", DotsTTSConfig) - -__all__ = ["DotsTTSConfig"]