Yuekai ZhangandClaude Fable 5
5be7541d01
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
86 lines
3.9 KiB
Python
86 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""NemotronVoiceChat pipeline: thinker (speech -> frame-locked text) -> talker
|
|
(text timeline -> 31-quantizer RVQ code stacks) -> code2wav (codes -> 22.05 kHz PCM).
|
|
|
|
``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` is an 11B full-duplex speech-to-speech
|
|
model. Offline it is served as a qwen3-omni-style 3-stage pipeline:
|
|
|
|
* Stage 0 (``thinker``, LLM_AR): vendored FastConformer perception + additive
|
|
channel fusion feeding the NemotronH backbone on vLLM; emits the frame-locked
|
|
agent text channel (also the final text output) and carries the full frame
|
|
timeline + metadata to the talker as a latent payload.
|
|
* Stage 1 (``talker``, LLM_AR): the vendored Gemma3-based EAR-TTS model stepped
|
|
per frame (classifier-free guidance + MoG sampling run inside the vendored
|
|
forward); emits one 31-code stack per frame as latents.
|
|
* Stage 2 (``code2wav``, LLM_GENERATION): the vendored RVQ-VAE codec decoder
|
|
(fp32) producing the final 22.05 kHz waveform.
|
|
|
|
The STT->TTS coupling is one-directional and token-ids-only (verified against
|
|
the NeMo reference), which is what makes the offline cascade exact.
|
|
"""
|
|
|
|
from vllm_omni.config.stage_config import (
|
|
PipelineConfig,
|
|
StageExecutionType,
|
|
StagePipelineConfig,
|
|
)
|
|
|
|
_PROC = "vllm_omni.model_executor.stage_input_processors.nemotron_voicechat"
|
|
|
|
NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig(
|
|
model_type="nemotron_voicechat",
|
|
model_arch="NemotronVoiceChatThinkerForConditionalGeneration",
|
|
# Named after the alias key so the deploy-yaml stem substring-matches the
|
|
# checkpoint directory name (NVIDIA-NemotronLabs-VoiceChat-11B) during
|
|
# bare-path auto-detection.
|
|
default_deploy_config_name="nemotron_labs_voicechat.yaml",
|
|
stages=(
|
|
StagePipelineConfig(
|
|
stage_id=0,
|
|
model_stage="thinker",
|
|
execution_type=StageExecutionType.LLM_AR,
|
|
input_sources=(),
|
|
owns_tokenizer=True,
|
|
final_output=True,
|
|
final_output_type="text",
|
|
engine_output_type="latent",
|
|
# thinker -> talker sync mode is token-path only: the frame-locked
|
|
# text timeline rides the talker's prompt_token_ids (built by
|
|
# thinker2talker_token_only on the talker stage); no connector
|
|
# payload and no full-payload wait for stage 1. With
|
|
# ``async_chunk: true`` in the deploy yaml, the timeline instead
|
|
# streams incrementally over the connector (cumulative ids.all per
|
|
# thinker step) so the talker starts before the thinker finishes.
|
|
async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk",
|
|
sampling_constraints={"detokenize": True},
|
|
),
|
|
StagePipelineConfig(
|
|
stage_id=1,
|
|
model_stage="talker",
|
|
execution_type=StageExecutionType.LLM_AR,
|
|
model_arch="NemotronVoiceChatTalkerForConditionalGeneration",
|
|
hf_config_name="talker_config",
|
|
input_sources=(0,),
|
|
engine_output_type="latent",
|
|
custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
|
|
async_chunk_process_next_stage_input_func=f"{_PROC}.talker2code2wav_async_chunk",
|
|
sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
|
|
sampling_constraints={"detokenize": False},
|
|
),
|
|
StagePipelineConfig(
|
|
stage_id=2,
|
|
model_stage="code2wav",
|
|
execution_type=StageExecutionType.LLM_GENERATION,
|
|
model_arch="NemotronVoiceChatCode2Wav",
|
|
hf_config_name="code2wav_config",
|
|
input_sources=(1,),
|
|
final_output=True,
|
|
final_output_type="audio",
|
|
engine_output_type="audio",
|
|
sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
|
|
sampling_constraints={"detokenize": False},
|
|
),
|
|
),
|
|
)
|