[New Model]: XiaomiMiMo/MiMo-Audio-7B-Instruct support (#750)
Signed-off-by: wangyu31577 <wangyu31577@hundsun.com> Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com> Signed-off-by: hsliu <liuhongsheng4@huawei.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: GG-li <3226868735@qq.com> Signed-off-by: Sihao Li <111170255+GG-li@users.noreply.github.com> Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com> Signed-off-by: mxuax <mxuax@connect.ust.hk> Signed-off-by: Baoyuan Qi <qibaoyuan@126.com> Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: wuzhongjian <wuzhongjian_yewu@cmss.chinamobile.com> Signed-off-by: dongbo910220 <1275604947@qq.com> Signed-off-by: dongbo910220 <32610838+dongbo910220@users.noreply.github.com> Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: baoyuan qi <qibaoyuan@126.com> Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Prajwal A <prajwalanagani@gmail.com> Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Signed-off-by: 丁宁 <nndding@gmail.com> Signed-off-by: SHIJIN ZHANG <75300765+Dovis01@users.noreply.github.com> Signed-off-by: dingning<dingning7@xiaomi.com> Signed-off-by: dingning <dingning7@xiaomi.com> Signed-off-by: dingning <dingning@xiaomi.com> Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com> Co-authored-by: wangyu31577 <wangyu31577@hundsun.com> Co-authored-by: Zhang Shijin <zhangshijin@xiaomi.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Sihao Li <111170255+GG-li@users.noreply.github.com> Co-authored-by: XU Mingshi <91017482+mxuax@users.noreply.github.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: JohnJan <wuzhongjian_yewu@cmss.chinamobile.com> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Co-authored-by: dongbo910220 <32610838+dongbo910220@users.noreply.github.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Co-authored-by: Junhong Liu <ljh_lbj@163.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: shijin zhang <zsj1364226740@gmail.com> Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk> Co-authored-by: root <root@hk01dgx028.cm.cluster> Co-authored-by: Prajwal A <34590600+LawJarp-A@users.noreply.github.com> Co-authored-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Co-authored-by: dingning <dingning7@xiaomi.com> Co-authored-by: ning ding <nndding@gmail.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
23 files changed
+8041
No files matched your search
@@ -43,6 +43,7 @@ th {
|
||||
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
|
||||
|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
|
||||
|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
|
||||
|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
|
||||
|
||||
|
||||
## List of Supported Models for NPU
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
# MiMo-Audio Offline Inference
|
||||
|
||||
This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally.
|
||||
|
||||
## Model Overview
|
||||
|
||||
MiMo-Audio provides multiple task variants for audio understanding and generation:
|
||||
|
||||
- **tts_sft**: Basic text-to-speech generation from text input.
|
||||
- **tts_sft_with_instruct**: TTS generation with explicit voice style instructions.
|
||||
- **tts_sft_with_audio**: TTS generation with audio reference for voice cloning.
|
||||
- **tts_sft_with_natural_instruction**: TTS generation from natural language descriptions embedded in text.
|
||||
- **audio_trancribing_sft**: Transcribe audio to text (speech-to-text).
|
||||
- **audio_understanding_sft**: Understand and analyze audio content with text queries.
|
||||
- **audio_understanding_sft_with_thinking**: Audio understanding with reasoning chain.
|
||||
- **spoken_dialogue_sft_multiturn**: Multi-turn spoken dialogue with audio input/output.
|
||||
- **speech2text_dialogue_sft_multiturn**: Multi-turn dialogue converting speech to text.
|
||||
- **text_dialogue_sft_multiturn**: Multi-turn text-only dialogue.
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
### Environment Variables
|
||||
|
||||
The `MIMO_AUDIO_TOKENIZER_PATH` environment variable is mandatory due to the specialized architecture:
|
||||
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run a single sample for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft
|
||||
```
|
||||
|
||||
Run batch samples for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--num-prompts {batch_size}
|
||||
```
|
||||
|
||||
When enabling multi-batch processing, if the total number of tokens passed to the next stage exceeds the `max_model_len` value in the `mimo_audio.yaml` configuration file, you must also synchronously update the `max_position_embeddings` value in `MiMo-Audio-7B-Instruct/config.json` to match the modified value.
|
||||
|
||||
Generated audio files are saved to `output_audio/` by default. `--num-prompts` also can be used to all tasks below.
|
||||
|
||||
## Task Usage
|
||||
|
||||
### tts_sft (Basic Text-to-Speech)
|
||||
|
||||
Generate speech from text input:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--text "The weather is so nice today."
|
||||
```
|
||||
|
||||
### tts_sft_with_instruct (TTS with Voice Instructions)
|
||||
|
||||
Generate speech with explicit voice style instructions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_instruct \
|
||||
--text "The weather is so nice today." \
|
||||
--instruct "Speak happily in a child's voice"
|
||||
```
|
||||
|
||||
### tts_sft_with_audio (TTS with Audio Reference)
|
||||
|
||||
Generate speech using an audio reference for voice cloning:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_audio \
|
||||
--text "The weather is so nice today." \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
```
|
||||
|
||||
### tts_sft_with_natural_instruction (Natural Language TTS)
|
||||
|
||||
Generate speech from text containing natural voice descriptions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_natural_instruction \
|
||||
--text "In a panting young male voice, he said: I can't run anymore, wait for me!"
|
||||
```
|
||||
|
||||
### audio_trancribing_sft (Speech-to-Text)
|
||||
|
||||
Transcribe audio to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_trancribing_sft \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
```
|
||||
|
||||
### audio_understanding_sft (Audio Understanding)
|
||||
|
||||
Understand and analyze audio content with text queries:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft \
|
||||
--text "Summarize the audio." \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
```
|
||||
|
||||
### audio_understanding_sft_with_thinking (Audio Understanding with Reasoning)
|
||||
|
||||
Audio understanding with reasoning chain:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft_with_thinking \
|
||||
--text "Summarize the audio." \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
```
|
||||
|
||||
### spoken_dialogue_sft_multiturn (Multi-turn Spoken Dialogue)
|
||||
|
||||
Multi-turn dialogue with audio input and output:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type spoken_dialogue_sft_multiturn \
|
||||
--audio-path "./prompt_speech_zh_m.wav"
|
||||
```
|
||||
|
||||
Note: This task uses hardcoded audio files in the script. The audio files used in examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
|
||||
|
||||
### speech2text_dialogue_sft_multiturn (Speech-to-Text Dialogue)
|
||||
|
||||
Multi-turn dialogue converting speech to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type speech2text_dialogue_sft_multiturn
|
||||
```
|
||||
|
||||
Note: This task uses hardcoded audio files and message lists in the script.
|
||||
|
||||
### text_dialogue_sft_multiturn (Text Dialogue)
|
||||
|
||||
Multi-turn text-only dialogue:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type text_dialogue_sft_multiturn
|
||||
```
|
||||
|
||||
Note: This task uses hardcoded message lists in the script.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Audio dependencies (soundfile, librosa)
|
||||
|
||||
This example depends on **soundfile** (read/write WAV) and **librosa** (load audio including MP3). Install the project requirements first:
|
||||
|
||||
```bash
|
||||
pip install -r requirements/common.txt
|
||||
# or at least: pip install soundfile>=0.13.1 librosa>=0.11.0
|
||||
```
|
||||
|
||||
- **`soundfile` / libsndfile not found**
|
||||
`soundfile` uses the C library **libsndfile**. On Linux, install the system package before pip:
|
||||
- Debian/Ubuntu: `sudo apt-get install libsndfile1`
|
||||
- For development builds: `sudo apt-get install libsndfile1-dev`
|
||||
- Then: `pip install soundfile`
|
||||
|
||||
- **`librosa` fails to load MP3 or reports "No backend available"**
|
||||
Loading MP3 (e.g. in `spoken_dialogue_sft_multiturn` with `.mp3` files) uses **ffmpeg** as the backend. Install ffmpeg:
|
||||
- Debian/Ubuntu: `sudo apt-get install ffmpeg`
|
||||
- macOS: `brew install ffmpeg`
|
||||
|
||||
- **`ImportError: No module named 'soundfile'` or `ModuleNotFoundError: ... librosa`**
|
||||
Ensure you are in the same Python environment where vLLM Omni and the example dependencies are installed, and that `requirements/common.txt` (or the packages above) are installed.
|
||||
|
||||
### Tokenizer path
|
||||
|
||||
- **`MIMO_AUDIO_TOKENIZER_PATH` not set or model fails to find tokenizer**
|
||||
Export the tokenizer path before running:
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
```
|
||||
See [Environment Variables](#environment-variables) in Setup.
|
||||
|
||||
### Other
|
||||
|
||||
- If the model or stage config fails to load, check [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings.
|
||||
- For errors when reading/writing WAV (e.g. unsupported format), ensure input files are standard WAV/MP3 and that `soundfile` is linked to a working libsndfile (see above).
|
||||
|
||||
## Notes
|
||||
|
||||
- The script uses default model paths and audio files embedded in `end2end.py`. Update them if your local cache path differs.
|
||||
- Use `--output-dir` to change the output folder (default: `./output_audio`).
|
||||
- Use `--num-prompts` to generate multiple prompts in one run (default: 1).
|
||||
- Audio files used in multi-turn dialogue examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
|
||||
- The script supports various configuration options for initialization timeouts, batch timeouts, and shared memory thresholds. See `--help` for details.
|
||||
@@ -0,0 +1,443 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
This example shows how to use vLLM for running offline inference
|
||||
with the correct prompt format on MiMo-Audio-Omni.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
from typing import NamedTuple
|
||||
|
||||
import soundfile as sf
|
||||
from message_convert import (
|
||||
get_audio_data,
|
||||
get_audio_understanding_sft_prompt,
|
||||
get_s2t_dialogue_sft_multiturn_prompt,
|
||||
get_spoken_dialogue_sft_multiturn_prompt,
|
||||
get_text_dialogue_sft_multiturn_prompt,
|
||||
get_tts_sft_prompt,
|
||||
to_prompt,
|
||||
)
|
||||
from vllm import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
|
||||
SEED = 42
|
||||
|
||||
MAX_CODE2WAV_TOKENS = 18192 # Maximum tokens supported by code2wav model
|
||||
|
||||
|
||||
class QueryResult(NamedTuple):
|
||||
inputs: dict
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
|
||||
|
||||
def get_codes_query_from_json(codes_path: str) -> QueryResult:
|
||||
with open(codes_path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
if isinstance(data, list):
|
||||
code_final = data
|
||||
elif isinstance(data, dict) and "code_final" in data:
|
||||
code_final = data["code_final"]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported codes json format in {codes_path}.\n"
|
||||
"Expect a JSON list[int] or {{'code_final': list[int]}}."
|
||||
)
|
||||
|
||||
if not isinstance(code_final, list) or not all(isinstance(x, int) for x in code_final):
|
||||
raise ValueError("code_final must be a list[int].")
|
||||
|
||||
if len(code_final) > MAX_CODE2WAV_TOKENS:
|
||||
print(f"[Warn] code_final len={len(code_final)} > {MAX_CODE2WAV_TOKENS}, truncating.")
|
||||
code_final = code_final[:MAX_CODE2WAV_TOKENS]
|
||||
|
||||
return QueryResult(
|
||||
inputs=OmniTokensPrompt(
|
||||
prompt_token_ids=code_final,
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
),
|
||||
limit_mm_per_prompt={},
|
||||
)
|
||||
|
||||
|
||||
def get_tts_sft(
|
||||
text="The weather is so nice today.",
|
||||
instruct=None,
|
||||
read_text_only=True,
|
||||
prompt_speech=None,
|
||||
audio_list=None,
|
||||
):
|
||||
res = get_tts_sft_prompt(
|
||||
text,
|
||||
instruct=instruct,
|
||||
read_text_only=read_text_only,
|
||||
prompt_speech=prompt_speech,
|
||||
)
|
||||
|
||||
prompt = to_prompt(res)
|
||||
final_prompt = {
|
||||
"prompt": prompt,
|
||||
}
|
||||
if audio_list is not None:
|
||||
final_prompt.update(
|
||||
{
|
||||
"multi_modal_data": {
|
||||
"audio": audio_list,
|
||||
},
|
||||
}
|
||||
)
|
||||
return final_prompt
|
||||
|
||||
|
||||
def get_audio_understanding_sft(audio_path, text="", thinking=False, use_sostm=False):
|
||||
audio_list = []
|
||||
audio_list.append(get_audio_data(audio_path))
|
||||
res = get_audio_understanding_sft_prompt(
|
||||
input_speech=audio_path, input_text=text, thinking=thinking, use_sostm=use_sostm
|
||||
)
|
||||
prompt = to_prompt(res)
|
||||
final_prompt = {
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"audio": audio_list,
|
||||
},
|
||||
}
|
||||
return final_prompt
|
||||
|
||||
|
||||
def get_spoken_dialogue_sft_multiturn(message_list, system_prompt=None, ref_audio_path=None, audio_list=None):
|
||||
res = get_spoken_dialogue_sft_multiturn_prompt(
|
||||
message_list, system_prompt=system_prompt, prompt_speech=ref_audio_path
|
||||
)
|
||||
prompt = to_prompt(res)
|
||||
final_prompt = {
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"audio": audio_list,
|
||||
},
|
||||
}
|
||||
return final_prompt
|
||||
|
||||
|
||||
def get_speech2text_dialogue_sft_multiturn(message_list, thinking=False, audio_list=None):
|
||||
res = get_s2t_dialogue_sft_multiturn_prompt(
|
||||
message_list,
|
||||
thinking=thinking,
|
||||
)
|
||||
prompt = to_prompt(res)
|
||||
final_prompt = {
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"audio": audio_list,
|
||||
},
|
||||
}
|
||||
return final_prompt
|
||||
|
||||
|
||||
def get_text_dialogue_sft_multiturn(
|
||||
message_list,
|
||||
):
|
||||
res = get_text_dialogue_sft_multiturn_prompt(
|
||||
message_list,
|
||||
)
|
||||
prompt = to_prompt(res)
|
||||
final_prompt = {
|
||||
"prompt": prompt,
|
||||
}
|
||||
return final_prompt
|
||||
|
||||
|
||||
query_map = {
|
||||
"tts_sft": get_tts_sft,
|
||||
"tts_sft_with_instruct": get_tts_sft,
|
||||
"tts_sft_with_audio": get_tts_sft,
|
||||
"tts_sft_with_natural_instruction": get_tts_sft,
|
||||
"audio_trancribing_sft": get_audio_understanding_sft,
|
||||
"audio_understanding_sft": get_audio_understanding_sft,
|
||||
"audio_understanding_sft_with_thinking": get_audio_understanding_sft,
|
||||
"spoken_dialogue_sft_multiturn": get_spoken_dialogue_sft_multiturn,
|
||||
"speech2text_dialogue_sft_multiturn": get_speech2text_dialogue_sft_multiturn,
|
||||
"text_dialogue_sft_multiturn": get_text_dialogue_sft_multiturn,
|
||||
}
|
||||
|
||||
|
||||
def main(args):
|
||||
model_name = args.model_name
|
||||
|
||||
# Get paths from args
|
||||
text = getattr(args, "text", None)
|
||||
audio_path = getattr(args, "audio_path", None)
|
||||
|
||||
instruct = getattr(args, "instruct", None)
|
||||
|
||||
# Get the query function and call it with appropriate parameters
|
||||
query_func = query_map[args.query_type]
|
||||
|
||||
omni_llm = Omni(
|
||||
model=model_name,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.enable_stats,
|
||||
log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
|
||||
init_sleep_seconds=args.init_sleep_seconds,
|
||||
batch_timeout=args.batch_timeout,
|
||||
init_timeout=args.init_timeout,
|
||||
shm_threshold_bytes=args.shm_threshold_bytes,
|
||||
)
|
||||
|
||||
thinker_sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=2048,
|
||||
seed=SEED,
|
||||
logit_bias={},
|
||||
repetition_penalty=1.1,
|
||||
)
|
||||
|
||||
code2wav_sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=4096 * 16,
|
||||
seed=SEED,
|
||||
detokenize=True,
|
||||
repetition_penalty=1.1,
|
||||
)
|
||||
|
||||
sampling_params_list = [
|
||||
thinker_sampling_params,
|
||||
code2wav_sampling_params,
|
||||
]
|
||||
|
||||
# Build query result based on query type
|
||||
# Notice: The audio files used in this example are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
|
||||
if args.query_type == "tts_sft":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft
|
||||
query_result = query_func(text=text, read_text_only=True)
|
||||
elif args.query_type == "tts_sft_with_instruct":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice"
|
||||
query_result = query_func(text=text, instruct=instruct, read_text_only=True)
|
||||
elif args.query_type == "tts_sft_with_audio":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_audio --audio_path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
audio_list = [get_audio_data(audio_path)]
|
||||
query_result = query_func(text=text, read_text_only=True, prompt_speech=audio_path, audio_list=audio_list)
|
||||
elif args.query_type == "tts_sft_with_natural_instruction":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
|
||||
query_result = query_func(text=text, read_text_only=False)
|
||||
elif args.query_type == "audio_trancribing_sft":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_trancribing_sft --audio_path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
audio_path = "spoken_dialogue_assistant_turn_1.wav"
|
||||
text = "Please transcribe this audio and repeat it once."
|
||||
query_result = query_func(text=text, audio_path=audio_path, use_sostm=True)
|
||||
elif args.query_type == "audio_understanding_sft":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_understanding_sft --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
query_result = query_func(text=text, audio_path=audio_path)
|
||||
elif args.query_type == "audio_understanding_sft_with_thinking":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
query_result = query_func(text=text, audio_path=audio_path, thinking=True)
|
||||
elif args.query_type == "spoken_dialogue_sft_multiturn":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type spoken_dialogue_sft_multiturn --audio_path "./prompt_speech_zh_m.wav"
|
||||
first_turn_text_response = "我没办法获取实时的天气信息。不过呢，你可以试试几个方法来查看今天的天气。首先，你可以用手机自带的天气功能，比如苹果手机的天气应用，或者直接在系统设置里查看。其次，你也可以用一些专业的天气服务，像是国外的AccuWeather、Weather.com，或者国内的中国天气网、墨迹天气等等。再有就是，你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市，我也可以帮你分析一下历史天气趋势，不过最新的数据还是需要你通过官方渠道去获取哦。"
|
||||
audio_list = []
|
||||
s1_audio_path = "weather_of_today.mp3"
|
||||
s2_audio_path = "spoken_dialogue_assistant_turn_1.wav"
|
||||
s3_audio_path = "beijing.mp3"
|
||||
audio_list.append(get_audio_data(audio_path))
|
||||
audio_list.append(get_audio_data(s1_audio_path))
|
||||
audio_list.append(get_audio_data(s2_audio_path))
|
||||
audio_list.append(get_audio_data(s3_audio_path))
|
||||
|
||||
message_list = [
|
||||
{"role": "user", "content": s1_audio_path},
|
||||
{"role": "assistant", "content": {"text": first_turn_text_response, "audio": s2_audio_path}},
|
||||
{"role": "user", "content": s3_audio_path},
|
||||
]
|
||||
query_result = query_func(message_list, system_prompt=None, ref_audio_path=audio_path, audio_list=audio_list)
|
||||
elif args.query_type == "speech2text_dialogue_sft_multiturn":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type speech2text_dialogue_sft_multiturn
|
||||
s1_audio_path = "weather_of_today.mp3"
|
||||
s2_audio_path = "beijing.mp3"
|
||||
audio_list = []
|
||||
audio_list.append(get_audio_data(s1_audio_path))
|
||||
audio_list.append(get_audio_data(s2_audio_path))
|
||||
message_list = [
|
||||
{"role": "user", "content": s1_audio_path},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Hello, I can't get real-time weather information. If you can tell me your city, I can help you analyze historical weather trends, but for the latest data, you'll need to get it through official channels.",
|
||||
},
|
||||
{"role": "user", "content": s2_audio_path},
|
||||
]
|
||||
query_result = query_func(message_list, thinking=True, audio_list=audio_list)
|
||||
elif args.query_type == "text_dialogue_sft_multiturn":
|
||||
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type text_dialogue_sft_multiturn
|
||||
message_list = [
|
||||
{"role": "user", "content": "Could you recommend some tourist attractions in China?"},
|
||||
{"role": "assistant", "content": "Hello, which city would you like to travel to?"},
|
||||
{"role": "user", "content": "Beijing"},
|
||||
]
|
||||
query_result = query_func(message_list=message_list)
|
||||
else:
|
||||
raise ValueError(f"Invalid query type: {args.query_type}")
|
||||
|
||||
prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
|
||||
|
||||
print("prompts", prompts)
|
||||
omni_outputs = omni_llm.generate(prompts, sampling_params_list)
|
||||
|
||||
output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
|
||||
if args.query_type is not None:
|
||||
output_dir = os.path.join(output_dir, args.query_type)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
for stage_outputs in omni_outputs:
|
||||
if stage_outputs.final_output_type == "text":
|
||||
for output in stage_outputs.request_output:
|
||||
request_id = output.request_id
|
||||
text_output = output.outputs[0].text
|
||||
# Save aligned text file per request
|
||||
prompt_text = output.prompt
|
||||
out_txt = os.path.join(output_dir, f"{request_id}.txt")
|
||||
lines = []
|
||||
lines.append("Prompt:\n")
|
||||
lines.append(str(prompt_text) + "\n")
|
||||
lines.append("vllm_text_output:\n")
|
||||
lines.append(str(text_output).strip() + "\n")
|
||||
try:
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
print("lines", lines)
|
||||
f.writelines(lines)
|
||||
except Exception as e:
|
||||
print(f"[Warn] Failed writing text file {out_txt}: {e}")
|
||||
print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
|
||||
elif stage_outputs.final_output_type == "audio":
|
||||
for output in stage_outputs.request_output:
|
||||
request_id = output.request_id
|
||||
audio_tensor = output.outputs[0].multimodal_output.get("audio")
|
||||
|
||||
if audio_tensor is None:
|
||||
continue
|
||||
|
||||
output_wav = os.path.join(output_dir, f"{request_id}.wav")
|
||||
|
||||
# Convert to numpy array and ensure correct format
|
||||
audio_numpy = audio_tensor.float().detach().cpu().numpy()
|
||||
|
||||
# Ensure audio is 1D (flatten if needed)
|
||||
if audio_numpy.ndim > 1:
|
||||
audio_numpy = audio_numpy.flatten()
|
||||
|
||||
# Save audio file with explicit WAV format
|
||||
sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
|
||||
print(f"Request ID: {request_id}, Audio saved to {output_wav}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model-name",
|
||||
"-m",
|
||||
type=str,
|
||||
default="XiaomiMiMo/MiMo-Audio-7B-Instruct",
|
||||
help="Backbone LLM path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
"-t",
|
||||
type=str,
|
||||
default="The weather is so nice today.",
|
||||
help="input text",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="tts_sft",
|
||||
choices=query_map.keys(),
|
||||
help="Query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
"-a",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local audio file. If not provided, uses default audio asset.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--instruct",
|
||||
type=str,
|
||||
default=None,
|
||||
help="instruct",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-stats",
|
||||
action="store_true",
|
||||
help="Enable writing detailed statistics (default: disabled)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-sleep-seconds",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-timeout",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Timeout for batching in seconds (default: 5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-timeout",
|
||||
type=int,
|
||||
default=5000,
|
||||
help="Timeout for initializing stages in seconds (default: 5000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shm-threshold-bytes",
|
||||
type=int,
|
||||
default=65536,
|
||||
help="Threshold for using shared memory in bytes (default: 65536)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default="./output_audio",
|
||||
help="Output audio wav directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-wav",
|
||||
default="output_audio",
|
||||
help="[Deprecated] Output wav directory (use --output-dir).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-prompts",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of prompts to generate.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--sampling-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sampling rate for audio.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default="../../../model_executor/stage_configs/mimo_audio.yaml",
|
||||
help="Path to a stage configs file.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,56 @@
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Your voice should be:"
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "data:audio/wav;base64,UklGRiRKAwBXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQBKAwAFAP3/CAD7/wkAAAD9/wgA+v8JAP//AQADAP//CAD+/wEABAABAP7/BgD8//7/BgD0/wYA/f/8/wMA9f8GAP///P8DAPz/AAAFAP3//v8CAP7////+////+/8IAP//AQAAAAMABgD6/wUA+v8DAAcA+/8AAAMAAAAAAP3//v8DAAMACgD8/wQACQD9/wUA/f8BAAcABQADAAAABAADAAEA+v8AAAIAAAAFAPz//f8DAAAA/f8AAAMA/v8FAP7/9v8GAP7//f8EAAAAAQAIAPv/+/8AAP3////+//3/AAAHAPr//v8BAPz/AwD///n/AAACAPr//f/8//7/AAD//wAA+f8GAPz/9v8CAPX/BgACAPv/AQD//wQA///7/wAAAAAHAAQA/f8FAAAAAgAAAPv/BAAGAAAABgACAP7/DQD7//3/BAD//wkA//8DAPz/AwAHAPf/AwAAAAAABAD7//f/AwD8/wAA/f/5/wQA/P8AAPr//P8AAAAA//8AAP////8DAPL/AQD9//r/CwD4/wEAAgD9//7//P/9//7/AAD//wAA/v8DAPr/+v8CAPr///8CAPz///8GAPr/AAAEAP//AgADAP//AAAHAPz/BgD//wUACwD//wYAAgABAAQABQD8/wgABgACAAUA//8CAAQAAAADAP7/AQAIAPr/AwD+//v/AgACAPv/AQAJAPf/AQD+//n/AgAAAP///f8CAAAA/P8DAPf/+/8HAPv///////7//f8AAP3//P8BAPr/AAD6//z/AAD5/wIAAAD7/wIAAgD4/wAA/v/6/wQAAQAAAAAABwAAAAAAAwD+/wIABQADAAEABAADAAAAAwABAAAABAAGAAEAAgAGAAAABQD///7/BQAAAAMAAQAAAAIABAD8/wMA/////wQA/f8AAAEAAQD//wMA+/8FAPr///////P/CwDz/wQA/P/6/wIA8/8AAPH//P/+//f//v/+//n//P/8//L//v/7//r////7//3//v/5//v/AAD5/wMAAAD8/wYA/P///////////wAABwD6/wwAAQD+/wUA/P8GAP7/AwD+/wUABwD+/wUA//8FAAMAAgD//wUAAAAEAAcA+/8KAAUAAgAEAAAA//8HAAAA+v8EAAIAAQAEAPz/+/8EAP7//////wAAAQABAP//+///////AAD9/wEA/f8DAAAA+P8AAPj/AQD+//v/AQD8/wIA/v/6/wAA/////wIA/P///wMA+/8BAPj///8CAAAAAwD+/wcAAAAFAP3/AAAFAP7/AwABAP7/AAACAPv/AwD+/wMABAD9/wUA/v8AAAEA/v8AAAIAAQAAAAAAAAD8/wIA/v/8/wMAAAD///3/AAD5//7/AwD2/wgAAAD8/wQA+f8CAP7/AAAAAP//AwD//wAA/v//////AAAAAAAAAwD9/wIA+//9/wQA+/8DAAAA//8AAP///f/+/wAAAAABAP//BgD8/wIAAgD3/wwA//8AAAkA/f8EAAAA//8BAAAAAwADAP//BQD///z/AwD8/wEAAwAAAAEAAQAAAAAA/v8BAAAAAQAGAP7/AAD///3//v8AAAAAAgAGAP//AgAAAP3////9/wAAAQACAAAAAQD///3/AQD9/wAABwD//wMABAD8/wEAAwD+/wIABQD+/wEAAwD6////BAD5/wYAAAD+/wgA+f8CAP3///8DAP7/AgAAAAAA/v8EAPv/AAAHAPv/CQACAP7/BgD//wIAAAADAAAAAgAEAPz/BQD8/wQA/P/+/wMA+f8EAP3//P8AAP////8AAPv/AgD8/wAAAgDz/wgA/f/7/wQA+P8BAAAA+v8BAP7/AAACAP3/AAACAP7/AQAEAPj/BwAAAAAABgD7/wcA+v8CAP7/+/8IAP3/AQACAAAAAQACAAAA/v8CAAcA/f8HAAIA/v8IAPv/AQAAAAAABAD+/wQAAgD//wMA/P8AAAAA//8BAPr/BQD9/wAA///4/wIA+/8AAP3//v8EAPr/AAD9//r/BQD3/wEA+v/6/wYA8/8BAPz/9v8CAPn/+f/+//3//v/+//3//v///wAA/P/7/wAAAAD//wAA//8AAP7/AgD4////BAD7/wYAAAADAAAAAwAAAPr/BQD6/wUABAD//wcA/f8EAPz//P8FAPv/BQADAP//BgAAAP//AQAAAAAAAgAAAAEAAwAAAAAA//8AAAAAAAAAAP7/BAD///7/AgD8/wAAAQD8////AQD9/wQA/v///wIA//8BAPv/AQD9/wMA/v/8/wQA+f8DAP3/9/8EAAAA/P8GAPz/AwAEAP//AAD8/wQAAAAGAAAAAgAFAAEABAD8/wEAAQADAAAAAgADAAAACAD7//7/BQD9/wUAAAD+/wcAAAABAAAA/f8DAAMAAAAAAAEA/v8AAAAA+f8AAAEA//8AAP//AAD7/wUA9//6/wQA9/8BAP3//f8AAP7/AAD7//3/AwD+//v//v/7//3/BAD3//r/CAD8/wIA/v/6//3/AQD+//j/AAAAAAgAAAD///v//f8BAP7//v/9/wEA//8EAAAA//8BAAEAAAD+/////f8AAP7/AgD9/wQABQD4/wIA/f/5/wAA/P/9/wEAAQABAPf/AAD8//7/AQD5/wAA/f8BAPf/+f/7//b/AwADAPz///8BAPn//P/4/////P8AAAAA+P8JAAAA/f/8//3/AAABAAAA8v8EAAMAAAACAPn/BAAOAAcA+//t////DwAEAAIA/f8HAAYADAD7//f/BgD//wwA//8BAAEADwAEAPv/BQD1/xEA//8DAP//9/8LAPn/AgABAAMABAAEAPz//f8FAPb/BAD+//r/CwAAAP7/BQD4////BAD4/wEA/v8AAAAA/v/+//z/BAABAP//BQD+//z/AQD7/wAA/P8DAAUAAAAHAP//AQD9////AAABAP//AAACAAQACAD+/wYA//8DAAcA/v8CAAEABQADAAIAAgACAAUABQAAAAAAAgAAAAAAAwABAAAABQAAAAYABgAAAAEA/v8EAAAA/P8GAP7/BAAGAPn/AQABAAIA/f8EAPn///8GAPT/BgD5/wAAAgD8/wQA+v8BAPv/9/8AAP3/AAAAAPz/AQD9//r////6////AAD9/////f8AAPz/AAD9//z/CQD5/wUA+v/7/wIA9P8FAPn/AQADAP3//v8AAP3/+/8DAPr/AQABAAAA/f8BAAAA9/8EAP7///8GAPr///8BAPz/AQD9/wAABQD//wEAAAD9/wMAAQD7/wAAAwAAAAMAAAD+/wMAAQAAAP3/AAABAP//AgD+/wEAAQAHAPv///8EAPb/CQD//wAAAAAAAAUA9v8JAPf///8FAPr/AwD9/wIA+v8DAP//+P8GAP////8EAPv/+/8EAPv/AAD+//7/AQD//wEA+/8AAAAA/v8CAP///P8GAP////8DAPv/BgAAAAAAAgD//wEA//8CAP7/AgAAAP//BgAAAAEAAgAAAP//BwD6/wIABAD6/woA+/8EAAAAAAAEAPv/BgD//wIAAAABAAMAAAAHAPf/BwAAAPr/BwD4/wQAAAAAAAAAAAD//wAAAAD7/wcA/v8AAAAA///+/wIA///7/wkA+/8GAAAA/f8BAP//AgD9/wYA//8FAAQA//8CAPz/AQAFAPr/AwAAAP//BgD5//7//v8CAPv/AAD9//b/CAD4//z//v/6/wAAAAD+//z/AAABAP///v////3/BwABAAAAAAAAAAUA/f8DAPb/AAAIAPv/AwD9////BAAAAPf/+/8AAPv/AQD7//v/AAABAP7/+v8BAPn/AQACAPr/AQACAAEA/f8FAPv/AAAFAPj/AAABAPv//P/+//n//v////r//f/+//r/AAD0//b//f/0/wMA9//+//7//f8BAPP/BAD6////AwD8/wEAAAAKAPf/BwACAPz/DAD9//3/AwAGAPn/BwD+//3/DgD7/wAAAAD7/wYA/f/5/wQA+v8GAAgA+P8IAAMAAwAKAP//AQAFAAMABQAAAAUABQAGAAcA//8IAPr/AgAAAPX/BAD+//3/BAAAAPX/AwD5//T/AwD3//n/BgD6//z/BAD3/wAAAAD+/wAA//8AAAAA/v/+//7///8BAAAAAAABAP7//v/+//b/AAD8/wAAAAD+/wEA/f8CAPr/AQD7/wAABQD7/wcA//8FAAcAAgAIAAUABAAGAAcAAwAGAAQABAAKAAkABQAIAAUABwADAAMAAAACAAMAAgAFAP7/DAABAAgA/v/+/wYA/P8KAPj/BAAGAP7/CQAAAAEACAAEAAUA/P///wIA+/8DAP7///8FAP//AAD7/wAA/f/2/wIA9P/9/wMA8v///wAA+/8DAAEA/P8CAP7//v/5//r/CAAAAAkABgAAAAUABwD+//3/CwD//w4ABwAAAAkAAgAJAPz/BwALAAEAEgD//wAACgABAAcAAgAFAAcABwAGAAAABAAEAAoAAwAIAAsABQAMAAAAAQACAAkABwAGAA4ACwAKABMAAgAAAAoA/f8JAAAABAADAAsACQD9/wQA+/8AAPf//v/z//r/AQDv//3/8f/w//r/8f/0/+//9P/u/+j/8v/g/+7/8f/r//X/8f/x/+7/7v/m/+v/7v/3//D/8//4/+3//v/z//b/9//7//n/9f8BAPT//P8AAPj/AAADAAAAAAD+/wIA+f/8/wQA+v8DAAcA/v8AAA8AAAAFABEA//8NAA4ABwAHABAAEgAKABgAFAARABYAFAARABIAHAAXABIAIAAYABQAIgAUABEAIQAXABUAGgAVABMAFwAVAAgADQAJAAkABAAJAAkAAAAKAPz/+f/3//z/8P/2/wAA8/8CAP7/8//0//j/7v/t/Line truncated
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "data:audio/wav;base64,UklGRgrVAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YebUAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA//8AAAAAAAAAAAAAAAAAAAAAAAAAAP////8AAAAAAAAAAAAAAAD//wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP////8AAAAAAAAAAAAA///p/9T/xv/G/8b/yf/K/8r/zP/M/83/zf/O/83/zf/P/8//z//O/87/z//P/8//z//P/8//z//P/8//z//P/8//z//Q/9D/0P/Q/9D/0P/R/9H/0f/R/9H/0f/R/9H/0f/R/9L/0v/S/9L/0v/S/9P/0//T/9P/0//T/9P/0//U/9T/1P/U/9T/1P/V/9T/1f/V/9X/1f/W/9b/1v/W/9b/1v/W/9b/1//X/9f/1//Y/9j/2P/Y/9j/2P/Z/9n/2f/Z/9n/2f/a/9r/2v/a/9v/2//b/9v/2//b/9v/3P/c/9z/3f/c/93/3f/d/93/3f/e/97/3v/e/9//3//f/9//4P/g/+D/4P/g/+D/4P/h/+H/4f/h/+H/4f/i/+L/4v/j/+P/4//j/+P/4//k/+P/5P/k/+T/5P/l/+X/5f/l/+X/5f/l/+X/5f/l/+X/5v/n/+f/5//o/+j/6P/o/+j/5v/n/+f/5//o/+j/6P/o/+n/6f/p/+j/6P/o/+j/6P/p/+n/6f/p/+n/6v/q/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//r/+r/6v/r/+r/6v/p/+r/6v/r/+r/6v/r/+z/7P/s/+v/7P/r/+v/6//s/+v/6//r/+v/6//t/+z/7f/t/+3/7f/r/+z/6//s/+v/7P/s/+z/7P/s/+z/7P/t/+3/7f/s/+z/6//s/+r/6//q/+z/7P/t/+3/7v/u/+3/7f/t/+z/7P/s/+z/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+z/7P/s/+z/7f/t/+3/7f/u/+3/7f/t/+3/7f/t/+3/7f/u/+3/7f/t/+3/7f/t/+3/7f/t/+3/7P/r/+z/6//t/+z/7f/t/+7/7f/u/+7/7f/t/+z/7P/t/+z/7P/t/+3/7P/t/+7/7v/v/+7/7v/t/+z/7P/t/+3/7v/u/+7/7v/v/+//7v/u/+3/7P/u/+z/7P/s/+z/7f/u/+7/7v/v/+7/7v/s/+z/7P/t/+3/7f/u/+7/7f/u/+3/7v/t/+3/7f/t/+3/7P/t/+z/7f/u/+3/7//u/+7/7f/t/+3/7f/u/+//7v/v/+7/7v/t/+7/7f/v/+7/7v/t/+7/7f/t/+3/7v/u/+7/7v/v/+3/7v/t/+z/7P/r/+z/6//t/+r/7P/t/+7/7v/u/+7/7f/s/+z/6//s/+v/7P/t/+3/7//w//D/7v/w/+7/7v/s/+3/6//t/+v/7f/u/+//7//u/+//7v/u/+r/6//q/+v/6//s/+7/7v/w/+//7//v/+3/7f/s/+v/7P/s/+3/7f/u/+7/7v/u/+3/7P/r/+z/7f/s/+z/7P/t/+3/7P/t/+3/7v/s/+3/7f/t/+z/6//r/+z/7P/s/+z/7P/t/+z/7f/t/+z/7v/v/+7/7P/s/+z/6//u/+3/7v/u/+7/7f/u/+3/7f/t/+v/7f/r/+v/6//s/+3/7P/u/+7/7v/v/+7/7v/u/+z/7P/r/+3/6//t/+3/7v/t/+7/7v/v/+7/7v/t/+3/7P/t/+7/7v/v/+3/7v/t/+3/7f/s/+z/7f/s/+3/7f/t/+7/7v/u/+3/7//u/+7/7f/s/+z/7P/s/+v/7P/t/+3/7f/v/+7/7v/u/+3/7f/t/+3/7f/u/+7/7v/u/+7/7//u/+7/7v/t/+3/7f/t/+3/7f/u/+//7//v/+//7v/u/+7/7v/u/+7/7v/u/+7/7v/t/+7/7v/u/+7/7v/u/+7/7f/t/+3/7f/t/+7/7f/t/+3/7P/u/+3/7v/t/+7/7v/u/+7/7v/v/+7/7v/v/+7/7v/u/+7/7v/u/+3/7f/u/+3/7v/u/+7/7f/u/+3/7v/u/+7/7f/t/+7/7v/v/+7/7v/u/+3/7v/t/+3/7f/t/+3/7f/u/+7/7v/u/+3/7v/t/+3/7P/t/+3/7f/u/+3/7v/u/+//7v/u/+7/7v/u/+7/7v/u/+7/7f/u/+3/7f/u/+7/7f/u/+3/7f/t/+3/7f/t/+3/7f/s/+z/7f/t/+3/7f/t/+3/7P/t/+3/7f/t/+z/7P/s/+3/7P/t/+3/7f/u/+3/7f/t/+3/7f/t/+7/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/s/+z/7f/t/+3/7f/t/+3/7f/t/+z/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+z/7P/s/+3/7f/t/+3/7f/t/+3/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+3/7P/t/+3/7f/t/+3/7f/t/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7f/s/+z/7P/s/+z/7P/s/+v/6//r/+z/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/6//s/+z/7P/s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+3/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/s/+3/7P/t/+3/7P/s/+z/7P/t/+z/7f/t/+3/7f/t/+3/7f/u/+z/7f/t/+3/7f/t/+3/7f/u/+3/7f/t/+z/7P/s/+z/7P/t/+z/7v/t/+7/7f/t/+7/7f/v/+z/7f/r/+3/7P/t/+3/7f/u/+7/7v/v/+7/7v/u/+z/7f/s/+7/7f/u/+3/7f/v/+7/7//u/+7/7//t/+//7v/u/+z/7P/s/+7/7f/w/+3/7f/u/+//7f/u/+7/7//s/+//6//w/+3/7//u/+3/7v/s/+7/7P/u/+//7v/x/+z/7//t/+3/7f/u/+//7//u/+7/6//w/+3/8f/t/+//7v/v/+7/7v/u/+//7v/w/+3/7v/t/+//7//t/+7/7v/w/+3/7P/v/+3/7//t/+//7v/w/+7/7//y/+3/7v/v/+3/7//t/+7/7P/v/+3/9P/w/+3/6//r//D/7//w/+3/8v/t/+7/7f/v/+3/7f/v//D/7f/t//D/7v/u/+3/8P/v/+7/7v/w/+z/7//r//L/7v/x/+3/7v/u/+7/7//t/+//7f/v/+7/7f/v/+7/7f/u/+7/7f/v/+//8P/u/+3/7//u/+//7f/v/+z/7f/v/+3/7v/p//D/7f/x/+7/7v/s/+3/8P/v/+//7f/r/+7/7P/t/+3/7v/t/+3/7f/u/+7/7//t/+7/7f/u/+3/6//s/+z/7f/t/+3/7//s/+//7f/y/+z/7v/s/+3/7P/v/+3/6//r/+3/7P/u/+//7f/t/+7/7v/v/+z/7f/r/+//7f/w/+z/7f/v/+v/7f/s/+7/7v/v/+z/7P/u/+z/7//s/+3/7f/u/+7/7f/w/+z/7v/t/+3/7v/t/+3/7P/u/+7/7v/s/+z/7f/v/+3/7v/t/+3/7f/t/+7/7P/t/+3/7v/v/+z/7f/u/+3/7v/v/+z/7v/t/+//7f/u/Line truncated
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "我没办法获取实时的天气信息。不过呢，你可以试试几个方法来查看今天的天气。首先，你可以用手机自带的天气功能，比如苹果手机的天气应用，或者直接在系统设置里查看。其次，你也可以用一些专业的天气服务，像是国外的AccuWeather、Weather.com，或者国内的中国天气网、墨迹天气等等。再有就是，你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市，我也可以帮你分析一下历史天气趋势，不过最新的数据还是需要你通过官方渠道去获取哦。"
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "data:audio/wav;base64,UklGRiQiKQBXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQAiKQAGAAcABwAIAAkACQAJAAoACwAMAAwADQAOAA4ADwAPABAAEAAQABAAEQAQABAAEQARABAAEQARABEAEAAQABEAEQAQABAAEAAPAA4ADwAOAA0ADQANAAwACwAJAAYABQAEAAIAAQAAAP7//f/7//r/+v/5//b/9v/0//X/9f/0//X/9P/0//L/9P/y//D/9P/z//T/8v/x//L/8v/x//H/8P/v/+//7v/u/+3/7P/t/+3/7P/t/+7/7v/u//D/8f/z//X/9f/3//n/+//7//3///8AAAEABAAGAAgACgAMAA4AEAASABMAFAATABQAFAAUABUAFQATABIAEgARABEAEAAOAAwACwAJAAcABgAEAAMAAgABAAEAAAAAAAAA//////7//f/8//v/+f/4//j/9//3//f/9v/2//b/9f/1//b/9f/1//b/9v/2//b/9f/1//X/9P/z//P/9P/z//P/8//0//T/9P/0//X/9v/2//f/+P/4//n/+v/7//z//f/+//7///8AAAEAAgACAAMABAAFAAYABwAJAAoADAANAA4ADwAPAA8AEAAQABEAEgASABMAFAAVABYAFwAXABYAFgAWABYAFgAVABQAEwATABMAEgASABIAEQARABAADwAOAA4ADQANAAsACgAIAAcABAACAAEA///9//v/+f/4//b/9f/z//H/8f/v/+3/7P/r/+n/5//m/+X/5f/k/+T/5P/l/+X/5//o/+r/7P/u//D/8v/1//f/+f/7//z//v8AAAEAAgAEAAQABQAHAAUABQAGAAQAAwACAAEAAgAAAAEAAQD+/wAA/v8AAP7//P////7//f/9/wAA/v/9//7//v///wAA//////////8AAAEAAwACAAIAAwADAAMAAgACAAIAAwACAAEAAAABAAEAAAD////////////////+/////v/9//3//P/8//z//f/9//3//v/+//7///8AAP////////////////////////7//v/9//3//f/8//r/+P/3//j/9//3//b/9f/2//b/9f/1//b/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//4//j/+P/5//n/+f/5//n/+f/5//n/+f/6//r/+//7//z//P/+////AAACAAMABQAGAAcACAAIAAkACQAJAAkACQAKAAoACQAJAAcABgAEAAMAAQD///3/+//6//n/+P/3//b/9v/1//b/9v/2//b/9v/2//b/9//4//n/+v/8//7//v///wEAAwAEAAUABwAJAAoACgALAA4ADAANAA0ADwAOAA0ADwAPAA8ADgARABEAEQARABEAEAAQABEAEQARABAAEAAQABEAEQAQABAADwAPAA4ADQAOAAwADQANAAwADQANAAsADAALAAsACgAJAAkACAAHAAcABQAFAAUAAgABAAEAAAD///7//P/6//n/9v/z//P/8v/w/+//7f/r/+v/6v/p/+r/6f/o/+j/6P/o/+j/6f/p/+n/6v/r/+z/7P/t/+7/7//w//H/8//1//b/+P/5//r//P/9////AAABAAIAAwAFAAYABwAIAAkACwALAAwADQAOAA4ADwAQABEAEgASABMAFAAVABUAFQAVABYAFQAVABUAFAATABIAEQAPAA4ADQAMAAsACgAJAAcABQADAAAA///8//r/+P/2//X/8//x/+//7v/t/+z/6//r/+r/6v/q/+v/6//s/+z/7f/t/+7/7//w//H/8v/z//T/9f/2//f/9//4//j/+f/6//r/+//7//z//P/9//7//v///wAAAAAAAAEAAQACAAMABAAGAAcACQAKAAwADQAPABAAEQASABIAEgATABUAFwAYABkAGgAbAB0AHgAfAB4AHgAdABwAGwAaABkAGAAWABQAEgAQAA0ACwAJAAcABQADAAEA/v/8//r/+P/2//T/8//x//D/7//t/+z/6//r/+v/6//r/+v/6//s/+3/7f/t/+7/7v/v/+//8P/w//D/8f/w//D/7//w//H/8f/w//D/8f/y//L/8v/y//P/8//1//f/+P/5//z//f/8/wAA//8DAAEAAAAEAAIAAQADAAMAAQAAAAEAAAABAAEA/////////v/9//7////9//z//P/9//3//P/8//z//f/+//7//v////////8AAAAAAQACAAIAAgACAAIAAgABAAEAAQABAAIAAQAAAAAAAAD///7//v/9//3/+//6//r/+v/6//r/+f/6//r/+v/5//n/+v/7//v/+v/6//v//f/+//7//v//////AAD///////////7//f/8//v/+//6//r/+f/5//n/+v/6//r/+v/6//z//P/9//3///8AAAEAAgADAAUABQAGAAYABwAIAAkACQAKAAsADAAMAAwADAAMAA0ADAAMAAsACwAKAAkACAAHAAUAAwACAAEAAAD+//3/+//6//j/9//1//T/8//y//L/8//z//T/9f/1//b/9v/2//f/+P/4//n/+f/5//r//P/9////AAACAAMABAAFAAYACAAJAAoACwANAA4ADQAOABEAEAARABEAFAATABMAFQAVABYAFQAXABYAFgAWABUAFAATABQAEwASABIAEQAQABAADwAPAA4ADQAMAAsACwAMAAoACwAKAAkACgAJAAgABwAFAAUAAwACAAEA///9//3/+v/5//j/9P/y//D/7v/s/+v/6P/m/+X/4v/h/+H/4f/g/+H/4P/g/+H/4v/j/+T/5f/m/+f/6P/q/+v/7f/u/+//8f/y//T/9f/3//j/+v/7//z//v///wEAAgADAAQABQAGAAgACQAKAAsADAAOAA8AEAASABMAFAAVABYAFgAXABgAGAAZABkAGgAaABoAGQAZABgAFgAVABMAEgAQAA8ADgAMAAoACQAHAAUABAACAAAA///9//v/+f/3//X/8//x//D/7//v/+7/7v/t/+3/7f/t/+7/7v/v//D/8P/x//L/8//0//X/9f/2//f/9//4//j/+f/5//r/+v/6//r/+v/6//r/+f/6//r/+v/7//v//P/8//3//v///wEAAgADAAUABgAIAAsADQAPABEAEwAUABYAFwAYABgAGQAZABkAGAAZABkAGQAZABgAFwAVABQAEgARAA4ACwAIAAUAAwABAP///v/8//v/+v/4//b/9f/0//P/8v/y//H/8P/w/+//7v/u/+3/7f/t/+3/7f/s/+z/7P/t/+3/7v/u/+7/7v/u/+//8P/v//D/8P/x//H/8f/y//P/9P/0//X/9f/1//f/+P/5//n/+//8//z//f/9//7//f///wAAAAAAAAEAAQD//wAA/f8AAPz/+//+//v/+v/8//z/+//7//z//f/+/////v///wEAAQAAAAIABAADAAIAAwADAAQAAwADAAMAAwADAAIAAgACAAEAAAAAAAAAAAD//////v/+//7//v/9//7//v/+//////////////////////8AAAAA////////AAAAAAAAAAAAAAEAAAD///7//////////v/9//3//////////v/+//7//v/+//7//v////////8AAAAAAQACAAQABQAGAAcACAAKAAoACgALAAwADAAMAAwADAAMAAwADAAMAAwACwAKAAkACQAJAAgACAAIAAgABwAHAAYABQAEAAMAAgAAAP///v/9//z/+v/5//f/9f/0//P/8v/x//D/7//u/+3/7f/s/+z/7P/s/+7/7//w//L/8//0//X/9v/3//j/+v/8//3///8BAAMABgAIAAoADAAOABAAEQARABIAEwAUABQAFQAVABYAFAAUABYAFAAUABQAFgAVABMAFQAVABUAFAAVABUAFQAVABQAFAATABQAFAAUABQAFAATABQAFAAUABMAEgASABAAEAAQAA4ADgAMAAkACAAGAAIAAAD8//r/9//0//L/7//t/+v/6f/n/+b/4//h/+D/3v/e/9//3v/e/97/3v/e/+H/4v/j/+b/5v/n/+n/6v/s/+7/7//w//L/8//0//b/9//4//n/+v/7//z//f/+////AAABAAIABAAFAAcACAAJAAoACwANAA4ADwARABIAEwAUABUAFgAXABcAFwAXABcAFgAWABUAFAATABIAEAAPAA0ACwAJAAYABAABAP///f/7//n/9//1//T/8v/w/+//7f/s/+r/6f/o/+f/5v/l/+T/5P/j/+T/5P/l/+X/5v/n/+j/6P/p/+r/6//s/+z/7f/u/+7/7v/v/+//7//v/+//8P/w//H/8f/y//P/9P/1//f/+P/5//v//P/+/wEAAwAFAAcACgALAA4AEAATABUAFwAZABsAHQAfACEAIwAkACUAJgAmACUAJQAkACMAIQAgAB0AGwAaABkAGAAWABQAEgAPAA4ADAALAAgABgAEAAIAAQAAAAAA///+//7//f/8//v/+f/5//j/+P/4//f/9//2//X/9f/0//T/9P/0//X/9f/1//X/9f/2//f/9//4//j/+f/6//v//P/8//3//v///wAAAAABAAIAAwAEAAUABQAGAAcACAAIAAcACAAJAAgACAAIAAgABwAIAAkACQAIAAoACgAIAAkABwAKAAcABgAKAAgABwAKAAoACQAJAAoACwALAAwACgAKAAsACQAIAAgACAAGAAQAAwACAAEA///+//3//P/7//r/+f/4//j/9//3//f/9//3//f/9//3//j/+P/4//j/+P/5//n/+f/5//n/+f/5//j/+f/5//n/+f/4//n/+f/6//r/+f/6//v/+//6//r/+//8//3//P/8//7/AAABAAIAAgADAAQABQAGAAcACAAJALine truncated
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "data:audio/wav;base64,UklGRkqYAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YSaYAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAP//AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD//wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP//AAAAAP//AAAAAAAA///o/9T/xv/G/8b/yP/K/8r/y//M/8z/zf/O/83/zf/P/87/zv/O/87/zv/P/8//z//P/8//z//P/8//z//P/8//z//P/8//0P/Q/8//z//Q/9D/0f/R/9D/0f/R/9H/0f/R/9H/0f/R/9H/0f/S/9L/0v/S/9L/0v/S/9L/0v/T/9P/0//T/9P/0//T/9P/0//T/9P/0//T/9P/0//U/9T/1f/V/9T/1P/U/9X/1P/V/9X/1f/V/9b/1v/W/9b/1v/W/9b/1v/W/9f/1//X/9f/1//X/9f/1//X/9j/2P/Y/9j/2P/Y/9j/2P/Y/9j/2f/Z/9r/2v/a/9r/2v/b/9v/2//b/9v/2//c/9v/3P/c/9z/3P/d/9z/3P/c/93/3f/e/93/3f/e/97/3v/e/97/3//f/+D/4P/g/9//3//g/+D/4P/g/+H/4f/g/+H/4P/h/+H/4f/h/+L/4v/i/+H/4f/i/+L/4v/i/+L/4v/j/+L/4//j/+P/4//j/+P/4//j/+P/5P/k/+T/4//l/+T/5P/k/+T/5P/l/+T/5f/l/+X/5f/m/+X/5v/m/+b/5v/m/+b/5v/m/+b/5v/n/+b/5//n/+b/5v/o/+f/5//n/+j/5//p/+f/5//o/+j/6P/n/+j/6P/p/+n/6f/o/+j/6P/p/+j/6P/o/+n/6P/p/+j/6P/p/+n/6f/q/+n/6v/p/+n/6f/q/+n/6f/p/+n/6f/p/+n/6f/p/+n/6v/p/+n/6f/p/+n/6f/p/+n/6f/r/+r/6v/q/+r/6v/r/+n/6v/p/+r/6v/r/+n/6v/q/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//q/+v/6//r/+r/6//r/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//q/+v/6//s/+r/6//q/+r/6v/s/+n/6v/q/+v/6v/r/+v/6//s/+v/6//s/+v/7P/r/+r/6//s/+z/6//r/+v/7P/r/+z/6//r/+v/6//r/+r/6//r/+v/6//r/+v/6//s/+v/6//r/+z/7P/r/+v/6//s/+v/7P/r/+v/6//r/+v/6//s/+v/6//r/+z/7P/s/+z/7P/r/+z/6//s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/s/+z/6//t/+v/7P/s/+v/7P/r/+z/6//s/+z/7P/r/+z/6//t/+v/7P/s/+v/7P/s/+z/6//s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/r/+z/7P/r/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/s/+z/6//r/+z/7P/s/+v/7P/s/+v/7P/s/+v/6//s/+v/7P/s/+z/6//s/+v/7P/q/+v/6//s/+v/6//r/+z/6//r/+z/7P/r/+z/6//s/+z/7P/r/+z/7P/s/+v/6//t/+r/7P/r/+v/6//s/+v/7P/s/+v/7P/r/+z/6//s/+v/6//r/+z/6//r/+v/6//r/+v/6//t/+v/7P/r/+z/6//r/+z/7P/s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/r/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/s/+z/7P/r/+v/6//s/+v/7P/r/+z/6//r/+v/6//s/+v/6//r/+z/7P/r/+z/6//s/+v/6//r/+v/6//s/+v/6//s/+r/6//r/+v/6//r/+v/6//r/+v/7P/r/+v/6//s/+v/6//r/+z/6//s/+v/6//q/+v/6//s/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/q/+z/6//s/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/7P/r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+r/6//r/+v/6//r/+r/6//r/+v/6//s/+v/6//q/+v/6v/s/+r/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+r/6//q/+v/6//r/+v/6//r/+v/6//r/+r/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+z/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+v/6//r/+z/6//r/+v/7P/r/+v/6//s/+v/6//q/+v/6//s/+r/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+v/7P/r/+z/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+z/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//s/+v/7P/r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6v/s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+z/7P/r/+v/7P/s/+v/6//s/+v/6//r/+v/6//r/+z/7P/s/+v/6//r/+z/6//s/+v/7P/r/+z/6//s/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/r/+v/6//s/+v/7P/r/+v/6//s/+v/7P/s/+v/7P/r/+z/7P/r/+z/6//s/+v/6//r/+z/6//s/+z/7P/s/+v/6//r/+v/6//r/+z/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+z/6//s/+v/6//s/+v/7P/s/+z/6//s/+v/6//s/+v/7P/s/+z/7P/s/+z/6//r/+v/6//s/+v/7P/s/+z/6//s/+v/6//r/+z/6//s/+v/6//r/+z/6//r/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/r/+v/7P/s/+z/7P/s/+v/7P/s/+z/6//s/+v/7P/s/+v/7P/s/+z/7P/s/+v/6//r/+v/6//s/+v/7P/r/+z/6//r/+z/6//s/+v/7P/r/+v/7P/r/+z/6//s/+v/6//s/+v/7P/r/+z/6//s/+z/7P/s/+v/7P/r/+z/7P/s/+v/6//r/+z/7P/s/+z/6//s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/6//t/+v/7P/s/+z/7P/s/Line truncated
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,714 @@
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchaudio
|
||||
from process_speechdata import InputSegment, StreamingInputSegment
|
||||
from torchaudio.transforms import MelSpectrogram
|
||||
|
||||
speech_zeroemb_idx = 151667
|
||||
empty_token = "<|empty|>"
|
||||
mimo_audio_tokenizer = None
|
||||
device = "cpu"
|
||||
|
||||
asr_zh_templates = [
|
||||
"请将这段语音转换为文字",
|
||||
"帮我识别这个音频文件中的内容",
|
||||
"把这段录音转成文本",
|
||||
"请转录这段语音",
|
||||
"将音频内容转换成文字格式",
|
||||
"识别并转写这段语音",
|
||||
"把语音内容写成文字",
|
||||
"转录这个音频片段",
|
||||
"将这段对话转换为文本",
|
||||
"麻烦帮我把这段录音整理成详细的文字记录",
|
||||
]
|
||||
|
||||
asr_en_templates = [
|
||||
"Please transcribe this audio file",
|
||||
"Convert this speech recording to text",
|
||||
"Transcribe the following voice message",
|
||||
"Turn this audio into readable text",
|
||||
"Please convert the recording to written format",
|
||||
"Transcribe what you hear in this audio",
|
||||
"Convert this spoken content to text",
|
||||
"Please write down what is said in this recording",
|
||||
"Transcribe this voice recording",
|
||||
"Could you please help me transcribe this important recording?",
|
||||
"Would you mind converting this voice message into a readable text format?",
|
||||
"I'd really appreciate it if you could turn this audio file into a written document",
|
||||
]
|
||||
|
||||
tts_zh_templates = [
|
||||
"请将这段文字转换为语音",
|
||||
"帮我把这个文本读出来",
|
||||
"将这些文字生成音频",
|
||||
"请朗读这段内容",
|
||||
"把这段话转换成语音文件",
|
||||
"生成这段文字的语音版本",
|
||||
"请用语音播报这些内容",
|
||||
"将文本转换为可听的音频",
|
||||
"帮我朗读这段文字",
|
||||
"把这些内容念出来",
|
||||
]
|
||||
|
||||
tts_en_templates = [
|
||||
"Please convert this text to speech",
|
||||
"Turn this writing into audio",
|
||||
"Generate speech from this text",
|
||||
"Read this content out loud",
|
||||
"Convert these words to voice",
|
||||
"Create an audio version of this text",
|
||||
"Please vocalize this content",
|
||||
"Turn this text into audible format",
|
||||
"Help me convert this writing to speech",
|
||||
"Make this text into spoken audio",
|
||||
]
|
||||
|
||||
|
||||
def detect_language(text):
|
||||
if re.search(r"[\u4e00-\u9fff]", text):
|
||||
return "zh"
|
||||
else:
|
||||
return "en"
|
||||
|
||||
|
||||
# ============================================
|
||||
# Common Helper Functions - InputSegment Creation
|
||||
# ============================================
|
||||
|
||||
|
||||
def create_segment(text: str = "", audio=None) -> InputSegment:
|
||||
"""Create a standard InputSegment with default zeroemb parameters"""
|
||||
return InputSegment(
|
||||
text=text,
|
||||
audio=audio,
|
||||
speech_zeroemb_idx=speech_zeroemb_idx,
|
||||
text_zeroemb_idx=empty_token,
|
||||
)
|
||||
|
||||
|
||||
def create_streaming_segment(
|
||||
text: str, audio, tokenizer, group_size: int, audio_channels: int
|
||||
) -> StreamingInputSegment:
|
||||
"""Create a StreamingInputSegment"""
|
||||
return StreamingInputSegment(
|
||||
text=text,
|
||||
audio=audio,
|
||||
tokenizer=tokenizer,
|
||||
group_size=group_size,
|
||||
audio_channels=audio_channels,
|
||||
speech_zeroemb_idx=speech_zeroemb_idx,
|
||||
text_zeroemb_idx=empty_token,
|
||||
)
|
||||
|
||||
|
||||
# ============================================
|
||||
# Common Helper Functions - Common Token Segment Creation
|
||||
# ============================================
|
||||
|
||||
|
||||
def create_user_start() -> InputSegment:
|
||||
"""Create user role start token"""
|
||||
return create_segment(text="<|im_start|>user\n")
|
||||
|
||||
|
||||
def create_user_end() -> InputSegment:
|
||||
"""Create role end token"""
|
||||
return create_segment(text="<|im_end|>\n")
|
||||
|
||||
|
||||
def create_assistant_start() -> InputSegment:
|
||||
"""Create assistant role start token"""
|
||||
return create_segment(text="<|im_start|>assistant\n")
|
||||
|
||||
|
||||
def create_system_start() -> InputSegment:
|
||||
"""Create system role start token"""
|
||||
return create_segment(text="<|im_start|>system\n")
|
||||
|
||||
|
||||
def create_thinking_segment(thinking: bool = False) -> InputSegment:
|
||||
"""Create thinking token, closed or open based on thinking parameter"""
|
||||
if thinking:
|
||||
return create_segment(text="<think>\n")
|
||||
else:
|
||||
return create_segment(text="<think>\n\n</think>\n")
|
||||
|
||||
|
||||
def create_sostm_segment() -> InputSegment:
|
||||
"""Create streaming output start token"""
|
||||
return create_segment(text="<|sostm|>")
|
||||
|
||||
|
||||
def create_assistant_start_with_sostm() -> InputSegment:
|
||||
"""Create assistant start token with sostm"""
|
||||
return create_segment(text="<|im_start|>assistant\n<|sostm|>")
|
||||
|
||||
|
||||
def create_assistant_start_with_think() -> InputSegment:
|
||||
"""Create assistant start token with think opening"""
|
||||
return create_segment(text="<|im_start|>assistant\n<think>\n")
|
||||
|
||||
|
||||
# ============================================
|
||||
# Common Helper Functions - Composite Segment Creation
|
||||
# ============================================
|
||||
|
||||
|
||||
def create_user_turn_with_audio(audio_tokenized, extra_text: str = None) -> list[InputSegment]:
|
||||
"""Create a user turn containing audio"""
|
||||
segments = [
|
||||
create_user_start(),
|
||||
create_segment(audio=audio_tokenized),
|
||||
]
|
||||
if extra_text:
|
||||
segments.append(create_segment(text=extra_text))
|
||||
segments.append(create_user_end())
|
||||
return segments
|
||||
|
||||
|
||||
def create_user_turn_with_text(text: str) -> list[InputSegment]:
|
||||
"""Create a text-only user turn"""
|
||||
return [
|
||||
create_user_start(),
|
||||
create_segment(text=text),
|
||||
create_user_end(),
|
||||
]
|
||||
|
||||
|
||||
def create_system_turn_with_voice_prompt(prompt_text: str, audio_token) -> list[InputSegment]:
|
||||
"""Create a system turn with voice prompt"""
|
||||
return [
|
||||
create_system_start(),
|
||||
create_segment(text=prompt_text),
|
||||
create_segment(text="", audio=audio_token),
|
||||
create_user_end(),
|
||||
]
|
||||
|
||||
|
||||
def create_system_turn_text_only(system_text: str) -> list[InputSegment]:
|
||||
"""Create a text-only system turn"""
|
||||
return [
|
||||
create_system_start(),
|
||||
create_segment(text=system_text),
|
||||
create_user_end(),
|
||||
]
|
||||
|
||||
|
||||
# ============================================
|
||||
# Common Helper Functions - Multi-turn Dialogue Processing
|
||||
# ============================================
|
||||
|
||||
|
||||
def process_multiturn_messages(
|
||||
message_list: list[dict],
|
||||
user_processor: Callable[[dict], list[InputSegment]],
|
||||
assistant_processor: Callable[[dict], list[InputSegment]],
|
||||
) -> list[InputSegment]:
|
||||
"""
|
||||
Generic multi-turn dialogue message processing function
|
||||
|
||||
Args:
|
||||
message_list: List of messages, each containing 'role' and 'content'
|
||||
user_processor: Function to process user messages
|
||||
assistant_processor: Function to process assistant messages
|
||||
|
||||
Returns:
|
||||
Processed list of InputSegments
|
||||
"""
|
||||
lm_prompt = []
|
||||
for message in message_list:
|
||||
role = message["role"]
|
||||
if role == "user":
|
||||
lm_prompt.extend(user_processor(message))
|
||||
elif role == "assistant":
|
||||
lm_prompt.extend(assistant_processor(message))
|
||||
else:
|
||||
raise ValueError(f"Invalid role: {role}")
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def create_text_user_message(message: dict) -> list[InputSegment]:
|
||||
"""Process a text-only user message"""
|
||||
return [
|
||||
create_user_start(),
|
||||
create_segment(text=message["content"]),
|
||||
create_user_end(),
|
||||
]
|
||||
|
||||
|
||||
def create_text_assistant_message(message: dict) -> list[InputSegment]:
|
||||
"""Process a text-only assistant message"""
|
||||
return [
|
||||
create_assistant_start(),
|
||||
create_segment(text=message["content"]),
|
||||
create_user_end(),
|
||||
]
|
||||
|
||||
|
||||
def create_audio_user_message(message: dict) -> list[InputSegment]:
|
||||
"""Process an audio user message"""
|
||||
return [
|
||||
create_user_start(),
|
||||
create_segment(audio=preprocess_input(message["content"])),
|
||||
create_user_end(),
|
||||
]
|
||||
|
||||
|
||||
def append_assistant_ending(
|
||||
lm_prompt: list[InputSegment], thinking: bool = False, use_sostm: bool = False
|
||||
) -> list[InputSegment]:
|
||||
"""
|
||||
Append assistant ending to the prompt
|
||||
|
||||
Args:
|
||||
lm_prompt: Existing prompt list
|
||||
thinking: Whether to use open thinking token
|
||||
use_sostm: Whether to use sostm token (for speech output)
|
||||
"""
|
||||
if use_sostm:
|
||||
lm_prompt.append(create_assistant_start_with_sostm())
|
||||
else:
|
||||
lm_prompt.append(create_assistant_start())
|
||||
lm_prompt.append(create_thinking_segment(thinking))
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def get_asr_sft_prompt(
|
||||
input: None | str = None,
|
||||
):
|
||||
"""Build prompt for ASR (Automatic Speech Recognition) task"""
|
||||
audio_tokenized = preprocess_input(input)
|
||||
template = random.choice(asr_zh_templates + asr_en_templates)
|
||||
|
||||
lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=template)
|
||||
lm_prompt = append_assistant_ending(lm_prompt, thinking=False)
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def resample_audio_if_needed(wav_tensor: torch.Tensor, original_sr: int):
|
||||
target_sr = 24000
|
||||
if original_sr != target_sr:
|
||||
wav_tensor = torchaudio.functional.resample(wav_tensor, original_sr, target_sr)
|
||||
return wav_tensor
|
||||
|
||||
|
||||
def wav2mel(wav, device="cpu"):
|
||||
mel_transform = MelSpectrogram(
|
||||
sample_rate=mimo_audio_tokenizer.config.sampling_rate,
|
||||
n_fft=mimo_audio_tokenizer.config.nfft,
|
||||
hop_length=mimo_audio_tokenizer.config.hop_length,
|
||||
win_length=mimo_audio_tokenizer.config.window_size,
|
||||
f_min=mimo_audio_tokenizer.config.fmin,
|
||||
f_max=mimo_audio_tokenizer.config.fmax,
|
||||
n_mels=mimo_audio_tokenizer.config.n_mels,
|
||||
power=1.0,
|
||||
center=True,
|
||||
).to(device)
|
||||
spec = mel_transform(wav[None, :])
|
||||
return torch.log(torch.clip(spec, min=1e-7)).squeeze()
|
||||
|
||||
|
||||
def group_by_length(features: torch.Tensor, lengths: torch.Tensor, max_length: int):
|
||||
if features.size(0) != lengths.sum().item():
|
||||
raise ValueError(f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}")
|
||||
|
||||
split_points = []
|
||||
current_sum = 0
|
||||
|
||||
for i, seq_len in enumerate(lengths):
|
||||
if current_sum + seq_len > max_length and current_sum > 0:
|
||||
split_points.append(i)
|
||||
current_sum = seq_len.item()
|
||||
else:
|
||||
current_sum += seq_len.item()
|
||||
|
||||
# Convert split points to group sizes
|
||||
group_sizes = []
|
||||
prev = 0
|
||||
for point in split_points:
|
||||
group_sizes.append(point - prev)
|
||||
prev = point
|
||||
if prev < len(lengths):
|
||||
group_sizes.append(len(lengths) - prev)
|
||||
|
||||
len_groups = torch.split(lengths, group_sizes)
|
||||
feature_sizes = [group.sum().item() for group in len_groups]
|
||||
feature_groups = torch.split(features, feature_sizes)
|
||||
|
||||
return feature_groups, len_groups
|
||||
|
||||
|
||||
def encode_batch(input_features: torch.Tensor, input_lens: torch.Tensor, max_length: int = 256000):
|
||||
feature_groups, len_groups = group_by_length(input_features, input_lens, max_length)
|
||||
|
||||
encoded_parts = []
|
||||
for features, lengths in zip(feature_groups, len_groups):
|
||||
with torch.no_grad():
|
||||
codes, _ = mimo_audio_tokenizer.encoder.encode(
|
||||
input_features=features.to(device), input_lens=lengths.to(device), return_codes_only=True
|
||||
)
|
||||
encoded_parts.append(codes)
|
||||
|
||||
return torch.cat(encoded_parts, dim=-1)
|
||||
|
||||
|
||||
def preprocess_input(input: None | str | torch.Tensor = None, device="cpu", audio_channels=4, group_size=8):
|
||||
if isinstance(input, torch.Tensor) or (isinstance(input, str) and os.path.isfile(input)):
|
||||
return "<|sosp|><|empty|><|eosp|>"
|
||||
|
||||
else:
|
||||
text = input
|
||||
if (
|
||||
text.isupper() or text.islower()
|
||||
): # If the text only contains upper-case or lower-case letters, capitalize it.
|
||||
text = text.capitalize()
|
||||
return text
|
||||
|
||||
|
||||
def _build_tts_system_prompt(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
|
||||
"""Build system prompt for TTS task"""
|
||||
if has_voice_prompt and voice_audio_token is not None:
|
||||
return [
|
||||
create_system_start(),
|
||||
create_segment(
|
||||
# text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be: "
|
||||
text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是："
|
||||
),
|
||||
create_segment(text="", audio=voice_audio_token),
|
||||
create_user_end(),
|
||||
]
|
||||
else:
|
||||
return create_system_turn_text_only("你需要根据指定的风格指令和文本内容来生成语音。")
|
||||
# return create_system_turn_text_only(
|
||||
# "You need to generate speech based on the specified style instructions and text content."
|
||||
# )
|
||||
|
||||
|
||||
def _build_tts_system_prompt_no_instruct(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
|
||||
"""Build system prompt for TTS task"""
|
||||
if has_voice_prompt and voice_audio_token is not None:
|
||||
return [
|
||||
create_system_start(),
|
||||
create_segment(
|
||||
# text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be:"
|
||||
text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是："
|
||||
),
|
||||
create_segment(text="", audio=voice_audio_token),
|
||||
create_user_end(),
|
||||
]
|
||||
else:
|
||||
return []
|
||||
|
||||
|
||||
def get_tts_sft_prompt(
|
||||
input: None | str = None,
|
||||
instruct=None,
|
||||
read_text_only=True,
|
||||
prompt_speech=None,
|
||||
):
|
||||
"""
|
||||
Build prompt for TTS (Text-to-Speech) task
|
||||
|
||||
Args:
|
||||
input: Input text
|
||||
instruct: Style instruction (e.g., "speak happily in a child's voice")
|
||||
read_text_only: Whether to read only plain text (False means text contains template)
|
||||
prompt_speech: Reference audio (for voice cloning)
|
||||
"""
|
||||
assistant_prompt_audio_token = preprocess_input(prompt_speech) if prompt_speech is not None else None
|
||||
|
||||
if not read_text_only:
|
||||
# Not just reading text, text contains template (template:text format)
|
||||
text = preprocess_input(input)
|
||||
lm_prompt = _build_tts_system_prompt(
|
||||
has_voice_prompt=assistant_prompt_audio_token is not None,
|
||||
voice_audio_token=assistant_prompt_audio_token,
|
||||
)
|
||||
lm_prompt.append(create_segment(text=f"<|im_start|>user\n{text}<|im_end|>\n"))
|
||||
lm_prompt.append(create_assistant_start_with_think())
|
||||
else:
|
||||
# Plain text (no instruction inside)
|
||||
language = detect_language(input)
|
||||
template = random.choice(tts_zh_templates if language == "zh" else tts_en_templates)
|
||||
text = preprocess_input(input)
|
||||
|
||||
if instruct is None:
|
||||
# No instruct instruction
|
||||
lm_prompt = _build_tts_system_prompt_no_instruct(
|
||||
has_voice_prompt=assistant_prompt_audio_token is not None,
|
||||
voice_audio_token=assistant_prompt_audio_token,
|
||||
)
|
||||
lm_prompt.extend(
|
||||
[
|
||||
create_segment(text=f"<|im_start|>user\n{template}: {text}<|im_end|>\n"),
|
||||
create_assistant_start_with_sostm(),
|
||||
]
|
||||
)
|
||||
else:
|
||||
# Has instruct instruction
|
||||
lm_prompt = _build_tts_system_prompt(
|
||||
has_voice_prompt=assistant_prompt_audio_token is not None,
|
||||
voice_audio_token=assistant_prompt_audio_token,
|
||||
)
|
||||
lm_prompt.append(create_segment(text=f"<|im_start|>user\n{template}: {text}({instruct})<|im_end|>\n"))
|
||||
lm_prompt.append(create_assistant_start_with_think())
|
||||
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def get_audio_understanding_sft_prompt(
|
||||
input_speech,
|
||||
input_text,
|
||||
thinking=False,
|
||||
use_sostm=False,
|
||||
):
|
||||
"""Build prompt for audio understanding task"""
|
||||
audio_tokenized = preprocess_input(input_speech)
|
||||
|
||||
lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=input_text)
|
||||
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking, use_sostm=use_sostm)
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def _build_voice_prompt_system(prompt_speech) -> list[InputSegment]:
|
||||
"""Build system prompt with voice prompt"""
|
||||
return create_system_turn_with_voice_prompt(
|
||||
prompt_text="Your voice should be：", audio_token=preprocess_input(prompt_speech)
|
||||
)
|
||||
|
||||
|
||||
def get_spoken_dialogue_sft_prompt(
|
||||
input_speech,
|
||||
system_prompt=None,
|
||||
prompt_speech=None,
|
||||
add_history=False,
|
||||
):
|
||||
"""
|
||||
Build prompt for spoken dialogue task
|
||||
|
||||
Args:
|
||||
input_speech: Input speech
|
||||
system_prompt: System prompt text
|
||||
prompt_speech: Reference audio (for voice cloning)
|
||||
add_history: Whether to add history (Note: history variable is undefined in original code)
|
||||
"""
|
||||
audio_tokenized = preprocess_input(input_speech)
|
||||
lm_prompt = []
|
||||
|
||||
# Note: history variable is undefined in original code, this branch may never execute
|
||||
# To use history feature, history should be passed as a parameter
|
||||
if add_history:
|
||||
# Simplified form of adding history
|
||||
lm_prompt = create_user_turn_with_audio(audio_tokenized)
|
||||
lm_prompt.append(create_assistant_start_with_sostm())
|
||||
else:
|
||||
# Add voice prompt (if available)
|
||||
if prompt_speech:
|
||||
lm_prompt.extend(_build_voice_prompt_system(prompt_speech))
|
||||
|
||||
# Add user turn
|
||||
lm_prompt.append(create_user_start())
|
||||
if system_prompt:
|
||||
lm_prompt.append(create_segment(text=system_prompt))
|
||||
lm_prompt.append(create_segment(audio=audio_tokenized))
|
||||
lm_prompt.append(create_user_end())
|
||||
lm_prompt.append(create_assistant_start_with_sostm())
|
||||
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def get_spoken_dialogue_sft_multiturn_prompt(
|
||||
message_list,
|
||||
system_prompt=None,
|
||||
prompt_speech=None,
|
||||
tokenizer=None,
|
||||
group_size=8,
|
||||
audio_channels=4,
|
||||
):
|
||||
"""
|
||||
Build prompt for multi-turn spoken dialogue task
|
||||
|
||||
Args:
|
||||
message_list: List of messages containing role and content
|
||||
system_prompt: System prompt text
|
||||
prompt_speech: Reference audio (for voice cloning)
|
||||
tokenizer: Tokenizer
|
||||
group_size: Group size
|
||||
audio_channels: Number of audio channels
|
||||
"""
|
||||
lm_prompt = []
|
||||
|
||||
# Add voice prompt (if available)
|
||||
if prompt_speech:
|
||||
lm_prompt.extend(
|
||||
create_system_turn_with_voice_prompt(
|
||||
prompt_text="Your voice should be:", audio_token=preprocess_input(prompt_speech)
|
||||
)
|
||||
)
|
||||
|
||||
# Define message processors
|
||||
def user_processor(msg):
|
||||
segments = [create_user_start()]
|
||||
if system_prompt:
|
||||
segments.append(create_segment(text=system_prompt))
|
||||
segments.append(create_segment(audio=preprocess_input(msg["content"])))
|
||||
segments.append(create_user_end())
|
||||
return segments
|
||||
|
||||
def assistant_processor(msg):
|
||||
return [
|
||||
create_assistant_start(),
|
||||
create_streaming_segment(
|
||||
text=msg["content"]["text"],
|
||||
audio=preprocess_input(msg["content"]["audio"]),
|
||||
tokenizer=tokenizer,
|
||||
group_size=group_size,
|
||||
audio_channels=audio_channels,
|
||||
),
|
||||
create_user_end(),
|
||||
]
|
||||
|
||||
# Process message list
|
||||
lm_prompt.extend(process_multiturn_messages(message_list, user_processor, assistant_processor))
|
||||
lm_prompt.append(create_assistant_start_with_sostm())
|
||||
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def get_s2t_dialogue_sft_prompt(
|
||||
input_speech,
|
||||
thinking=False,
|
||||
):
|
||||
"""Build prompt for speech-to-text dialogue task"""
|
||||
audio_tokenized = preprocess_input(input_speech)
|
||||
|
||||
lm_prompt = create_user_turn_with_audio(audio_tokenized)
|
||||
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def get_s2t_dialogue_sft_multiturn_prompt(message_list, thinking=False):
|
||||
"""Build prompt for multi-turn speech-to-text dialogue task"""
|
||||
lm_prompt = process_multiturn_messages(
|
||||
message_list, user_processor=create_audio_user_message, assistant_processor=create_text_assistant_message
|
||||
)
|
||||
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def get_text_dialogue_sft_prompt(
|
||||
input_text,
|
||||
thinking=False,
|
||||
):
|
||||
"""Build prompt for text-only dialogue task"""
|
||||
lm_prompt = create_user_turn_with_text(input_text)
|
||||
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def get_text_dialogue_sft_multiturn_prompt(
|
||||
message_list,
|
||||
thinking=False,
|
||||
):
|
||||
"""Build prompt for multi-turn text-only dialogue task"""
|
||||
lm_prompt = process_multiturn_messages(
|
||||
message_list, user_processor=create_text_user_message, assistant_processor=create_text_assistant_message
|
||||
)
|
||||
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
|
||||
return lm_prompt
|
||||
|
||||
|
||||
def get_in_context_learning_s2s_prompt(
|
||||
instruction,
|
||||
prompt_examples,
|
||||
audio,
|
||||
tokenizer=None,
|
||||
group_size=8,
|
||||
audio_channels=4,
|
||||
):
|
||||
"""
|
||||
Build prompt for In-Context Learning speech-to-speech task
|
||||
|
||||
Args:
|
||||
instruction: Instruction text
|
||||
prompt_examples: List of examples, each containing input_audio, output_transcription, output_audio
|
||||
audio: Input audio to be processed
|
||||
tokenizer: Tokenizer
|
||||
group_size: Group size
|
||||
audio_channels: Number of audio channels
|
||||
"""
|
||||
prompt = [create_segment(text=f"[Int]:{instruction}\n")]
|
||||
|
||||
# Add examples
|
||||
for example in prompt_examples:
|
||||
prompt.extend(
|
||||
[
|
||||
create_segment(audio=preprocess_input(example["input_audio"])),
|
||||
create_segment(text="\n"),
|
||||
create_streaming_segment(
|
||||
text=example["output_transcription"],
|
||||
audio=preprocess_input(example["output_audio"]),
|
||||
tokenizer=tokenizer,
|
||||
group_size=group_size,
|
||||
audio_channels=audio_channels,
|
||||
),
|
||||
create_segment(text=" \n\n"),
|
||||
]
|
||||
)
|
||||
|
||||
# Add input audio to be processed
|
||||
prompt.extend(
|
||||
[
|
||||
create_segment(audio=preprocess_input(audio)),
|
||||
create_segment(text="\n"),
|
||||
create_sostm_segment(),
|
||||
]
|
||||
)
|
||||
|
||||
return prompt
|
||||
|
||||
|
||||
def get_audio_data(audio_url):
|
||||
if audio_url.startswith("data:"):
|
||||
header, b64_data = audio_url.split(",", 1)
|
||||
audio_bytes = base64.b64decode(b64_data.strip())
|
||||
audio_file = io.BytesIO(audio_bytes)
|
||||
else:
|
||||
# File path
|
||||
audio_file = audio_url
|
||||
|
||||
audio_signal, sr = librosa.load(audio_file, sr=24000)
|
||||
audio_data = (audio_signal.astype(np.float32), sr)
|
||||
return audio_data
|
||||
|
||||
|
||||
def to_prompt(input_segs):
|
||||
out_put = []
|
||||
|
||||
for input_seg in input_segs:
|
||||
if isinstance(input_seg, StreamingInputSegment) and input_seg.text:
|
||||
out_put.append("<|sostm|>")
|
||||
if input_seg.audio is not None and isinstance(input_seg.audio, str):
|
||||
out_put.append(input_seg.text)
|
||||
out_put.append("<|eot|>")
|
||||
out_put.append("<|empty|>")
|
||||
else:
|
||||
out_put.append(input_seg.text)
|
||||
out_put.append("<|eot|>")
|
||||
out_put.append("<|eostm|>")
|
||||
|
||||
else:
|
||||
out_put.append(input_seg.text)
|
||||
if input_seg.audio is not None:
|
||||
out_put.append(input_seg.audio)
|
||||
|
||||
prompt = "".join(out_put)
|
||||
return prompt
|
||||
@@ -0,0 +1,261 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright 2025 Xiaomi Corporation.
|
||||
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class InputSegment:
|
||||
def __init__(
|
||||
self,
|
||||
text: str = "",
|
||||
audio: torch.Tensor = None,
|
||||
tokenized_text: torch.Tensor = None,
|
||||
speech_zeroemb_idx: int | list[int] = 1024,
|
||||
text_zeroemb_idx: int = 152067,
|
||||
add_sosp_eosp=True,
|
||||
) -> None:
|
||||
has_text = text is not None
|
||||
has_tokenized_text = tokenized_text is not None
|
||||
assert has_text or has_tokenized_text, "Text or tokenized text must be provided"
|
||||
|
||||
self.audio = audio
|
||||
self.text = text
|
||||
self.tokenized_text = tokenized_text
|
||||
self.speech_zeroemb_idx = speech_zeroemb_idx
|
||||
self.text_zeroemb_idx = text_zeroemb_idx
|
||||
self.add_sosp_eosp = add_sosp_eosp
|
||||
|
||||
@staticmethod
|
||||
def insert_between(tensor, i, value=-1):
|
||||
return torch.scatter(
|
||||
torch.full(
|
||||
(1, tensor.shape[1] + (tensor.shape[1] - 1) * i + i),
|
||||
value,
|
||||
dtype=tensor.dtype,
|
||||
),
|
||||
1,
|
||||
torch.arange(0, tensor.shape[1], dtype=torch.int64)[None] * (i + 1),
|
||||
tensor,
|
||||
)
|
||||
|
||||
def to_input_id(
|
||||
self,
|
||||
tokenizer,
|
||||
group_size: int,
|
||||
audio_channels: int = 8,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if (
|
||||
self.audio is None
|
||||
): ## If no audio is provided, tokenize the text with (group_size-1) tokens inserted between each, fill with -100; and set audio to empty indices
|
||||
if self.tokenized_text is None:
|
||||
tokenized_text = tokenizer(
|
||||
self.text,
|
||||
return_tensors="pt",
|
||||
truncation=True,
|
||||
max_length=999999,
|
||||
padding=False,
|
||||
add_special_tokens=False,
|
||||
)["input_ids"].int()
|
||||
else:
|
||||
tokenized_text = self.tokenized_text.unsqueeze(0)
|
||||
|
||||
if group_size > 1:
|
||||
tokenized_text = self.insert_between(tokenized_text, group_size - 1, value=-100)
|
||||
|
||||
if isinstance(self.speech_zeroemb_idx, list):
|
||||
audio_part_input_id = torch.zeros((audio_channels, tokenized_text.shape[1]), dtype=torch.int)
|
||||
for i, idx in enumerate(self.speech_zeroemb_idx):
|
||||
audio_part_input_id[i, :] = idx
|
||||
else:
|
||||
audio_part_input_id = torch.full(
|
||||
(audio_channels, tokenized_text.shape[1]), self.speech_zeroemb_idx, dtype=torch.int
|
||||
)
|
||||
|
||||
else: # If audio is provided, add sosp/eosp markers and convert audio: fill with empty tokens based on audio length
|
||||
sosp_token = tokenizer.convert_tokens_to_ids("<|sosp|>") if self.add_sosp_eosp else None
|
||||
eosp_token = tokenizer.convert_tokens_to_ids("<|eosp|>") if self.add_sosp_eosp else None
|
||||
audio_part = self.audio.reshape(-1, audio_channels).T # [audio_channels, seqlen]
|
||||
|
||||
assert audio_part.shape[1] % group_size == 0, (
|
||||
f"Audio shape {audio_part.shape} is not divisible by group_size {group_size}"
|
||||
)
|
||||
|
||||
text_len = audio_part.shape[1] // group_size
|
||||
empty_token = self.text_zeroemb_idx
|
||||
if empty_token is None:
|
||||
empty_token = tokenizer.eod
|
||||
tokenized_text = torch.full((1, text_len), empty_token, dtype=torch.int)
|
||||
|
||||
tokenized_text = (
|
||||
torch.cat(
|
||||
[
|
||||
torch.tensor([[sosp_token]], dtype=torch.int),
|
||||
tokenized_text,
|
||||
torch.tensor([[eosp_token]], dtype=torch.int),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
if self.add_sosp_eosp
|
||||
else tokenized_text
|
||||
)
|
||||
tokenized_text = self.insert_between(tokenized_text, group_size - 1, value=-100)
|
||||
|
||||
if self.add_sosp_eosp:
|
||||
if isinstance(self.speech_zeroemb_idx, list):
|
||||
sosp_part = torch.zeros((audio_channels, group_size), dtype=torch.int)
|
||||
eosp_part = torch.zeros((audio_channels, group_size), dtype=torch.int)
|
||||
for i, idx in enumerate(self.speech_zeroemb_idx):
|
||||
sosp_part[i, :] = idx
|
||||
eosp_part[i, :] = idx
|
||||
audio_part_input_id = torch.cat([sosp_part, audio_part, eosp_part], dim=1)
|
||||
else:
|
||||
audio_part_input_id = torch.cat(
|
||||
[
|
||||
torch.full((audio_channels, group_size), self.speech_zeroemb_idx, dtype=torch.int),
|
||||
audio_part,
|
||||
torch.full((audio_channels, group_size), self.speech_zeroemb_idx, dtype=torch.int),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
else:
|
||||
audio_part_input_id = audio_part
|
||||
|
||||
input_ids = torch.cat([tokenized_text, audio_part_input_id], dim=0) # [n_rvq + 1, seqlen]
|
||||
|
||||
return input_ids
|
||||
|
||||
|
||||
class StreamingInputSegment:
|
||||
def __init__(
|
||||
self,
|
||||
text: str = "",
|
||||
audio: torch.Tensor = None,
|
||||
tokenized_text: torch.Tensor = None,
|
||||
speech_zeroemb_idx: int | list[int] = 1024,
|
||||
text_zeroemb_idx: int = 152067,
|
||||
text_segment_size: int = 5,
|
||||
audio_segment_size: int = 5,
|
||||
tokenizer=None,
|
||||
group_size=None,
|
||||
audio_channels=None,
|
||||
) -> None:
|
||||
has_text = text is not None
|
||||
has_tokenized_text = tokenized_text is not None
|
||||
assert has_text or has_tokenized_text, "Text or tokenized text must be provided"
|
||||
|
||||
self.audio = audio
|
||||
self.text = text
|
||||
self.tokenized_text = tokenized_text
|
||||
self.speech_zeroemb_idx = speech_zeroemb_idx
|
||||
self.text_zeroemb_idx = text_zeroemb_idx
|
||||
self.text_segment_size = text_segment_size
|
||||
self.audio_segment_size = audio_segment_size
|
||||
self.tokenizer = tokenizer
|
||||
self.group_size = group_size
|
||||
self.audio_channels = audio_channels
|
||||
|
||||
def to_input_id(
|
||||
self,
|
||||
tokenizer,
|
||||
group_size: int,
|
||||
audio_channels: int = 8,
|
||||
):
|
||||
if self.tokenized_text is None:
|
||||
tokenized_text = tokenizer(
|
||||
self.text,
|
||||
return_tensors="pt",
|
||||
truncation=True,
|
||||
max_length=999999,
|
||||
padding=False,
|
||||
add_special_tokens=False,
|
||||
)["input_ids"].int() # [1, seqlen]
|
||||
else:
|
||||
tokenized_text = self.tokenized_text.unsqueeze(0)
|
||||
|
||||
tokenized_text = tokenized_text.squeeze(0)
|
||||
|
||||
text_segments = tokenized_text.split(self.text_segment_size, dim=0)
|
||||
audio_segments = self.audio.split(self.audio_segment_size * group_size * audio_channels, dim=0)
|
||||
|
||||
tokenized_segments = []
|
||||
tokenized_segments.append(
|
||||
InputSegment(
|
||||
text="<|sostm|>",
|
||||
speech_zeroemb_idx=self.speech_zeroemb_idx,
|
||||
text_zeroemb_idx=self.text_zeroemb_idx,
|
||||
),
|
||||
)
|
||||
|
||||
eot_tokens = tokenizer(
|
||||
"<|eot|>",
|
||||
return_tensors="pt",
|
||||
truncation=True,
|
||||
max_length=999999,
|
||||
padding=False,
|
||||
add_special_tokens=False,
|
||||
)["input_ids"][0].to(text_segments[-1])
|
||||
|
||||
text_segments = text_segments[:-1] + (torch.cat([text_segments[-1], eot_tokens], dim=0),)
|
||||
|
||||
length = min(len(text_segments), len(audio_segments)) # The number of empty tokens here may vary
|
||||
for i in range(length):
|
||||
text_segment = text_segments[i]
|
||||
audio_segment = audio_segments[i]
|
||||
|
||||
tokenized_segments.append(
|
||||
InputSegment(
|
||||
tokenized_text=text_segment,
|
||||
speech_zeroemb_idx=self.speech_zeroemb_idx,
|
||||
text_zeroemb_idx=self.text_zeroemb_idx,
|
||||
),
|
||||
)
|
||||
tokenized_segments.append(
|
||||
InputSegment(
|
||||
audio=audio_segment,
|
||||
add_sosp_eosp=False,
|
||||
speech_zeroemb_idx=self.speech_zeroemb_idx,
|
||||
text_zeroemb_idx=self.text_zeroemb_idx,
|
||||
),
|
||||
)
|
||||
|
||||
for j in range(length, len(text_segments)):
|
||||
tokenized_segments.append(
|
||||
InputSegment(
|
||||
tokenized_text=text_segments[j],
|
||||
speech_zeroemb_idx=self.speech_zeroemb_idx,
|
||||
text_zeroemb_idx=self.text_zeroemb_idx,
|
||||
),
|
||||
)
|
||||
|
||||
for j in range(length, len(audio_segments)):
|
||||
tokenized_segments.append(
|
||||
InputSegment(
|
||||
audio=audio_segments[j],
|
||||
add_sosp_eosp=False,
|
||||
speech_zeroemb_idx=self.speech_zeroemb_idx,
|
||||
text_zeroemb_idx=self.text_zeroemb_idx,
|
||||
),
|
||||
)
|
||||
|
||||
tokenized_segments.append(
|
||||
InputSegment(
|
||||
text="<|eostm|>",
|
||||
speech_zeroemb_idx=self.speech_zeroemb_idx,
|
||||
text_zeroemb_idx=self.text_zeroemb_idx,
|
||||
),
|
||||
)
|
||||
|
||||
input_ids = [
|
||||
seg.to_input_id(
|
||||
self.tokenizer,
|
||||
self.group_size,
|
||||
self.audio_channels,
|
||||
)
|
||||
for seg in tokenized_segments
|
||||
]
|
||||
|
||||
input_ids = torch.cat(input_ids, dim=1).type(torch.int64) # [n_rvq + 1, seqlen]
|
||||
|
||||
return input_ids
|
||||
@@ -0,0 +1,58 @@
|
||||
# Online serving Example of vLLM-Omni for MiMo-Audio
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
Please refer to [README.md](../../../README.md)
|
||||
|
||||
## Run examples (MiMo-Audio)
|
||||
|
||||
### Launch the Server
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
|
||||
vllm-omni serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
|
||||
--served-model-name "MiMo-Audio-7B-Instruct" \
|
||||
--port 18091 --stage-configs-path ./vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
|
||||
```
|
||||
> ⚠️ **Important**
|
||||
> **MiMo-Audio is not compatible with the default chat template.**
|
||||
> The provided `chat_template.jinja` implements MiMo-specific role, audio token, and instruction formatting and **must be used for all inference**.
|
||||
|
||||
|
||||
### Send Multi-modal Request
|
||||
|
||||
Get into the example folder
|
||||
```bash
|
||||
cd examples/online_serving/mimo_audio
|
||||
```
|
||||
|
||||
#### Send request via python
|
||||
|
||||
```bash
|
||||
# Audio dialogue task
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type multi_audios \
|
||||
--message-json ../../offline_inference/mimo_audio/message_base64_wav.json
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
|
||||
- `--query-type` (or `-q`): Query type (default: `multi_audios`)
|
||||
- Options: `multi_audios`, `text`
|
||||
- `--message-json` (or `-m`): Path to `base64` multi rounds audio messages json file
|
||||
- Do not pass any value for "text" query type
|
||||
- Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs, only for "Are these two audio clips the same?" task
|
||||
- Example: `---message-json ./examples/offline_inference/mimo_audio/message_base64_wav.json`
|
||||
- `--prompt` (or `-p`): Custom text prompt/question, only for query type is "text"(TTS task)
|
||||
- Attention! Do not pass any value for "multi_audios" query type
|
||||
- Example: `--prompt "What are the main activities shown in this video?"`
|
||||
|
||||
|
||||
For example, to use multi rounds audios with local files:
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type multi_audios \
|
||||
--message-json ../../offline_inference/mimo_audio/message_base64_wav.json
|
||||
```
|
||||
@@ -0,0 +1,86 @@
|
||||
{%- if tools %}
|
||||
{{- '<|im_start|>system\n' }}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{{- messages[0]['content'] }}
|
||||
{%- else %}
|
||||
{{- 'You are a helpful assistant.' }}
|
||||
{%- endif %}
|
||||
{{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
|
||||
{%- for tool in tools %}
|
||||
{{- "\n" }}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
|
||||
{%- else %}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{%- set _m = '<|sosp|><|empty|><|eosp|>' -%}
|
||||
{%- set _raw0 = messages[0]['content'] if messages[0]['content'] is string else '' -%}
|
||||
{%- if _m in _raw0 %}
|
||||
{%- set _t0 = (_raw0 | replace(_m ~ '\n', '') | replace(_m, '') | trim) -%}
|
||||
{{- '<|im_start|>system\n' + (_t0 ~ _m if _t0 else _m) + '<|im_end|>\n' }}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>system\n' + _raw0 + '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
|
||||
{%- for message in messages %}
|
||||
{%- if message['role'] == 'assistant' %}
|
||||
{{- '<|im_start|>assistant' }}
|
||||
{%- set _sosp = '<|sosp|><|empty|><|eosp|>' -%}
|
||||
{%- set _text = message['content'] if message['content'] is string else '' -%}
|
||||
{%- if _sosp in _text %}
|
||||
{%- set _clean = _text | replace(_sosp, '') -%}
|
||||
{%- set _body = _clean[1:] if (_clean and _clean[0] == '\n') else _clean -%}
|
||||
{{- '\n<|sostm|>' + _body + '<|eot|><|empty|><|eostm|>' }}
|
||||
{%- else %}
|
||||
{%- set _body = _text[1:] if (_text and _text[0] == '\n') else _text -%}
|
||||
{{- '\n<|sostm|>' + _body + '<|eot|><|eostm|>' }}
|
||||
{%- endif %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
|
||||
{%- elif message['role'] == 'user' %}
|
||||
{%- set _m = '<|sosp|><|empty|><|eosp|>' -%}
|
||||
{%- set _raw = message['content'] if message['content'] is string else '' -%}
|
||||
{%- if _m in _raw %}
|
||||
{%- set _t = (_raw | replace(_m ~ '\n', '') | replace(_m, '') | trim) -%}
|
||||
{{- '<|im_start|>user\n' + (_t ~ _m if _t else _m) + '<|im_end|>\n' }}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>user\n' + _raw + '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
|
||||
{%- elif message['role'] == 'system' %}
|
||||
{%- if not loop.first %}
|
||||
{%- set _m = '<|sosp|><|empty|><|eosp|>' -%}
|
||||
{%- set _raw = message['content'] if message['content'] is string else '' -%}
|
||||
{%- if _m in _raw %}
|
||||
{%- set _t = (_raw | replace(_m ~ '\n', '') | replace(_m, '') | trim) -%}
|
||||
{{- '<|im_start|>system\n' + (_t ~ _m if _t else _m) + '<|im_end|>\n' }}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>system\n' + _raw + '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
|
||||
{%- elif message['role'] == 'tool' %}
|
||||
{%- if (loop.index0 == 0) or (messages[loop.index0 - 1]['role'] != 'tool') %}
|
||||
{{- '<|im_start|>user' }}
|
||||
{%- endif %}
|
||||
{%- set _m = '<|sosp|><|empty|><|eosp|>' -%}
|
||||
{%- set _raw = message['content'] if message['content'] is string else '' -%}
|
||||
{%- if _m in _raw %}
|
||||
{%- set _t = (_raw | replace(_m ~ '\n', '') | replace(_m, '') | trim) -%}
|
||||
{{- '\n<tool_response>\n' + (_t ~ _m if _t else _m) + '\n</tool_response>' }}
|
||||
{%- else %}
|
||||
{{- '\n<tool_response>\n' + _raw + '\n</tool_response>' }}
|
||||
{%- endif %}
|
||||
{%- if loop.last or (messages[loop.index0 + 1]['role'] != 'tool') %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n<|sostm|>' }}
|
||||
{%- endif %}
|
||||
+545
@@ -0,0 +1,545 @@
|
||||
import base64
|
||||
import concurrent.futures
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from openai import OpenAI
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:18091/v1"
|
||||
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
|
||||
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
SEED = 42
|
||||
|
||||
|
||||
def encode_base64_content_from_url(content_url: str) -> str:
|
||||
"""Encode a content retrieved from a remote url to base64 format."""
|
||||
|
||||
with requests.get(content_url) as response:
|
||||
response.raise_for_status()
|
||||
result = base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def encode_base64_content_from_file(file_path: str) -> str:
|
||||
"""Encode a local file to base64 format."""
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
result = base64.b64encode(content).decode("utf-8")
|
||||
return result
|
||||
|
||||
|
||||
def get_video_url_from_path(video_path: str | None) -> str:
|
||||
"""Convert a video path (local file or URL) to a video URL format for the API.
|
||||
|
||||
If video_path is None or empty, returns the default URL.
|
||||
If video_path is a local file path, encodes it to base64 data URL.
|
||||
If video_path is a URL, returns it as-is.
|
||||
"""
|
||||
if not video_path:
|
||||
# Default video URL
|
||||
return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
|
||||
# Check if it's a URL (starts with http:// or https://)
|
||||
if video_path.startswith(("http://", "https://")):
|
||||
return video_path
|
||||
|
||||
# Otherwise, treat it as a local file path
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"Video file not found: {video_path}")
|
||||
|
||||
# Detect video MIME type from file extension
|
||||
video_path_lower = video_path.lower()
|
||||
if video_path_lower.endswith(".mp4"):
|
||||
mime_type = "video/mp4"
|
||||
elif video_path_lower.endswith(".webm"):
|
||||
mime_type = "video/webm"
|
||||
elif video_path_lower.endswith(".mov"):
|
||||
mime_type = "video/quicktime"
|
||||
elif video_path_lower.endswith(".avi"):
|
||||
mime_type = "video/x-msvideo"
|
||||
elif video_path_lower.endswith(".mkv"):
|
||||
mime_type = "video/x-matroska"
|
||||
else:
|
||||
# Default to mp4 if extension is unknown
|
||||
mime_type = "video/mp4"
|
||||
|
||||
video_base64 = encode_base64_content_from_file(video_path)
|
||||
return f"data:{mime_type};base64,{video_base64}"
|
||||
|
||||
|
||||
def get_image_url_from_path(image_path: str | None) -> str:
|
||||
"""Convert an image path (local file or URL) to an image URL format for the API.
|
||||
|
||||
If image_path is None or empty, returns the default URL.
|
||||
If image_path is a local file path, encodes it to base64 data URL.
|
||||
If image_path is a URL, returns it as-is.
|
||||
"""
|
||||
if not image_path:
|
||||
# Default image URL
|
||||
return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
|
||||
|
||||
# Check if it's a URL (starts with http:// or https://)
|
||||
if image_path.startswith(("http://", "https://")):
|
||||
return image_path
|
||||
|
||||
# Otherwise, treat it as a local file path
|
||||
if not os.path.exists(image_path):
|
||||
raise FileNotFoundError(f"Image file not found: {image_path}")
|
||||
|
||||
# Detect image MIME type from file extension
|
||||
image_path_lower = image_path.lower()
|
||||
if image_path_lower.endswith((".jpg", ".jpeg")):
|
||||
mime_type = "image/jpeg"
|
||||
elif image_path_lower.endswith(".png"):
|
||||
mime_type = "image/png"
|
||||
elif image_path_lower.endswith(".gif"):
|
||||
mime_type = "image/gif"
|
||||
elif image_path_lower.endswith(".webp"):
|
||||
mime_type = "image/webp"
|
||||
else:
|
||||
# Default to jpeg if extension is unknown
|
||||
mime_type = "image/jpeg"
|
||||
|
||||
image_base64 = encode_base64_content_from_file(image_path)
|
||||
return f"data:{mime_type};base64,{image_base64}"
|
||||
|
||||
|
||||
def get_audio_url_from_path(audio_path: str | None) -> str:
|
||||
"""Convert an audio path (local file or URL) to an audio URL format for the API.
|
||||
|
||||
If audio_path is None or empty, returns the default URL.
|
||||
If audio_path is already a base64 data URL (starts with "data:"), returns it as-is.
|
||||
If audio_path is a URL (starts with http:// or https://), returns it as-is.
|
||||
If audio_path is a local file path, encodes it to base64 data URL.
|
||||
"""
|
||||
if not audio_path:
|
||||
# Default audio URL
|
||||
return AudioAsset("mary_had_lamb").url
|
||||
|
||||
# Check if it's already a base64 data URL
|
||||
if audio_path.startswith("data:"):
|
||||
return audio_path
|
||||
|
||||
# Check if it's a URL (starts with http:// or https://)
|
||||
if audio_path.startswith(("http://", "https://")):
|
||||
return audio_path
|
||||
|
||||
# Otherwise, treat it as a local file path
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
|
||||
# Detect audio MIME type from file extension
|
||||
audio_path_lower = audio_path.lower()
|
||||
if audio_path_lower.endswith((".mp3", ".mpeg")):
|
||||
mime_type = "audio/mpeg"
|
||||
elif audio_path_lower.endswith(".wav"):
|
||||
mime_type = "audio/wav"
|
||||
elif audio_path_lower.endswith(".ogg"):
|
||||
mime_type = "audio/ogg"
|
||||
elif audio_path_lower.endswith(".flac"):
|
||||
mime_type = "audio/flac"
|
||||
elif audio_path_lower.endswith(".m4a"):
|
||||
mime_type = "audio/mp4"
|
||||
else:
|
||||
# Default to wav if extension is unknown
|
||||
mime_type = "audio/wav"
|
||||
|
||||
audio_base64 = encode_base64_content_from_file(audio_path)
|
||||
return f"data:{mime_type};base64,{audio_base64}"
|
||||
|
||||
|
||||
def load_messages_from_json(message_json_path: str) -> dict[str, Any]:
|
||||
"""Load messages from a JSON file."""
|
||||
with open(message_json_path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return data
|
||||
|
||||
|
||||
def process_audio_url_in_content(content: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Process audio_url in content, handling both file paths and base64 URLs."""
|
||||
if content.get("type") == "audio_url":
|
||||
audio_url = content.get("audio_url", {}).get("url")
|
||||
if audio_url:
|
||||
# Process the audio URL (handles base64, file paths, and URLs)
|
||||
processed_url = get_audio_url_from_path(audio_url)
|
||||
content = content.copy()
|
||||
content["audio_url"] = {"url": processed_url}
|
||||
return content
|
||||
|
||||
|
||||
def get_system_prompt(message_json_path: str | None = None):
|
||||
"""Get system prompt, optionally from message.json file."""
|
||||
if message_json_path and os.path.exists(message_json_path):
|
||||
data = load_messages_from_json(message_json_path)
|
||||
messages = data.get("messages", [])
|
||||
# Find the first system message
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
# Process audio URLs in the content
|
||||
processed_content = []
|
||||
for content_item in msg.get("content", []):
|
||||
processed_item = process_audio_url_in_content(content_item)
|
||||
processed_content.append(processed_item)
|
||||
return {
|
||||
"role": "system",
|
||||
"content": processed_content,
|
||||
}
|
||||
|
||||
|
||||
def get_text_query(custom_prompt: str | None = None):
|
||||
question = f"Please convert this text to speech: {custom_prompt}"
|
||||
prompt = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{question}",
|
||||
}
|
||||
],
|
||||
}
|
||||
return prompt
|
||||
|
||||
|
||||
def get_multi_audios_query(
|
||||
audio_path: str | None = None,
|
||||
custom_prompt: str | None = None,
|
||||
message_json_path: str | None = None,
|
||||
) -> list[dict[str, Any]] | dict[str, Any]:
|
||||
"""Get multi-audios query, optionally from message.json file."""
|
||||
if message_json_path and os.path.exists(message_json_path):
|
||||
data = load_messages_from_json(message_json_path)
|
||||
messages = data.get("messages", [])
|
||||
|
||||
# Skip the first system role and return all other messages
|
||||
rebuild_prompt_messages = []
|
||||
skipped_first_system = False
|
||||
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system" and not skipped_first_system:
|
||||
skipped_first_system = True
|
||||
continue
|
||||
|
||||
# Process all content items in the message
|
||||
processed_content = []
|
||||
for content_item in msg.get("content", []):
|
||||
processed_item = process_audio_url_in_content(content_item)
|
||||
processed_content.append(processed_item)
|
||||
|
||||
rebuild_prompt_messages.append(
|
||||
{
|
||||
"role": msg.get("role"),
|
||||
"content": processed_content,
|
||||
}
|
||||
)
|
||||
|
||||
return rebuild_prompt_messages
|
||||
|
||||
# Original behavior when message_json_path is not provided
|
||||
question = custom_prompt or "Are these two audio clips the same?"
|
||||
audio_url = get_audio_url_from_path(audio_path)
|
||||
prompt = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": audio_url},
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": AudioAsset("winning_call").url},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{question}",
|
||||
},
|
||||
],
|
||||
}
|
||||
return prompt
|
||||
|
||||
|
||||
query_map = {
|
||||
"multi_audios": get_multi_audios_query,
|
||||
"text": get_text_query,
|
||||
}
|
||||
|
||||
|
||||
def run_multimodal_generation(args) -> None:
|
||||
model_name = "MiMo-Audio-7B-Instruct"
|
||||
thinker_sampling_params = {
|
||||
"temperature": 0.0, # Deterministic - no randomness
|
||||
"top_p": 1.0, # Disable nucleus sampling
|
||||
"top_k": -1, # Disable top-k sampling
|
||||
"max_tokens": 3200,
|
||||
"seed": SEED, # Fixed seed for sampling
|
||||
"detokenize": True,
|
||||
"repetition_penalty": 1.1,
|
||||
}
|
||||
code2wav_sampling_params = {
|
||||
"temperature": 0.0, # Deterministic - no randomness
|
||||
"top_p": 1.0, # Disable nucleus sampling
|
||||
"top_k": -1, # Disable top-k sampling
|
||||
"max_tokens": 4096 * 16,
|
||||
"seed": SEED, # Fixed seed for sampling
|
||||
"detokenize": True,
|
||||
"repetition_penalty": 1.1,
|
||||
}
|
||||
|
||||
sampling_params_list = [
|
||||
thinker_sampling_params,
|
||||
code2wav_sampling_params,
|
||||
]
|
||||
|
||||
# Get paths and custom prompt from args
|
||||
audio_path = getattr(args, "audio_path", None)
|
||||
custom_prompt = getattr(args, "prompt", None)
|
||||
message_json_path = getattr(args, "message_json", None)
|
||||
output_audio_path = getattr(args, "output_audio_path", None)
|
||||
|
||||
# Get the query function and call it with appropriate parameters
|
||||
query_func = query_map[args.query_type]
|
||||
if args.query_type == "multi_audios":
|
||||
prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt, message_json_path=message_json_path)
|
||||
elif args.query_type == "text":
|
||||
prompt = query_func(custom_prompt=custom_prompt)
|
||||
else:
|
||||
prompt = query_func()
|
||||
|
||||
extra_body = {
|
||||
"sampling_params_list": sampling_params_list
|
||||
# Optional, it has a default setting in stage_configs of the corresponding model.
|
||||
}
|
||||
|
||||
# Build messages list (filter None so concurrent tasks get valid structure)
|
||||
system = get_system_prompt(message_json_path=message_json_path)
|
||||
if args.query_type == "multi_audios" and isinstance(prompt, list):
|
||||
messages = ([system] if system else []) + prompt
|
||||
elif args.query_type == "text":
|
||||
messages = [prompt]
|
||||
|
||||
num_concurrent_requests = args.num_concurrent_requests
|
||||
start_time = time.perf_counter()
|
||||
|
||||
def run_one_request(req_id: int):
|
||||
"""Submit one request and return (req_id, response) to preserve request order."""
|
||||
return req_id, client.chat.completions.create(
|
||||
messages=messages,
|
||||
model=model_name,
|
||||
extra_body=extra_body,
|
||||
stream=args.stream,
|
||||
)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
|
||||
futures = [executor.submit(run_one_request, req_id) for req_id in range(num_concurrent_requests)]
|
||||
|
||||
if not args.stream:
|
||||
# Collect by req_id so chat_completions[i] = response for request i
|
||||
results_by_req = [None] * num_concurrent_requests
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
req_id, completion = future.result()
|
||||
results_by_req[req_id] = completion
|
||||
chat_completions = results_by_req
|
||||
|
||||
assert len(chat_completions) == num_concurrent_requests
|
||||
count = 0
|
||||
for req_id, chat_completion in enumerate(chat_completions):
|
||||
chat_completion_id = getattr(chat_completion, "id", "")
|
||||
for choice in chat_completion.choices:
|
||||
if choice.message.audio:
|
||||
audio_data = base64.b64decode(choice.message.audio.data)
|
||||
audio_file_path = f"{output_audio_path}/{args.query_type}/audio_{count}.wav"
|
||||
os.makedirs(os.path.dirname(audio_file_path), exist_ok=True)
|
||||
with open(audio_file_path, "wb") as f:
|
||||
f.write(audio_data)
|
||||
print(f"[req {req_id}_{chat_completion_id}] Audio saved to {audio_file_path}")
|
||||
count += 1
|
||||
elif choice.message.content:
|
||||
print(
|
||||
f"[req {req_id}_{chat_completion_id}] Chat completion output from text:",
|
||||
choice.message.content,
|
||||
)
|
||||
else:
|
||||
# Stream mode: process chunks from all requests in real-time,
|
||||
# displaying them in the order they arrive.
|
||||
chunk_queue = queue.Queue()
|
||||
audio_counters = {}
|
||||
chat_completion_id_by_req = {}
|
||||
chat_completion_id_lock = threading.Lock()
|
||||
last_text_req_id = None
|
||||
|
||||
def _stream_reader(req_id: int, stream):
|
||||
"""Read one stream and relay every chunk to the shared queue."""
|
||||
try:
|
||||
for chunk in stream:
|
||||
chunk_queue.put((req_id, time.perf_counter(), chunk))
|
||||
chunk_queue.put((req_id, time.perf_counter(), None))
|
||||
except Exception as e:
|
||||
print(f"\n[Request {req_id}] stream error: {e}", file=sys.stderr, flush=True)
|
||||
chunk_queue.put((req_id, time.time(), None))
|
||||
|
||||
# Kick off a reader thread per request
|
||||
reader_threads = []
|
||||
for req_id, future in enumerate(futures):
|
||||
req_id_from_future, stream = future.result()
|
||||
assert req_id == req_id_from_future, f"Request ID mismatch: {req_id} != {req_id_from_future}"
|
||||
audio_counters[req_id] = 0
|
||||
t = threading.Thread(target=_stream_reader, args=(req_id, stream), daemon=True)
|
||||
t.start()
|
||||
reader_threads.append(t)
|
||||
|
||||
# Main loop – consume chunks in arrival order
|
||||
active_streams = set(range(num_concurrent_requests))
|
||||
while active_streams:
|
||||
try:
|
||||
request_id, ts, chunk = chunk_queue.get(timeout=2.0)
|
||||
except queue.Empty:
|
||||
if all(not t.is_alive() for t in reader_threads):
|
||||
break
|
||||
continue
|
||||
|
||||
if chunk is None:
|
||||
elapsed = ts - start_time
|
||||
print(f" ({elapsed:.2f}s)", flush=True)
|
||||
with chat_completion_id_lock:
|
||||
chat_completion_id = chat_completion_id_by_req.get(request_id, "")
|
||||
print(
|
||||
f" [req {request_id}_{chat_completion_id}] Time finished for streaming: ",
|
||||
datetime.datetime.now(),
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
active_streams.discard(request_id)
|
||||
continue
|
||||
|
||||
with chat_completion_id_lock:
|
||||
if request_id not in chat_completion_id_by_req:
|
||||
chat_completion_id_by_req[request_id] = getattr(chunk, "id", "")
|
||||
chat_completion_id = chat_completion_id_by_req[request_id]
|
||||
|
||||
modality = getattr(chunk, "modality", None)
|
||||
elapsed = ts - start_time
|
||||
for choice in chunk.choices:
|
||||
content = getattr(choice.delta, "content", None) if hasattr(choice, "delta") else None
|
||||
|
||||
if modality == "audio" and content:
|
||||
audio_data = base64.b64decode(content)
|
||||
audio_dir = (
|
||||
f"{output_audio_path}/{args.query_type}/{num_concurrent_requests}/{chat_completion_id}"
|
||||
)
|
||||
os.makedirs(audio_dir, exist_ok=True)
|
||||
audio_file_path = f"{audio_dir}/audio_{audio_counters[request_id]}.wav"
|
||||
with open(audio_file_path, "wb") as f:
|
||||
f.write(audio_data)
|
||||
print(
|
||||
f"\n[{elapsed:7.2f}s][req {request_id}_{chat_completion_id}] Audio saved to {audio_file_path}",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
audio_counters[request_id] += 1
|
||||
|
||||
elif modality == "text" and content:
|
||||
if last_text_req_id != request_id:
|
||||
if last_text_req_id is not None:
|
||||
print(flush=True)
|
||||
print(
|
||||
f"\n[{elapsed:7.2f}s][req {request_id}_{chat_completion_id}] content:",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
last_text_req_id = request_id
|
||||
print(
|
||||
f"\n[{elapsed:7.2f}s][req {request_id}_{chat_completion_id}] {content}", end="", flush=True
|
||||
)
|
||||
|
||||
# Final newline if the last output was streaming text
|
||||
if last_text_req_id is not None:
|
||||
print(flush=True)
|
||||
|
||||
for t in reader_threads:
|
||||
t.join(timeout=1.0)
|
||||
|
||||
elapsed = time.perf_counter() - start_time
|
||||
print(
|
||||
f"num_concurrent_requests_{num_concurrent_requests}>>>>>>>>>Time finished for streaming<<<<<<<<: ",
|
||||
elapsed,
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
timing_audio_folder = f"{output_audio_path}/{args.query_type}/{num_concurrent_requests}"
|
||||
os.makedirs(timing_audio_folder, exist_ok=True)
|
||||
timing_file = os.path.join(timing_audio_folder, "streaming_finish_time.txt")
|
||||
with open(timing_file, "w") as f:
|
||||
f.write(f"num_concurrent_requests_{num_concurrent_requests} elapsed_seconds: {elapsed}\n")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="multi_audios",
|
||||
choices=query_map.keys(),
|
||||
help="Query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
"-p",
|
||||
type=str,
|
||||
default="Who are you?",
|
||||
help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--message-json",
|
||||
"-m",
|
||||
type=str,
|
||||
# default="../../offline_inference/mimo_audio/message_base64_wav.json",
|
||||
default="../../offline_inference/mimo_audio/message_base64_wav_tts.json",
|
||||
help="Path to message.json file containing conversation history. When provided, "
|
||||
"system prompt and multi_audios query will be loaded from this file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-audio-path",
|
||||
"-o",
|
||||
type=str,
|
||||
default="./",
|
||||
help="Path to save the generated audio files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stream",
|
||||
action="store_true",
|
||||
help="Stream the response.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-concurrent-requests",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of concurrent requests to send. Default is 1.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_multimodal_generation(args)
|
||||
@@ -29,6 +29,15 @@ class DummyReqState:
|
||||
pass
|
||||
|
||||
|
||||
class MiMoAudioForConditionalGeneration(torch.nn.Module):
|
||||
"""Dummy model whose class name must exactly match the production check."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
# No real forward needed for these tests.
|
||||
|
||||
|
||||
class DummyTalkerMTP(torch.nn.Module):
|
||||
"""A fake talker_mtp module for deterministic CPU testing."""
|
||||
|
||||
@@ -82,6 +91,30 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
|
||||
return runner
|
||||
|
||||
|
||||
def _make_runner_for_mimo(req_id="r_mimo"):
|
||||
"""Create a minimal runner with MiMoAudio-like model and request state."""
|
||||
runner = object.__new__(OmniGPUModelRunner)
|
||||
runner.model = MiMoAudioForConditionalGeneration()
|
||||
|
||||
# Minimal vllm_config / model_config used by helper.
|
||||
class _DummyModelConfig:
|
||||
async_chunk = False
|
||||
|
||||
class _DummyVllmConfig:
|
||||
model_config = _DummyModelConfig()
|
||||
|
||||
runner.vllm_config = _DummyVllmConfig()
|
||||
|
||||
# Attach a single request state with mm_features and additional_information_cpu.
|
||||
req_state = DummyReqState()
|
||||
req_state.mm_features = ["mm_feature_obj"]
|
||||
req_state.additional_information_cpu = {"some_key": "some_value"}
|
||||
|
||||
runner.requests = {req_id: req_state}
|
||||
|
||||
return runner
|
||||
|
||||
|
||||
def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
|
||||
# Patch the module-level `set_forward_context` symbol used inside
|
||||
# OmniGPUModelRunner._talker_mtp_forward.
|
||||
@@ -132,3 +165,32 @@ def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
|
||||
|
||||
# Ensure no changes were made
|
||||
assert torch.allclose(inputs_embeds, before)
|
||||
|
||||
|
||||
def test_maybe_attach_mimo_audio_req_infos_enriches_dict():
|
||||
runner = _make_runner_for_mimo()
|
||||
req_id = "r_mimo"
|
||||
req_state = runner.requests[req_id]
|
||||
|
||||
# Existing req_infos should be copied and enriched, not mutated in place.
|
||||
original_req_infos = {"existing": 1}
|
||||
enriched = OmniGPUModelRunner._maybe_attach_mimo_audio_req_infos(runner, req_state, original_req_infos, req_id)
|
||||
|
||||
assert enriched is not original_req_infos
|
||||
assert enriched["existing"] == 1
|
||||
# mm_features should be filled from req_state when missing
|
||||
assert enriched["mm_features"] == req_state.mm_features
|
||||
# req_id should always be attached
|
||||
assert enriched["req_id"] == req_id
|
||||
|
||||
|
||||
def test_maybe_attach_mimo_audio_req_infos_no_req_state_returns_input():
|
||||
runner = _make_runner_for_mimo()
|
||||
req_id = "missing"
|
||||
req_state = None
|
||||
req_infos = {"k": "v"}
|
||||
|
||||
result = OmniGPUModelRunner._maybe_attach_mimo_audio_req_infos(runner, req_state, req_infos, req_id)
|
||||
|
||||
# When no req_state, helper should be a no-op.
|
||||
assert result is req_infos
|
||||
Whitespace-only changes.
@@ -0,0 +1,206 @@
|
||||
# Copyright 2025 Xiaomi Corporation.
|
||||
import copy
|
||||
from dataclasses import dataclass
|
||||
|
||||
from transformers import PretrainedConfig, Qwen2Config
|
||||
|
||||
SPAN_CODEC_START_TOKEN_ID = 151670
|
||||
SPAN_CODEC_END_TOKEN_ID = 151672
|
||||
TALKER_CODEC_PAD_TOKEN_ID = 151667
|
||||
TEXT_GROUP_SIZE = 5
|
||||
PAD_GROUP_SIZE = 5
|
||||
NO_INTERLEAVE_NEXT_TOKEN_ID = 151671
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiMoAudioConfig(Qwen2Config):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
speech_vocab_size: str | int = "1025-1025-129-129-129-129-129-129",
|
||||
speech_zeroemb_idx: str | int = "1024-1024-128-128-128-128-128-128",
|
||||
delay_pattern: str = "0-1-2-3-4-5-6-7",
|
||||
head_dim: int = 128,
|
||||
group_size: int = 4,
|
||||
audio_channels: int = 8,
|
||||
local_dim: int = 1024,
|
||||
local_layers: int = 16,
|
||||
local_attn_heads: int = 64,
|
||||
local_ffn_dim: int = 4096,
|
||||
local_attn_dropout: float = 0.1,
|
||||
input_local_layers: int = 6,
|
||||
input_local_dim: int | None = None,
|
||||
input_full_attention: bool | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
**kwargs,
|
||||
)
|
||||
self.speech_vocab_size = speech_vocab_size
|
||||
self.speech_zeroemb_idx = speech_zeroemb_idx
|
||||
self.delay_pattern = delay_pattern
|
||||
|
||||
self.head_dim = head_dim
|
||||
|
||||
self.group_size = group_size
|
||||
self.audio_channels = audio_channels
|
||||
|
||||
self.local_dim = local_dim
|
||||
self.local_layers = local_layers
|
||||
self.local_attn_heads = local_attn_heads
|
||||
self.local_ffn_dim = local_ffn_dim
|
||||
self.local_attn_dropout = local_attn_dropout
|
||||
|
||||
self.input_local_layers = input_local_layers
|
||||
self.input_local_dim = input_local_dim or local_dim
|
||||
|
||||
self.input_full_attention = input_full_attention
|
||||
|
||||
def _parse_maybe_list(self, value: str | int, length: int) -> list[int]:
|
||||
if isinstance(value, str) and "-" in value:
|
||||
return [int(s) for s in value.split("-")]
|
||||
return [int(value)] * length
|
||||
|
||||
def parsed_speech_empty_ids(self):
|
||||
return self._parse_maybe_list(self.speech_zeroemb_idx, self.audio_channels)
|
||||
|
||||
def parsed_speech_vocab_sizes(self):
|
||||
return self._parse_maybe_list(self.speech_vocab_size, self.audio_channels)
|
||||
|
||||
def parsed_delay_pattern(self):
|
||||
return self._parse_maybe_list(self.delay_pattern, self.audio_channels)
|
||||
|
||||
def local_config(self):
|
||||
config = copy.deepcopy(self)
|
||||
|
||||
config.hidden_size = self.local_dim
|
||||
config.num_hidden_layers = self.local_layers
|
||||
config.num_attention_heads = self.local_attn_heads
|
||||
config.num_key_value_heads = self.local_attn_heads
|
||||
config.head_dim = config.hidden_size // self.local_attn_heads
|
||||
config.intermediate_size = self.local_ffn_dim
|
||||
config.attention_dropout = self.local_attn_dropout
|
||||
|
||||
return config
|
||||
|
||||
def input_local_config(self):
|
||||
config = copy.deepcopy(self)
|
||||
|
||||
config.hidden_size = self.input_local_dim
|
||||
config.num_hidden_layers = self.input_local_layers
|
||||
config.num_attention_heads = self.local_attn_heads
|
||||
config.num_key_value_heads = self.local_attn_heads
|
||||
config.head_dim = config.hidden_size // self.local_attn_heads
|
||||
config.intermediate_size = config.hidden_size * 4
|
||||
config.attention_dropout = self.local_attn_dropout
|
||||
|
||||
return config
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"speech_vocab_size={self.speech_vocab_size!r}, "
|
||||
f"speech_zeroemb_idx={self.speech_zeroemb_idx!r}, "
|
||||
f"delay_pattern={self.delay_pattern!r}, "
|
||||
f"head_dim={self.head_dim}, "
|
||||
f"group_size={self.group_size}, "
|
||||
f"audio_channels={self.audio_channels}, "
|
||||
f"local_dim={self.local_dim}, "
|
||||
f"local_layers={self.local_layers}, "
|
||||
f"local_attn_heads={self.local_attn_heads}, "
|
||||
f"local_ffn_dim={self.local_ffn_dim}, "
|
||||
f"local_attn_dropout={self.local_attn_dropout}, "
|
||||
f"input_local_layers={self.input_local_layers}, "
|
||||
f"input_local_dim={self.input_local_dim}, "
|
||||
f"input_full_attention={self.input_full_attention})"
|
||||
)
|
||||
|
||||
|
||||
class MiMoAudioTokenizerConfig(PretrainedConfig):
|
||||
model_type = "mimo_audio_tokenizer"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_audio_seconds: int = 1800,
|
||||
stride_size: int = 2,
|
||||
avg_pooler: int = 1,
|
||||
d_model: int = 768,
|
||||
scale_embedding: bool = True,
|
||||
kernel_size: int = 3,
|
||||
activation_function: str = "gelu",
|
||||
encoder_layers: int = 8,
|
||||
encoder_skip_layer_id: int = None,
|
||||
encoder_attention_heads: int = 12,
|
||||
encoder_ffn_dim: int = 3072,
|
||||
encoder_causal: bool = False,
|
||||
encoder_attn_window_size: list[int] = None,
|
||||
decoder_layers: int = 8,
|
||||
decoder_attention_heads: int = 12,
|
||||
decoder_ffn_dim: int = 3072,
|
||||
decoder_kernel_size: int = 3,
|
||||
decoder_stride_size: int = 2,
|
||||
decoder_causal: bool = True,
|
||||
decoder_attn_window_size: list[int] = None,
|
||||
nfft: int = 1024,
|
||||
vocoder_dim: int = 512,
|
||||
vocoder_intermediate_dim: int = 4096,
|
||||
vocoder_num_layers: int = 30,
|
||||
n_mels: int = 80,
|
||||
sampling_rate: int = 24000,
|
||||
hop_length: int = 240,
|
||||
window_size: int = 1024,
|
||||
vocoder_padding: str = "same",
|
||||
fmin: int = 0,
|
||||
fmax: int = None,
|
||||
num_quantizers: int = 12,
|
||||
codebook_size: list[int] = None,
|
||||
threshold_ema_dead_code: int = 10,
|
||||
position_embedding_type: str = "rope",
|
||||
rope_theta: int = 10000,
|
||||
rope_type: str = "default",
|
||||
ln_type: str = "LayerNorm",
|
||||
vocoder_attention_heads: int = 4,
|
||||
vocoder_attn_window_size: list[int] = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.max_audio_seconds = max_audio_seconds
|
||||
self.stride_size = stride_size
|
||||
self.avg_pooler = avg_pooler
|
||||
self.d_model = d_model
|
||||
self.scale_embedding = scale_embedding
|
||||
self.kernel_size = kernel_size
|
||||
self.activation_function = activation_function
|
||||
self.encoder_layers = encoder_layers
|
||||
self.encoder_skip_layer_id = encoder_skip_layer_id
|
||||
self.encoder_attention_heads = encoder_attention_heads
|
||||
self.encoder_ffn_dim = encoder_ffn_dim
|
||||
self.encoder_causal = encoder_causal
|
||||
self.encoder_attn_window_size = encoder_attn_window_size if encoder_attn_window_size is not None else [-1, -1]
|
||||
self.decoder_layers = decoder_layers
|
||||
self.decoder_attention_heads = decoder_attention_heads
|
||||
self.decoder_ffn_dim = decoder_ffn_dim
|
||||
self.decoder_kernel_size = decoder_kernel_size
|
||||
self.decoder_stride_size = decoder_stride_size
|
||||
self.decoder_causal = decoder_causal
|
||||
self.decoder_attn_window_size = decoder_attn_window_size if decoder_attn_window_size is not None else [-1, -1]
|
||||
self.nfft = nfft
|
||||
self.vocoder_dim = vocoder_dim
|
||||
self.vocoder_intermediate_dim = vocoder_intermediate_dim
|
||||
self.vocoder_num_layers = vocoder_num_layers
|
||||
self.n_mels = n_mels
|
||||
self.sampling_rate = sampling_rate
|
||||
self.hop_length = hop_length
|
||||
self.window_size = window_size
|
||||
self.vocoder_padding = vocoder_padding
|
||||
self.fmin = fmin
|
||||
self.fmax = fmax
|
||||
self.num_quantizers = num_quantizers
|
||||
self.codebook_size = codebook_size if codebook_size is not None else [1024]
|
||||
self.threshold_ema_dead_code = threshold_ema_dead_code
|
||||
self.position_embedding_type = position_embedding_type
|
||||
self.rope_theta = rope_theta
|
||||
self.rope_type = rope_type
|
||||
self.ln_type = ln_type
|
||||
self.vocoder_attention_heads = vocoder_attention_heads
|
||||
self.vocoder_attn_window_size = vocoder_attn_window_size if vocoder_attn_window_size is not None else [40, 10]
|
||||
@@ -0,0 +1,914 @@
|
||||
# Copyright 2025 Xiaomi Corporation.
|
||||
import os
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from functools import cached_property
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers import BatchFeature, Qwen2Config
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config.multimodal import BaseDummyOptions
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.models import SupportsPP
|
||||
from vllm.model_executor.models.interfaces import SupportsMultiModal
|
||||
from vllm.model_executor.models.utils import init_vllm_registered_model
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.inputs import (
|
||||
AudioItem,
|
||||
ModalityData,
|
||||
MultiModalBatchedField,
|
||||
MultiModalDataDict,
|
||||
MultiModalFieldConfig,
|
||||
MultiModalFieldElem,
|
||||
MultiModalKwargsItem,
|
||||
)
|
||||
from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, MultiModalDataParser
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
BaseMultiModalProcessor,
|
||||
BaseProcessingInfo,
|
||||
ProcessorInputs,
|
||||
PromptReplacement,
|
||||
PromptUpdate,
|
||||
PromptUpdateDetails,
|
||||
)
|
||||
from vllm.multimodal.utils import group_mm_kwargs_by_modality
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.utils.cache import LRUCache
|
||||
from vllm.utils.collection_utils import is_list_of
|
||||
from vllm.utils.platform_utils import is_pin_memory_available
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
|
||||
from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import (
|
||||
NO_INTERLEAVE_NEXT_TOKEN_ID,
|
||||
PAD_GROUP_SIZE,
|
||||
SPAN_CODEC_END_TOKEN_ID,
|
||||
SPAN_CODEC_START_TOKEN_ID,
|
||||
TALKER_CODEC_PAD_TOKEN_ID,
|
||||
TEXT_GROUP_SIZE,
|
||||
MiMoAudioConfig,
|
||||
)
|
||||
from vllm_omni.model_executor.models.mimo_audio.mimo_audio_code2wav import get_tokenizer_worker
|
||||
from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni import OmniOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def interleave_5_and_5_in_span(
|
||||
input_ids: list[int],
|
||||
*,
|
||||
span_start_token: int = SPAN_CODEC_START_TOKEN_ID,
|
||||
span_end_token: int = SPAN_CODEC_END_TOKEN_ID,
|
||||
pad_token_id: int = TALKER_CODEC_PAD_TOKEN_ID,
|
||||
text_group_size: int = TEXT_GROUP_SIZE,
|
||||
pad_group_size: int = PAD_GROUP_SIZE,
|
||||
no_interleave_next_token: int = NO_INTERLEAVE_NEXT_TOKEN_ID,
|
||||
) -> list[int]:
|
||||
"""
|
||||
Interleave text tokens and padding tokens within spans marked by special tokens.
|
||||
|
||||
Finds spans delimited by span_start_token and span_end_token, then reorganizes
|
||||
tokens inside each span by grouping text tokens (excluding padding) into groups
|
||||
of text_group_size, followed by pad_group_size padding tokens. Interleaving is
|
||||
skipped if the token after span_end_token is no_interleave_next_token.
|
||||
|
||||
Args:
|
||||
input_ids: Input token ID list
|
||||
span_start_token: Token marking span start (default: 151670)
|
||||
span_end_token: Token marking span end (default: 151672)
|
||||
pad_token_id: Padding token ID (default: 151667)
|
||||
text_group_size: Number of text tokens per group (default: 5)
|
||||
pad_group_size: Number of padding tokens per group (default: 5)
|
||||
no_interleave_next_token: Skip interleaving if this token follows span_end (default: 151671)
|
||||
|
||||
Returns:
|
||||
Processed token list with same length as input
|
||||
"""
|
||||
if not input_ids:
|
||||
return input_ids
|
||||
|
||||
original_len = len(input_ids)
|
||||
output_ids: list[int] = []
|
||||
cursor = 0
|
||||
|
||||
while cursor < original_len:
|
||||
# Non-span starting point: Copy as is
|
||||
if input_ids[cursor] != span_start_token:
|
||||
output_ids.append(input_ids[cursor])
|
||||
cursor += 1
|
||||
continue
|
||||
|
||||
# ===== Enter span processing =====
|
||||
output_ids.append(input_ids[cursor])
|
||||
cursor += 1
|
||||
|
||||
span_end_pos = cursor
|
||||
while span_end_pos < original_len and input_ids[span_end_pos] != span_end_token:
|
||||
span_end_pos += 1
|
||||
|
||||
# If span_end_token is not found: Output the remaining as is and end
|
||||
if span_end_pos >= original_len:
|
||||
output_ids.extend(input_ids[cursor:])
|
||||
break
|
||||
|
||||
next_pos = span_end_pos + 1
|
||||
next_token = input_ids[next_pos] if next_pos < original_len else None
|
||||
|
||||
# Rule:
|
||||
# After "-end" is 151667(PAD) -> for interlacing
|
||||
# After "-end" is 151671 -> Do not interlace
|
||||
do_interleave = next_token == pad_token_id
|
||||
if next_token == no_interleave_next_token:
|
||||
do_interleave = False
|
||||
|
||||
if not do_interleave:
|
||||
output_ids.extend(input_ids[cursor:span_end_pos])
|
||||
output_ids.append(span_end_token)
|
||||
cursor = span_end_pos + 1
|
||||
continue
|
||||
|
||||
span_content = input_ids[cursor:span_end_pos]
|
||||
span_original_len = len(span_content)
|
||||
span_text_tokens = [t for t in span_content if t != pad_token_id]
|
||||
|
||||
rebuilt_span: list[int] = []
|
||||
text_cursor = 0
|
||||
total_text = len(span_text_tokens)
|
||||
|
||||
while text_cursor < total_text:
|
||||
take_text = min(text_group_size, total_text - text_cursor)
|
||||
rebuilt_span.extend(span_text_tokens[text_cursor : text_cursor + take_text])
|
||||
text_cursor += take_text
|
||||
|
||||
if text_cursor < total_text:
|
||||
rebuilt_span.extend([pad_token_id] * pad_group_size)
|
||||
|
||||
output_ids.extend(rebuilt_span)
|
||||
output_ids.append(span_end_token)
|
||||
|
||||
extra_tokens_added = len(rebuilt_span) - span_original_len
|
||||
cursor = span_end_pos + 1
|
||||
|
||||
while extra_tokens_added > 0 and cursor < original_len:
|
||||
if input_ids[cursor] == pad_token_id:
|
||||
extra_tokens_added -= 1
|
||||
cursor += 1
|
||||
else:
|
||||
output_ids.append(input_ids[cursor])
|
||||
cursor += 1
|
||||
|
||||
if extra_tokens_added > 0:
|
||||
return output_ids[:original_len]
|
||||
|
||||
if len(output_ids) > original_len:
|
||||
output_ids = output_ids[:original_len]
|
||||
elif len(output_ids) < original_len:
|
||||
output_ids.extend([pad_token_id] * (original_len - len(output_ids)))
|
||||
|
||||
return output_ids
|
||||
|
||||
|
||||
class MiMoAudioLLMProcessingInfo(
|
||||
BaseProcessingInfo,
|
||||
):
|
||||
def get_hf_config(self):
|
||||
return self.ctx.get_hf_config(Qwen2Config)
|
||||
|
||||
def get_hf_processor(
|
||||
self,
|
||||
*,
|
||||
# Ignored in initialization
|
||||
sampling_rate: int | None = None,
|
||||
**kwargs: object,
|
||||
):
|
||||
# MiMoAudio doesn't use a standard HF processor like Qwen2Audio
|
||||
# Instead it uses MiMoAudioTokenizer directly
|
||||
# This method may return None or a dummy processor
|
||||
# depending on the actual implementation requirements
|
||||
return None
|
||||
|
||||
def get_data_parser(self) -> MultiModalDataParser:
|
||||
sampling_rate = 24000
|
||||
return MiMoAudioDataParser(target_sr=sampling_rate)
|
||||
|
||||
def get_feature_extractor(
|
||||
self,
|
||||
*,
|
||||
# Ignored in initialization
|
||||
sampling_rate: int | None = None,
|
||||
):
|
||||
# MiMoAudio uses MiMoAudioTokenizer instead of WhisperFeatureExtractor
|
||||
# This method may need to return the tokenizer or None
|
||||
# depending on the actual implementation requirements
|
||||
return None
|
||||
|
||||
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
|
||||
return {"audio": None}
|
||||
|
||||
def get_mm_max_tokens_per_item(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
) -> Mapping[str, int]:
|
||||
# MiMoAudio has different tokenization than Qwen2Audio
|
||||
# This is a placeholder - adjust based on actual MiMoAudio tokenizer limits
|
||||
# For now, return a large value or None
|
||||
# This should be adjusted based on MiMoAudio's actual constraints
|
||||
return {"audio": 1}
|
||||
|
||||
|
||||
class MiMoAudioLLMDummyInputsBuilder(BaseDummyInputsBuilder[MiMoAudioLLMProcessingInfo]):
|
||||
_processor_inputs_cache: LRUCache = LRUCache(capacity=1024)
|
||||
|
||||
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
|
||||
num_audios = mm_counts.get("audio", 0)
|
||||
return (" <|empty|>" * num_audios).strip()
|
||||
|
||||
def get_dummy_mm_data(
|
||||
self, seq_len: int, mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions] | None = None
|
||||
) -> MultiModalDataDict:
|
||||
num_audios = mm_counts.get("audio", 0)
|
||||
if num_audios == 0:
|
||||
return {}
|
||||
# Return dummy raw audio data (not encoded codes)
|
||||
# This will be processed by _parse_audio_data like real audio
|
||||
# Use 1 second of audio at target_sr (24000 Hz)
|
||||
dummy_audio_length = mm_options.get("audio").length if mm_options else 24000 # 1 second at 24kHz
|
||||
dummy_audio = np.zeros((dummy_audio_length,), dtype=np.float32)
|
||||
return {"audio": [(dummy_audio, 24000)] * num_audios}
|
||||
|
||||
def get_dummy_processor_inputs(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
mm_options: Mapping[str, BaseDummyOptions] | None = None,
|
||||
) -> ProcessorInputs:
|
||||
dummy_text = self.get_dummy_text(mm_counts)
|
||||
dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
|
||||
|
||||
dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
|
||||
|
||||
dummy_processor_inputs = ProcessorInputs(
|
||||
prompt=dummy_text,
|
||||
mm_items=dummy_mm_items,
|
||||
)
|
||||
|
||||
return dummy_processor_inputs
|
||||
|
||||
|
||||
class MiMoAudioDataParser(MultiModalDataParser):
|
||||
def __init__(self, target_sr: int, use_mono_channel: bool = True):
|
||||
super().__init__(target_sr=target_sr)
|
||||
self.use_mono_channel = use_mono_channel
|
||||
self.target_sr = target_sr
|
||||
|
||||
tokenizer_device = os.environ.get("MIMO_AUDIO_TOKENIZER_DEVICE", None)
|
||||
|
||||
if tokenizer_device:
|
||||
if tokenizer_device.lower() == "cpu":
|
||||
self.device = torch.device("cpu")
|
||||
else:
|
||||
# Support formats like "cuda", "cuda:0", "cuda:6", etc.
|
||||
self.device = torch.device(tokenizer_device)
|
||||
else:
|
||||
# Default to cuda (will use current GPU)
|
||||
self.device = torch.device(f"cuda:{torch.cuda.current_device()}" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
self.audio_tokenizer_path = os.environ.get("MIMO_AUDIO_TOKENIZER_PATH", None)
|
||||
if not self.audio_tokenizer_path:
|
||||
raise ValueError(
|
||||
"Audio tokenizer path is not set. Provide "
|
||||
"`model_config.audio_tokenizer_path` in the stage config "
|
||||
"or export MIMO_AUDIO_TOKENIZER_PATH."
|
||||
)
|
||||
|
||||
if not os.path.exists(self.audio_tokenizer_path):
|
||||
raise ValueError(
|
||||
"Audio tokenizer not exists. Provide "
|
||||
"`model_config.audio_tokenizer_path` in the stage config "
|
||||
"or export MIMO_AUDIO_TOKENIZER_PATH."
|
||||
)
|
||||
|
||||
self.tokenizer_config_path = os.environ.get("MIMO_AUDIO_TOKENIZER_PATH", None)
|
||||
|
||||
self.mimo_tokenizer = get_tokenizer_worker(
|
||||
device=self.device,
|
||||
config_path=self.tokenizer_config_path,
|
||||
audio_tokenizer_path=self.audio_tokenizer_path,
|
||||
)
|
||||
|
||||
def _parse_audio_data(
|
||||
self,
|
||||
data: ModalityData[AudioItem],
|
||||
):
|
||||
if data is None:
|
||||
return AudioProcessorItems(None)
|
||||
|
||||
# also check single audio item with sampling rate
|
||||
if self._is_empty(data) or (isinstance(data, tuple) and self._is_empty(data[0])):
|
||||
return None
|
||||
|
||||
if (
|
||||
is_list_of(data, float)
|
||||
or isinstance(data, (np.ndarray, torch.Tensor))
|
||||
and data.ndim == 1
|
||||
or isinstance(data, tuple)
|
||||
):
|
||||
data_items = [data]
|
||||
elif isinstance(data, (np.ndarray, torch.Tensor)):
|
||||
data_items = [elem for elem in data]
|
||||
else:
|
||||
data_items = data
|
||||
|
||||
new_audios = list[np.ndarray]()
|
||||
for data_item in data_items:
|
||||
if not isinstance(data_item, tuple) or not isinstance(data_item[0], np.ndarray):
|
||||
new_audios.append(data_item)
|
||||
continue
|
||||
|
||||
audio, orig_sr = self._get_audio_with_sr(data_item)
|
||||
if orig_sr is None:
|
||||
new_audio = audio
|
||||
else:
|
||||
new_audio = self.audio_resampler.resample(audio, orig_sr=orig_sr)
|
||||
wav_tensor = torch.from_numpy(new_audio.astype(np.float32))
|
||||
if wav_tensor.dim() == 1:
|
||||
wav_tensor = wav_tensor.unsqueeze(0)
|
||||
if self.use_mono_channel:
|
||||
if wav_tensor.size(0) > 1:
|
||||
wav_mono = wav_tensor.mean(dim=0) # [samples]
|
||||
else:
|
||||
wav_mono = wav_tensor[0]
|
||||
audio_codes = self.mimo_tokenizer.encode(audio=(wav_mono, self.target_sr))
|
||||
new_audio = audio_codes.detach().cpu().numpy().astype(np.int64).tolist()
|
||||
new_audios.append(new_audio)
|
||||
|
||||
return AudioProcessorItems(new_audios)
|
||||
|
||||
|
||||
class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProcessingInfo]):
|
||||
def _call_hf_processor(
|
||||
self,
|
||||
prompt: str,
|
||||
mm_data: Mapping[str, object],
|
||||
mm_kwargs: Mapping[str, Any],
|
||||
tok_kwargs: Mapping[str, object],
|
||||
) -> BatchFeature:
|
||||
sampling_rate = 24000
|
||||
audios = mm_data.pop("audios", [])
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
if audios:
|
||||
mm_data["audio"] = audios
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt_ids = tokenizer.encode(prompt)
|
||||
else:
|
||||
prompt_ids = prompt
|
||||
|
||||
# Text-only input not supported in composite processor
|
||||
if not mm_data.get("audio", []):
|
||||
prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
|
||||
|
||||
return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
|
||||
|
||||
# feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
|
||||
mm_kwargs = dict(
|
||||
**mm_kwargs,
|
||||
sampling_rate=sampling_rate,
|
||||
)
|
||||
|
||||
audio_lengths = [len(x) for x in audios]
|
||||
if all(audio_length == 8 for audio_length in audio_lengths):
|
||||
audio_lengths = [len(x[0]) for x in audios]
|
||||
|
||||
return BatchFeature(
|
||||
data={
|
||||
"input_ids": torch.tensor([prompt_ids], dtype=torch.int64),
|
||||
"audio_embeds": audios,
|
||||
"audio_lengths": audio_lengths,
|
||||
"mm_data": mm_data,
|
||||
"mm_kwargs": mm_kwargs,
|
||||
"tok_kwargs": tok_kwargs,
|
||||
},
|
||||
tensor_type=None,
|
||||
)
|
||||
|
||||
def _get_mm_fields_config(
|
||||
self,
|
||||
hf_inputs: BatchFeature,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
# MiMoAudio may have different field names than Qwen2Audio
|
||||
# Adjust based on actual MiMoAudio model requirements
|
||||
# For now, return empty dict or adjust based on actual needs
|
||||
return {
|
||||
"audio_embeds": MultiModalFieldConfig.batched("audio"),
|
||||
"audio_lengths": MultiModalFieldConfig.batched("audio"),
|
||||
}
|
||||
|
||||
def _get_prompt_updates(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
out_mm_kwargs,
|
||||
) -> Sequence[PromptUpdate]:
|
||||
# MiMoAudio doesn't use a standard HF processor like Qwen2Audio
|
||||
# Get tokenizer and vocabulary directly
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
vocab = tokenizer.get_vocab()
|
||||
|
||||
# MiMoAudio may use different audio tokens
|
||||
# Adjust these based on actual tokenizer vocabulary
|
||||
# Common tokens in MiMoAudio might be different from Qwen2Audio
|
||||
audio_token = "<|empty|>"
|
||||
|
||||
audio_bos_id = None
|
||||
audio_eos_id = None
|
||||
audio_token_id = vocab.get(audio_token, None)
|
||||
|
||||
# If tokens don't exist, try alternative tokens or use None
|
||||
if audio_token_id is None:
|
||||
# Try alternative token names or skip audio token replacement
|
||||
return []
|
||||
|
||||
audio_kwargs = out_mm_kwargs.get("audio", [])
|
||||
# feature_attention_mask = out_mm_kwargs.get("feature_attention_mask", None)
|
||||
|
||||
if audio_kwargs:
|
||||
audio_output_lengths = [item["audio_lengths"].data for item in audio_kwargs]
|
||||
else:
|
||||
audio_output_lengths = []
|
||||
if mm_items and "audio" in mm_items:
|
||||
try:
|
||||
audios = mm_items.get_items("audio", AudioProcessorItems)
|
||||
if audios is not None:
|
||||
num_audios = audios.get_count()
|
||||
audio_output_lengths = [audios.get_audio_length(i) for i in range(num_audios)]
|
||||
except Exception as e:
|
||||
logger.error(f"_get_prompt_updates failed: {e}")
|
||||
|
||||
def get_replacement_mimo_audio(item_idx: int):
|
||||
if item_idx >= len(audio_output_lengths):
|
||||
# If no length info, use a default or skip
|
||||
return PromptUpdateDetails.select_token_id(
|
||||
[audio_token_id],
|
||||
embed_token_id=audio_token_id,
|
||||
)
|
||||
|
||||
num_features = audio_output_lengths[item_idx] // 4
|
||||
if num_features == 0:
|
||||
try:
|
||||
audios = mm_items.get_items("audio", AudioProcessorItems)
|
||||
audio_len = audios.get_audio_length(item_idx)
|
||||
raise ValueError(f"The audio (len={audio_len}) is too short to be represented inside the model")
|
||||
except (AttributeError, KeyError):
|
||||
# If AudioProcessorItems is not available, use default
|
||||
num_features = 1
|
||||
|
||||
audio_tokens = [audio_token_id] * num_features
|
||||
|
||||
# Build replacement tokens following Qwen2Audio pattern
|
||||
# [audio_bos_id] + audio_tokens + [audio_eos_id]
|
||||
replacement_tokens = []
|
||||
if audio_bos_id is not None:
|
||||
replacement_tokens.append(audio_bos_id)
|
||||
replacement_tokens.extend(audio_tokens)
|
||||
if audio_eos_id is not None:
|
||||
replacement_tokens.append(audio_eos_id)
|
||||
|
||||
# If no bos/eos tokens, just use audio tokens
|
||||
if not replacement_tokens:
|
||||
replacement_tokens = audio_tokens
|
||||
|
||||
return PromptUpdateDetails.select_token_id(
|
||||
replacement_tokens,
|
||||
embed_token_id=audio_token_id,
|
||||
)
|
||||
|
||||
return [
|
||||
PromptReplacement(
|
||||
modality="audio",
|
||||
target=audio_token,
|
||||
replacement=get_replacement_mimo_audio,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
MiMoAudioLLMMultiModalProcessor,
|
||||
info=MiMoAudioLLMProcessingInfo,
|
||||
dummy_inputs=MiMoAudioLLMDummyInputsBuilder,
|
||||
)
|
||||
class MiMoAudioForConditionalGeneration(
|
||||
nn.Module,
|
||||
SupportsPP,
|
||||
SupportsMultiModal,
|
||||
CustomProcessMixin,
|
||||
):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
|
||||
if modality.startswith("image"):
|
||||
return "<|vision_start|><|IMAGE|><|vision_end|>"
|
||||
if modality.startswith("video"):
|
||||
return "<|vision_start|><|VIDEO|><|vision_end|>"
|
||||
if modality.startswith("audio"):
|
||||
return "<|sosp|><|empty|><|eosp|>"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.has_preprocess = False
|
||||
self.have_multimodal_outputs = True
|
||||
self.requires_raw_input_tokens = True
|
||||
config = vllm_config.model_config.hf_config
|
||||
config = MiMoAudioConfig(**vars(config)) if isinstance(config, Qwen2Config) else config
|
||||
quant_config = vllm_config.quant_config
|
||||
lora_config = vllm_config.lora_config
|
||||
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
|
||||
self.config = config
|
||||
self.lora_config = lora_config
|
||||
|
||||
self.quant_config = quant_config
|
||||
self.vllm_config = vllm_config
|
||||
|
||||
self.llm_prev_hidden_state = None
|
||||
self.device = vllm_config.device_config.device
|
||||
self.pin_memory = is_pin_memory_available()
|
||||
|
||||
if "model_stage" in os.environ:
|
||||
self.model_stage = os.environ["model_stage"]
|
||||
else:
|
||||
self.model_stage = vllm_config.model_config.model_stage
|
||||
|
||||
if self.model_stage == "fused_thinker_talker":
|
||||
self.has_preprocess = True
|
||||
self.set_custom_preprocess(self.fused_thinker_talker_preprocess)
|
||||
# Initialize fused_thinker_talker llm model (multimodal processing)
|
||||
self.fused_thinker_talker = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
# prefix=maybe_prefix(prefix, "model"),
|
||||
hf_config=vllm_config.model_config.hf_config,
|
||||
# Use registry architecture key
|
||||
architectures=["MiMoAudioLLMModel"],
|
||||
)
|
||||
self.token2wav = None
|
||||
self.model = self.fused_thinker_talker
|
||||
|
||||
elif self.model_stage == "code2wav":
|
||||
self.fused_thinker_talker = None
|
||||
# Initialize token2wav (code->mel->wav) like thinker/talker
|
||||
self.token2wav = init_vllm_registered_model(
|
||||
vllm_config=vllm_config,
|
||||
# prefix=maybe_prefix(prefix, "token2wav"),
|
||||
# hf_config=self.token2wav_config,
|
||||
architectures=["MiMoAudioToken2WavModel"],
|
||||
)
|
||||
|
||||
self.model = self.token2wav
|
||||
else:
|
||||
raise ValueError("Invalid model stage")
|
||||
|
||||
# Set up intermediate tensors
|
||||
self.make_empty_intermediate_tensors = (
|
||||
(self.fused_thinker_talker.make_empty_intermediate_tensors)
|
||||
if self.model_stage == "fused_thinker_talker"
|
||||
else lambda: None
|
||||
)
|
||||
|
||||
def fused_thinker_talker_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
|
||||
# Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported.
|
||||
|
||||
# Ensure we have base embeddings when only ids are provided
|
||||
if input_embeds is None and input_ids is not None:
|
||||
input_embeds = self.fused_thinker_talker.embed_input_ids(input_ids)
|
||||
|
||||
span_len = input_ids.shape[0]
|
||||
|
||||
if span_len > 1:
|
||||
# prefill
|
||||
return self.fused_thinker_talker_prefill(input_ids, input_embeds, **info_dict)
|
||||
else:
|
||||
# decode
|
||||
return input_ids, input_embeds, info_dict
|
||||
|
||||
def fused_thinker_talker_prefill(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
|
||||
empty_token_id = self.fused_thinker_talker.empty_token_id
|
||||
|
||||
mm_kwargs = list[tuple[str, MultiModalKwargsItem]]()
|
||||
mm_features = info_dict.get("mm_features", [])
|
||||
mm_embeddings = []
|
||||
prompt_ids = torch.tensor(
|
||||
interleave_5_and_5_in_span(input_ids.tolist()),
|
||||
dtype=torch.int64,
|
||||
device=input_ids.device,
|
||||
)
|
||||
|
||||
for mm_feature in mm_features:
|
||||
mm_item = mm_feature.data
|
||||
|
||||
if mm_item is not None:
|
||||
mm_item["mm_offset"] = MultiModalFieldElem(
|
||||
data=mm_feature.mm_position.offset,
|
||||
field=MultiModalBatchedField(),
|
||||
)
|
||||
mm_kwargs.append((mm_feature.modality, mm_item))
|
||||
|
||||
for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
|
||||
mm_kwargs,
|
||||
device=self.device,
|
||||
pin_memory=self.pin_memory,
|
||||
):
|
||||
# MiMoAudio is only supported audio modality
|
||||
mm_kwargs_group["prompt_ids"] = prompt_ids
|
||||
mm_kwargs_group["modality_preprocess"] = True
|
||||
mm_kwargs_group["mm_offset"] = torch.tensor(
|
||||
mm_kwargs_group["mm_offset"], dtype=torch.long, device=self.device
|
||||
)
|
||||
mm_embeddings = self.fused_thinker_talker.embed_multimodal(**mm_kwargs_group)
|
||||
|
||||
input_embeds = self.fused_thinker_talker.embed_input_ids(
|
||||
prompt_ids, multimodal_embeddings=mm_embeddings, is_multimodal=(prompt_ids == empty_token_id)
|
||||
)
|
||||
|
||||
return prompt_ids, input_embeds, info_dict
|
||||
|
||||
@staticmethod
|
||||
def _module_device(module: nn.Module) -> torch.device:
|
||||
try:
|
||||
return next(module.parameters()).device
|
||||
except StopIteration:
|
||||
# No parameters; fall back to buffers or cpu
|
||||
for _, buf in module.named_buffers(recurse=True):
|
||||
return buf.device
|
||||
return torch.device("cpu")
|
||||
|
||||
def move_submodules_to_devices(
|
||||
self,
|
||||
*,
|
||||
llm_device: str | torch.device | None = None,
|
||||
token2wav_device: str | torch.device | None = None,
|
||||
) -> None:
|
||||
"""Optionally move thinker/talker/token2wav to different devices.
|
||||
|
||||
Example:
|
||||
model.move_submodules_to_devices(
|
||||
thinker_device='cuda:0',
|
||||
talker_device='cuda:1',
|
||||
token2wav_device='cpu',
|
||||
)
|
||||
"""
|
||||
if llm_device is not None and self.fused_thinker_talker is not None:
|
||||
self.fused_thinker_talker.to(llm_device)
|
||||
if token2wav_device is not None and self.token2wav is not None:
|
||||
self.token2wav.to(token2wav_device)
|
||||
|
||||
@cached_property
|
||||
def sampler(self):
|
||||
if hasattr(self.model, "sampler"):
|
||||
return self.model.sampler
|
||||
return Sampler()
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings=None,
|
||||
is_multimodal: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if self.model_stage == "code2wav":
|
||||
return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
|
||||
return self.model.embed_input_ids(input_ids, multimodal_embeddings, is_multimodal=is_multimodal)
|
||||
|
||||
def embed_multimodal(self, **kwargs):
|
||||
# Delegate to thinker model for multimodal processing
|
||||
return self.model.embed_multimodal(**kwargs)
|
||||
|
||||
def last_index_of(self, list, value):
|
||||
return len(list) - 1 - list[::-1].index(value)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:
|
||||
"""Load weights for all components of the omni model."""
|
||||
loaded_weights = set()
|
||||
llm_weights = []
|
||||
token2wav_weights = []
|
||||
for k, v in weights:
|
||||
if (
|
||||
k.startswith("model.")
|
||||
or k.startswith("lm_head.")
|
||||
or k.startswith("input_local_transformer.")
|
||||
or k.startswith("hidden_states_downcast.")
|
||||
or k.startswith("local_transformer")
|
||||
or k.startswith("speech_embeddings")
|
||||
or k.startswith("speech_group_downcast")
|
||||
):
|
||||
if self.fused_thinker_talker:
|
||||
llm_weights.append((k, v))
|
||||
elif k.startswith("mel_transform."):
|
||||
if self.token2wav:
|
||||
token2wav_weights.append((k, v))
|
||||
else:
|
||||
pass
|
||||
|
||||
# Load fused_thinker_talker llm weights
|
||||
if self.fused_thinker_talker:
|
||||
if llm_weights:
|
||||
llm_loaded = self.fused_thinker_talker.load_weights(llm_weights)
|
||||
else:
|
||||
llm_loaded = set([k for k, v in llm_weights])
|
||||
# thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker")
|
||||
loaded_weights.update(llm_loaded)
|
||||
|
||||
## load code2wav
|
||||
if self.token2wav:
|
||||
token2wav_loaded = self.token2wav.load_weights(
|
||||
token2wav_weights,
|
||||
os.environ["MIMO_AUDIO_TOKENIZER_PATH"],
|
||||
)
|
||||
|
||||
loaded_weights.update(token2wav_loaded)
|
||||
|
||||
@staticmethod
|
||||
def insert_between(input_ids: torch.Tensor, group_size: int, value: int = -100):
|
||||
if group_size < 0:
|
||||
raise ValueError("group_size must be non-negative")
|
||||
|
||||
if not isinstance(input_ids, torch.Tensor):
|
||||
input_ids = torch.tensor(input_ids)
|
||||
if input_ids.dim() != 1:
|
||||
raise ValueError("input_ids must be 1-D tensor or list")
|
||||
|
||||
L = input_ids.numel()
|
||||
if group_size == 0 or L == 0:
|
||||
return input_ids.clone()
|
||||
|
||||
new_len = L * (group_size + 1)
|
||||
# Create output tensor filled with value, dtype & device consistent with input_ids
|
||||
out = input_ids.new_full((new_len,), value)
|
||||
|
||||
# Place original elements at the first position of each block: 0, group_size+1, 2*(group_size+1), ...
|
||||
positions = torch.arange(L, dtype=torch.long, device=input_ids.device) * (group_size + 1)
|
||||
out[positions] = input_ids
|
||||
|
||||
return out
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
generate_audio: bool = True,
|
||||
codec: torch.Tensor | None = None,
|
||||
sampling_metadata: SamplingMetadata | None = None,
|
||||
logits_index: int | None = None,
|
||||
sampler=None,
|
||||
additional_information: dict[str, object] | None = None,
|
||||
**kwargs: object,
|
||||
) -> torch.Tensor | IntermediateTensors | OmniOutput:
|
||||
"""
|
||||
Workflow:
|
||||
1) llm: multimodal understanding → text hidden states.
|
||||
2) If audio requested and codec not provided, use talker to derive codec.
|
||||
3) If audio requested (or codec provided), use token2wav to synthesize waveform.
|
||||
4) Return text hidden states (and audio when applicable).
|
||||
"""
|
||||
if self.model_stage == "fused_thinker_talker":
|
||||
next_speech_tokens, text_hidden_states = self.generate_codes(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
inputs_embeds=inputs_embeds,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=text_hidden_states.reshape(-1, text_hidden_states.shape[-1]),
|
||||
multimodal_outputs={"code_predictor_codes": next_speech_tokens},
|
||||
)
|
||||
|
||||
if self.model_stage == "code2wav":
|
||||
code = (
|
||||
input_ids # tensor [seq_len]
|
||||
if input_ids is not None
|
||||
else torch.zeros(
|
||||
inputs_embeds.shape[0],
|
||||
dtype=torch.long,
|
||||
device=inputs_embeds.device,
|
||||
)
|
||||
)
|
||||
|
||||
audio_tensor = self.generate_audio(code)
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"model_outputs": audio_tensor.reshape(1, -1) if audio_tensor is not None else audio_tensor
|
||||
},
|
||||
)
|
||||
|
||||
def generate_codes(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
inputs_embeds: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
**kwargs: object,
|
||||
):
|
||||
llm_dev = self._module_device(self.fused_thinker_talker)
|
||||
|
||||
# if input_ids is None, set it to a zero tensor, in the length of the
|
||||
# same as the embedding seq length
|
||||
if input_ids is None:
|
||||
input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=llm_dev).unsqueeze(0) # (1, 0)
|
||||
|
||||
# 1) Thinker (ensure inputs on thinker's device)
|
||||
if input_ids is not None and input_ids.device != llm_dev:
|
||||
input_ids = input_ids.to(llm_dev)
|
||||
if positions is not None and positions.device != llm_dev:
|
||||
positions = positions.to(llm_dev)
|
||||
if inputs_embeds is not None and inputs_embeds.device != llm_dev:
|
||||
inputs_embeds = inputs_embeds.to(llm_dev)
|
||||
|
||||
# Run llm
|
||||
llm_output = self.fused_thinker_talker(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
next_speech_tokens = None
|
||||
if isinstance(llm_output, tuple):
|
||||
if len(llm_output) == 2:
|
||||
next_speech_tokens, text_hidden_states = llm_output
|
||||
elif len(llm_output) == 3:
|
||||
ids, next_speech_tokens, text_hidden_states = llm_output
|
||||
else:
|
||||
text_hidden_states = llm_output
|
||||
|
||||
return next_speech_tokens, text_hidden_states
|
||||
|
||||
def generate_audio(self, code: torch.Tensor):
|
||||
token2wav_dev = self._module_device(self.token2wav)
|
||||
# Check if in CUDA graph capture phase
|
||||
is_capturing = torch.cuda.is_current_stream_capturing()
|
||||
|
||||
if isinstance(code, torch.Tensor):
|
||||
if is_capturing:
|
||||
# During CUDA graph capture, avoid device movement operations
|
||||
# Assume tensor is already on the correct device, only convert dtype
|
||||
code_tensor = code.to(dtype=torch.long)
|
||||
else:
|
||||
# During non-capture phase, normally perform device movement and type conversion
|
||||
code_tensor = code.to(dtype=torch.long, device=token2wav_dev)
|
||||
else:
|
||||
# If code is not a Tensor, should avoid creating new tensors during capture
|
||||
# This case should be rare during capture, as code usually comes from input_ids
|
||||
if is_capturing:
|
||||
# During capture, if code is not a Tensor, try to use current stream device
|
||||
# But this case should theoretically not happen
|
||||
code_tensor = torch.as_tensor(code, dtype=torch.long)
|
||||
else:
|
||||
code_tensor = torch.as_tensor(code, dtype=torch.long, device=token2wav_dev)
|
||||
if code_tensor.ndim == 2 and code_tensor.shape[0] == 1:
|
||||
code_tensor = code_tensor.squeeze(0)
|
||||
|
||||
with torch.inference_mode():
|
||||
audio_tensor = self.token2wav(codes=code_tensor) # code_tensor tensor [seq_len]
|
||||
|
||||
return audio_tensor
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata | None = None,
|
||||
) -> torch.Tensor | None:
|
||||
# logits = self.logits_processor(self.lm_head, hidden_states,
|
||||
# sampling_metadata)
|
||||
# Handle OmniOutput type
|
||||
if isinstance(hidden_states, OmniOutput):
|
||||
hidden_states = hidden_states.text_hidden_states
|
||||
|
||||
# Use thinker model for logits computation
|
||||
return self.model.compute_logits(hidden_states)
|
||||
@@ -0,0 +1,647 @@
|
||||
# Copyright 2025 Xiaomi Corporation.
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import torchaudio
|
||||
from torch import nn
|
||||
from torchaudio.transforms import MelSpectrogram
|
||||
from transformers import AutoTokenizer, Qwen2Config
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.models import SupportsPP
|
||||
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.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID, MiMoAudioConfig
|
||||
from vllm_omni.model_executor.models.mimo_audio.modeling_audio_tokenizer import MiMoAudioTokenizer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
DUMMY_CODE_SHAPE = 36
|
||||
|
||||
|
||||
class MiMoAudioTokenizerWorker:
|
||||
def __init__(
|
||||
self,
|
||||
device_str: str,
|
||||
config_path: str,
|
||||
audio_tokenizer_path: str,
|
||||
):
|
||||
self.device = device_str
|
||||
logger.info("[tokenizer worker] Loading MiMoAudioConfig from %s", config_path)
|
||||
load_cfg_start = time.monotonic()
|
||||
self.config: MiMoAudioConfig = MiMoAudioConfig.from_pretrained(config_path)
|
||||
logger.info(
|
||||
"[tokenizer worker] Config loaded in %.2fs",
|
||||
time.monotonic() - load_cfg_start,
|
||||
)
|
||||
|
||||
if device_str == "cpu":
|
||||
model_dtype = torch.float32 # CPU must use float32
|
||||
else:
|
||||
model_dtype = torch.bfloat16 # GPU uses bfloat16
|
||||
|
||||
logger.info(
|
||||
"[tokenizer worker] Loading MiMo-Audio Tokenizer from %s (device=%s, dtype=%s)",
|
||||
audio_tokenizer_path,
|
||||
self.device,
|
||||
model_dtype,
|
||||
)
|
||||
start_loading_slm_tokenizer_time = time.monotonic()
|
||||
logger.info(f"MiMoAudioTokenizer,device:{self.device}")
|
||||
self.audio_tokenizer = MiMoAudioTokenizer.from_pretrained(
|
||||
audio_tokenizer_path,
|
||||
dtype=model_dtype,
|
||||
device_map={"": self.device},
|
||||
)
|
||||
# Move to target device and only cast to bfloat16 on non-CPU devices.
|
||||
self.audio_tokenizer.to(self.device)
|
||||
if self.device != "cpu" and model_dtype == torch.bfloat16:
|
||||
self.audio_tokenizer.to(dtype=torch.bfloat16)
|
||||
self.audio_tokenizer.eval()
|
||||
|
||||
logger.info(
|
||||
f"Audio Tokenizers loaded in "
|
||||
f"{time.monotonic() - start_loading_slm_tokenizer_time:.2f} seconds, "
|
||||
f"device: {self.device}"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[tokenizer worker] Building MelSpectrogram transform (sr=%s, n_fft=%s)",
|
||||
self.audio_tokenizer.config.sampling_rate,
|
||||
self.audio_tokenizer.config.nfft,
|
||||
)
|
||||
mel_start = time.monotonic()
|
||||
self.mel_transform = (
|
||||
MelSpectrogram(
|
||||
sample_rate=self.audio_tokenizer.config.sampling_rate,
|
||||
n_fft=self.audio_tokenizer.config.nfft,
|
||||
hop_length=self.audio_tokenizer.config.hop_length,
|
||||
win_length=self.audio_tokenizer.config.window_size,
|
||||
f_min=self.audio_tokenizer.config.fmin,
|
||||
f_max=self.audio_tokenizer.config.fmax,
|
||||
n_mels=self.audio_tokenizer.config.n_mels,
|
||||
power=1.0,
|
||||
center=True,
|
||||
)
|
||||
.to(self.device)
|
||||
.to(torch.float32)
|
||||
)
|
||||
logger.info(
|
||||
"[tokenizer worker] MelSpectrogram ready in %.2fs",
|
||||
time.monotonic() - mel_start,
|
||||
)
|
||||
|
||||
self.group_size = self.config.group_size
|
||||
self.audio_channels = self.config.audio_channels
|
||||
self.sample_rate = self.audio_tokenizer.config.sampling_rate
|
||||
|
||||
# Warmup (skip for CPU due to potential shape mismatch issues)
|
||||
if device_str != "cpu":
|
||||
logger.info("[tokenizer worker] Running warmup encode/decode...")
|
||||
warmup_start = time.monotonic()
|
||||
warmup_dtype = torch.float32 if device_str == "cpu" else torch.bfloat16
|
||||
try:
|
||||
self.encode_wav_base(torch.zeros(self.sample_rate, dtype=warmup_dtype))
|
||||
self.decode(torch.zeros(self.audio_channels, self.group_size, dtype=torch.long))
|
||||
logger.info(
|
||||
"[tokenizer worker] Warmup finished in %.2fs",
|
||||
time.monotonic() - warmup_start,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("[tokenizer worker] Warmup failed (non-critical): %s", str(e))
|
||||
else:
|
||||
logger.info("[tokenizer worker] Skipping warmup for CPU device")
|
||||
|
||||
def resample_audio_if_needed(self, wav_tensor: torch.Tensor, original_sr: int):
|
||||
"""Resample audio if sample rate doesn't match config"""
|
||||
target_sr = self.sample_rate
|
||||
if original_sr != target_sr:
|
||||
wav_tensor = torchaudio.functional.resample(wav_tensor, original_sr, target_sr)
|
||||
return wav_tensor
|
||||
|
||||
def wav2mel(self, wav: torch.Tensor):
|
||||
"""Convert waveform to mel spectrogram using consistent processing"""
|
||||
wav = wav.to(torch.float32)
|
||||
spec = self.mel_transform(wav[None, :])
|
||||
return torch.log(torch.clip(spec, min=1e-7)).squeeze()
|
||||
|
||||
def group_by_length(self, features: torch.Tensor, lengths: torch.Tensor, max_length: int):
|
||||
if features.size(0) != lengths.sum().item():
|
||||
raise ValueError(f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}")
|
||||
|
||||
split_points = []
|
||||
current_sum = 0
|
||||
|
||||
for i, seq_len in enumerate(lengths):
|
||||
if current_sum + seq_len > max_length and current_sum > 0:
|
||||
split_points.append(i)
|
||||
current_sum = seq_len.item()
|
||||
else:
|
||||
current_sum += seq_len.item()
|
||||
|
||||
# Convert split points to group sizes
|
||||
group_sizes = []
|
||||
prev = 0
|
||||
for point in split_points:
|
||||
group_sizes.append(point - prev)
|
||||
prev = point
|
||||
if prev < len(lengths):
|
||||
group_sizes.append(len(lengths) - prev)
|
||||
|
||||
len_groups = torch.split(lengths, group_sizes)
|
||||
feature_sizes = [group.sum().item() for group in len_groups]
|
||||
feature_groups = torch.split(features, feature_sizes)
|
||||
|
||||
return feature_groups, len_groups
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode_batch_base(
|
||||
self,
|
||||
feature_groups: list[torch.Tensor],
|
||||
len_groups: list[torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
"""Run this in cuda stream if available"""
|
||||
encoded_parts = []
|
||||
for features, lengths in zip(feature_groups, len_groups):
|
||||
codes, _ = self.audio_tokenizer.encoder.encode(
|
||||
input_features=features.to(self.device),
|
||||
input_lens=lengths.to(self.device),
|
||||
return_codes_only=True,
|
||||
)
|
||||
encoded_parts.append(codes)
|
||||
codes_packed = torch.cat(encoded_parts, dim=1)
|
||||
|
||||
codes = codes_packed.transpose(0, 1).detach()
|
||||
audio_codes = codes[:, : self.audio_channels]
|
||||
|
||||
# Pad the sequence to be a multiple of group_size by repeating the last frame
|
||||
T = audio_codes.shape[0]
|
||||
if T % self.group_size != 0:
|
||||
pad = self.group_size - (T % self.group_size)
|
||||
last_tokens = audio_codes[-1, :] # Keep dim for repeat
|
||||
padding_tokens = last_tokens.expand(pad, -1)
|
||||
audio_codes = torch.cat([audio_codes, padding_tokens], dim=0)
|
||||
|
||||
audio_codes = audio_codes.transpose(0, 1).cpu()
|
||||
|
||||
return audio_codes # [audio_channels, T] cpu
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode(self, audio: tuple[torch.Tensor, int], max_length: float | None = 256000) -> torch.Tensor:
|
||||
"""wav: [samples] cpu, sample_rate = 24000"""
|
||||
wav, original_sr = audio
|
||||
wav = self.resample_audio_if_needed(wav, original_sr)
|
||||
wav = wav.to(self.device)
|
||||
|
||||
mel = self.wav2mel(wav).transpose(0, 1) # (seq_len, n_mels)
|
||||
input_len = mel.size(0)
|
||||
input_features = mel
|
||||
segment_size = 6000
|
||||
input_len_seg = [segment_size] * (input_len // segment_size)
|
||||
|
||||
if input_len % segment_size > 0:
|
||||
input_len_seg.append(input_len % segment_size)
|
||||
|
||||
input_lens = torch.tensor(input_len_seg, device=self.device)
|
||||
|
||||
feature_groups, len_groups = self.group_by_length(input_features, input_lens, max_length)
|
||||
|
||||
audio_codes = self.encode_batch_base(feature_groups, len_groups)
|
||||
return audio_codes # [audio_channels, T] cpu
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode_wav_base(
|
||||
self,
|
||||
wav: torch.Tensor, # [samples] cpu
|
||||
) -> torch.Tensor:
|
||||
"""Run this in cuda stream if available"""
|
||||
wav = wav.to(self.device)
|
||||
|
||||
mel = self.wav2mel(wav).transpose(0, 1) # (seq_len, n_mels)
|
||||
input_features = mel # [seq_len, n_mels]
|
||||
input_lens = torch.tensor([mel.shape[0]], device=self.device)
|
||||
codes_packed, _ = self.audio_tokenizer.encoder.encode(
|
||||
input_features=input_features,
|
||||
input_lens=input_lens,
|
||||
return_codes_only=True,
|
||||
)
|
||||
codes = codes_packed.transpose(0, 1).detach()
|
||||
audio_codes = codes[:, : self.audio_channels]
|
||||
|
||||
# Pad the sequence to be a multiple of group_size by repeating the last frame
|
||||
T = audio_codes.shape[0]
|
||||
if T % self.group_size != 0:
|
||||
pad = self.group_size - (T % self.group_size)
|
||||
last_tokens = audio_codes[-1, :] # Keep dim for repeat
|
||||
padding_tokens = last_tokens.expand(pad, -1)
|
||||
audio_codes = torch.cat([audio_codes, padding_tokens], dim=0)
|
||||
|
||||
audio_codes = audio_codes.transpose(0, 1).cpu()
|
||||
|
||||
return audio_codes # [audio_channels, T] cpu
|
||||
|
||||
@torch.inference_mode()
|
||||
def decode(
|
||||
self,
|
||||
tokens: torch.Tensor, # [audio_channels, T] cpu
|
||||
) -> torch.Tensor:
|
||||
"""Decode audio tokens to waveform using the tokenizer's decoder"""
|
||||
tokens = tokens.to(self.device)
|
||||
with torch.no_grad():
|
||||
decoded_audio: torch.Tensor = self.audio_tokenizer.decode(tokens)
|
||||
decoded_audio = decoded_audio.float().reshape(-1).detach().cpu()
|
||||
return decoded_audio # [samples] cpu
|
||||
|
||||
|
||||
@dataclass
|
||||
class AudioStreamerConfig:
|
||||
group_size: int
|
||||
audio_channels: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class MiMoAudioCodes:
|
||||
sosp: int
|
||||
eosp: int
|
||||
sostm: int
|
||||
eostm: int
|
||||
im_end: int
|
||||
pad: int
|
||||
eot: int
|
||||
empty: int
|
||||
|
||||
@classmethod
|
||||
def from_tokenizer(cls, tokenizer: AutoTokenizer) -> "MiMoAudioCodes":
|
||||
def idx(token: str) -> int:
|
||||
token_id = tokenizer.convert_tokens_to_ids(token)
|
||||
if not isinstance(token_id, int):
|
||||
raise ValueError(f"Token {token} is not mapped to a single id.")
|
||||
return token_id
|
||||
|
||||
return cls(
|
||||
sosp=idx("<|sosp|>"),
|
||||
eosp=idx("<|eosp|>"),
|
||||
sostm=idx("<|sostm|>"),
|
||||
eostm=idx("<|eostm|>"),
|
||||
im_end=idx("<|im_end|>"),
|
||||
pad=idx("<|endoftext|>"),
|
||||
eot=idx("<|eot|>"),
|
||||
empty=idx("<|empty|>"),
|
||||
)
|
||||
|
||||
|
||||
def extract_audio_code_tensor(
|
||||
flat_codes: torch.Tensor,
|
||||
group_size: int,
|
||||
audio_channels: int,
|
||||
codes: MiMoAudioCodes,
|
||||
) -> torch.Tensor | None:
|
||||
"""Convert flattened talker output into [audio_channels, T] codes."""
|
||||
if flat_codes.numel() == 0:
|
||||
return None
|
||||
|
||||
group_width = group_size * (audio_channels + 1)
|
||||
usable = (flat_codes.numel() // group_width) * group_width
|
||||
if usable == 0:
|
||||
return None
|
||||
|
||||
groups = flat_codes[:usable].view(-1, group_size, audio_channels + 1)
|
||||
audio_buffer: list[torch.Tensor] = []
|
||||
|
||||
for group in groups:
|
||||
text_token = int(group[0, 0].item())
|
||||
if text_token == codes.empty:
|
||||
audio_buffer.append(group[:, 1:])
|
||||
elif text_token == codes.eostm:
|
||||
break
|
||||
|
||||
if not audio_buffer:
|
||||
return None
|
||||
|
||||
audio_tokens = torch.cat(audio_buffer, dim=0).transpose(0, 1).contiguous()
|
||||
return audio_tokens # [audio_channels, T]
|
||||
|
||||
|
||||
def _normalize_tokenizer_worker_cache_key(
|
||||
device: torch.device,
|
||||
config_path: str | None,
|
||||
audio_tokenizer_path: str,
|
||||
) -> tuple[str, str, str]:
|
||||
"""Normalize cache key so that same tokenizer always hits the same cache entry."""
|
||||
device_type = device.type if isinstance(device, torch.device) else str(device).split(":")[0]
|
||||
# Use realpath so symlinks / trailing slash don't create duplicate entries
|
||||
ap = audio_tokenizer_path or ""
|
||||
if ap and os.path.exists(ap):
|
||||
ap = os.path.realpath(ap)
|
||||
cp = config_path or ""
|
||||
if cp and os.path.exists(cp):
|
||||
cp = os.path.realpath(cp)
|
||||
|
||||
if not cp and ap:
|
||||
cp = os.path.dirname(ap)
|
||||
return (device_type, cp, ap)
|
||||
|
||||
|
||||
_TOKENIZER_WORKER_CACHE: dict[tuple[str, str, str], MiMoAudioTokenizerWorker] = {}
|
||||
|
||||
|
||||
def get_tokenizer_worker(
|
||||
device: torch.device,
|
||||
config_path: str,
|
||||
audio_tokenizer_path: str,
|
||||
) -> MiMoAudioTokenizerWorker:
|
||||
key = _normalize_tokenizer_worker_cache_key(device, config_path, audio_tokenizer_path)
|
||||
if key not in _TOKENIZER_WORKER_CACHE:
|
||||
device_type = key[0]
|
||||
_TOKENIZER_WORKER_CACHE[key] = MiMoAudioTokenizerWorker(
|
||||
device_str=device_type,
|
||||
config_path=config_path,
|
||||
audio_tokenizer_path=audio_tokenizer_path,
|
||||
)
|
||||
logger.info(
|
||||
"[tokenizer cache worker] cache size=%s pid=%s",
|
||||
len(_TOKENIZER_WORKER_CACHE),
|
||||
os.getpid(),
|
||||
)
|
||||
return _TOKENIZER_WORKER_CACHE[key]
|
||||
|
||||
|
||||
class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
|
||||
"""Decode MiMo audio codes to waveform for the code2wav stage."""
|
||||
|
||||
def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
config = vllm_config.model_config.hf_config
|
||||
config = MiMoAudioConfig(**vars(config)) if isinstance(config, Qwen2Config) else config
|
||||
self.config = config
|
||||
self.vllm_config = vllm_config
|
||||
self.quant_config = vllm_config.quant_config
|
||||
self.lora_config = vllm_config.lora_config
|
||||
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
|
||||
device_str = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
self.device = torch.device(device_str)
|
||||
self.sample_rate = getattr(config, "audio_sample_rate", 24000)
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path, trust_remote_code=True)
|
||||
self.codes = MiMoAudioCodes.from_tokenizer(self.tokenizer)
|
||||
self.streamer_config = AudioStreamerConfig(
|
||||
group_size=self.config.group_size,
|
||||
audio_channels=self.config.audio_channels,
|
||||
)
|
||||
|
||||
self.audio_tokenizer_path = getattr(vllm_config.model_config, "audio_tokenizer_path", None) or os.environ.get(
|
||||
"MIMO_AUDIO_TOKENIZER_PATH"
|
||||
)
|
||||
if not self.audio_tokenizer_path:
|
||||
raise ValueError(
|
||||
"Audio tokenizer path is not set. Provide "
|
||||
"`model_config.audio_tokenizer_path` in the stage config "
|
||||
"or export MIMO_AUDIO_TOKENIZER_PATH."
|
||||
)
|
||||
|
||||
self.tokenizer_config_path = (
|
||||
getattr(vllm_config.model_config, "audio_tokenizer_config_path", None)
|
||||
or os.environ.get("MIMO_AUDIO_TOKENIZER_PATH")
|
||||
or self.config.name_or_path
|
||||
)
|
||||
|
||||
self._tokenizer_service: MiMoAudioTokenizerWorker | None = get_tokenizer_worker(
|
||||
device=self.device,
|
||||
config_path=self.tokenizer_config_path,
|
||||
audio_tokenizer_path=self.audio_tokenizer_path,
|
||||
)
|
||||
# samples per codec frame for streaming context strip (same as tokenizer frames_per_token)
|
||||
audio_tokenizer_config = self._tokenizer_service.audio_tokenizer.config
|
||||
self.total_upsample = (
|
||||
getattr(audio_tokenizer_config, "avg_pooler", 2)
|
||||
* getattr(audio_tokenizer_config, "stride_size", 2)
|
||||
* getattr(audio_tokenizer_config, "hop_length", 240)
|
||||
) * self.config.group_size
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
audio__dict_path: str | None = None,
|
||||
**kwargs,
|
||||
) -> set[str]:
|
||||
# Decoder has no trainable weights to load in this stage.
|
||||
return set()
|
||||
|
||||
def chunked_decode_streaming(
|
||||
self,
|
||||
codes: torch.Tensor,
|
||||
chunk_size: int = 10,
|
||||
left_context_size: int = 10,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Decode one chunk of codes and return waveform with left context removed.
|
||||
|
||||
Used when async_chunk is True: each chunk is decoded and we strip
|
||||
context_size * total_upsample samples from the start (left context).
|
||||
"""
|
||||
wav_chunk = self._decode_waveform_from_codes(codes)
|
||||
if wav_chunk.numel() == 0:
|
||||
return wav_chunk
|
||||
# Infer number of codec frames: flat format has 36 elements per frame (9*4 with prepend)
|
||||
num_flat = codes.numel() if codes.ndim == 1 else codes.shape[-1]
|
||||
num_chunks = num_flat // 36
|
||||
if num_chunks <= chunk_size:
|
||||
context_size = 0
|
||||
else:
|
||||
context_size = left_context_size
|
||||
drop = context_size * self.total_upsample
|
||||
if drop >= wav_chunk.numel():
|
||||
return wav_chunk[:0].clone()
|
||||
return wav_chunk[drop:].contiguous()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
codes: torch.Tensor | None = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
# Support both input_ids (from vLLM) and codes (from mimo_audio.py)
|
||||
code_tensor = codes if codes is not None else input_ids
|
||||
|
||||
if code_tensor is None:
|
||||
raise ValueError("Either input_ids or codes must be provided")
|
||||
|
||||
if code_tensor.numel() == 0:
|
||||
raise ValueError("code_tensor is empty.")
|
||||
|
||||
if getattr(self.vllm_config.model_config, "async_chunk", False):
|
||||
waveform = self.chunked_decode_streaming(code_tensor, chunk_size=10, left_context_size=10)
|
||||
else:
|
||||
waveform = self._decode_waveform_from_codes(code_tensor)
|
||||
|
||||
# For warmup/dummy run: if waveform is empty, return dummy hidden_states
|
||||
# _dummy_run expects a tensor that can be processed by extract_multimodal_outputs
|
||||
if waveform.numel() == 0:
|
||||
# Return dummy hidden_states with shape [seq_len, hidden_size]
|
||||
# This is needed for warmup to avoid NoneType errors
|
||||
hidden_size = getattr(self.config, "hidden_size", None)
|
||||
if hidden_size is None:
|
||||
# Fallback: try to get from vllm_config.model_config
|
||||
hidden_size = getattr(
|
||||
self.vllm_config.model_config,
|
||||
"hidden_size",
|
||||
4096, # Default fallback
|
||||
)
|
||||
|
||||
# Return [seq_len=1, hidden_size] tensor for warmup
|
||||
dummy_hidden = torch.zeros(
|
||||
(1, hidden_size),
|
||||
dtype=torch.bfloat16,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
return dummy_hidden
|
||||
|
||||
if waveform.numel() == 1 and waveform.item() == 0:
|
||||
# Designed for Text output only
|
||||
return None
|
||||
|
||||
return waveform
|
||||
|
||||
def _get_full_code_sequence(
|
||||
self,
|
||||
code_tensor: torch.Tensor | None,
|
||||
kwargs: dict,
|
||||
) -> torch.Tensor | None:
|
||||
"""Gather the full prompt token ids (code sequence) if available."""
|
||||
if code_tensor is not None and code_tensor.numel() > 1:
|
||||
return code_tensor
|
||||
|
||||
sampling_metadata = kwargs.get("sampling_metadata")
|
||||
prompt_token_ids = getattr(sampling_metadata, "prompt_token_ids", None) if sampling_metadata else None
|
||||
|
||||
if prompt_token_ids is None:
|
||||
return code_tensor
|
||||
|
||||
if isinstance(prompt_token_ids, torch.Tensor):
|
||||
return prompt_token_ids.detach().to(torch.long).view(-1)
|
||||
|
||||
if isinstance(prompt_token_ids, list):
|
||||
if len(prompt_token_ids) == 0:
|
||||
return code_tensor
|
||||
if isinstance(prompt_token_ids[0], list):
|
||||
flat = prompt_token_ids[0]
|
||||
else:
|
||||
flat = prompt_token_ids
|
||||
if len(flat) == 0:
|
||||
return code_tensor
|
||||
return torch.tensor(flat, dtype=torch.long)
|
||||
|
||||
return code_tensor
|
||||
|
||||
def _check_dummy_code_tensor(self, code_tensor: torch.Tensor) -> bool:
|
||||
if code_tensor is not None and code_tensor.numel() == DUMMY_CODE_SHAPE:
|
||||
code_groups = code_tensor.view(self.config.group_size, self.config.audio_channels + 1)
|
||||
return (
|
||||
(code_groups[:, 0] == TALKER_CODEC_PAD_TOKEN_ID).all() and (code_groups[:, 1:].sum() == 0).all()
|
||||
).item()
|
||||
return False
|
||||
|
||||
def _decode_waveform_from_codes(self, code_tensor: torch.Tensor) -> torch.Tensor:
|
||||
# Check if in CUDA graph capture phase
|
||||
is_capturing = torch.cuda.is_current_stream_capturing()
|
||||
|
||||
# During CUDA graph capture, return dummy tensor to avoid operations like .cpu() which are not allowed
|
||||
if is_capturing:
|
||||
# Return an empty dummy tensor with shape and dtype consistent with normal output
|
||||
return torch.zeros(0, dtype=torch.float32, device=self.device)
|
||||
|
||||
if code_tensor is None or code_tensor.numel() == 0:
|
||||
return torch.zeros(0, dtype=torch.float32, device=self.device)
|
||||
|
||||
if self._check_dummy_code_tensor(code_tensor):
|
||||
return torch.zeros(1, dtype=torch.float32, device=self.device)
|
||||
|
||||
if code_tensor.ndim > 1:
|
||||
if code_tensor.shape[0] == 1:
|
||||
code_tensor = code_tensor.squeeze(0)
|
||||
|
||||
elif code_tensor.shape[0] == self.streamer_config.audio_channels + 1:
|
||||
code_tensor = code_tensor.transpose(0, 1).contiguous().view(-1)
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"code2wav expects shape [L], [1, L] or [audio_channels+1, T], got {code_tensor.shape}"
|
||||
)
|
||||
|
||||
flat_codes = code_tensor.detach().to(torch.long).cpu()
|
||||
|
||||
audio_codes = extract_audio_code_tensor(
|
||||
flat_codes,
|
||||
self.streamer_config.group_size,
|
||||
self.streamer_config.audio_channels,
|
||||
self.codes,
|
||||
)
|
||||
|
||||
if audio_codes is None or audio_codes.numel() == 0:
|
||||
return torch.zeros(0, dtype=torch.float32, device=self.device)
|
||||
|
||||
decoded_audio = self._tokenizer_service.decode(audio_codes)
|
||||
|
||||
result = decoded_audio.to(self.device)
|
||||
|
||||
return result
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings=None,
|
||||
is_multimodal: bool = False,
|
||||
) -> torch.Tensor:
|
||||
return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size())
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor | None:
|
||||
# code2wav stage does not need actual logits, but needs to return a valid tensor
|
||||
# Return a dummy logits tensor with shape [batch_size, vocab_size]
|
||||
if hidden_states.numel() == 0:
|
||||
return None
|
||||
|
||||
batch_size = hidden_states.shape[0] if hidden_states.ndim > 0 else 1
|
||||
vocab_size = self.config.vocab_size
|
||||
|
||||
# Create an all-zero logits tensor (sampler will choose the first token, usually pad token)
|
||||
# This is safe for the code2wav stage as we do not rely on sampling results
|
||||
dummy_logits = torch.zeros(
|
||||
(batch_size, vocab_size),
|
||||
dtype=hidden_states.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
return dummy_logits
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> SamplerOutput | None:
|
||||
# code2wav stage does not need sampling, but needs to return a valid SamplerOutput to avoid CUDA errors
|
||||
# Use Sampler to handle sampling, even though the result will not be used
|
||||
if logits is None or logits.numel() == 0:
|
||||
return None
|
||||
|
||||
# Use Sampler for sampling (although the result will not be used)
|
||||
# This ensures the sampling process can complete normally and avoids CUDA errors
|
||||
sampler = Sampler()
|
||||
sampler_output = sampler(
|
||||
logits=logits,
|
||||
sampling_metadata=sampling_metadata,
|
||||
)
|
||||
return sampler_output
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,824 @@
|
||||
# Copyright 2025 Xiaomi Corporation.
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
from transformers.activations import ACT2FN
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from .config_mimo_audio import MiMoAudioTokenizerConfig
|
||||
from .modeling_rope_utils import ROPE_INIT_FUNCTIONS, apply_rotary_pos_emb, dynamic_rope_update
|
||||
from .quantization import ResidualVectorQuantizer
|
||||
|
||||
logger = init_logger(__name__)
|
||||
is_flash_atth_available = False
|
||||
try:
|
||||
from flash_attn import flash_attn_varlen_func
|
||||
|
||||
is_flash_atth_available = True
|
||||
except Exception:
|
||||
logger.warning("flash_attn not installed")
|
||||
|
||||
|
||||
def get_sequence_mask(inputs, inputs_length):
|
||||
if inputs.dim() == 3:
|
||||
bsz, tgt_len, _ = inputs.size()
|
||||
else:
|
||||
bsz, tgt_len = inputs_length.shape[0], torch.max(inputs_length)
|
||||
sequence_mask = torch.arange(0, tgt_len).to(inputs.device)
|
||||
sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view(bsz, tgt_len, 1)
|
||||
unpacking_index = torch.cumsum(sequence_mask.to(torch.int64).view(-1), dim=0) - 1
|
||||
return sequence_mask, unpacking_index
|
||||
|
||||
|
||||
def unpack_hidden_states(hidden_states, lengths, sequence_mask=None, unpacking_index=None):
|
||||
bsz = lengths.shape[0]
|
||||
if sequence_mask is None or unpacking_index is None:
|
||||
sequence_mask, unpacking_index = get_sequence_mask(hidden_states, lengths)
|
||||
hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(
|
||||
bsz, torch.max(lengths), hidden_states.shape[-1]
|
||||
)
|
||||
hidden_states = torch.where(sequence_mask, hidden_states, 0)
|
||||
return hidden_states
|
||||
|
||||
|
||||
def get_position_ids(lengths):
|
||||
total_len = lengths.sum()
|
||||
offset = torch.cat([torch.zeros(1).to(lengths), lengths[:-1].cumsum(dim=0)])
|
||||
offset = torch.repeat_interleave(offset, lengths)
|
||||
position_ids = torch.arange(0, total_len).to(offset) - offset
|
||||
return position_ids
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamingConfig:
|
||||
seg_point: int = field(default=60 * 25)
|
||||
process_seg_point: bool = field(default=True)
|
||||
left_overlap: int = field(default=10 * 25)
|
||||
right_overlap: int = field(default=40)
|
||||
seg_point_left_overlap: int = field(default=0)
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamingCache:
|
||||
hidden_states: list[torch.Tensor] = field(default=None)
|
||||
processed_lengths: list[int] = field(default=None)
|
||||
|
||||
|
||||
class ISTFT(nn.Module):
|
||||
"""
|
||||
Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
|
||||
windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
|
||||
See issue: https://github.com/pytorch/pytorch/issues/62323
|
||||
Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
|
||||
The NOLA constraint is met as we trim padded samples anyway.
|
||||
|
||||
Args:
|
||||
n_fft (int): Size of Fourier transform.
|
||||
hop_length (int): The distance between neighboring sliding window frames.
|
||||
win_length (int): The size of window frame and STFT filter.
|
||||
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
|
||||
"""
|
||||
|
||||
def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
|
||||
super().__init__()
|
||||
if padding not in ["center", "same"]:
|
||||
raise ValueError("Padding must be 'center' or 'same'.")
|
||||
self.padding = padding
|
||||
self.n_fft = n_fft
|
||||
self.hop_length = hop_length
|
||||
self.win_length = win_length
|
||||
window = torch.hann_window(win_length)
|
||||
self.register_buffer("window", window)
|
||||
|
||||
def forward(self, spec: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
|
||||
|
||||
Args:
|
||||
spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
|
||||
N is the number of frequency bins, and T is the number of time frames.
|
||||
|
||||
Returns:
|
||||
Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
|
||||
"""
|
||||
if self.padding == "center":
|
||||
# Fallback to pytorch native implementation
|
||||
return torch.istft(
|
||||
spec,
|
||||
self.n_fft,
|
||||
self.hop_length,
|
||||
self.win_length,
|
||||
self.window,
|
||||
center=True,
|
||||
)
|
||||
elif self.padding == "same":
|
||||
pad = (self.win_length - self.hop_length) // 2
|
||||
else:
|
||||
raise ValueError("Padding must be 'center' or 'same'.")
|
||||
|
||||
assert spec.dim() == 3, "Expected a 3D tensor as input"
|
||||
B, N, T = spec.shape
|
||||
|
||||
# Inverse FFT
|
||||
ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
|
||||
ifft = ifft * self.window[None, :, None]
|
||||
|
||||
# Overlap and Add
|
||||
output_size = (T - 1) * self.hop_length + self.win_length
|
||||
y = torch.nn.functional.fold(
|
||||
ifft,
|
||||
output_size=(1, output_size),
|
||||
kernel_size=(1, self.win_length),
|
||||
stride=(1, self.hop_length),
|
||||
)[:, 0, 0, pad:-pad]
|
||||
|
||||
# Window envelope
|
||||
window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
|
||||
window_envelope = torch.nn.functional.fold(
|
||||
window_sq,
|
||||
output_size=(1, output_size),
|
||||
kernel_size=(1, self.win_length),
|
||||
stride=(1, self.hop_length),
|
||||
).squeeze()[pad:-pad]
|
||||
|
||||
# Normalize
|
||||
assert (window_envelope > 1e-11).all()
|
||||
y = y / window_envelope
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class ISTFTHead(nn.Module):
|
||||
"""
|
||||
ISTFT Head module for predicting STFT complex coefficients.
|
||||
|
||||
Args:
|
||||
dim (int): Hidden dimension of the model.
|
||||
n_fft (int): Size of Fourier transform.
|
||||
hop_length (int): The distance between neighboring sliding window frames, which should align with
|
||||
the resolution of the input features.
|
||||
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
|
||||
"""
|
||||
|
||||
def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
|
||||
super().__init__()
|
||||
out_dim = n_fft + 2
|
||||
self.out = torch.nn.Linear(dim, out_dim)
|
||||
self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass of the ISTFTHead module.
|
||||
|
||||
Args:
|
||||
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
|
||||
L is the sequence length, and H denotes the model dimension.
|
||||
|
||||
Returns:
|
||||
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
|
||||
"""
|
||||
x = self.out(x).transpose(1, 2)
|
||||
mag, p = x.chunk(2, dim=1)
|
||||
mag = torch.exp(mag)
|
||||
mag = torch.clip(mag, max=1e2) # safeguard to prevent excessively large magnitudes
|
||||
# wrapping happens here. These two lines produce real and imaginary value
|
||||
x = torch.cos(p)
|
||||
y = torch.sin(p)
|
||||
# recalculating phase here does not produce anything new
|
||||
# only costs time
|
||||
# phase = torch.atan2(y, x)
|
||||
# S = mag * torch.exp(phase * 1j)
|
||||
# better directly produce the complex value
|
||||
original_dtype = x.dtype
|
||||
S = mag.float() * (x.float() + 1j * y.float())
|
||||
audio = self.istft(S)
|
||||
audio = audio.to(original_dtype)
|
||||
return audio
|
||||
|
||||
|
||||
class RotaryEmbedding(nn.Module):
|
||||
def __init__(self, base, dim, max_seq_len, rope_type="default", device=None):
|
||||
super().__init__()
|
||||
self.max_seq_len = max_seq_len
|
||||
self.rope_type = rope_type
|
||||
|
||||
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
||||
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(device=device, base=base, dim=dim)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self.original_inv_freq = self.inv_freq
|
||||
|
||||
@torch.no_grad()
|
||||
@dynamic_rope_update
|
||||
def forward(self, x, position_ids):
|
||||
inv_freq_expanded = self.inv_freq[:, None].float().expand(-1, 1).to(x.device)
|
||||
position_ids_expanded = position_ids[None, :].float()
|
||||
|
||||
device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
|
||||
with torch.autocast(device_type=device_type, enabled=False): # Force float32
|
||||
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(0, 1)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
cos = emb.cos() * self.attention_scaling
|
||||
sin = emb.sin() * self.attention_scaling
|
||||
|
||||
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
"""
|
||||
RMSNorm is equivalent to T5LayerNorm
|
||||
"""
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states):
|
||||
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
|
||||
if self.weight.dtype in [torch.float16, torch.bfloat16]:
|
||||
hidden_states = hidden_states.to(self.weight.dtype)
|
||||
|
||||
return self.weight * hidden_states
|
||||
|
||||
|
||||
LAYER_NORM = {"LayerNorm": nn.LayerNorm, "RMSNorm": RMSNorm}
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, embed_dim, num_heads, window_size=(-1, -1), causal=False):
|
||||
super().__init__()
|
||||
self.embed_dim = embed_dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = embed_dim // num_heads
|
||||
self.window_size = window_size
|
||||
|
||||
self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
|
||||
self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
|
||||
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
|
||||
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
|
||||
|
||||
self.causal = causal
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
seq_len: torch.Tensor,
|
||||
rope_position_embeddings=None,
|
||||
):
|
||||
total_seq_len, _ = hidden_states.size()
|
||||
|
||||
query_states = self.q_proj(hidden_states).view(total_seq_len, self.num_heads, self.head_dim)
|
||||
key_states = self.k_proj(hidden_states).view(total_seq_len, self.num_heads, self.head_dim)
|
||||
value_states = self.v_proj(hidden_states).view(total_seq_len, self.num_heads, self.head_dim)
|
||||
|
||||
if rope_position_embeddings is not None:
|
||||
cos, sin = rope_position_embeddings
|
||||
query_states = apply_rotary_pos_emb(query_states, cos, sin)
|
||||
key_states = apply_rotary_pos_emb(key_states, cos, sin)
|
||||
|
||||
if hidden_states.is_cuda and is_flash_atth_available is True:
|
||||
# === Use flash-attn in GPU mode ===
|
||||
cu_len = F.pad(torch.cumsum(seq_len, dim=0), (1, 0), "constant", 0).to(torch.int32)
|
||||
max_seqlen = torch.max(seq_len).to(torch.int32).detach()
|
||||
attn_output = flash_attn_varlen_func(
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
cu_len,
|
||||
cu_len,
|
||||
max_seqlen,
|
||||
max_seqlen,
|
||||
causal=self.causal,
|
||||
window_size=self.window_size,
|
||||
)
|
||||
attn_output = attn_output.reshape(total_seq_len, self.embed_dim)
|
||||
|
||||
else:
|
||||
# === Fallback implementation in CPU / Eager mode ===
|
||||
cu_len = F.pad(torch.cumsum(seq_len, dim=0), (1, 0), "constant", 0).to(torch.long)
|
||||
attn_output = torch.zeros_like(hidden_states)
|
||||
|
||||
for i, slen in enumerate(seq_len.tolist()):
|
||||
start_idx = cu_len[i].item()
|
||||
end_idx = cu_len[i + 1].item()
|
||||
|
||||
q = query_states[start_idx:end_idx] # [slen, num_heads, head_dim]
|
||||
k = key_states[start_idx:end_idx]
|
||||
v = value_states[start_idx:end_idx]
|
||||
|
||||
q = q.transpose(0, 1) # [num_heads, slen, head_dim]
|
||||
k = k.transpose(0, 1)
|
||||
v = v.transpose(0, 1)
|
||||
|
||||
attn_scores = torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim**0.5)
|
||||
if self.causal:
|
||||
mask = torch.tril(torch.ones_like(attn_scores))
|
||||
attn_scores = attn_scores.masked_fill(mask == 0, float("-inf"))
|
||||
attn_probs = F.softmax(attn_scores, dim=-1)
|
||||
attn_out = torch.matmul(attn_probs, v) # [num_heads, slen, head_dim]
|
||||
|
||||
attn_out = attn_out.transpose(0, 1).reshape(slen, self.embed_dim)
|
||||
attn_output[start_idx:end_idx] = attn_out
|
||||
|
||||
attn_output = self.out_proj(attn_output)
|
||||
return attn_output
|
||||
|
||||
|
||||
class TransformerLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
act,
|
||||
d_model,
|
||||
encoder_attention_heads,
|
||||
encoder_ffn_dim,
|
||||
causal,
|
||||
ln_type="LayerNorm",
|
||||
attn_window_size=(-1, -1),
|
||||
):
|
||||
super().__init__()
|
||||
self.embed_dim = d_model
|
||||
self.self_attn = Attention(self.embed_dim, encoder_attention_heads, attn_window_size, causal)
|
||||
|
||||
self.self_attn_layer_norm = LAYER_NORM[ln_type](self.embed_dim)
|
||||
|
||||
self.activation_fn = act
|
||||
self.fc1 = nn.Linear(self.embed_dim, encoder_ffn_dim)
|
||||
self.fc2 = nn.Linear(encoder_ffn_dim, self.embed_dim)
|
||||
|
||||
self.final_layer_norm = LAYER_NORM[ln_type](self.embed_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
seq_len: torch.Tensor,
|
||||
rope_position_embeddings: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.self_attn_layer_norm(hidden_states)
|
||||
hidden_states = self.self_attn(hidden_states, seq_len, rope_position_embeddings=rope_position_embeddings)
|
||||
hidden_states = residual + hidden_states
|
||||
residual = hidden_states
|
||||
hidden_states = self.final_layer_norm(hidden_states)
|
||||
hidden_states = self.activation_fn(self.fc1(hidden_states))
|
||||
hidden_states = self.fc2(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
if (hidden_states.dtype == torch.float16 or hidden_states.dtype == torch.bfloat16) and (
|
||||
torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
|
||||
):
|
||||
clamp_value = torch.finfo(hidden_states.dtype).max - 1000
|
||||
hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class TransformerVocos(nn.Module):
|
||||
def __init__(self, config: MiMoAudioTokenizerConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.max_source_positions = self.config.max_audio_seconds * self.config.sampling_rate // self.config.hop_length
|
||||
self.embeddings = nn.Linear(config.n_mels, config.vocoder_dim, bias=False)
|
||||
|
||||
self.position_embedding = RotaryEmbedding(
|
||||
config.rope_theta,
|
||||
config.vocoder_dim // config.vocoder_attention_heads,
|
||||
self.max_source_positions,
|
||||
self.config.rope_type,
|
||||
)
|
||||
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
TransformerLayer(
|
||||
ACT2FN[self.config.activation_function],
|
||||
self.config.vocoder_dim,
|
||||
self.config.vocoder_attention_heads,
|
||||
self.config.vocoder_intermediate_dim,
|
||||
causal=False,
|
||||
ln_type=self.config.ln_type,
|
||||
attn_window_size=self.config.vocoder_attn_window_size,
|
||||
)
|
||||
for _ in range(self.config.vocoder_num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.layer_norm = LAYER_NORM[self.config.ln_type](self.config.vocoder_dim)
|
||||
self.hop_size = self.config.hop_length
|
||||
self.head = ISTFTHead(
|
||||
self.config.vocoder_dim,
|
||||
self.config.nfft,
|
||||
self.config.hop_length,
|
||||
self.config.vocoder_padding,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, input_length):
|
||||
x = x.transpose(1, 2)
|
||||
attention_mask, unpacking_index = get_sequence_mask(x, input_length)
|
||||
x = torch.masked_select(x, attention_mask).view(torch.sum(input_length), self.config.n_mels)
|
||||
x = self.embeddings(x)
|
||||
position_ids = torch.arange(0, x.size(0), device=x.device, dtype=torch.long)
|
||||
rope_position_embeddings = self.position_embedding(x, position_ids)
|
||||
for idx, layer in enumerate(self.layers):
|
||||
x = layer(x, input_length, rope_position_embeddings=rope_position_embeddings)
|
||||
|
||||
x = self.layer_norm(x)
|
||||
x = unpack_hidden_states(x, input_length, attention_mask, unpacking_index)
|
||||
x = self.head(x)
|
||||
output_length = input_length * self.hop_size
|
||||
return x[:, None, :], output_length
|
||||
|
||||
|
||||
class AudioEncoder(nn.Module):
|
||||
def __init__(self, config: MiMoAudioTokenizerConfig):
|
||||
super().__init__()
|
||||
config._attn_implementation = "flash_attention_2"
|
||||
self.config = config
|
||||
self.max_source_positions = (
|
||||
config.max_audio_seconds * config.sampling_rate // config.hop_length
|
||||
) // config.stride_size
|
||||
self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
|
||||
|
||||
self.skip_layer_idx = config.encoder_skip_layer_id
|
||||
self.conv1 = nn.Conv1d(config.n_mels, config.d_model, kernel_size=config.kernel_size, padding=1)
|
||||
self.conv2 = nn.Conv1d(
|
||||
config.d_model,
|
||||
config.d_model,
|
||||
kernel_size=config.kernel_size,
|
||||
stride=config.stride_size,
|
||||
padding=1,
|
||||
)
|
||||
|
||||
self.position_embedding = RotaryEmbedding(
|
||||
config.rope_theta,
|
||||
config.d_model // config.encoder_attention_heads,
|
||||
self.max_source_positions,
|
||||
config.rope_type,
|
||||
)
|
||||
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
TransformerLayer(
|
||||
ACT2FN[config.activation_function],
|
||||
config.d_model,
|
||||
config.encoder_attention_heads,
|
||||
config.encoder_ffn_dim,
|
||||
causal=self.config.encoder_causal,
|
||||
ln_type=self.config.ln_type,
|
||||
attn_window_size=self.config.encoder_attn_window_size,
|
||||
)
|
||||
for _ in range(config.encoder_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.layer_norm = LAYER_NORM[config.ln_type](config.d_model)
|
||||
|
||||
if self.config.avg_pooler != 1:
|
||||
self.down_sample_layer = nn.Sequential(
|
||||
nn.Conv1d(
|
||||
config.d_model,
|
||||
config.d_model,
|
||||
config.avg_pooler,
|
||||
config.avg_pooler,
|
||||
bias=False,
|
||||
),
|
||||
nn.GELU(),
|
||||
)
|
||||
self.down_sample_norm = LAYER_NORM[config.ln_type](config.d_model)
|
||||
else:
|
||||
self.down_sample_layer = None
|
||||
|
||||
if self.config.num_quantizers != 0:
|
||||
self.quantizer = ResidualVectorQuantizer(
|
||||
dimension=self.config.d_model,
|
||||
n_q=self.config.num_quantizers,
|
||||
bins=self.config.codebook_size,
|
||||
threshold_ema_dead_code=self.config.threshold_ema_dead_code,
|
||||
)
|
||||
else:
|
||||
self.quantizer = None
|
||||
|
||||
def get_features(self, input_features, output_length):
|
||||
input_features = input_features.to(self.conv1.weight)
|
||||
inputs_embeds = nn.functional.gelu(self.conv1(input_features))
|
||||
inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))
|
||||
inputs_embeds = inputs_embeds.permute(0, 2, 1)
|
||||
bsz, tgt_len, _ = inputs_embeds.size()
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
position_ids = get_position_ids(output_length).long().to(input_features.device)
|
||||
rope_position_embeddings = self.position_embedding(input_features, position_ids)
|
||||
|
||||
attention_mask, unpacking_index = get_sequence_mask(hidden_states, output_length)
|
||||
|
||||
hidden_states = torch.masked_select(hidden_states, attention_mask).view(
|
||||
torch.sum(output_length), self.config.d_model
|
||||
)
|
||||
|
||||
skip_connect_hidden_states = 0.0
|
||||
for idx, encoder_layer in enumerate(self.layers):
|
||||
hidden_states = encoder_layer(
|
||||
hidden_states,
|
||||
output_length,
|
||||
rope_position_embeddings=rope_position_embeddings,
|
||||
)
|
||||
if (self.skip_layer_idx is not None) and idx == self.skip_layer_idx - 1:
|
||||
skip_connect_hidden_states = hidden_states.clone()
|
||||
|
||||
hidden_states += skip_connect_hidden_states
|
||||
hidden_states = self.layer_norm(hidden_states)
|
||||
|
||||
if self.down_sample_layer is not None:
|
||||
hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(
|
||||
bsz, tgt_len, self.config.d_model
|
||||
)
|
||||
if hidden_states.size(1) % self.config.avg_pooler:
|
||||
pad_len = self.config.avg_pooler - hidden_states.size(1) % self.config.avg_pooler
|
||||
hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, pad_len), mode="constant", value=0.0)
|
||||
tgt_len += pad_len
|
||||
tgt_len = tgt_len // self.config.avg_pooler
|
||||
hidden_states = self.down_sample_layer(hidden_states.transpose(1, 2))
|
||||
output_length = (
|
||||
output_length // self.config.avg_pooler + (output_length % self.config.avg_pooler != 0).int()
|
||||
)
|
||||
hidden_states = hidden_states.transpose(1, 2)
|
||||
attention_mask, unpacking_index = get_sequence_mask(hidden_states, output_length)
|
||||
hidden_states = torch.masked_select(hidden_states, attention_mask).view(
|
||||
torch.sum(output_length), self.config.d_model
|
||||
)
|
||||
hidden_states = self.down_sample_norm(hidden_states)
|
||||
|
||||
return (
|
||||
hidden_states,
|
||||
output_length,
|
||||
attention_mask,
|
||||
unpacking_index,
|
||||
tgt_len,
|
||||
bsz,
|
||||
)
|
||||
|
||||
def get_output_length(self, mel_len):
|
||||
tgt_len = mel_len + 3 - self.config.kernel_size
|
||||
return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(
|
||||
self,
|
||||
input_features,
|
||||
input_lens=None,
|
||||
output_length=None,
|
||||
return_codes_only=False,
|
||||
n_q=None,
|
||||
use_quantizer=True,
|
||||
):
|
||||
if output_length is None:
|
||||
output_length = self.get_output_length(input_lens)
|
||||
input_features = unpack_hidden_states(input_features, input_lens)
|
||||
hidden_states, output_length, attention_mask, unpacking_index, tgt_len, bsz = self.get_features(
|
||||
input_features=input_features.transpose(1, 2),
|
||||
output_length=output_length,
|
||||
)
|
||||
|
||||
dtype = hidden_states.dtype
|
||||
|
||||
if use_quantizer and self.quantizer is not None:
|
||||
self.quantizer.float()
|
||||
|
||||
codes = self.quantizer.encode(hidden_states.float(), n_q=n_q)
|
||||
if return_codes_only:
|
||||
return codes, output_length
|
||||
hidden_states = self.quantizer.decode(codes)
|
||||
hidden_states = hidden_states.to(dtype)
|
||||
else:
|
||||
codes = None
|
||||
|
||||
hidden_states_packed = hidden_states.clone()
|
||||
|
||||
# unpacking
|
||||
hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(bsz, tgt_len, self.config.d_model)
|
||||
hidden_states = torch.where(attention_mask, hidden_states, 0)
|
||||
return hidden_states, hidden_states_packed, output_length, codes
|
||||
|
||||
@torch.no_grad()
|
||||
def decode_vq(self, codes):
|
||||
self.quantizer.float()
|
||||
hidden_states = self.quantizer.decode(codes)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CausalConvTranspose1d(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride):
|
||||
super().__init__()
|
||||
self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride)
|
||||
self.norm = nn.GroupNorm(1, out_channels)
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
def forward(self, hidden_states, input_length, output_dim=None):
|
||||
kernel_size = self.conv.kernel_size[0]
|
||||
stride = self.conv.stride[0]
|
||||
bsz = input_length.shape[0]
|
||||
|
||||
if output_dim is None:
|
||||
output_dim = hidden_states.dim()
|
||||
if hidden_states.dim() <= 2: # unpack sequence to 3d
|
||||
sequence_mask, unpacking_index = get_sequence_mask(hidden_states, input_length)
|
||||
hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(
|
||||
bsz, torch.max(input_length), self.in_channels
|
||||
)
|
||||
hidden_states = torch.where(sequence_mask, hidden_states, 0)
|
||||
|
||||
hidden_states = hidden_states.transpose(2, 1) # (N, L, C) -> (N, C, L)
|
||||
hidden_states = self.conv(hidden_states)
|
||||
hidden_states = self.norm(hidden_states)
|
||||
hidden_states = hidden_states.transpose(2, 1) # (N, C, L) -> (N, L, C)
|
||||
|
||||
casual_padding_right = max(0, kernel_size - stride)
|
||||
hidden_states = hidden_states[:, : hidden_states.shape[1] - casual_padding_right, :]
|
||||
output_length = (input_length - 1) * stride + kernel_size - casual_padding_right
|
||||
sequence_mask, _ = get_sequence_mask(hidden_states, output_length)
|
||||
if output_dim <= 2:
|
||||
hidden_states = torch.masked_select(hidden_states, sequence_mask).view(-1, self.out_channels)
|
||||
else:
|
||||
hidden_states = torch.where(sequence_mask, hidden_states, 0)
|
||||
hidden_states = hidden_states[:, : torch.max(output_length), :]
|
||||
return hidden_states, output_length
|
||||
|
||||
|
||||
class AudioDecoder(nn.Module):
|
||||
def __init__(self, config: MiMoAudioTokenizerConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.max_source_positions = self.config.max_audio_seconds * self.config.sampling_rate // self.config.hop_length
|
||||
|
||||
if self.config.avg_pooler != 1:
|
||||
self.dconv1 = CausalConvTranspose1d(
|
||||
self.config.d_model,
|
||||
self.config.d_model,
|
||||
self.config.avg_pooler,
|
||||
self.config.avg_pooler,
|
||||
)
|
||||
else:
|
||||
self.dconv1 = None
|
||||
|
||||
self.position_embedding = RotaryEmbedding(
|
||||
config.rope_theta,
|
||||
config.d_model // config.decoder_attention_heads,
|
||||
self.max_source_positions,
|
||||
config.rope_type,
|
||||
)
|
||||
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
TransformerLayer(
|
||||
ACT2FN[self.config.activation_function],
|
||||
self.config.d_model,
|
||||
self.config.decoder_attention_heads,
|
||||
self.config.decoder_ffn_dim,
|
||||
causal=self.config.decoder_causal,
|
||||
ln_type=self.config.ln_type,
|
||||
attn_window_size=self.config.decoder_attn_window_size,
|
||||
)
|
||||
for _ in range(self.config.decoder_layers)
|
||||
]
|
||||
)
|
||||
self.layer_norm = LAYER_NORM[config.ln_type](self.config.d_model)
|
||||
self.dconv2 = CausalConvTranspose1d(
|
||||
self.config.d_model,
|
||||
self.config.n_mels,
|
||||
self.config.decoder_kernel_size,
|
||||
self.config.decoder_stride_size,
|
||||
)
|
||||
self.vocoder = TransformerVocos(config)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
audio_embed,
|
||||
input_length,
|
||||
):
|
||||
assert audio_embed.shape[-1] == self.config.d_model
|
||||
audio_embed = audio_embed.to(self.layer_norm.weight)
|
||||
|
||||
if self.dconv1 is not None:
|
||||
audio_embed, output_length = self.dconv1(audio_embed, input_length, output_dim=3)
|
||||
else:
|
||||
output_length = input_length
|
||||
|
||||
hidden_states = audio_embed
|
||||
|
||||
position_ids = get_position_ids(output_length).long().to(hidden_states.device)
|
||||
rope_position_embeddings = self.position_embedding(hidden_states, position_ids)
|
||||
|
||||
# packing hidden states
|
||||
attention_mask, _ = get_sequence_mask(hidden_states, output_length)
|
||||
hidden_states = torch.masked_select(hidden_states, attention_mask).view(
|
||||
torch.sum(output_length), self.config.d_model
|
||||
)
|
||||
|
||||
for idx, encoder_layer in enumerate(self.layers):
|
||||
hidden_states = encoder_layer(
|
||||
hidden_states,
|
||||
output_length,
|
||||
rope_position_embeddings=rope_position_embeddings,
|
||||
)
|
||||
|
||||
hidden_states = self.layer_norm(hidden_states)
|
||||
|
||||
coarse_mel, output_length = self.dconv2(hidden_states, output_length, output_dim=3)
|
||||
|
||||
recon_wav, wav_length = self.vocoder(
|
||||
x=coarse_mel.transpose(1, 2),
|
||||
input_length=output_length,
|
||||
)
|
||||
|
||||
return recon_wav
|
||||
|
||||
|
||||
class MiMoAudioTokenizer(PreTrainedModel):
|
||||
config_class = MiMoAudioTokenizerConfig
|
||||
|
||||
def __init__(self, config: MiMoAudioTokenizerConfig):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
self.sampling_rate = config.sampling_rate
|
||||
self.encoder = AudioEncoder(config=config)
|
||||
self.decoder = AudioDecoder(config=config)
|
||||
self.downsample_rate = int(self.config.hop_length * 2 * self.config.avg_pooler)
|
||||
|
||||
def get_output_length(self, mel_len):
|
||||
tgt_len = mel_len + 3 - self.config.kernel_size
|
||||
return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(self, mels, input_lens, use_quantizer=True):
|
||||
input_features = mels
|
||||
hidden_states, hidden_states_packed, encoder_output_length, codes = self.encoder.encode(
|
||||
input_features, input_lens=input_lens, use_quantizer=use_quantizer
|
||||
)
|
||||
return hidden_states, hidden_states_packed, encoder_output_length, codes
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, codes):
|
||||
hidden_states = self.encoder.decode_vq(codes)
|
||||
output = self.decoder(
|
||||
hidden_states,
|
||||
torch.tensor([hidden_states.size(0)], device=hidden_states.device),
|
||||
)
|
||||
return output
|
||||
|
||||
@torch.no_grad()
|
||||
def streaming_decode(
|
||||
self,
|
||||
codes_chunks,
|
||||
chunk_input_lengths,
|
||||
history_cache=StreamingCache(),
|
||||
streaming_config=StreamingConfig(),
|
||||
last_chunk=False,
|
||||
):
|
||||
hidden_states = self.encoder.decode_vq(codes_chunks)
|
||||
input_lengths = []
|
||||
input_hidden_states = []
|
||||
start_idx = 0
|
||||
cache_hidden_states = []
|
||||
for i, input_length in enumerate(chunk_input_lengths):
|
||||
sample_hidden_states = hidden_states[start_idx : start_idx + input_length]
|
||||
start_idx += input_length
|
||||
if history_cache.hidden_states is not None:
|
||||
sample_hidden_states = torch.cat([history_cache.hidden_states[i], sample_hidden_states], dim=0)
|
||||
input_length += history_cache.hidden_states[i].size(0)
|
||||
input_hidden_states.append(sample_hidden_states)
|
||||
cache_hidden_states.append(sample_hidden_states.clone())
|
||||
input_lengths.append(input_length)
|
||||
input_hidden_states = torch.cat(input_hidden_states, dim=0)
|
||||
input_lengths = torch.tensor(input_lengths, device=hidden_states.device)
|
||||
output = self.decoder(input_hidden_states, input_lengths)
|
||||
return_wavs = []
|
||||
frames_per_token = self.config.avg_pooler * self.config.stride_size * self.config.hop_length
|
||||
processed_lengths = []
|
||||
for i, wav in enumerate(output):
|
||||
wav = wav.float().detach().cpu()
|
||||
start_idx = history_cache.processed_lengths[i] if history_cache.processed_lengths is not None else 0
|
||||
if last_chunk:
|
||||
return_wavs.append(wav[:, start_idx * frames_per_token :])
|
||||
new_processed_length = input_lengths[i].item()
|
||||
elif input_lengths[i].item() <= streaming_config.right_overlap:
|
||||
return_wavs.append(None)
|
||||
new_processed_length = 0
|
||||
else:
|
||||
end_idx = input_lengths[i].item() - streaming_config.right_overlap
|
||||
wav = wav[:, start_idx * frames_per_token : end_idx * frames_per_token]
|
||||
return_wavs.append(wav)
|
||||
new_processed_length = end_idx
|
||||
if input_lengths[i].item() > streaming_config.left_overlap:
|
||||
cache_hidden_states[i] = cache_hidden_states[i][-streaming_config.left_overlap :]
|
||||
new_processed_length -= input_lengths[i].item() - streaming_config.left_overlap
|
||||
processed_lengths.append(new_processed_length)
|
||||
history_cache.hidden_states = cache_hidden_states
|
||||
history_cache.processed_lengths = processed_lengths
|
||||
|
||||
return return_wavs, history_cache
|
||||
@@ -0,0 +1,703 @@
|
||||
# Copyright 2025 Xiaomi Corporation.
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from typing import Optional
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import is_torch_available, logging
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
if is_torch_available():
|
||||
import torch
|
||||
|
||||
|
||||
def dynamic_rope_update(rope_forward: Callable) -> Callable:
|
||||
"""
|
||||
Decorator function to update the RoPE parameters in the forward pass, if the model is using a dynamic RoPE
|
||||
(i.e. a RoPE implementation that may recompute its frequencies in the forward pass).
|
||||
|
||||
Args:
|
||||
rope_forward (Callable):
|
||||
The forward pass of the RoPE implementation.
|
||||
|
||||
Returns:
|
||||
The decorated forward pass.
|
||||
"""
|
||||
|
||||
def longrope_frequency_update(self, position_ids, device):
|
||||
"""Longrope uses long factor if sequence is larger than original pretraining length, short otherwise."""
|
||||
seq_len = torch.max(position_ids) + 1
|
||||
if hasattr(self.config, "original_max_position_embeddings"):
|
||||
original_max_position_embeddings = self.config.original_max_position_embeddings
|
||||
else:
|
||||
original_max_position_embeddings = self.config.max_position_embeddings
|
||||
if seq_len > original_max_position_embeddings:
|
||||
if not hasattr(self, "long_inv_freq"):
|
||||
self.long_inv_freq, _ = self.rope_init_fn(
|
||||
self.config, device, seq_len=original_max_position_embeddings + 1
|
||||
)
|
||||
self.register_buffer("inv_freq", self.long_inv_freq, persistent=False)
|
||||
else:
|
||||
# This .to() is needed if the model has been moved to a device after being initialized (because
|
||||
# the buffer is automatically moved, but not the original copy)
|
||||
self.original_inv_freq = self.original_inv_freq.to(device)
|
||||
self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
|
||||
|
||||
def dynamic_frequency_update(self, position_ids, device):
|
||||
"""
|
||||
dynamic RoPE layers should recompute `inv_freq` in the following situations:
|
||||
1 - growing beyond the cached sequence length (allow scaling)
|
||||
2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
|
||||
"""
|
||||
seq_len = torch.max(position_ids) + 1
|
||||
if seq_len > self.max_seq_len_cached: # growth
|
||||
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation
|
||||
self.max_seq_len_cached = seq_len
|
||||
|
||||
if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset
|
||||
# This .to() is needed if the model has been moved to a device after being initialized (because
|
||||
# the buffer is automatically moved, but not the original copy)
|
||||
self.original_inv_freq = self.original_inv_freq.to(device)
|
||||
self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
|
||||
self.max_seq_len_cached = self.original_max_seq_len
|
||||
|
||||
@wraps(rope_forward)
|
||||
def wrapper(self, x, position_ids):
|
||||
if "dynamic" in self.rope_type:
|
||||
dynamic_frequency_update(self, position_ids, device=x.device)
|
||||
elif self.rope_type == "longrope":
|
||||
longrope_frequency_update(self, position_ids, device=x.device)
|
||||
return rope_forward(self, x, position_ids)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def _compute_default_rope_parameters(
|
||||
config: PretrainedConfig | None = None,
|
||||
device: Optional["torch.device"] = None,
|
||||
seq_len: int | None = None,
|
||||
**rope_kwargs,
|
||||
) -> tuple["torch.Tensor", float]:
|
||||
"""
|
||||
Computes the inverse frequencies according to the original RoPE implementation
|
||||
Args:
|
||||
config ([`~transformers.PretrainedConfig`]):
|
||||
The model configuration.
|
||||
device (`torch.device`):
|
||||
The device to use for initialization of the inverse frequencies.
|
||||
seq_len (`int`, *optional*):
|
||||
The current sequence length. Unused for this type of RoPE.
|
||||
rope_kwargs (`Dict`, *optional*):
|
||||
BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
|
||||
Returns:
|
||||
Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
|
||||
post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
|
||||
"""
|
||||
if config is not None and len(rope_kwargs) > 0:
|
||||
raise ValueError(
|
||||
"Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in "
|
||||
f"`_compute_default_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}"
|
||||
)
|
||||
if len(rope_kwargs) > 0:
|
||||
base = rope_kwargs["base"]
|
||||
dim = rope_kwargs["dim"]
|
||||
elif config is not None:
|
||||
base = config.rope_theta
|
||||
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
|
||||
head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
|
||||
dim = int(head_dim * partial_rotary_factor)
|
||||
|
||||
attention_factor = 1.0 # Unused in this type of RoPE
|
||||
|
||||
# Compute the inverse frequencies
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
|
||||
return inv_freq, attention_factor
|
||||
|
||||
|
||||
def _compute_linear_scaling_rope_parameters(
|
||||
config: PretrainedConfig | None = None,
|
||||
device: Optional["torch.device"] = None,
|
||||
seq_len: int | None = None,
|
||||
**rope_kwargs,
|
||||
) -> tuple["torch.Tensor", float]:
|
||||
"""
|
||||
Computes the inverse frequencies with linear scaling. Credits to the Reddit user /u/kaiokendev
|
||||
Args:
|
||||
config ([`~transformers.PretrainedConfig`]):
|
||||
The model configuration.
|
||||
device (`torch.device`):
|
||||
The device to use for initialization of the inverse frequencies.
|
||||
seq_len (`int`, *optional*):
|
||||
The current sequence length. Unused for this type of RoPE.
|
||||
rope_kwargs (`Dict`, *optional*):
|
||||
BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
|
||||
Returns:
|
||||
Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
|
||||
post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
|
||||
"""
|
||||
if config is not None and len(rope_kwargs) > 0:
|
||||
raise ValueError(
|
||||
"Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in "
|
||||
f"`_compute_linear_scaling_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}"
|
||||
)
|
||||
if len(rope_kwargs) > 0:
|
||||
factor = rope_kwargs["factor"]
|
||||
elif config is not None:
|
||||
factor = config.rope_scaling["factor"]
|
||||
|
||||
# Gets the default RoPE parameters
|
||||
inv_freq, attention_factor = _compute_default_rope_parameters(config, device, seq_len, **rope_kwargs)
|
||||
|
||||
# Then applies linear scaling to the frequencies.
|
||||
# NOTE: originally, scaling was applied to the position_ids. However, we get `embs = inv_freq @ position_ids`, so
|
||||
# applying scaling to the inverse frequencies is equivalent.
|
||||
inv_freq /= factor
|
||||
return inv_freq, attention_factor
|
||||
|
||||
|
||||
def _compute_dynamic_ntk_parameters(
|
||||
config: PretrainedConfig | None = None,
|
||||
device: Optional["torch.device"] = None,
|
||||
seq_len: int | None = None,
|
||||
**rope_kwargs,
|
||||
) -> tuple["torch.Tensor", float]:
|
||||
"""
|
||||
Computes the inverse frequencies with NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla
|
||||
Args:
|
||||
config ([`~transformers.PretrainedConfig`]):
|
||||
The model configuration.
|
||||
device (`torch.device`):
|
||||
The device to use for initialization of the inverse frequencies.
|
||||
seq_len (`int`, *optional*):
|
||||
The current sequence length, used to update the dynamic RoPE at inference time.
|
||||
rope_kwargs (`Dict`, *optional*):
|
||||
BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
|
||||
Returns:
|
||||
Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
|
||||
post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
|
||||
"""
|
||||
# TODO (joao): use the new `original_max_position_embeddings` from rope_scaling
|
||||
if config is not None and len(rope_kwargs) > 0:
|
||||
raise ValueError(
|
||||
"Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in "
|
||||
f"`_compute_dynamic_ntk_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}"
|
||||
)
|
||||
if len(rope_kwargs) > 0:
|
||||
base = rope_kwargs["base"]
|
||||
dim = rope_kwargs["dim"]
|
||||
max_position_embeddings = rope_kwargs["max_position_embeddings"]
|
||||
factor = rope_kwargs["factor"]
|
||||
elif config is not None:
|
||||
base = config.rope_theta
|
||||
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
|
||||
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
||||
dim = int(head_dim * partial_rotary_factor)
|
||||
max_position_embeddings = config.max_position_embeddings
|
||||
factor = config.rope_scaling["factor"]
|
||||
|
||||
attention_factor = 1.0 # Unused in this type of RoPE
|
||||
|
||||
# seq_len: default to max_position_embeddings, e.g. at init time
|
||||
seq_len = seq_len if seq_len is not None and seq_len > max_position_embeddings else max_position_embeddings
|
||||
|
||||
# Compute the inverse frequencies
|
||||
base = base * ((factor * seq_len / max_position_embeddings) - (factor - 1)) ** (dim / (dim - 2))
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
|
||||
return inv_freq, attention_factor
|
||||
|
||||
|
||||
def _compute_yarn_parameters(
|
||||
config: PretrainedConfig,
|
||||
device: "torch.device",
|
||||
seq_len: int | None = None,
|
||||
**rope_kwargs,
|
||||
) -> tuple["torch.Tensor", float]:
|
||||
"""
|
||||
Computes the inverse frequencies with NTK scaling. Please refer to the
|
||||
[original paper](https://huggingface.co/papers/2309.00071)
|
||||
Args:
|
||||
config ([`~transformers.PretrainedConfig`]):
|
||||
The model configuration.
|
||||
device (`torch.device`):
|
||||
The device to use for initialization of the inverse frequencies.
|
||||
seq_len (`int`, *optional*):
|
||||
The current sequence length. Unused for this type of RoPE.
|
||||
rope_kwargs (`Dict`, *optional*):
|
||||
BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
|
||||
Returns:
|
||||
Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
|
||||
post-processing scaling factor applied to the computed cos/sin.
|
||||
"""
|
||||
# No need to keep BC with yarn, unreleased when this new pattern was created.
|
||||
if len(rope_kwargs) > 0:
|
||||
raise ValueError(
|
||||
f"Unexpected arguments: `**rope_kwargs` should be unset in `_compute_yarn_parameters`, got {rope_kwargs}"
|
||||
)
|
||||
|
||||
base = config.rope_theta
|
||||
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
|
||||
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
||||
dim = int(head_dim * partial_rotary_factor)
|
||||
factor = config.rope_scaling["factor"]
|
||||
attention_factor = config.rope_scaling.get("attention_factor")
|
||||
mscale = config.rope_scaling.get("mscale")
|
||||
mscale_all_dim = config.rope_scaling.get("mscale_all_dim")
|
||||
|
||||
# NOTE: DeekSeek-V3 (and potentially other models) modify `max_position_embeddings` and have a
|
||||
# `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two
|
||||
# values to compute the default attention scaling factor, instead of using `factor`.
|
||||
if "original_max_position_embeddings" in config.rope_scaling:
|
||||
original_max_position_embeddings = config.rope_scaling["original_max_position_embeddings"]
|
||||
factor = config.max_position_embeddings / original_max_position_embeddings
|
||||
else:
|
||||
original_max_position_embeddings = config.max_position_embeddings
|
||||
|
||||
def get_mscale(scale, mscale=1):
|
||||
if scale <= 1:
|
||||
return 1.0
|
||||
return 0.1 * mscale * math.log(scale) + 1.0
|
||||
|
||||
# Sets the attention factor as suggested in the paper
|
||||
if attention_factor is None:
|
||||
if mscale and mscale_all_dim:
|
||||
attention_factor = float(get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim))
|
||||
else:
|
||||
attention_factor = get_mscale(factor)
|
||||
|
||||
# Optional config options
|
||||
# beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly)
|
||||
beta_fast = config.rope_scaling.get("beta_fast") or 32
|
||||
beta_slow = config.rope_scaling.get("beta_slow") or 1
|
||||
|
||||
# Compute the inverse frequencies
|
||||
def find_correction_dim(num_rotations, dim, base, max_position_embeddings):
|
||||
"""Inverse dimension formula to find the dimension based on the number of rotations"""
|
||||
return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
|
||||
|
||||
def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings):
|
||||
"""Find dimension range bounds based on rotations"""
|
||||
low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings))
|
||||
high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings))
|
||||
return max(low, 0), min(high, dim - 1)
|
||||
|
||||
def linear_ramp_factor(min, max, dim):
|
||||
if min == max:
|
||||
max += 0.001 # Prevent singularity
|
||||
|
||||
linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
|
||||
ramp_func = torch.clamp(linear_func, 0, 1)
|
||||
return ramp_func
|
||||
|
||||
# Note on variable naming: "interpolation" comes from the original technique, where we interpolate the position IDs
|
||||
# to expand the possible context length. In other words, interpolation = apply scaling factor.
|
||||
pos_freqs = base ** (torch.arange(0, dim, 2).to(device=device, dtype=torch.float) / dim)
|
||||
inv_freq_extrapolation = 1.0 / pos_freqs
|
||||
inv_freq_interpolation = 1.0 / (factor * pos_freqs)
|
||||
|
||||
low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_max_position_embeddings)
|
||||
|
||||
# Get n-dimensional rotational scaling corrected for extrapolation
|
||||
inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).to(device=device, dtype=torch.float)
|
||||
inv_freq = (
|
||||
inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)
|
||||
+ inv_freq_extrapolation * inv_freq_extrapolation_factor
|
||||
)
|
||||
return inv_freq, attention_factor
|
||||
|
||||
|
||||
def _compute_longrope_parameters(
|
||||
config: PretrainedConfig,
|
||||
device: "torch.device",
|
||||
seq_len: int | None = None,
|
||||
**rope_kwargs,
|
||||
) -> tuple["torch.Tensor", float]:
|
||||
"""
|
||||
Computes the inverse frequencies with LongRoPE scaling. Please refer to the
|
||||
[original implementation](https://github.com/microsoft/LongRoPE)
|
||||
Args:
|
||||
config ([`~transformers.PretrainedConfig`]):
|
||||
The model configuration.
|
||||
device (`torch.device`):
|
||||
The device to use for initialization of the inverse frequencies.
|
||||
seq_len (`int`, *optional*):
|
||||
The current sequence length.
|
||||
rope_kwargs (`Dict`, *optional*):
|
||||
BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
|
||||
Returns:
|
||||
Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
|
||||
post-processing scaling factor applied to the computed cos/sin.
|
||||
"""
|
||||
# TODO (joao): use the new `original_max_position_embeddings` from rope_scaling
|
||||
# No need to keep BC with longrope, unreleased when this new pattern was created.
|
||||
if len(rope_kwargs) > 0:
|
||||
raise ValueError(
|
||||
"Unexpected arguments: `**rope_kwargs` should be unset in `_compute_longrope_parameters`, got "
|
||||
f"{rope_kwargs}"
|
||||
)
|
||||
|
||||
base = config.rope_theta
|
||||
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
|
||||
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
||||
dim = int(head_dim * partial_rotary_factor)
|
||||
long_factor = config.rope_scaling["long_factor"]
|
||||
short_factor = config.rope_scaling["short_factor"]
|
||||
factor = config.rope_scaling.get("factor")
|
||||
attention_factor = config.rope_scaling.get("attention_factor")
|
||||
|
||||
# NOTE: Phi3 (and potentially other models) modify `max_position_embeddings` and have a
|
||||
# `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two
|
||||
# values to compute the default attention scaling factor, instead of using `factor`.
|
||||
if hasattr(config, "original_max_position_embeddings"):
|
||||
original_max_position_embeddings = config.original_max_position_embeddings
|
||||
factor = config.max_position_embeddings / config.original_max_position_embeddings
|
||||
else:
|
||||
original_max_position_embeddings = config.max_position_embeddings
|
||||
|
||||
# Sets the attention factor as suggested in the paper
|
||||
if attention_factor is None:
|
||||
if factor <= 1.0:
|
||||
attention_factor = 1.0
|
||||
else:
|
||||
attention_factor = math.sqrt(1 + math.log(factor) / math.log(original_max_position_embeddings))
|
||||
|
||||
# Compute the inverse frequencies -- scaled based on the target sequence length
|
||||
if seq_len and seq_len > original_max_position_embeddings:
|
||||
ext_factors = torch.tensor(long_factor, dtype=torch.float32, device=device)
|
||||
else:
|
||||
ext_factors = torch.tensor(short_factor, dtype=torch.float32, device=device)
|
||||
inv_freq_shape = torch.arange(0, dim, 2, dtype=torch.int64, device=device).float() / dim
|
||||
inv_freq = 1.0 / (ext_factors * base**inv_freq_shape)
|
||||
|
||||
return inv_freq, attention_factor
|
||||
|
||||
|
||||
def _compute_llama3_parameters(
|
||||
config: PretrainedConfig,
|
||||
device: "torch.device",
|
||||
seq_len: int | None = None,
|
||||
**rope_kwargs,
|
||||
) -> tuple["torch.Tensor", float]:
|
||||
"""
|
||||
Computes the inverse frequencies for llama 3.1.
|
||||
|
||||
Args:
|
||||
config ([`~transformers.PretrainedConfig`]):
|
||||
The model configuration.
|
||||
device (`torch.device`):
|
||||
The device to use for initialization of the inverse frequencies.
|
||||
seq_len (`int`, *optional*):
|
||||
The current sequence length. Unused for this type of RoPE.
|
||||
rope_kwargs (`Dict`, *optional*):
|
||||
BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
|
||||
Returns:
|
||||
Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
|
||||
post-processing scaling factor applied to the computed cos/sin.
|
||||
"""
|
||||
# Gets the default RoPE parameters
|
||||
inv_freq, attention_factor = _compute_default_rope_parameters(config, device, seq_len, **rope_kwargs)
|
||||
|
||||
factor = config.rope_scaling["factor"] # `8` in the original implementation
|
||||
low_freq_factor = config.rope_scaling["low_freq_factor"] # `1` in the original implementation
|
||||
high_freq_factor = config.rope_scaling["high_freq_factor"] # `4` in the original implementation
|
||||
old_context_len = config.rope_scaling["original_max_position_embeddings"] # `8192` in the original implementation
|
||||
|
||||
low_freq_wavelen = old_context_len / low_freq_factor
|
||||
high_freq_wavelen = old_context_len / high_freq_factor
|
||||
|
||||
wavelen = 2 * math.pi / inv_freq
|
||||
# wavelen < high_freq_wavelen: do nothing
|
||||
# wavelen > low_freq_wavelen: divide by factor
|
||||
inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq)
|
||||
# otherwise: interpolate between the two, using a smooth factor
|
||||
smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
|
||||
smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
|
||||
is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
|
||||
inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
|
||||
|
||||
return inv_freq_llama, attention_factor
|
||||
|
||||
|
||||
# This maps the "rope_type" string field in rope config to the corresponding function to compute the RoPE parameters
|
||||
# from the model config. You can append new {'rope_type': callable} pairs to this dictionary to enable custom RoPE
|
||||
# parameterizations, as long as the callable has the same signature.
|
||||
ROPE_INIT_FUNCTIONS = {
|
||||
"default": _compute_default_rope_parameters,
|
||||
"linear": _compute_linear_scaling_rope_parameters,
|
||||
"dynamic": _compute_dynamic_ntk_parameters,
|
||||
"yarn": _compute_yarn_parameters,
|
||||
"longrope": _compute_longrope_parameters,
|
||||
"llama3": _compute_llama3_parameters,
|
||||
}
|
||||
|
||||
|
||||
def _check_received_keys(
|
||||
rope_type: str,
|
||||
received_keys: set,
|
||||
required_keys: set,
|
||||
optional_keys: set | None = None,
|
||||
ignore_keys: set | None = None,
|
||||
):
|
||||
"""Compare the received keys in `config.rope_scaling` against the expected and optional keys"""
|
||||
# BC: "rope_type" was originally "type" -- let's check for "rope_type" when "type" is present
|
||||
if "type" in received_keys:
|
||||
received_keys -= {"type"}
|
||||
required_keys.add("rope_type")
|
||||
|
||||
# Some models need to store model-specific keys, and we don't want to throw warning at them
|
||||
if ignore_keys is not None:
|
||||
received_keys -= ignore_keys
|
||||
|
||||
missing_keys = required_keys - received_keys
|
||||
if missing_keys:
|
||||
raise KeyError(f"Missing required keys in `rope_scaling` for 'rope_type'='{rope_type}': {missing_keys}")
|
||||
|
||||
if optional_keys is not None:
|
||||
unused_keys = received_keys - required_keys - optional_keys
|
||||
else:
|
||||
unused_keys = received_keys - required_keys
|
||||
if unused_keys:
|
||||
logger.warning(f"Unrecognized keys in `rope_scaling` for 'rope_type'='{rope_type}': {unused_keys}")
|
||||
|
||||
|
||||
def _validate_default_rope_parameters(config: PretrainedConfig, ignore_keys: set | None = None):
|
||||
rope_scaling = config.rope_scaling
|
||||
rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type"
|
||||
required_keys = {"rope_type"}
|
||||
received_keys = set(rope_scaling.keys())
|
||||
_check_received_keys(rope_type, received_keys, required_keys, ignore_keys=ignore_keys)
|
||||
|
||||
|
||||
def _validate_linear_scaling_rope_parameters(config: PretrainedConfig, ignore_keys: set | None = None):
|
||||
rope_scaling = config.rope_scaling
|
||||
rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type"
|
||||
required_keys = {"rope_type", "factor"}
|
||||
received_keys = set(rope_scaling.keys())
|
||||
_check_received_keys(rope_type, received_keys, required_keys, ignore_keys=ignore_keys)
|
||||
|
||||
factor = rope_scaling["factor"]
|
||||
if factor is None or not isinstance(factor, float) or factor < 1.0:
|
||||
logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}")
|
||||
|
||||
|
||||
def _validate_dynamic_scaling_rope_parameters(config: PretrainedConfig, ignore_keys: set | None = None):
|
||||
rope_scaling = config.rope_scaling
|
||||
rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type"
|
||||
required_keys = {"rope_type", "factor"}
|
||||
# TODO (joao): update logic for the inclusion of `original_max_position_embeddings`
|
||||
optional_keys = {"original_max_position_embeddings"}
|
||||
received_keys = set(rope_scaling.keys())
|
||||
_check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys)
|
||||
|
||||
factor = rope_scaling["factor"]
|
||||
if factor is None or not isinstance(factor, float) or factor < 1.0:
|
||||
logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}")
|
||||
|
||||
|
||||
def _validate_yarn_parameters(config: PretrainedConfig, ignore_keys: set | None = None):
|
||||
rope_scaling = config.rope_scaling
|
||||
rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type"
|
||||
required_keys = {"rope_type", "factor"}
|
||||
optional_keys = {
|
||||
"attention_factor",
|
||||
"beta_fast",
|
||||
"beta_slow",
|
||||
"original_max_position_embeddings",
|
||||
"mscale",
|
||||
"mscale_all_dim",
|
||||
}
|
||||
received_keys = set(rope_scaling.keys())
|
||||
_check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys)
|
||||
|
||||
factor = rope_scaling["factor"]
|
||||
if factor is None or not isinstance(factor, float) or factor < 1.0:
|
||||
logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}")
|
||||
|
||||
attention_factor = rope_scaling.get("attention_factor")
|
||||
if attention_factor is not None and (not isinstance(attention_factor, float) or attention_factor < 0):
|
||||
logger.warning(
|
||||
f"`rope_scaling`'s attention_factor field must be a float greater than 0, got {attention_factor}"
|
||||
)
|
||||
beta_fast = rope_scaling.get("beta_fast")
|
||||
if beta_fast is not None and not isinstance(beta_fast, float):
|
||||
logger.warning(f"`rope_scaling`'s beta_fast field must be a float, got {beta_fast}")
|
||||
beta_slow = rope_scaling.get("beta_slow")
|
||||
if beta_slow is not None and not isinstance(beta_slow, float):
|
||||
logger.warning(f"`rope_scaling`'s beta_slow field must be a float, got {beta_slow}")
|
||||
|
||||
if (beta_fast or 32) < (beta_slow or 1):
|
||||
logger.warning(
|
||||
f"`rope_scaling`'s beta_fast field must be greater than beta_slow, got beta_fast={beta_fast} "
|
||||
f"(defaults to 32 if None) and beta_slow={beta_slow} (defaults to 1 if None)"
|
||||
)
|
||||
|
||||
|
||||
def _validate_longrope_parameters(config: PretrainedConfig, ignore_keys: set | None = None):
|
||||
rope_scaling = config.rope_scaling
|
||||
rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type"
|
||||
required_keys = {"rope_type", "short_factor", "long_factor"}
|
||||
# TODO (joao): update logic for the inclusion of `original_max_position_embeddings`
|
||||
optional_keys = {"attention_factor", "factor", "original_max_position_embeddings"}
|
||||
received_keys = set(rope_scaling.keys())
|
||||
_check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys)
|
||||
|
||||
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
|
||||
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
||||
dim = int(head_dim * partial_rotary_factor)
|
||||
|
||||
short_factor = rope_scaling.get("short_factor")
|
||||
if not isinstance(short_factor, list) and all(isinstance(x, (int, float)) for x in short_factor):
|
||||
logger.warning(f"`rope_scaling`'s short_factor field must be a list of numbers, got {short_factor}")
|
||||
if not len(short_factor) == dim // 2:
|
||||
logger.warning(f"`rope_scaling`'s short_factor field must have length {dim // 2}, got {len(short_factor)}")
|
||||
|
||||
long_factor = rope_scaling.get("long_factor")
|
||||
if not isinstance(long_factor, list) and all(isinstance(x, (int, float)) for x in long_factor):
|
||||
logger.warning(f"`rope_scaling`'s long_factor field must be a list of numbers, got {long_factor}")
|
||||
if not len(long_factor) == dim // 2:
|
||||
logger.warning(f"`rope_scaling`'s long_factor field must have length {dim // 2}, got {len(long_factor)}")
|
||||
|
||||
# Handle Phi3 divergence: prefer the use of `attention_factor` and/or `factor` over
|
||||
# `original_max_position_embeddings` to compute internal variables. The latter lives outside `rope_scaling` and is
|
||||
# unique to longrope (= undesirable)
|
||||
if hasattr(config, "original_max_position_embeddings"):
|
||||
logger.warning_once(
|
||||
"This model has set a `original_max_position_embeddings` field, to be used together with "
|
||||
"`max_position_embeddings` to determine a scaling factor. Please set the `factor` field of `rope_scaling`"
|
||||
"with this ratio instead -- we recommend the use of this field over `original_max_position_embeddings`, "
|
||||
"as it is compatible with most model architectures."
|
||||
)
|
||||
else:
|
||||
factor = rope_scaling.get("factor")
|
||||
if factor is None:
|
||||
logger.warning("Missing required keys in `rope_scaling`: 'factor'")
|
||||
elif not isinstance(factor, float) or factor < 1.0:
|
||||
logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}")
|
||||
|
||||
attention_factor = rope_scaling.get("attention_factor")
|
||||
if attention_factor is not None:
|
||||
if not isinstance(attention_factor, float) or attention_factor < 0.0:
|
||||
logger.warning(
|
||||
f"`rope_scaling`'s attention_factor field must be a float greater than 0, got {attention_factor}"
|
||||
)
|
||||
|
||||
|
||||
def _validate_llama3_parameters(config: PretrainedConfig, ignore_keys: set | None = None):
|
||||
rope_scaling = config.rope_scaling
|
||||
rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type"
|
||||
required_keys = {
|
||||
"rope_type",
|
||||
"factor",
|
||||
"original_max_position_embeddings",
|
||||
"low_freq_factor",
|
||||
"high_freq_factor",
|
||||
}
|
||||
received_keys = set(rope_scaling.keys())
|
||||
_check_received_keys(rope_type, received_keys, required_keys, ignore_keys=ignore_keys)
|
||||
|
||||
factor = rope_scaling["factor"]
|
||||
if factor is None or not isinstance(factor, float) or factor < 1.0:
|
||||
logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}")
|
||||
|
||||
low_freq_factor = rope_scaling["low_freq_factor"]
|
||||
high_freq_factor = rope_scaling["high_freq_factor"]
|
||||
if low_freq_factor is None or not isinstance(low_freq_factor, float):
|
||||
logger.warning(f"`rope_scaling`'s low_freq_factor field must be a float, got {low_freq_factor}")
|
||||
if high_freq_factor is None or not isinstance(high_freq_factor, float):
|
||||
logger.warning(f"`rope_scaling`'s high_freq_factor field must be a float, got {high_freq_factor}")
|
||||
if high_freq_factor <= low_freq_factor:
|
||||
logger.warning(
|
||||
"`rope_scaling`'s high_freq_factor field must be greater than low_freq_factor, got high_freq_factor="
|
||||
f"{high_freq_factor} and low_freq_factor={low_freq_factor}"
|
||||
)
|
||||
|
||||
original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
|
||||
if original_max_position_embeddings is None or not isinstance(original_max_position_embeddings, int):
|
||||
logger.warning(
|
||||
"`rope_scaling`'s original_max_position_embeddings field must be an integer, got "
|
||||
f"{original_max_position_embeddings}"
|
||||
)
|
||||
if original_max_position_embeddings >= config.max_position_embeddings:
|
||||
logger.warning(
|
||||
"`rope_scaling`'s original_max_position_embeddings field must be less than max_position_embeddings, got "
|
||||
f"{original_max_position_embeddings} and max_position_embeddings={config.max_position_embeddings}"
|
||||
)
|
||||
|
||||
|
||||
# Like `ROPE_INIT_FUNCTIONS`, this validation function mapping can be dynamically updated for custom RoPE types.
|
||||
ROPE_VALIDATION_FUNCTIONS = {
|
||||
"default": _validate_default_rope_parameters,
|
||||
"linear": _validate_linear_scaling_rope_parameters,
|
||||
"dynamic": _validate_dynamic_scaling_rope_parameters,
|
||||
"yarn": _validate_yarn_parameters,
|
||||
"longrope": _validate_longrope_parameters,
|
||||
"llama3": _validate_llama3_parameters,
|
||||
}
|
||||
|
||||
|
||||
def rope_config_validation(config: PretrainedConfig, ignore_keys: set | None = None):
|
||||
"""
|
||||
Validate the RoPE config arguments, given a `PretrainedConfig` object
|
||||
"""
|
||||
rope_scaling = getattr(config, "rope_scaling", None) # not a default parameter in `PretrainedConfig`
|
||||
if rope_scaling is None:
|
||||
return
|
||||
|
||||
# BC: "rope_type" was originally "type"
|
||||
rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", "default"))
|
||||
validation_fn = ROPE_VALIDATION_FUNCTIONS.get(rope_type)
|
||||
if validation_fn is not None:
|
||||
validation_fn(config, ignore_keys=ignore_keys)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Missing validation function mapping in `ROPE_VALIDATION_FUNCTIONS` for 'rope_type'='{rope_type}'"
|
||||
)
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
"""Rotates half the hidden dims of the input."""
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(x, cos, sin, position_ids=None, unsqueeze_dim=1):
|
||||
"""Applies Rotary Position Embedding to the query and key tensors.
|
||||
|
||||
Args:
|
||||
x (`torch.Tensor`): The input tensor.
|
||||
cos (`torch.Tensor`): The cosine part of the rotary embedding.
|
||||
sin (`torch.Tensor`): The sine part of the rotary embedding.
|
||||
position_ids (`torch.Tensor`, *optional*):
|
||||
Deprecated and unused.
|
||||
unsqueeze_dim (`int`, *optional*, defaults to 1):
|
||||
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
|
||||
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
|
||||
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
|
||||
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
|
||||
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
|
||||
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
|
||||
Returns:
|
||||
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
|
||||
"""
|
||||
cos = cos.unsqueeze(unsqueeze_dim)
|
||||
sin = sin.unsqueeze(unsqueeze_dim)
|
||||
x_embed = (x * cos) + (rotate_half(x) * sin)
|
||||
return x_embed
|
||||
@@ -0,0 +1,450 @@
|
||||
# Copyright 2025 Xiaomi Corporation.
|
||||
import typing as tp
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange, repeat
|
||||
from torch import nn
|
||||
|
||||
|
||||
def rank():
|
||||
if dist.is_initialized():
|
||||
return dist.get_rank()
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def world_size():
|
||||
if dist.is_initialized():
|
||||
return dist.get_world_size()
|
||||
else:
|
||||
return 1
|
||||
|
||||
|
||||
def default(val: tp.Any, d: tp.Any) -> tp.Any:
|
||||
return val if val is not None else d
|
||||
|
||||
|
||||
def ema_inplace(moving_avg, new, decay: float):
|
||||
if dist.is_initialized():
|
||||
dist.all_reduce(new, op=dist.ReduceOp.SUM)
|
||||
moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
|
||||
|
||||
|
||||
def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
|
||||
return (x + epsilon) / (x.sum() + n_categories * epsilon)
|
||||
|
||||
|
||||
def uniform_init(*shape: int):
|
||||
t = torch.empty(shape)
|
||||
nn.init.kaiming_uniform_(t)
|
||||
return t
|
||||
|
||||
|
||||
def sample_vectors(samples, num: int):
|
||||
num_samples, device = samples.shape[0], samples.device
|
||||
|
||||
if num_samples >= num:
|
||||
indices = torch.randperm(num_samples, device=device)[:num]
|
||||
else:
|
||||
indices = torch.randint(0, num_samples, (num,), device=device)
|
||||
|
||||
selected_samples = samples[indices]
|
||||
|
||||
if dist.is_initialized():
|
||||
dist.broadcast(selected_samples, src=0)
|
||||
|
||||
return selected_samples
|
||||
|
||||
|
||||
def kmeans(samples, num_clusters: int, num_iters: int = 10):
|
||||
dim, dtype = samples.shape[-1], samples.dtype
|
||||
|
||||
means = sample_vectors(samples, num_clusters)
|
||||
|
||||
for _ in range(num_iters):
|
||||
dists = -(samples.pow(2).sum(1, keepdim=True) - 2 * samples @ means.t() + means.t().pow(2).sum(0, keepdim=True))
|
||||
|
||||
buckets = dists.max(dim=-1).indices
|
||||
bins = torch.bincount(buckets, minlength=num_clusters)
|
||||
|
||||
new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
|
||||
new_means = new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
|
||||
|
||||
if dist.is_initialized():
|
||||
dist.all_reduce(bins, op=dist.ReduceOp.SUM)
|
||||
dist.all_reduce(new_means, op=dist.ReduceOp.SUM)
|
||||
|
||||
zero_mask = bins == 0
|
||||
bins_min_clamped = bins.masked_fill(zero_mask, 1)
|
||||
|
||||
new_means = new_means / bins_min_clamped[..., None]
|
||||
|
||||
means = torch.where(zero_mask[..., None], means, new_means)
|
||||
|
||||
return means, bins
|
||||
|
||||
|
||||
class EuclideanCodebook(nn.Module):
|
||||
"""Codebook with Euclidean distance.
|
||||
Args:
|
||||
dim (int): Dimension.
|
||||
codebook_size (int): Codebook size.
|
||||
kmeans_init (bool): Whether to use k-means to initialize the codebooks.
|
||||
If set to true, run the k-means algorithm on the first training batch and use
|
||||
the learned centroids as initialization.
|
||||
kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
|
||||
decay (float): Decay for exponential moving average over the codebooks.
|
||||
epsilon (float): Epsilon value for numerical stability.
|
||||
threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
|
||||
that have an exponential moving average cluster size less than the specified threshold with
|
||||
randomly selected vector from the current batch.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
codebook_size: int,
|
||||
kmeans_init: int = False,
|
||||
kmeans_iters: int = 10,
|
||||
decay: float = 0.99,
|
||||
epsilon: float = 1e-5,
|
||||
threshold_ema_dead_code: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
self.decay = decay
|
||||
init_fn: tp.Callable[..., torch.Tensor] | tp.Any = uniform_init if not kmeans_init else torch.zeros
|
||||
embed = init_fn(codebook_size, dim)
|
||||
|
||||
self.codebook_size = codebook_size
|
||||
|
||||
self.kmeans_iters = kmeans_iters
|
||||
self.epsilon = epsilon
|
||||
self.threshold_ema_dead_code = threshold_ema_dead_code
|
||||
|
||||
self.register_buffer("inited", torch.Tensor([not kmeans_init]))
|
||||
self.register_buffer("cluster_size", torch.zeros(codebook_size))
|
||||
self.register_buffer("embed", embed)
|
||||
self.register_buffer("embed_avg", embed.clone())
|
||||
|
||||
@torch.jit.ignore
|
||||
def init_embed_(self, data):
|
||||
if self.inited:
|
||||
return
|
||||
|
||||
embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
|
||||
self.embed.data.copy_(embed)
|
||||
self.embed_avg.data.copy_(embed.clone())
|
||||
self.cluster_size.data.copy_(cluster_size)
|
||||
self.inited.data.copy_(torch.Tensor([True]))
|
||||
|
||||
def replace_(self, samples, mask):
|
||||
# modified_codebook = torch.where(
|
||||
# mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
|
||||
# )
|
||||
replace_num = mask.sum()
|
||||
modified_codebook = self.embed.clone()
|
||||
modified_codebook[mask] = sample_vectors(samples, replace_num)
|
||||
self.embed.data.copy_(modified_codebook)
|
||||
|
||||
def expire_codes_(self, batch_samples):
|
||||
if self.threshold_ema_dead_code == 0:
|
||||
return
|
||||
|
||||
expired_codes = self.cluster_size < self.threshold_ema_dead_code
|
||||
if not torch.any(expired_codes):
|
||||
return
|
||||
|
||||
batch_samples = rearrange(batch_samples, "... d -> (...) d")
|
||||
self.replace_(batch_samples, mask=expired_codes)
|
||||
|
||||
def preprocess(self, x):
|
||||
x = rearrange(x, "... d -> (...) d")
|
||||
return x
|
||||
|
||||
def quantize(self, x):
|
||||
embed = self.embed.t()
|
||||
dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed + embed.pow(2).sum(0, keepdim=True))
|
||||
embed_ind = dist.max(dim=-1).indices
|
||||
return embed_ind
|
||||
|
||||
def postprocess_emb(self, embed_ind, shape):
|
||||
return embed_ind.view(*shape[:-1])
|
||||
|
||||
def dequantize(self, embed_ind):
|
||||
quantize = F.embedding(embed_ind, self.embed)
|
||||
return quantize
|
||||
|
||||
def encode(self, x):
|
||||
shape = x.shape
|
||||
# pre-process
|
||||
x = self.preprocess(x)
|
||||
# quantize
|
||||
embed_ind = self.quantize(x)
|
||||
# post-process
|
||||
embed_ind = self.postprocess_emb(embed_ind, shape)
|
||||
return embed_ind
|
||||
|
||||
def decode(self, embed_ind):
|
||||
quantize = self.dequantize(embed_ind)
|
||||
return quantize
|
||||
|
||||
def forward(self, x):
|
||||
shape, dtype = x.shape, x.dtype
|
||||
x = self.preprocess(x)
|
||||
|
||||
self.init_embed_(x)
|
||||
|
||||
embed_ind = self.quantize(x)
|
||||
embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
|
||||
embed_ind = self.postprocess_emb(embed_ind, shape)
|
||||
quantize = self.dequantize(embed_ind)
|
||||
|
||||
if self.training:
|
||||
# We do the expiry of code at that point as buffers are in sync
|
||||
# and all the workers will take the same decision.
|
||||
self.expire_codes_(x)
|
||||
ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
|
||||
embed_sum = x.t() @ embed_onehot
|
||||
ema_inplace(self.embed_avg, embed_sum.t().contiguous(), self.decay)
|
||||
cluster_size = (
|
||||
laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon) * self.cluster_size.sum()
|
||||
)
|
||||
embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
|
||||
self.embed.data.copy_(embed_normalized)
|
||||
|
||||
return quantize, embed_ind
|
||||
|
||||
|
||||
class VectorQuantization(nn.Module):
|
||||
"""Vector quantization implementation.
|
||||
Currently supports only euclidean distance.
|
||||
Args:
|
||||
dim (int): Dimension
|
||||
codebook_size (int): Codebook size
|
||||
codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
|
||||
decay (float): Decay for exponential moving average over the codebooks.
|
||||
epsilon (float): Epsilon value for numerical stability.
|
||||
kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
|
||||
kmeans_iters (int): Number of iterations used for kmeans initialization.
|
||||
threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
|
||||
that have an exponential moving average cluster size less than the specified threshold with
|
||||
randomly selected vector from the current batch.
|
||||
commitment_weight (float): Weight for commitment loss.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
codebook_size: int,
|
||||
codebook_dim: int | None = None,
|
||||
decay: float = 0.99,
|
||||
epsilon: float = 1e-5,
|
||||
kmeans_init: bool = True,
|
||||
kmeans_iters: int = 50,
|
||||
threshold_ema_dead_code: int = 2,
|
||||
commitment_weight: float = 1.0,
|
||||
):
|
||||
super().__init__()
|
||||
_codebook_dim: int = default(codebook_dim, dim)
|
||||
|
||||
requires_projection = _codebook_dim != dim
|
||||
self.project_in = nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity()
|
||||
self.project_out = nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity()
|
||||
|
||||
self.epsilon = epsilon
|
||||
self.commitment_weight = commitment_weight
|
||||
|
||||
self._codebook = EuclideanCodebook(
|
||||
dim=_codebook_dim,
|
||||
codebook_size=codebook_size,
|
||||
kmeans_init=kmeans_init,
|
||||
kmeans_iters=kmeans_iters,
|
||||
decay=decay,
|
||||
epsilon=epsilon,
|
||||
threshold_ema_dead_code=threshold_ema_dead_code,
|
||||
)
|
||||
self.codebook_size = codebook_size
|
||||
|
||||
@property
|
||||
def codebook(self):
|
||||
return self._codebook.embed
|
||||
|
||||
def encode(self, x):
|
||||
# x = rearrange(x, "b d n -> b n d")
|
||||
x = self.project_in(x)
|
||||
embed_in = self._codebook.encode(x)
|
||||
return embed_in
|
||||
|
||||
def decode(self, embed_ind):
|
||||
quantize = self._codebook.decode(embed_ind)
|
||||
quantize = self.project_out(quantize)
|
||||
# quantize = rearrange(quantize, "b n d -> b d n")
|
||||
return quantize
|
||||
|
||||
def forward(self, x):
|
||||
device = x.device
|
||||
x = self.project_in(x)
|
||||
|
||||
quantize, embed_ind = self._codebook(x)
|
||||
|
||||
if self.training:
|
||||
quantize = x + (quantize - x).detach()
|
||||
|
||||
loss = torch.tensor([0.0], device=device, requires_grad=self.training)
|
||||
|
||||
if self.training:
|
||||
if self.commitment_weight > 0:
|
||||
commit_loss = F.mse_loss(quantize.detach(), x)
|
||||
loss = loss + commit_loss * self.commitment_weight
|
||||
|
||||
quantize = self.project_out(quantize)
|
||||
# quantize = rearrange(quantize, "b n d -> b d n")
|
||||
return quantize, embed_ind, loss
|
||||
|
||||
|
||||
class ResidualVectorQuantization(nn.Module):
|
||||
"""Residual vector quantization implementation.
|
||||
Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
|
||||
"""
|
||||
|
||||
def __init__(self, *, num_quantizers, codebook_size, **kwargs):
|
||||
super().__init__()
|
||||
if isinstance(codebook_size, int):
|
||||
codebook_size = [codebook_size] * num_quantizers
|
||||
elif len(codebook_size) < num_quantizers:
|
||||
codebook_size += [codebook_size[-1]] * (num_quantizers - len(codebook_size))
|
||||
self.layers = nn.ModuleList(
|
||||
[VectorQuantization(codebook_size=codebook_size[i], **kwargs) for i in range(num_quantizers)]
|
||||
)
|
||||
|
||||
def forward(self, x, n_q: int | None = None, layers: list | None = None):
|
||||
quantized_out = 0.0
|
||||
residual = x
|
||||
|
||||
all_losses = []
|
||||
all_indices = []
|
||||
out_quantized = []
|
||||
|
||||
n_q = n_q or len(self.layers)
|
||||
|
||||
for i, layer in enumerate(self.layers[:n_q]):
|
||||
quantized, indices, loss = layer(residual)
|
||||
residual = residual - quantized
|
||||
quantized_out = quantized_out + quantized
|
||||
|
||||
all_indices.append(indices)
|
||||
all_losses.append(loss)
|
||||
if layers and i in layers:
|
||||
out_quantized.append(quantized_out)
|
||||
|
||||
out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
|
||||
return quantized_out, out_indices, out_losses, out_quantized
|
||||
|
||||
def encode(self, x: torch.Tensor, n_q: int | None = None, st: int | None = None) -> torch.Tensor:
|
||||
residual = x
|
||||
all_indices = []
|
||||
n_q = len(self.layers) if n_q is None else n_q
|
||||
st = 0 if st is None else st
|
||||
for layer in self.layers[st:n_q]:
|
||||
indices = layer.encode(residual)
|
||||
quantized = layer.decode(indices)
|
||||
residual = residual - quantized
|
||||
all_indices.append(indices)
|
||||
out_indices = torch.stack(all_indices)
|
||||
return out_indices
|
||||
|
||||
def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor:
|
||||
quantized_out = torch.tensor(0.0, device=q_indices.device)
|
||||
for i, indices in enumerate(q_indices):
|
||||
layer = self.layers[st + i]
|
||||
quantized = layer.decode(indices)
|
||||
quantized_out = quantized_out + quantized
|
||||
return quantized_out
|
||||
|
||||
|
||||
class ResidualVectorQuantizer(nn.Module):
|
||||
"""Residual Vector Quantizer.
|
||||
Args:
|
||||
dimension (int): Dimension of the codebooks.
|
||||
n_q (int): Number of residual vector quantizers used.
|
||||
bins (int): Codebook size.
|
||||
decay (float): Decay for exponential moving average over the codebooks.
|
||||
kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
|
||||
kmeans_iters (int): Number of iterations used for kmeans initialization.
|
||||
threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
|
||||
that have an exponential moving average cluster size less than the specified threshold with
|
||||
randomly selected vector from the current batch.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dimension: int = 256,
|
||||
n_q: int = 8,
|
||||
bins: int | list = 1024,
|
||||
decay: float = 0.99,
|
||||
kmeans_init: bool = True,
|
||||
kmeans_iters: int = 50,
|
||||
threshold_ema_dead_code: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
self.n_q = n_q
|
||||
self.dimension = dimension
|
||||
self.bins = bins
|
||||
self.decay = decay
|
||||
self.kmeans_init = kmeans_init
|
||||
self.kmeans_iters = kmeans_iters
|
||||
self.threshold_ema_dead_code = threshold_ema_dead_code
|
||||
self.vq = ResidualVectorQuantization(
|
||||
dim=self.dimension,
|
||||
codebook_size=self.bins,
|
||||
num_quantizers=self.n_q,
|
||||
decay=self.decay,
|
||||
kmeans_init=self.kmeans_init,
|
||||
kmeans_iters=self.kmeans_iters,
|
||||
threshold_ema_dead_code=self.threshold_ema_dead_code,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
n_q: int | None = None,
|
||||
layers: list | None = None,
|
||||
):
|
||||
"""Residual vector quantization on the given input tensor.
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor.
|
||||
n_q (int): Number of quantizer used to quantize. Default: All quantizers.
|
||||
layers (list): Layer that need to return quantized. Default: None.
|
||||
Returns:
|
||||
QuantizedResult:
|
||||
The quantized (or approximately quantized) representation with
|
||||
the associated number quantizers and layer quantized required to return.
|
||||
"""
|
||||
n_q = n_q if n_q else self.n_q
|
||||
quantized, codes, commit_loss, quantized_list = self.vq(x, n_q=n_q, layers=layers)
|
||||
return quantized, codes, torch.mean(commit_loss), quantized_list
|
||||
|
||||
def encode(self, x: torch.Tensor, n_q: int | None = None, st: int | None = None) -> torch.Tensor:
|
||||
"""Encode a given input tensor with the specified sample rate at the given bandwidth.
|
||||
The RVQ encode method sets the appropriate number of quantizer to use
|
||||
and returns indices for each quantizer.
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor.
|
||||
n_q (int): Number of quantizer used to quantize. Default: All quantizers.
|
||||
st (int): Start to encode input from which layers. Default: 0.
|
||||
"""
|
||||
n_q = n_q if n_q else self.n_q
|
||||
st = st or 0
|
||||
codes = self.vq.encode(x, n_q=n_q, st=st)
|
||||
return codes
|
||||
|
||||
def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor:
|
||||
"""Decode the given codes to the quantized representation.
|
||||
Args:
|
||||
codes (torch.Tensor): Input indices for each quantizer.
|
||||
st (int): Start to decode input codes from which layers. Default: 0.
|
||||
"""
|
||||
quantized = self.vq.decode(codes, st=st)
|
||||
return quantized
|
||||
@@ -63,6 +63,23 @@ _OMNI_MODELS = {
|
||||
"qwen3_tts_code2wav",
|
||||
"Qwen3TTSCode2Wav",
|
||||
),
|
||||
## mimo_audio
|
||||
"MiMoAudioModel": (
|
||||
"mimo_audio",
|
||||
"mimo_audio",
|
||||
"MiMoAudioForConditionalGeneration",
|
||||
),
|
||||
"MiMoAudioLLMModel": (
|
||||
"mimo_audio",
|
||||
"mimo_audio_llm",
|
||||
"MiMoAudioLLMForConditionalGeneration",
|
||||
),
|
||||
"MiMoAudioToken2WavModel": (
|
||||
"mimo_audio",
|
||||
"mimo_audio_code2wav",
|
||||
"MiMoAudioToken2WavForConditionalGenerationVLLM",
|
||||
),
|
||||
## glm_image
|
||||
"GlmImageForConditionalGeneration": (
|
||||
"glm_image",
|
||||
"glm_image_ar",
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
# stage config for running mimo-audio with architecture of OmniLLM.
|
||||
|
||||
# The following config has been verified on 1x H20-96G GPU.
|
||||
async_chunk: false
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
dtype: bfloat16
|
||||
model_stage: fused_thinker_talker
|
||||
model_arch: MiMoAudioForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: true # need to discuss
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent # change the param name,such as pooling_output
|
||||
max_model_len: 8192
|
||||
max_num_batched_tokens: 8192
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.6
|
||||
top_p: 0.95
|
||||
top_k: 50
|
||||
max_tokens: 18192
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: MiMoAudioForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
gpu_memory_utilization: 0.2
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
max_model_len: 18192
|
||||
max_num_batched_tokens: 18192
|
||||
async_scheduling: false
|
||||
engine_input_source: [ 0 ]
|
||||
is_comprehension: false
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 18192
|
||||
seed: 42
|
||||
detokenize: false
|
||||
@@ -0,0 +1,67 @@
|
||||
# stage config for running mimo-audio with architecture of OmniLLM.
|
||||
|
||||
# The following config has been verified on 1x H20-96G GPU.
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
dtype: bfloat16
|
||||
model_stage: fused_thinker_talker
|
||||
model_arch: MiMoAudioForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: true # need to discuss
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent # change the param name,such as pooling_output
|
||||
max_model_len: 8192
|
||||
max_num_batched_tokens: 8192
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav_async_chunk
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.6
|
||||
top_p: 0.95
|
||||
top_k: 50
|
||||
max_tokens: 18192
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.1
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: MiMoAudioForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
gpu_memory_utilization: 0.2
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
max_model_len: 8192
|
||||
max_num_batched_tokens: 8192
|
||||
engine_input_source: [ 0 ]
|
||||
is_comprehension: false
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 18192
|
||||
seed: 42
|
||||
detokenize: false
|
||||
@@ -0,0 +1,206 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.inputs import TextPrompt
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def prepend_and_flatten_colmajor(x: torch.Tensor, pad_vec: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Prepend a padding vector to the input tensor and flatten in column-major order.
|
||||
|
||||
This function expands the padding vector to match the batch dimensions of the input
|
||||
tensor, prepends it to the row dimension, and then flattens the result in column-major
|
||||
order (transposing before flattening).
|
||||
|
||||
Args:
|
||||
x: Input tensor with shape (..., R, C) where R is the row dimension and C is
|
||||
the column dimension. Example: (B, 1, 8, 4) where B is batch size.
|
||||
pad_vec: Padding vector with shape (C,) to be prepended to x. The vector will
|
||||
be expanded to match the batch dimensions of x.
|
||||
|
||||
Returns:
|
||||
A flattened 1D tensor in column-major order with shape (-1,). The result
|
||||
contains the padded row followed by all rows of x, flattened column by column.
|
||||
"""
|
||||
pad_row = pad_vec.view(1, -1)
|
||||
|
||||
# Expand pad_row to the front of x, keeping other batch dimensions consistent
|
||||
# Example: x shape = (B,1,R,C) → pad shape = (B,1,1,C)
|
||||
pad_expand = pad_row.view(*([1] * (x.dim() - 2)), 1, x.size(-1)).expand(*x.shape[:-2], 1, x.size(-1))
|
||||
|
||||
# Prepend to the row dimension
|
||||
y = torch.cat([pad_expand, x], dim=-2) # (..., R+1, C)
|
||||
|
||||
# Flatten in column-major order:
|
||||
# First transpose (..., R+1, C) → (..., C, R+1)
|
||||
# Then flatten
|
||||
y_col_major = y.permute(*range(y.dim() - 2), -1, -2).reshape(-1)
|
||||
|
||||
return y_col_major
|
||||
|
||||
|
||||
def _make_finished_sentinel() -> dict[str, Any]:
|
||||
"""Return a minimal payload with finished=True so Stage-1 can end the request."""
|
||||
return {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
|
||||
|
||||
|
||||
def llm2code2wav_async_chunk(
|
||||
transfer_manager: Any, pooling_output: dict[str, Any], request: Any
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Async chunk version: convert stage-0 pooling_output to code2wav payload (pooling / connector accumulation).
|
||||
|
||||
Accumulates codes in connector per request_id,
|
||||
returns payload only when chunk_size is full or request is finished; returns None when waiting.
|
||||
"""
|
||||
if "code_predictor_codes" not in pooling_output:
|
||||
if request.is_finished():
|
||||
return _make_finished_sentinel()
|
||||
return None
|
||||
|
||||
code_predictor_codes = pooling_output["code_predictor_codes"]
|
||||
|
||||
if code_predictor_codes is None:
|
||||
if request.is_finished():
|
||||
return _make_finished_sentinel()
|
||||
return None
|
||||
if isinstance(code_predictor_codes, torch.Tensor):
|
||||
if code_predictor_codes.numel() == 0:
|
||||
if request.is_finished():
|
||||
return _make_finished_sentinel()
|
||||
return None
|
||||
elif hasattr(code_predictor_codes, "__len__"):
|
||||
if len(code_predictor_codes) == 0:
|
||||
return None
|
||||
|
||||
if isinstance(code_predictor_codes, torch.Tensor):
|
||||
if not code_predictor_codes.any():
|
||||
if request.is_finished():
|
||||
return _make_finished_sentinel()
|
||||
return None
|
||||
code_tensor = code_predictor_codes.to(torch.long)
|
||||
else:
|
||||
code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long)
|
||||
if not code_tensor.any():
|
||||
if request.is_finished():
|
||||
return _make_finished_sentinel()
|
||||
return None
|
||||
|
||||
if code_tensor.ndim == 3:
|
||||
code_tensor = code_tensor.unsqueeze(0)
|
||||
if code_tensor.ndim != 4 or code_tensor.shape[-2:] != (8, 4):
|
||||
return None
|
||||
|
||||
pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4, device=code_tensor.device, dtype=code_tensor.dtype)
|
||||
code_final = prepend_and_flatten_colmajor(code_tensor, pad_vec)
|
||||
code_list = code_final.tolist()
|
||||
if sum(code_list) == 0:
|
||||
if request.is_finished():
|
||||
return _make_finished_sentinel()
|
||||
return None
|
||||
|
||||
request_id = getattr(request, "external_req_id", None)
|
||||
if request_id is None:
|
||||
return None
|
||||
|
||||
chunk_size = left_context_size = 10
|
||||
transfer_manager.code_prompt_token_ids[request_id].append(code_list)
|
||||
length = len(transfer_manager.code_prompt_token_ids[request_id])
|
||||
chunk_length = length % chunk_size
|
||||
if chunk_length != 0 and not request.is_finished():
|
||||
return None
|
||||
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size
|
||||
end_index = min(length, left_context_size + context_length)
|
||||
|
||||
info = {
|
||||
"code_predictor_codes": (
|
||||
torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]).reshape(-1).tolist()
|
||||
),
|
||||
"finished": torch.tensor(request.is_finished(), dtype=torch.bool),
|
||||
}
|
||||
return info
|
||||
|
||||
|
||||
def llm2code2wav(
|
||||
stage_list: list[Any],
|
||||
engine_input_source: list[int],
|
||||
prompt: OmniTokensPrompt | TextPrompt | None = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[OmniTokensPrompt]:
|
||||
"""
|
||||
Process talker outputs to create code2wav inputs.
|
||||
|
||||
Workflow:
|
||||
1. Extract talker's codec code outputs (8-layer RVQ codes)
|
||||
2. Flatten codes for code2wav input
|
||||
3. Package for code2wav stage
|
||||
|
||||
Args:
|
||||
stage_list: List of stage objects
|
||||
engine_input_source: Source stage IDs (typically [1] for talker)
|
||||
prompt: Original prompt data
|
||||
requires_multimodal_data: Whether multimodal data is required
|
||||
|
||||
Returns:
|
||||
List of OmniTokensPrompt for code2wav stage
|
||||
"""
|
||||
if not engine_input_source:
|
||||
raise ValueError("engine_input_source cannot be empty")
|
||||
|
||||
source_stage_id = engine_input_source[0]
|
||||
if source_stage_id >= len(stage_list):
|
||||
raise IndexError(f"Invalid stage_id: {source_stage_id}")
|
||||
|
||||
if stage_list[source_stage_id].engine_outputs is None:
|
||||
raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
|
||||
|
||||
talker_outputs = stage_list[source_stage_id].engine_outputs
|
||||
code2wav_inputs = []
|
||||
|
||||
# Process each talker output
|
||||
for i, talker_output in enumerate(talker_outputs):
|
||||
output = talker_output.outputs[0]
|
||||
|
||||
# Extract codec codes from talker output
|
||||
# Expected shape: [8, seq_len] (8-layer RVQ codes)
|
||||
if "code_predictor_codes" in output.multimodal_output:
|
||||
codec_codes = output.multimodal_output["code_predictor_codes"].to(torch.long) # [seq_batch_size, 1, 8, 4]
|
||||
is_all_zero = (codec_codes == 0).all(dim=(1, 2, 3))
|
||||
non_zero_indices = (~is_all_zero).nonzero(as_tuple=True)[0]
|
||||
if len(non_zero_indices) == 0:
|
||||
# All codec codes are zero - skip this request with a warning
|
||||
request_id = getattr(talker_output, "request_id", f"unknown_{i}")
|
||||
logger.warning(
|
||||
"Skipping request %s: all codec codes are zero (empty output from Stage-0). "
|
||||
"This may indicate the model failed to generate valid audio codes.",
|
||||
request_id,
|
||||
)
|
||||
else:
|
||||
if len(non_zero_indices) < codec_codes.shape[0]:
|
||||
codec_codes = codec_codes[non_zero_indices]
|
||||
elif "latent" in output.multimodal_output and "code_predictor_codes" not in output.multimodal_output:
|
||||
codec_codes = torch.zeros(1, 1, 8, 4, dtype=torch.long)
|
||||
else:
|
||||
raise ValueError(f"Invalid multimodal_output: {output.multimodal_output}")
|
||||
|
||||
pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4)
|
||||
|
||||
code_final = prepend_and_flatten_colmajor(codec_codes, pad_vec)
|
||||
code_final = code_final.tolist()
|
||||
|
||||
code2wav_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=code_final,
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
)
|
||||
)
|
||||
|
||||
return code2wav_inputs
|
||||
@@ -1013,6 +1013,29 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
for req_id, req_infos in cached_infos.items():
|
||||
self._update_request_information(req_id, req_infos)
|
||||
|
||||
def _maybe_attach_mimo_audio_req_infos(
|
||||
self,
|
||||
req_state: CachedRequestState | None,
|
||||
req_infos: dict | None,
|
||||
req_id: str,
|
||||
) -> dict | None:
|
||||
"""Attach MiMoAudio-specific fields into req_infos if applicable.
|
||||
|
||||
This helper is intentionally small and self-contained so that it can be
|
||||
unit-tested to prevent regressions when updating MiMoAudio handling.
|
||||
"""
|
||||
if req_state is None or self.model.__class__.__name__ != "MiMoAudioForConditionalGeneration":
|
||||
return req_infos
|
||||
|
||||
# Always operate on a dict copy to avoid mutating shared instances.
|
||||
req_infos = dict(req_infos) if isinstance(req_infos, dict) else {}
|
||||
mm_features = getattr(req_state, "mm_features", None)
|
||||
if mm_features and (not req_infos.get("mm_features")):
|
||||
req_infos["mm_features"] = mm_features
|
||||
req_infos["req_id"] = req_id
|
||||
|
||||
return req_infos
|
||||
|
||||
def _preprocess(
|
||||
self,
|
||||
scheduler_output: "SchedulerOutput",
|
||||
@@ -1137,6 +1160,10 @@ class OmniGPUModelRunner(GPUModelRunner):
|
||||
for req_index, req_id in enumerate(self.input_batch.req_ids):
|
||||
req_state = self.requests.get(req_id)
|
||||
req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
|
||||
|
||||
# mimo-audio check
|
||||
req_infos = self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id)
|
||||
|
||||
start_offset = int(self.query_start_loc.cpu[req_index])
|
||||
sched_tokens = int(num_scheduled_tokens_np[req_index])
|
||||
s, e = start_offset, start_offset + sched_tokens
|
||||
|
||||
Reference in new issue
Block a user