104 lines
4.8 KiB
Python
104 lines
4.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""MiniMax Music 3 pipeline topology: AR talker -> acoustic decoder.
|
|
|
|
Stage 0 ``minimax_music3_ar``: a Qwen3 backbone (the repo's ``language_model/``
|
|
subfolder is a plain ``Qwen3ForCausalLM``) plus a depth decoder, wrapped by a
|
|
class that owns its sampler. It emits 200-frame windows of frame hidden states
|
|
and the matching 8-codebook RVQ codes.
|
|
|
|
Stage 1 ``minimax_music3_acoustic``: flow-matching DiT plus DAV vocoder. Pure
|
|
compute, no KV cache, weights loaded from the repo's ``condition_encoder/``,
|
|
``transformer/`` and ``vocoder/`` subfolders. Output is 32 kHz stereo.
|
|
|
|
Classifier-free guidance is mandatory for this checkpoint: every request
|
|
decodes as a cond/uncond pair, and the uncond row is produced by
|
|
``prompt_expand_func``. The deploy YAMLs therefore carry
|
|
``default_sampling_params.extra_args`` with ``cfg_role`` and ``cfg_scale`` so
|
|
that ``has_sampling_extra_args`` is on for the stage and per-row extra args
|
|
reach the model's forward.
|
|
"""
|
|
|
|
from vllm_omni.config.stage_config import (
|
|
PipelineConfig,
|
|
StageExecutionType,
|
|
StagePipelineConfig,
|
|
)
|
|
|
|
_PROC = "vllm_omni.model_executor.stage_input_processors.minimax_music3"
|
|
|
|
# ``<|audio_end|>`` in the checkpoint tokenizer (pinned by
|
|
# ``models/minimax_music3/prompt.SPECIAL_TOKEN_IDS``). The talker's own sampler
|
|
# forces this at the frame budget; the stop id is the engine-side safety net.
|
|
MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID = 151670
|
|
|
|
# MiniMax captions carry the full style description and run to roughly 5000
|
|
# characters, an order of magnitude past the 500-character serving default.
|
|
# ``tts_args`` is a pipeline-level stage extra, not a deploy-YAML field: it
|
|
# reaches the serving layer through ``StageConfig.yaml_extras``, which is where
|
|
# ``ServingSpeech._compute_max_instructions_length`` reads it from.
|
|
_MAX_INSTRUCTIONS_LENGTH = 20_000
|
|
|
|
MINIMAX_MUSIC3_PIPELINE = PipelineConfig(
|
|
model_type="minimax_music3",
|
|
model_arch="MiniMaxMusic3TalkerForConditionalGeneration",
|
|
hf_architectures=("MiniMaxMusic3ForConditionalGeneration",),
|
|
default_deploy_config_name="minimax_music3.yaml",
|
|
stages=(
|
|
StagePipelineConfig(
|
|
stage_id=0,
|
|
model_stage="minimax_music3_ar",
|
|
execution_type=StageExecutionType.LLM_AR,
|
|
input_sources=(),
|
|
owns_tokenizer=True,
|
|
engine_output_type="latent",
|
|
# The repo root config.json only names the composite architecture;
|
|
# the AR weights and the tokenizer live in their own subfolders.
|
|
model_subdir="language_model",
|
|
tokenizer_subdir="tokenizer",
|
|
# detokenize: the sampled stream is audio codes, never text.
|
|
# stop_token_ids: <|audio_end|> closes the song.
|
|
sampling_constraints={
|
|
"detokenize": False,
|
|
"stop_token_ids": [MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID],
|
|
},
|
|
# CFG is mandatory here, so this expands on essentially every
|
|
# request rather than only on an opt-in flag.
|
|
prompt_expand_func=f"{_PROC}.expand_cfg_prompts",
|
|
async_chunk_process_next_stage_input_func=f"{_PROC}.ar2acoustic_async_chunk",
|
|
# retains_state_across_chunks stays False. It is read only by
|
|
# OmniGenerationScheduler, so it is a no-op on an LLM_AR stage;
|
|
# more to the point, stage 0 produces chunks rather than parking on
|
|
# them, so it never holds a runner slot while waiting for one.
|
|
extras={"tts_args": {"max_instructions_length": _MAX_INSTRUCTIONS_LENGTH}},
|
|
),
|
|
StagePipelineConfig(
|
|
stage_id=1,
|
|
model_stage="minimax_music3_acoustic",
|
|
execution_type=StageExecutionType.LLM_GENERATION,
|
|
input_sources=(0,),
|
|
final_output=True,
|
|
final_output_type="audio",
|
|
# "audio" (not "latent"): the accumulator concatenates audio over
|
|
# time, whereas the latent path would stack the two stereo channels.
|
|
engine_output_type="audio",
|
|
model_arch="MiniMaxMusic3AcousticForConditionalGeneration",
|
|
# This stage reads the repo ROOT so it can reach condition_encoder/,
|
|
# transformer/ and vocoder/, and the root ships no tokenizer files.
|
|
tokenizer_subdir="tokenizer",
|
|
# The acoustic stage carries the previous window's latent forward as
|
|
# the next window's flow-matching prompt, so a request parked between
|
|
# chunks still owns its model state and must keep counting against
|
|
# max_num_seqs.
|
|
retains_state_across_chunks=True,
|
|
sync_process_input_func=f"{_PROC}.ar2acoustic",
|
|
sampling_constraints={"detokenize": True},
|
|
),
|
|
),
|
|
)
|
|
|
|
__all__ = [
|
|
"MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID",
|
|
"MINIMAX_MUSIC3_PIPELINE",
|
|
]
|